mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
Merge remote-tracking branch 'origin/main' into litellm_lite_cli_reasoning_efforts
This commit is contained in:
commit
00b2f11196
79 changed files with 5076 additions and 417 deletions
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: database-access
|
||||
description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Database access
|
||||
|
|
|
|||
4
.github/workflows/spend-logs-guard.yml
vendored
4
.github/workflows/spend-logs-guard.yml
vendored
|
|
@ -4,10 +4,6 @@ on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the P
|
|||
pull_request_target:
|
||||
branches:
|
||||
- main
|
||||
pull_request_review:
|
||||
types:
|
||||
- submitted
|
||||
- dismissed
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ so a path rule would put every migration on the owner. This gate reads the diff
|
|||
added line in the watched files mentions the table it passes at once, otherwise it passes only once
|
||||
the owner's latest review on the pull request is an approval that still covers every spend-log line
|
||||
in the head (approvals survive pushes, so a later commit may not add spend-log lines the owner never
|
||||
saw). The workflow re-runs on every review.
|
||||
saw). Reviews do not trigger the workflow, so the check is re-run by hand after approving.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -27,7 +27,8 @@ WATCHED: Final = (
|
|||
"litellm-proxy-extras/litellm_proxy_extras/schema.prisma",
|
||||
)
|
||||
MIGRATIONS_PREFIX: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/"
|
||||
TABLE: Final = re.compile(r"(?<![A-Za-z0-9])LiteLLM_SpendLogs", re.IGNORECASE)
|
||||
TABLE: Final = re.compile(r"(?<![A-Za-z0-9])LiteLLM_SpendLogs(?![A-Za-z0-9])", re.IGNORECASE)
|
||||
DECIDING_STATES: Final = frozenset({"APPROVED", "CHANGES_REQUESTED", "DISMISSED"})
|
||||
|
||||
|
||||
class Reviewer(BaseModel):
|
||||
|
|
@ -73,8 +74,12 @@ def _path_of(section: str) -> str:
|
|||
|
||||
|
||||
def approved_by(reviews: Sequence[Review], login: str) -> bool:
|
||||
"""True when `login`'s most recent review on the pull request is an approval."""
|
||||
states: Final = tuple(review.state for review in reviews if review.user.login == login)
|
||||
"""True when `login`'s most recent review that changes GitHub's approval state is an approval.
|
||||
|
||||
COMMENTED and PENDING reviews leave an earlier approval in force on GitHub, so they are skipped here too.
|
||||
"""
|
||||
own: Final = tuple(review.state for review in reviews if review.user.login == login)
|
||||
states: Final = tuple(state for state in own if state in DECIDING_STATES)
|
||||
return bool(states) and states[-1] == "APPROVED"
|
||||
|
||||
|
||||
|
|
@ -148,7 +153,7 @@ def main(argv: Sequence[str]) -> int:
|
|||
else f"LiteLLM_SpendLogs lines changed after @{args.approver} approved {approval.commit_id}, "
|
||||
f"a new approving review is needed:\n"
|
||||
)
|
||||
sys.stdout.write(f"{reason}{listed}The check re-runs when a review is submitted.\n")
|
||||
sys.stdout.write(f"{reason}{listed}After approving, re-run this job.\n")
|
||||
return 1
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -343,11 +343,16 @@ def btree_max_entry_bytes(block_size: int) -> int:
|
|||
|
||||
|
||||
def _value_bytes(column: ColumnType, identifier: "sql.Identifier") -> "sql.Composed":
|
||||
"""The bytes a value takes in the index entry: its length, or its stored size when
|
||||
Postgres compressed it, since `index_form_tuple` keeps the compressed bytes."""
|
||||
from psycopg import sql
|
||||
|
||||
if column.type_name in _OCTET_LENGTH_TYPES:
|
||||
return sql.SQL("octet_length({})").format(identifier)
|
||||
return sql.SQL("octet_length({}::text)").format(identifier)
|
||||
length: Final = (
|
||||
sql.SQL("octet_length({})").format(identifier)
|
||||
if column.type_name in _OCTET_LENGTH_TYPES
|
||||
else sql.SQL("octet_length({}::text)").format(identifier)
|
||||
)
|
||||
return sql.SQL("least({}, pg_column_size({}))").format(length, identifier)
|
||||
|
||||
|
||||
def _aligned(offset: "sql.Composable", align: str) -> "sql.Composed":
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: rust-error
|
||||
description: Define, wrap, convert, classify, report or test Rust error types in litellm-rust with thiserror. Use when adding an error type or variant, mapping one crate's error into another, turning an error into a gateway response or Python exception, deciding retry or fault from an error, or replacing call sites that build errors from strings
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust errors
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: rust-gateway
|
||||
description: Add or change an HTTP route, extractor, response envelope, stream or router mount in the litellm-rust gateway crates (crates/gateway and crates/gateway-*). Use when adding an endpoint, parsing a request body, choosing an error envelope for an API surface, relaying SSE, or writing a router-level test
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust gateway
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: rust-string-enums
|
||||
description: Define or refactor Rust string-valued enums and their Serde adapters in litellm-rust, using Strum and serde_with while preserving parsing, wire values, and schemas
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust string enums
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: rust-tracing
|
||||
description: Add or change Rust diagnostic tracing in litellm-rust, including route spans, subscriber layers, and Python logger delivery
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust tracing
|
||||
|
|
|
|||
|
|
@ -436,6 +436,7 @@ _POSTGRES_OPERATION_BY_CALL_TYPE: Final[Mapping[str, PostgresOperation]] = Mappi
|
|||
"reset_budget_windows": PostgresOperation("select", None),
|
||||
"write_budget_windows": PostgresOperation("update", None),
|
||||
"roll_window_spend_row": PostgresOperation("update", None),
|
||||
"open_window_spend_row": PostgresOperation("upsert", None),
|
||||
"seed_window_spend": PostgresOperation("select", None),
|
||||
"select_window_spend_rows": PostgresOperation("select", None),
|
||||
"commit_window_spend_updates": PostgresOperation("upsert", None),
|
||||
|
|
|
|||
|
|
@ -2758,7 +2758,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if entry is None or not is_str_keyed_dict(litellm_params):
|
||||
return
|
||||
metadata: Final[object] = litellm_params.get("metadata")
|
||||
bucket: Final = metadata if is_str_keyed_dict(metadata) else {}
|
||||
bucket: Final = {**metadata} if is_str_keyed_dict(metadata) else {}
|
||||
litellm_params["metadata"] = bucket
|
||||
with_guardrail_entry(bucket, entry)
|
||||
|
||||
|
|
|
|||
|
|
@ -268,6 +268,7 @@ async def caller_within_model_budget(
|
|||
user_id=user_id,
|
||||
user_model_max_budget=user_model_max_budget,
|
||||
model=model,
|
||||
team_id=user_api_key_auth.team_id,
|
||||
)
|
||||
except litellm.BudgetExceededError:
|
||||
return "denied"
|
||||
|
|
@ -309,6 +310,7 @@ async def caller_within_model_budget(
|
|||
end_user_id=end_user_id,
|
||||
end_user_model_max_budget=end_user_model_max_budget,
|
||||
model=model,
|
||||
team_id=user_api_key_auth.team_id,
|
||||
)
|
||||
except litellm.BudgetExceededError:
|
||||
return "denied"
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import ssl
|
|||
import sys
|
||||
import typing
|
||||
import urllib.request
|
||||
import weakref
|
||||
from collections.abc import Callable, Generator
|
||||
from typing import ClassVar, Final
|
||||
|
||||
|
|
@ -109,6 +110,24 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]:
|
|||
raise mapped_exc(message) from exc
|
||||
|
||||
|
||||
def _release_transport_anchor(_transport: object) -> None:
|
||||
"""Release a transport anchored to a response stream. See ``_anchor_transport_to``."""
|
||||
|
||||
|
||||
def _anchor_transport_to(stream: "AiohttpResponseStream", transport: object) -> None:
|
||||
"""Keep the transport alive for as long as the stream can still read.
|
||||
|
||||
The transport finalizer closes the session it owns, and nothing on the path from a
|
||||
live response back to that session references the transport, so a client dropped
|
||||
mid-stream would close the session under the body. ``weakref.finalize`` holds the
|
||||
transport in its own registry rather than on the stream, which keeps the transport
|
||||
out of the response's reference cycle (an exception carrying the response, for one)
|
||||
so it is finalized by refcount once the anchor drops and can still schedule an async
|
||||
close, instead of being collected in the same cyclic pass as the session it closes.
|
||||
"""
|
||||
weakref.finalize(stream, _release_transport_anchor, transport)
|
||||
|
||||
|
||||
class AiohttpResponseStream(httpx.AsyncByteStream):
|
||||
CHUNK_SIZE = 1024 * 16
|
||||
|
||||
|
|
@ -463,10 +482,12 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
|||
# Re-raise if it's a different RuntimeError
|
||||
raise
|
||||
|
||||
stream: Final = AiohttpResponseStream(response)
|
||||
_anchor_transport_to(stream, self)
|
||||
return httpx.Response(
|
||||
status_code=response.status,
|
||||
headers=response.headers,
|
||||
stream=AiohttpResponseStream(response),
|
||||
stream=stream,
|
||||
request=request,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -2743,10 +2743,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
chunk_indices = support.get("groundingChunkIndices", [])
|
||||
|
||||
if start_index is not None and end_index is not None and chunk_indices:
|
||||
# Use the first chunk's URL for the annotation
|
||||
first_chunk_idx = chunk_indices[0]
|
||||
if first_chunk_idx in chunk_to_uri_map:
|
||||
uri_info = chunk_to_uri_map[first_chunk_idx]
|
||||
for chunk_index in chunk_indices:
|
||||
if chunk_index not in chunk_to_uri_map:
|
||||
continue
|
||||
uri_info = chunk_to_uri_map[chunk_index]
|
||||
|
||||
url_citation: ChatCompletionAnnotationURLCitation = {
|
||||
"start_index": start_index,
|
||||
|
|
|
|||
|
|
@ -21589,7 +21589,13 @@
|
|||
],
|
||||
"supports_vision": true,
|
||||
"supports_audio_input": true,
|
||||
"supports_video_input": true
|
||||
"supports_video_input": true,
|
||||
"input_cost_per_token_priority": 5.6252e-07,
|
||||
"output_cost_per_token_priority": 3.37498e-06,
|
||||
"cache_read_input_token_cost_priority": 5.628e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.6252e-07,
|
||||
"input_cost_per_audio_token_priority": 1.12497e-06,
|
||||
"cache_read_input_audio_token_cost_priority": 1.1249e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-1-flash-image": {
|
||||
"cache_creation_input_token_cost": 6.2503e-07,
|
||||
|
|
@ -21676,7 +21682,14 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.50002e-06,
|
||||
"output_cost_per_token_priority": 2.699998e-05,
|
||||
"cache_read_input_token_cost_priority": 4.5003e-07,
|
||||
"cache_creation_input_token_cost_priority": 4.50002e-06,
|
||||
"input_cost_per_token_above_200k_tokens_priority": 8.99997e-06,
|
||||
"output_cost_per_token_above_200k_tokens_priority": 4.049997e-05,
|
||||
"cache_read_input_token_cost_above_200k_tokens_priority": 8.9999e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-flash": {
|
||||
"cache_creation_input_token_cost": 6.2503e-07,
|
||||
|
|
@ -21710,7 +21723,13 @@
|
|||
],
|
||||
"supports_vision": true,
|
||||
"supports_audio_input": true,
|
||||
"supports_video_input": true
|
||||
"supports_video_input": true,
|
||||
"input_cost_per_token_priority": 1.12497e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.1249e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.12497e-06,
|
||||
"input_cost_per_audio_token_priority": 2.25001e-06,
|
||||
"cache_read_input_audio_token_cost_priority": 2.2498e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-pro": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -21757,7 +21776,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-7-flash": {
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
|
|
@ -21785,10 +21808,14 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-6-flash": {
|
||||
"cache_creation_input_token_cost": 1.87502e-06,
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
"cache_read_input_token_cost": 7.497e-08,
|
||||
"input_cost_per_token": 7.4998e-07,
|
||||
"input_dbu_cost_per_token": 1.0714e-05,
|
||||
|
|
@ -21816,7 +21843,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-5-flash": {
|
||||
"cache_creation_input_token_cost": 1.87502e-06,
|
||||
|
|
@ -21846,7 +21877,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 3.37498e-06,
|
||||
"output_cost_per_token_priority": 2.025002e-05,
|
||||
"cache_read_input_token_cost_priority": 3.3747e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.37498e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-5-flash-lite": {
|
||||
"cache_creation_input_token_cost": 3.7499e-07,
|
||||
|
|
@ -21878,7 +21913,11 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 6.7501e-07,
|
||||
"output_cost_per_token_priority": 5.62499e-06,
|
||||
"cache_read_input_token_cost_priority": 6.748e-08,
|
||||
"cache_creation_input_token_cost_priority": 6.7501e-07
|
||||
},
|
||||
"databricks/databricks-gemma-3-12b": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -21921,7 +21960,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
"supports_vision": false,
|
||||
"input_cost_per_token_priority": 2.45e-06,
|
||||
"output_cost_per_token_priority": 7.7e-06,
|
||||
"cache_read_input_token_cost_priority": 4.55e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.45e-06
|
||||
},
|
||||
"databricks/databricks-glm-5-3": {
|
||||
"cache_creation_input_token_cost": 1.4e-06,
|
||||
|
|
@ -21951,7 +21994,11 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 2.45e-06,
|
||||
"output_cost_per_token_priority": 7.7e-06,
|
||||
"cache_read_input_token_cost_priority": 4.55e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.45e-06
|
||||
},
|
||||
"databricks/databricks-glm-5-3-flash": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -21980,7 +22027,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 2.625e-07,
|
||||
"output_cost_per_token_priority": 8.75e-07,
|
||||
"cache_read_input_token_cost_priority": 5.25e-08,
|
||||
"cache_creation_input_token_cost_priority": 2.625e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -21999,7 +22050,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 2.49998e-06,
|
||||
"output_cost_per_token_priority": 1.999998e-05,
|
||||
"cache_read_input_token_cost_priority": 2.4997e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-1": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -22018,7 +22073,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 2.49998e-06,
|
||||
"output_cost_per_token_priority": 1.999998e-05,
|
||||
"cache_read_input_token_cost_priority": 2.4997e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-2": {
|
||||
"cache_creation_input_token_cost": 1.75e-06,
|
||||
|
|
@ -22037,7 +22096,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 3.5e-06,
|
||||
"output_cost_per_token_priority": 2.8e-05,
|
||||
"cache_read_input_token_cost_priority": 3.5e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.5e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-3-codex": {
|
||||
"cache_creation_input_token_cost": 1.75e-06,
|
||||
|
|
@ -22056,7 +22119,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 3.5e-06,
|
||||
"output_cost_per_token_priority": 2.8e-05,
|
||||
"cache_read_input_token_cost_priority": 3.5e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.5e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -22087,7 +22154,14 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.00003e-06,
|
||||
"output_cost_per_token_priority": 2.999997e-05,
|
||||
"cache_read_input_token_cost_priority": 5.0001e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 5.00003e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.999997e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 5.0001e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-mini": {
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
|
|
@ -22115,7 +22189,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.50003e-06,
|
||||
"output_cost_per_token_priority": 8.99997e-06,
|
||||
"cache_read_input_token_cost_priority": 1.5001e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.50003e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-nano": {
|
||||
"cache_creation_input_token_cost": 1.9999e-07,
|
||||
|
|
@ -22175,7 +22253,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 8.00002e-06,
|
||||
"output_cost_per_token_priority": 4.000003e-05,
|
||||
"cache_read_input_token_cost_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_priority": 9.99999e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 1.599997e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 6.000001e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.59999e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.999998e-05
|
||||
},
|
||||
"databricks/databricks-gpt-5-6-terra": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -22207,7 +22293,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 2.399999e-05,
|
||||
"cache_read_input_token_cost_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 3.600002e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-6-luna": {
|
||||
"cache_creation_input_token_cost": 2.4997e-07,
|
||||
|
|
@ -22239,7 +22333,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 3.9998e-07,
|
||||
"output_cost_per_token_priority": 2.40002e-06,
|
||||
"cache_read_input_token_cost_priority": 3.997e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.0001e-07,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.0003e-07,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 3.60003e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.001e-08,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.00002e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-5": {
|
||||
"cache_creation_input_token_cost": 5.00003e-06,
|
||||
|
|
@ -22270,7 +22372,14 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.249997e-05,
|
||||
"output_cost_per_token_priority": 7.50001e-05,
|
||||
"cache_read_input_token_cost_priority": 1.24999e-06,
|
||||
"cache_creation_input_token_cost_priority": 1.249997e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 1.249997e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 7.50001e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.24999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-5-pro": {
|
||||
"cache_creation_input_token_cost": 2.999997e-05,
|
||||
|
|
@ -22319,7 +22428,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 4.5003e-07,
|
||||
"output_cost_per_token_priority": 3.60003e-06,
|
||||
"cache_read_input_token_cost_priority": 4.501e-08,
|
||||
"cache_creation_input_token_cost_priority": 4.5003e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5-nano": {
|
||||
"cache_creation_input_token_cost": 4.998e-08,
|
||||
|
|
@ -22469,7 +22582,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.25e-06,
|
||||
"output_cost_per_token_priority": 2.625e-05,
|
||||
"cache_read_input_token_cost_priority": 5.25e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.25e-06
|
||||
},
|
||||
"databricks/system.ai.kimi-k3": {
|
||||
"litellm_provider": "databricks",
|
||||
|
|
@ -22551,7 +22668,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 3.85e-07,
|
||||
"output_cost_per_token_priority": 3.85e-06,
|
||||
"cache_read_input_token_cost_priority": 3.85e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.85e-07
|
||||
},
|
||||
"databricks/databricks-qwen3-next-80b-a3b-instruct": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -77774,7 +77895,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.25e-07,
|
||||
"output_cost_per_token_priority": 2.1e-06,
|
||||
"cache_read_input_token_cost_priority": 5.25e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.25e-07
|
||||
},
|
||||
"databricks/databricks-deepseek-v4-1-flash-1017": {
|
||||
"cache_creation_input_token_cost": 3.0002e-07,
|
||||
|
|
@ -77897,7 +78022,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.50002e-06,
|
||||
"output_cost_per_token_priority": 2.250003e-05,
|
||||
"cache_read_input_token_cost_priority": 2.2498e-07,
|
||||
"cache_creation_input_token_cost_priority": 4.50002e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-pro": {
|
||||
"cache_creation_input_token_cost": 2.999997e-05,
|
||||
|
|
@ -77959,7 +78088,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 1.99997e-05,
|
||||
"cache_read_input_token_cost_priority": 1.9999e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-6-1-terra": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -77991,7 +78128,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.99997e-06,
|
||||
"output_cost_per_token_priority": 9.99999e-06,
|
||||
"cache_read_input_token_cost_priority": 1.0003e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4.00001e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 1.500002e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06
|
||||
},
|
||||
"databricks/databricks-gpt-6-astra": {
|
||||
"cache_creation_input_token_cost": 1.249997e-05,
|
||||
|
|
@ -78023,7 +78168,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.999998e-05,
|
||||
"output_cost_per_token_priority": 9.99999e-05,
|
||||
"cache_read_input_token_cost_priority": 1.99997e-06,
|
||||
"cache_creation_input_token_cost_priority": 2.500001e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4.000003e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 0.0001500002,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05
|
||||
},
|
||||
"databricks/databricks-gpt-6-luna": {
|
||||
"cache_creation_input_token_cost": 1.2502e-07,
|
||||
|
|
@ -78055,7 +78208,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.9999e-07,
|
||||
"output_cost_per_token_priority": 1.0003e-06,
|
||||
"cache_read_input_token_cost_priority": 2.002e-08,
|
||||
"cache_creation_input_token_cost_priority": 2.4997e-07,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 3.9998e-07,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 1.5001e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07
|
||||
},
|
||||
"databricks/databricks-gpt-6-sol": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -78087,7 +78248,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 1.99997e-05,
|
||||
"cache_read_input_token_cost_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-image-1": {
|
||||
"cache_creation_input_token_cost": 5.00003e-06,
|
||||
|
|
|
|||
|
|
@ -3487,6 +3487,13 @@
|
|||
"title": "Entities",
|
||||
"type": "object"
|
||||
},
|
||||
"guardrails": {
|
||||
"additionalProperties": {
|
||||
"$ref": "#/components/schemas/GuardrailMetric"
|
||||
},
|
||||
"title": "Guardrails",
|
||||
"type": "object"
|
||||
},
|
||||
"mcp_servers": {
|
||||
"additionalProperties": {
|
||||
"$ref": "#/components/schemas/MetricWithMetadata"
|
||||
|
|
@ -3802,6 +3809,74 @@
|
|||
"title": "ExportType",
|
||||
"type": "string"
|
||||
},
|
||||
"GuardrailKeyMetric": {
|
||||
"properties": {
|
||||
"metadata": {
|
||||
"$ref": "#/components/schemas/KeyMetadata"
|
||||
},
|
||||
"metrics": {
|
||||
"$ref": "#/components/schemas/GuardrailSpendMetrics"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"metrics",
|
||||
"metadata"
|
||||
],
|
||||
"title": "GuardrailKeyMetric",
|
||||
"type": "object"
|
||||
},
|
||||
"GuardrailMetadata": {
|
||||
"properties": {
|
||||
"guardrail_name": {
|
||||
"title": "Guardrail Name",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"guardrail_name"
|
||||
],
|
||||
"title": "GuardrailMetadata",
|
||||
"type": "object"
|
||||
},
|
||||
"GuardrailMetric": {
|
||||
"properties": {
|
||||
"api_key_breakdown": {
|
||||
"additionalProperties": {
|
||||
"$ref": "#/components/schemas/GuardrailKeyMetric"
|
||||
},
|
||||
"title": "Api Key Breakdown",
|
||||
"type": "object"
|
||||
},
|
||||
"metadata": {
|
||||
"$ref": "#/components/schemas/GuardrailMetadata"
|
||||
},
|
||||
"metrics": {
|
||||
"$ref": "#/components/schemas/GuardrailSpendMetrics"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"metrics",
|
||||
"metadata"
|
||||
],
|
||||
"title": "GuardrailMetric",
|
||||
"type": "object"
|
||||
},
|
||||
"GuardrailSpendMetrics": {
|
||||
"properties": {
|
||||
"spend": {
|
||||
"default": 0.0,
|
||||
"title": "Spend",
|
||||
"type": "number"
|
||||
},
|
||||
"untracked_units": {
|
||||
"default": 0,
|
||||
"title": "Untracked Units",
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"title": "GuardrailSpendMetrics",
|
||||
"type": "object"
|
||||
},
|
||||
"HTTPAuthSecurityScheme": {
|
||||
"description": "Defines a security scheme using HTTP authentication.",
|
||||
"properties": {
|
||||
|
|
|
|||
|
|
@ -370,7 +370,7 @@ class _KeyModelBudgetLimiter(Protocol):
|
|||
|
||||
class _UserModelBudgetLimiter(Protocol):
|
||||
async def is_user_within_model_budget(
|
||||
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str
|
||||
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None
|
||||
) -> bool: ...
|
||||
|
||||
|
||||
|
|
@ -388,11 +388,15 @@ class _EveryModelBudgetLimiter(Protocol):
|
|||
async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ...
|
||||
|
||||
async def is_user_within_model_budget(
|
||||
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str
|
||||
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None
|
||||
) -> bool: ...
|
||||
|
||||
async def is_end_user_within_model_budget(
|
||||
self, end_user_id: str, end_user_model_max_budget: Mapping[str, object], model: str
|
||||
self,
|
||||
end_user_id: str,
|
||||
end_user_model_max_budget: Mapping[str, object],
|
||||
model: str,
|
||||
team_id: str | None = None,
|
||||
) -> bool: ...
|
||||
|
||||
async def is_team_within_model_budget(
|
||||
|
|
@ -503,6 +507,7 @@ async def _check_user_model_budget(
|
|||
user_id=valid_token.user_id,
|
||||
user_model_max_budget=user_model_max_budget,
|
||||
model=model_name,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -547,6 +552,7 @@ async def _is_within_every_model_budget(
|
|||
user_id=valid_token.user_id,
|
||||
user_model_max_budget=user_model_max_budget,
|
||||
model=model,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
if (
|
||||
valid_token.end_user_id is not None
|
||||
|
|
@ -557,6 +563,7 @@ async def _is_within_every_model_budget(
|
|||
end_user_id=valid_token.end_user_id,
|
||||
end_user_model_max_budget=end_user_model_max_budget,
|
||||
model=model,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
if valid_token.team_id is not None and isinstance(team_model_max_budget, Mapping) and team_model_max_budget:
|
||||
await model_max_budget_limiter.is_team_within_model_budget(
|
||||
|
|
@ -3154,6 +3161,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
|
|||
end_user_id=valid_token.end_user_id,
|
||||
end_user_model_max_budget=end_user_mmb,
|
||||
model=model_name,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
|
||||
# Check 6: Additional Common Checks across jwt + key auth
|
||||
|
|
@ -4733,6 +4741,7 @@ async def _run_post_custom_auth_checks(
|
|||
end_user_id=valid_token.end_user_id,
|
||||
end_user_model_max_budget=end_user_mmb,
|
||||
model=model_name,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
|
||||
# team / user / end_user / project context objects are fetched by
|
||||
|
|
|
|||
|
|
@ -64,6 +64,16 @@ _ROLL_WINDOW_SPEND_SQL: Final = (
|
|||
"AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
||||
_OPEN_WINDOW_SPEND_SQL: Final = (
|
||||
'INSERT INTO "LiteLLM_BudgetWindowSpend" '
|
||||
"(entity_type, entity_id, window_duration, window_start, spend, created_at, updated_at) "
|
||||
"VALUES ($1, $2, $3, ($4::timestamptz AT TIME ZONE 'UTC'), 0, "
|
||||
"($4::timestamptz AT TIME ZONE 'UTC'), ($4::timestamptz AT TIME ZONE 'UTC')) "
|
||||
"ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET "
|
||||
"window_start = EXCLUDED.window_start, spend = 0, updated_at = EXCLUDED.updated_at "
|
||||
"WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
||||
_SEED_FROM_SPEND_LOGS_KEY_SQL: Final = (
|
||||
"SELECT COALESCE(SUM(spend), 0.0) AS total, "
|
||||
"COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch "
|
||||
|
|
@ -389,3 +399,29 @@ async def roll_window_spend_row(
|
|||
to_naive_utc(new_window_start),
|
||||
to_naive_utc(datetime.now(timezone.utc)),
|
||||
)
|
||||
|
||||
|
||||
async def open_window_spend_row(
|
||||
prisma_client: "PrismaClient",
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_duration: str,
|
||||
window_start: datetime,
|
||||
replaces_window_start: datetime | None,
|
||||
) -> None:
|
||||
"""Start a window at `window_start` unless another writer already started one.
|
||||
|
||||
The row is created when it does not exist, and moved onto the new window only
|
||||
while it still holds `replaces_window_start`, the ended window the caller saw.
|
||||
A writer that lost the race leaves the winner's window in place, so reading
|
||||
the row back gives every writer the same window start.
|
||||
"""
|
||||
async with db_span("open_window_spend_row", "LiteLLM_BudgetWindowSpend"):
|
||||
await prisma_client.db.execute_raw(
|
||||
_OPEN_WINDOW_SPEND_SQL,
|
||||
entity_type,
|
||||
entity_id,
|
||||
window_duration,
|
||||
to_naive_utc(window_start),
|
||||
None if replaces_window_start is None else to_naive_utc(replaces_window_start),
|
||||
)
|
||||
|
|
|
|||
176
litellm/proxy/hooks/model_budget_shared_windows.py
Normal file
176
litellm/proxy/hooks/model_budget_shared_windows.py
Normal file
|
|
@ -0,0 +1,176 @@
|
|||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy.db.budget_window_spend_writer import open_window_spend_row
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
WindowSpendTransaction,
|
||||
build_window_spend_transaction,
|
||||
)
|
||||
from litellm.repositories.table_repositories import BudgetWindowSpendRepository
|
||||
from litellm.router_strategy.budget_limiter import budget_window_seconds
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
MODEL_BUDGET_WINDOW_ENTITY_TYPE: Final = "model_max_budget"
|
||||
_READ_TTL_SECONDS: Final = 5
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SharedWindow:
|
||||
started_at: float
|
||||
spend: float
|
||||
|
||||
def is_open(self, budget_duration: str, now: float) -> bool:
|
||||
return now < self.started_at + budget_window_seconds(self.started_at, budget_duration)
|
||||
|
||||
|
||||
class LocalTTLCache(Protocol):
|
||||
def get_cache(self, key: str) -> object: ...
|
||||
|
||||
def set_cache(self, key: str, value: object, *, ttl: float) -> None: ...
|
||||
|
||||
def delete_cache(self, key: str) -> None: ...
|
||||
|
||||
|
||||
class SharedWindows(Protocol):
|
||||
async def read(self, spend_key: str, budget_duration: str) -> SharedWindow | None: ...
|
||||
|
||||
async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None: ...
|
||||
|
||||
async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None: ...
|
||||
|
||||
|
||||
PrismaClientSource: TypeAlias = Callable[[], "PrismaClient | None"]
|
||||
WindowSpendSink: TypeAlias = Callable[[WindowSpendTransaction], Awaitable[None]]
|
||||
|
||||
|
||||
def _proxy_prisma_client() -> "PrismaClient | None":
|
||||
try:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
except ImportError:
|
||||
return None
|
||||
return prisma_client
|
||||
|
||||
|
||||
async def _queue_on_proxy_spend_writer(transaction: WindowSpendTransaction) -> None:
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
await proxy_logging_obj.db_spend_update_writer.window_spend_update_queue.add_update(transaction)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Reading:
|
||||
window: SharedWindow | None
|
||||
|
||||
|
||||
class DatabaseSharedWindows:
|
||||
def __init__(
|
||||
self,
|
||||
prisma_client_source: PrismaClientSource = _proxy_prisma_client,
|
||||
sink: WindowSpendSink = _queue_on_proxy_spend_writer,
|
||||
read_ttl_seconds: int = _READ_TTL_SECONDS,
|
||||
) -> None:
|
||||
self._prisma_client_source = prisma_client_source
|
||||
self._sink = sink
|
||||
self._read_ttl_seconds = read_ttl_seconds
|
||||
self._readings: LocalTTLCache = InMemoryCache(max_size_in_memory=10_000, default_ttl=read_ttl_seconds)
|
||||
|
||||
async def read(self, spend_key: str, budget_duration: str) -> SharedWindow | None:
|
||||
cached: Final = self._readings.get_cache(key=spend_key)
|
||||
if isinstance(cached, _Reading):
|
||||
return cached.window
|
||||
window: Final = await self._read_row(spend_key=spend_key, budget_duration=budget_duration)
|
||||
self._readings.set_cache(key=spend_key, value=_Reading(window=window), ttl=self._read_ttl_seconds)
|
||||
return window
|
||||
|
||||
async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None:
|
||||
"""When the window every replica shares started, opening one at `now` if none is open.
|
||||
|
||||
Without a database every window is this replica's own, so it starts at `now`. None means the database did
|
||||
not answer, so this replica cannot know which window the others are in yet.
|
||||
"""
|
||||
prisma_client: Final = self._prisma_client_source()
|
||||
if prisma_client is None:
|
||||
return now
|
||||
try:
|
||||
seen: Final = await _fetch_row(
|
||||
prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=True
|
||||
)
|
||||
if seen is not None and seen.is_open(budget_duration=budget_duration, now=now):
|
||||
return seen.started_at
|
||||
await open_window_spend_row(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=MODEL_BUDGET_WINDOW_ENTITY_TYPE,
|
||||
entity_id=spend_key,
|
||||
window_duration=budget_duration,
|
||||
window_start=datetime.fromtimestamp(now, tz=timezone.utc),
|
||||
replaces_window_start=None
|
||||
if seen is None
|
||||
else datetime.fromtimestamp(seen.started_at, tz=timezone.utc),
|
||||
)
|
||||
opened: Final = await _fetch_row(
|
||||
prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=True
|
||||
)
|
||||
except Exception: # noqa: BLE001 # an unreachable database leaves this replica on its own window
|
||||
verbose_proxy_logger.exception("Unable to open the shared model budget window")
|
||||
return None
|
||||
if opened is None or not opened.is_open(budget_duration=budget_duration, now=now):
|
||||
return None
|
||||
return opened.started_at
|
||||
|
||||
async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None:
|
||||
if self._prisma_client_source() is None:
|
||||
return
|
||||
try:
|
||||
await self._sink(
|
||||
build_window_spend_transaction(
|
||||
entity_type=MODEL_BUDGET_WINDOW_ENTITY_TYPE,
|
||||
entity_id=spend_key,
|
||||
window_duration=budget_duration,
|
||||
window_start=datetime.fromtimestamp(started_at, tz=timezone.utc),
|
||||
spend=spend,
|
||||
)
|
||||
)
|
||||
except Exception: # noqa: BLE001 # spend tracking must never fail the success callback
|
||||
verbose_proxy_logger.debug("Unable to queue shared model budget spend", exc_info=True)
|
||||
|
||||
async def _read_row(self, spend_key: str, budget_duration: str) -> SharedWindow | None:
|
||||
prisma_client: Final = self._prisma_client_source()
|
||||
if prisma_client is None:
|
||||
return None
|
||||
try:
|
||||
return await _fetch_row(
|
||||
prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=False
|
||||
)
|
||||
except Exception: # noqa: BLE001 # an unreadable row leaves enforcement on this replica's own counter
|
||||
verbose_proxy_logger.exception("Unable to read shared model budget spend")
|
||||
return None
|
||||
|
||||
|
||||
async def _fetch_row(
|
||||
prisma_client: "PrismaClient", spend_key: str, budget_duration: str, use_writer: bool
|
||||
) -> SharedWindow | None:
|
||||
row: Final = await BudgetWindowSpendRepository(prisma_client, use_writer=use_writer).table.find_unique(
|
||||
where={
|
||||
"entity_type_entity_id_window_duration": {
|
||||
"entity_type": MODEL_BUDGET_WINDOW_ENTITY_TYPE,
|
||||
"entity_id": spend_key,
|
||||
"window_duration": budget_duration,
|
||||
}
|
||||
}
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
window_start: Final = row.window_start
|
||||
started_at: Final = (
|
||||
window_start if window_start.tzinfo is not None else window_start.replace(tzinfo=timezone.utc)
|
||||
).timestamp()
|
||||
return SharedWindow(started_at=started_at, spend=float(row.spend or 0.0))
|
||||
|
||||
|
||||
PROXY_SHARED_WINDOWS: Final = DatabaseSharedWindows()
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -3,6 +3,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence
|
|||
from collections.abc import Set as AbstractSet
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import date, datetime, timedelta
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, NoReturn, Protocol
|
||||
|
||||
|
|
@ -28,6 +29,10 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
|||
DailySpendData,
|
||||
DailySpendMetadata,
|
||||
GroupedData,
|
||||
GuardrailKeyMetric,
|
||||
GuardrailMetadata,
|
||||
GuardrailMetric,
|
||||
GuardrailSpendMetrics,
|
||||
KeyMetadata,
|
||||
KeyMetricWithMetadata,
|
||||
MetricWithMetadata,
|
||||
|
|
@ -39,6 +44,7 @@ from litellm.types.repositories.daily_activity import (
|
|||
DailyActivityTable,
|
||||
EntityRollupRow,
|
||||
GroupingSetsRow,
|
||||
GuardrailSpendRow,
|
||||
KeyMetadataRow,
|
||||
RollupMetricsRow,
|
||||
SpendLogsWindow,
|
||||
|
|
@ -1013,17 +1019,68 @@ def _fold_entity_rollups_sync(
|
|||
)
|
||||
|
||||
|
||||
def _guardrail_metric(day: DailySpendData, rows: Sequence[GuardrailSpendRow]) -> GuardrailMetric:
|
||||
return GuardrailMetric(
|
||||
metrics=GuardrailSpendMetrics(
|
||||
spend=sum(row.spend for row in rows),
|
||||
untracked_units=sum(row.untracked_units for row in rows),
|
||||
),
|
||||
metadata=GuardrailMetadata(guardrail_name=rows[0].guardrail_name),
|
||||
api_key_breakdown={
|
||||
row.api_key: GuardrailKeyMetric(
|
||||
metrics=GuardrailSpendMetrics(spend=row.spend, untracked_units=row.untracked_units),
|
||||
metadata=day.breakdown.api_keys[row.api_key].metadata,
|
||||
)
|
||||
for row in rows
|
||||
if row.api_key in day.breakdown.api_keys
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _day_with_guardrails(day: DailySpendData, rows: Sequence[GuardrailSpendRow]) -> DailySpendData:
|
||||
by_guardrail: Final = groupby(sorted(rows, key=lambda row: row.guardrail_id), key=lambda row: row.guardrail_id)
|
||||
guardrails: Final = {guardrail_id: _guardrail_metric(day, tuple(group)) for guardrail_id, group in by_guardrail}
|
||||
return day.model_copy(update={"breakdown": day.breakdown.model_copy(update={"guardrails": guardrails})})
|
||||
|
||||
|
||||
def _with_guardrail_breakdown(
|
||||
results: Sequence[DailySpendData], rows: Sequence[GuardrailSpendRow]
|
||||
) -> tuple[DailySpendData, ...]:
|
||||
rows_by_date: Final = MappingProxyType(
|
||||
{day: tuple(group) for day, group in groupby(sorted(rows, key=lambda row: row.date), key=lambda row: row.date)}
|
||||
)
|
||||
return tuple(_day_with_guardrails(day, rows_by_date.get(day.date.strftime("%Y-%m-%d"), ())) for day in results)
|
||||
|
||||
|
||||
async def _guardrail_spend_or_empty(
|
||||
repository: DailyActivityRepository, scope: DailyActivityScope
|
||||
) -> tuple[GuardrailSpendRow, ...]:
|
||||
from prisma.errors import PrismaError
|
||||
|
||||
try:
|
||||
return await repository.guardrail_spend(scope)
|
||||
except PrismaError as exc:
|
||||
verbose_proxy_logger.warning("Returning daily activity without its guardrail breakdown: %s", exc)
|
||||
return ()
|
||||
|
||||
|
||||
async def get_daily_activity_aggregated(
|
||||
repository: DailyActivityRepository,
|
||||
scope: DailyActivityScope,
|
||||
*,
|
||||
entity_metadata_field: Mapping[str, dict[str, object]] | None = None,
|
||||
include_entity_breakdown: bool = False,
|
||||
include_guardrail_breakdown: bool = False,
|
||||
api_key_limit: int = constants.USAGE_TOP_API_KEYS_DEFAULT,
|
||||
) -> SpendAnalyticsPaginatedResponse:
|
||||
try:
|
||||
aggregated_rows: Final = await repository.aggregated(
|
||||
scope, include_entity_breakdown=include_entity_breakdown, api_key_limit=api_key_limit
|
||||
aggregated_rows, guardrail_rows = await asyncio.gather(
|
||||
repository.aggregated(
|
||||
scope, include_entity_breakdown=include_entity_breakdown, api_key_limit=api_key_limit
|
||||
),
|
||||
_guardrail_spend_or_empty(repository, scope)
|
||||
if include_guardrail_breakdown and not scope.model
|
||||
else asyncio.sleep(0, result=()),
|
||||
)
|
||||
records: Final = aggregated_rows.grouping_rows
|
||||
aggregated: Final = await _aggregate_grouping_sets_records(
|
||||
|
|
@ -1058,8 +1115,13 @@ async def get_daily_activity_aggregated(
|
|||
api_key_metadata=entity_key_metadata,
|
||||
entity_metadata_field=entity_metadata_field,
|
||||
)
|
||||
results: Final = (
|
||||
list(await asyncio.to_thread(_with_guardrail_breakdown, aggregated["results"], guardrail_rows))
|
||||
if guardrail_rows
|
||||
else aggregated["results"]
|
||||
)
|
||||
return SpendAnalyticsPaginatedResponse(
|
||||
results=aggregated["results"],
|
||||
results=results,
|
||||
metadata=DailySpendMetadata(
|
||||
total_spend=aggregated["totals"].spend,
|
||||
total_flat_cost=aggregated["totals"].flat_cost,
|
||||
|
|
|
|||
|
|
@ -300,6 +300,7 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver,
|
|||
resolved.scope,
|
||||
entity_metadata_field=resolved.entity_metadata,
|
||||
include_entity_breakdown=resolver.include_entity_breakdown,
|
||||
include_guardrail_breakdown=resolver.include_guardrail_breakdown,
|
||||
api_key_limit=api_key_limit,
|
||||
)
|
||||
if resolver.entity != "team":
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ class EntityScopeResolver:
|
|||
resolve: EntityScopeResolve
|
||||
include_entity_breakdown: bool
|
||||
operation_names: OperationNames
|
||||
include_guardrail_breakdown: bool = False
|
||||
|
||||
|
||||
def _query_ids(value: str | None) -> tuple[str, ...] | None:
|
||||
|
|
@ -364,6 +365,7 @@ USER_RESOLVER = EntityScopeResolver(
|
|||
query=_user_query,
|
||||
resolve=_resolve_user,
|
||||
include_entity_breakdown=False,
|
||||
include_guardrail_breakdown=True,
|
||||
operation_names=MappingProxyType(
|
||||
{
|
||||
"aggregated": "get_user_daily_activity_aggregated",
|
||||
|
|
@ -385,6 +387,7 @@ TEAM_RESOLVER = EntityScopeResolver(
|
|||
query=_team_query,
|
||||
resolve=_resolve_team,
|
||||
include_entity_breakdown=True,
|
||||
include_guardrail_breakdown=True,
|
||||
operation_names=MappingProxyType(
|
||||
{
|
||||
"aggregated": "get_team_daily_activity_aggregated",
|
||||
|
|
@ -423,6 +426,7 @@ ORGANIZATION_RESOLVER = EntityScopeResolver(
|
|||
query=_organization_query,
|
||||
resolve=_resolve_organization,
|
||||
include_entity_breakdown=True,
|
||||
include_guardrail_breakdown=True,
|
||||
operation_names=MappingProxyType(
|
||||
{
|
||||
"aggregated": "get_organization_daily_activity_aggregated",
|
||||
|
|
|
|||
|
|
@ -499,7 +499,7 @@ async def new_user(
|
|||
- permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
|
||||
- metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
- mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
|
||||
|
|
@ -1702,7 +1702,7 @@ async def user_update(
|
|||
- permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
|
||||
- metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
- mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
|
||||
|
|
|
|||
|
|
@ -2328,7 +2328,7 @@ async def generate_key_fn(
|
|||
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
|
||||
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only.
|
||||
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||||
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||||
|
|
@ -2538,7 +2538,7 @@ async def generate_service_account_key_fn(
|
|||
- metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
- guardrails: Optional[List[str]] - List of active guardrails for the key
|
||||
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||||
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||||
|
|
@ -3855,7 +3855,7 @@ async def update_key_fn(
|
|||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
|
||||
- spend: Optional[float] - Amount spent by key
|
||||
- max_budget: Optional[float] - Max budget for key
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
- soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget.
|
||||
|
|
@ -6303,7 +6303,7 @@ async def regenerate_key_fn(
|
|||
- tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
|
||||
- spend: Optional[float] - Amount spent by key
|
||||
- max_budget: Optional[float] - Max budget for key
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
- soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ All /team management endpoints
|
|||
|
||||
import asyncio
|
||||
import copy
|
||||
import itertools
|
||||
import json
|
||||
import math
|
||||
import traceback
|
||||
|
|
@ -139,7 +140,12 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
from litellm.proxy.db.db_span import db_span
|
||||
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
RoutedModels,
|
||||
build_model_max_budget_usage,
|
||||
entries_outranked_by,
|
||||
entry_names_model,
|
||||
is_enforced_budget_entry,
|
||||
proxy_routed_models,
|
||||
resolve_model_budget,
|
||||
)
|
||||
from litellm.proxy.management.teams.authz import TEAM_OR_ORG_ADMIN, TeamRole, is_team_admin, team_access_denied
|
||||
|
|
@ -1467,8 +1473,10 @@ def _check_team_model_budget_update_authority(
|
|||
data: UpdateTeamRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
existing_model_max_budget: Mapping[str, object] | None,
|
||||
routed_models_source: Callable[[str | None], RoutedModels] = proxy_routed_models,
|
||||
) -> None:
|
||||
"""Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap."""
|
||||
"""Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap, or
|
||||
add an entry that takes precedence over one."""
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return
|
||||
if "model_max_budget" not in data.model_fields_set or not existing_model_max_budget:
|
||||
|
|
@ -1487,29 +1495,97 @@ def _check_team_model_budget_update_authority(
|
|||
)
|
||||
},
|
||||
)
|
||||
for model_name, proposed in requested.items():
|
||||
governing = resolve_model_budget(model=model_name, model_max_budget=existing_model_max_budget)
|
||||
if governing is None:
|
||||
continue
|
||||
cap = governing.budget_config
|
||||
if cap.max_budget is None:
|
||||
continue
|
||||
if (
|
||||
proposed.max_budget is None
|
||||
or proposed.max_budget > cap.max_budget
|
||||
or proposed.budget_duration != cap.budget_duration
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
f"Only a proxy admin can raise a team's model_max_budget for {model_name!r} or change its "
|
||||
f"budget_duration. Current max_budget={cap.max_budget} per {cap.budget_duration} "
|
||||
f"(entry {governing.budget_model!r}), requested={proposed.max_budget} per "
|
||||
f"{proposed.budget_duration}."
|
||||
)
|
||||
},
|
||||
routed_models: Final = routed_models_source(data.team_id)
|
||||
for entry, proposed in requested.items():
|
||||
existing = _existing_model_cap(existing_model_max_budget.get(entry))
|
||||
if existing is not None:
|
||||
_check_existing_model_cap_kept(entry=entry, cap_entry=entry, proposed=proposed, existing=existing)
|
||||
if is_enforced_budget_entry(existing) or not is_enforced_budget_entry(proposed):
|
||||
continue
|
||||
governing = resolve_model_budget(model=entry, model_max_budget=existing_model_max_budget, team_id=data.team_id)
|
||||
if governing is not None and entry_names_model(entry=governing.budget_model, model=entry):
|
||||
_check_existing_model_cap_kept(
|
||||
entry=entry, cap_entry=governing.budget_model, proposed=proposed, existing=governing.budget_config
|
||||
)
|
||||
_check_new_entry_outranks_no_cap(
|
||||
entry=entry,
|
||||
existing_model_max_budget=existing_model_max_budget,
|
||||
routed_models=routed_models,
|
||||
team_id=data.team_id,
|
||||
)
|
||||
|
||||
|
||||
def _check_existing_model_cap_kept(entry: str, cap_entry: str, proposed: BudgetConfig, existing: BudgetConfig) -> None:
|
||||
if existing.max_budget is None or not _looser_model_cap(proposed=proposed, cap=existing):
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
f"Only a proxy admin can raise a team's model_max_budget for {entry!r} or change its "
|
||||
f"budget_duration. Current max_budget={existing.max_budget} per {existing.budget_duration} "
|
||||
f"(entry {cap_entry!r}), requested={proposed.max_budget} per {proposed.budget_duration}."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _check_new_entry_outranks_no_cap(
|
||||
entry: str,
|
||||
existing_model_max_budget: Mapping[str, object],
|
||||
routed_models: RoutedModels,
|
||||
team_id: str | None,
|
||||
) -> None:
|
||||
"""A new entry gets its own counter, so the models it takes from a capped entry stop drawing from that cap."""
|
||||
outranked: Final = next(
|
||||
_capped_entries_outranked_by(
|
||||
entry=entry,
|
||||
existing_model_max_budget=existing_model_max_budget,
|
||||
routed_models=routed_models,
|
||||
team_id=team_id,
|
||||
),
|
||||
None,
|
||||
)
|
||||
if outranked is None:
|
||||
return
|
||||
outranked_name, cap = outranked
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
f"Only a proxy admin can add a team model_max_budget entry {entry!r} that takes precedence over "
|
||||
f"{outranked_name!r}, since the models they share would stop drawing from its max_budget="
|
||||
f"{cap.max_budget} per {cap.budget_duration}."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _capped_entries_outranked_by(
|
||||
entry: str,
|
||||
existing_model_max_budget: Mapping[str, object],
|
||||
routed_models: RoutedModels,
|
||||
team_id: str | None,
|
||||
) -> Iterator[tuple[str, BudgetConfig]]:
|
||||
governing: Final = resolve_model_budget(model=entry, model_max_budget=existing_model_max_budget, team_id=team_id)
|
||||
outranked_names: Final = itertools.chain(
|
||||
()
|
||||
if governing is None or entry_names_model(entry=governing.budget_model, model=entry)
|
||||
else (governing.budget_model,),
|
||||
entries_outranked_by(entry=entry, model_max_budget=existing_model_max_budget, routed_models=routed_models),
|
||||
)
|
||||
outranked_caps: Final = ((name, _existing_model_cap(existing_model_max_budget[name])) for name in outranked_names)
|
||||
yield from (
|
||||
(name, cap) for name, cap in outranked_caps if name != entry and cap is not None and cap.max_budget is not None
|
||||
)
|
||||
|
||||
|
||||
def _looser_model_cap(proposed: BudgetConfig, cap: BudgetConfig) -> bool:
|
||||
return (
|
||||
proposed.max_budget is None
|
||||
or (cap.max_budget is not None and proposed.max_budget > cap.max_budget)
|
||||
or (is_enforced_budget_entry(cap) and proposed.budget_duration != cap.budget_duration)
|
||||
)
|
||||
|
||||
|
||||
def _should_auto_add_team_creator(
|
||||
|
|
@ -1572,7 +1648,7 @@ async def new_team(
|
|||
- prompts: Optional[List[str]] - List of prompts that the team is allowed to use.
|
||||
- organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
|
||||
- model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
|
||||
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
- guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
|
||||
- policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
|
||||
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
|
||||
|
|
@ -2350,7 +2426,7 @@ async def update_team(
|
|||
- tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
|
||||
- organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
|
||||
- model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
|
||||
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
- model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
- guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
|
||||
- policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
|
||||
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm.repositories.daily_activity_sql import (
|
|||
build_cache_leakage_keys_sql,
|
||||
build_entity_rollup_sql,
|
||||
build_export_sql,
|
||||
build_guardrail_spend_sql,
|
||||
build_key_page_sql,
|
||||
build_key_search_sql,
|
||||
build_model_top_keys_sql,
|
||||
|
|
@ -36,6 +37,7 @@ from litellm.types.repositories.daily_activity import (
|
|||
ExportRow,
|
||||
ExportType,
|
||||
GroupingSetsRow,
|
||||
GuardrailSpendRow,
|
||||
KeyMetadataRow,
|
||||
KeyPage,
|
||||
KeySpendRow,
|
||||
|
|
@ -100,6 +102,7 @@ class DailyActivityDatabase(Protocol):
|
|||
|
||||
_GROUPING_ADAPTER: Final = TypeAdapter(tuple[GroupingSetsRow, ...])
|
||||
_ENTITY_ADAPTER: Final = TypeAdapter(tuple[EntityRollupRow, ...])
|
||||
_GUARDRAIL_SPEND_ADAPTER: Final = TypeAdapter(tuple[GuardrailSpendRow, ...])
|
||||
_KEY_SPEND_ADAPTER: Final = TypeAdapter(tuple[KeySpendRow, ...])
|
||||
_KEY_PAGE_TOTAL_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(int)
|
||||
_USER_SPEND_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...])
|
||||
|
|
@ -191,6 +194,9 @@ class DailyActivityRepository:
|
|||
distinct_api_keys=distinct_api_keys,
|
||||
)
|
||||
|
||||
async def guardrail_spend(self, scope: DailyActivityScope) -> tuple[GuardrailSpendRow, ...]:
|
||||
return _GUARDRAIL_SPEND_ADAPTER.validate_python(await self._query(build_guardrail_spend_sql(scope)))
|
||||
|
||||
async def search_keys(self, scope: DailyActivityScope, *, search: str, limit: int) -> tuple[str, ...]:
|
||||
if not 1 <= limit <= constants.USAGE_KEY_SEARCH_MAX:
|
||||
raise ValueError(f"limit must be between 1 and {constants.USAGE_KEY_SEARCH_MAX}")
|
||||
|
|
|
|||
|
|
@ -261,6 +261,25 @@ def build_entity_rollup_sql(scope: DailyActivityScope, *, api_key_limit: int) ->
|
|||
return SqlQuery(sql=sql, params=(*where_params, PTU_SENTINEL_API_KEY, api_key_limit))
|
||||
|
||||
|
||||
def build_guardrail_spend_sql(scope: DailyActivityScope) -> SqlQuery:
|
||||
pg_table: Final = PRISMA_TO_PG_TABLE[scope.table]
|
||||
where_clause, where_params = build_where_clause(scope)
|
||||
sql: Final = f"""
|
||||
SELECT g.date, g.guardrail_id, COALESCE(t.guardrail_name, g.guardrail_id) AS guardrail_name, g.api_key,
|
||||
COALESCE(SUM(g.cost), 0)::float AS spend,
|
||||
COALESCE(SUM(CASE WHEN g.cost IS NULL THEN g.units ELSE g.untracked_units END), 0)::bigint AS untracked_units
|
||||
FROM "LiteLLM_DailyGuardrailUsageUnits" g
|
||||
LEFT JOIN "LiteLLM_GuardrailsTable" t ON t.guardrail_id = g.guardrail_id
|
||||
WHERE g.date >= $1 AND g.date <= $2
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM "{pg_table}"
|
||||
WHERE {where_clause} AND "{pg_table}".date = g.date AND "{pg_table}".api_key = g.api_key
|
||||
)
|
||||
GROUP BY g.date, g.guardrail_id, t.guardrail_name, g.api_key
|
||||
"""
|
||||
return SqlQuery(sql=sql, params=where_params)
|
||||
|
||||
|
||||
def _key_spend_select() -> str:
|
||||
return """
|
||||
COALESCE(SUM(spend), 0)::float AS spend,
|
||||
|
|
|
|||
|
|
@ -11850,6 +11850,11 @@ class Router:
|
|||
indices: Final = self.model_name_to_deployment_indices.get(model_name)
|
||||
return None if indices is None else tuple(self.model_list[idx] for idx in indices)
|
||||
|
||||
def team_public_model_names_of(self, team_id: str) -> frozenset[str]:
|
||||
"""The public names `team_id` serves its own deployments under."""
|
||||
with self._deployment_index_lock:
|
||||
return frozenset(name for owner, name in self.team_model_to_deployment_indices if owner == team_id)
|
||||
|
||||
def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[DeploymentTypedDict, ...] | None:
|
||||
"""The deployments `team_id` serves under `public_model_name`, or None when that pair is not indexed."""
|
||||
with self._deployment_index_lock:
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from datetime import datetime
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import service_target
|
||||
|
|
@ -417,17 +417,35 @@ def _should_run_cooldown_logic(
|
|||
return True
|
||||
|
||||
|
||||
_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, list[str | Mapping[str, object]]] | str, ...])
|
||||
_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[object, ...])
|
||||
_FALLBACK_ENTRY_ADAPTER: Final = TypeAdapter(Mapping[str, list[str | Mapping[str, object]] | str])
|
||||
|
||||
|
||||
def _fallback_target_name(target: str | Mapping[str, object]) -> str:
|
||||
return target if isinstance(target, str) else str(target.get("model", ""))
|
||||
|
||||
|
||||
def _fallback_target_names(
|
||||
targets: Sequence[str | Mapping[str, object]] | str,
|
||||
) -> list[str]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
|
||||
return [targets] if isinstance(targets, str) else [_fallback_target_name(target) for target in targets]
|
||||
|
||||
|
||||
def _fallback_chain_names(
|
||||
entry: Mapping[str, Sequence[str | Mapping[str, object]]],
|
||||
) -> Mapping[str, list[str]]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
|
||||
return {key: [_fallback_target_name(target) for target in targets] for key, targets in entry.items()}
|
||||
entry: object,
|
||||
) -> Mapping[str, list[str]] | None: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
|
||||
try:
|
||||
validated: Final = _FALLBACK_ENTRY_ADAPTER.validate_python(entry)
|
||||
except ValidationError:
|
||||
return None
|
||||
return {key: _fallback_target_names(targets) for key, targets in validated.items()}
|
||||
|
||||
|
||||
def _fallback_entries(configured: object) -> tuple[object, ...]:
|
||||
try:
|
||||
return _FALLBACK_ENTRIES_ADAPTER.validate_python(configured)
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool:
|
||||
|
|
@ -437,9 +455,7 @@ def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_
|
|||
if requested_model_group is None or not configured:
|
||||
return False
|
||||
chains: Final = tuple(
|
||||
_fallback_chain_names(entry)
|
||||
for entry in _FALLBACK_ENTRIES_ADAPTER.validate_python(configured)
|
||||
if isinstance(entry, Mapping)
|
||||
chain for chain in map(_fallback_chain_names, _fallback_entries(configured)) if chain is not None
|
||||
)
|
||||
fallback_group, _ = get_fallback_model_group(fallbacks=chains, model_group=requested_model_group)
|
||||
return bool(fallback_group)
|
||||
|
|
|
|||
|
|
@ -65,6 +65,34 @@ class MetricWithMetadata(MetricBase):
|
|||
api_key_breakdown: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata}
|
||||
|
||||
|
||||
class GuardrailSpendMetrics(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
spend: float = 0.0
|
||||
untracked_units: int = 0
|
||||
|
||||
|
||||
class GuardrailKeyMetric(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
metrics: GuardrailSpendMetrics
|
||||
metadata: KeyMetadata
|
||||
|
||||
|
||||
class GuardrailMetadata(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
guardrail_name: str
|
||||
|
||||
|
||||
class GuardrailMetric(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
metrics: GuardrailSpendMetrics
|
||||
metadata: GuardrailMetadata
|
||||
api_key_breakdown: Mapping[str, GuardrailKeyMetric] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class BreakdownMetrics(LiteLLMBaseModel):
|
||||
"""Breakdown of spend by different dimensions"""
|
||||
|
||||
|
|
@ -75,6 +103,7 @@ class BreakdownMetrics(LiteLLMBaseModel):
|
|||
endpoints: dict[str, MetricWithMetadata] = Field(default_factory=dict) # endpoint -> {metrics, metadata}
|
||||
api_keys: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata}
|
||||
entities: dict[str, MetricWithMetadata] = Field(default_factory=dict) # entity -> {metrics, metadata}
|
||||
guardrails: Mapping[str, GuardrailMetric] = Field(default_factory=dict) # guardrail_id -> spend
|
||||
|
||||
|
||||
class DailySpendData(LiteLLMBaseModel):
|
||||
|
|
|
|||
|
|
@ -168,6 +168,16 @@ class EntityRollupRow(RollupMetricsRow):
|
|||
distinct_api_keys: int | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GuardrailSpendRow:
|
||||
date: str
|
||||
guardrail_id: str
|
||||
guardrail_name: str
|
||||
api_key: str
|
||||
spend: float
|
||||
untracked_units: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AggregatedRows:
|
||||
grouping_rows: tuple[GroupingSetsRow, ...]
|
||||
|
|
|
|||
|
|
@ -21589,7 +21589,13 @@
|
|||
],
|
||||
"supports_vision": true,
|
||||
"supports_audio_input": true,
|
||||
"supports_video_input": true
|
||||
"supports_video_input": true,
|
||||
"input_cost_per_token_priority": 5.6252e-07,
|
||||
"output_cost_per_token_priority": 3.37498e-06,
|
||||
"cache_read_input_token_cost_priority": 5.628e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.6252e-07,
|
||||
"input_cost_per_audio_token_priority": 1.12497e-06,
|
||||
"cache_read_input_audio_token_cost_priority": 1.1249e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-1-flash-image": {
|
||||
"cache_creation_input_token_cost": 6.2503e-07,
|
||||
|
|
@ -21676,7 +21682,14 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.50002e-06,
|
||||
"output_cost_per_token_priority": 2.699998e-05,
|
||||
"cache_read_input_token_cost_priority": 4.5003e-07,
|
||||
"cache_creation_input_token_cost_priority": 4.50002e-06,
|
||||
"input_cost_per_token_above_200k_tokens_priority": 8.99997e-06,
|
||||
"output_cost_per_token_above_200k_tokens_priority": 4.049997e-05,
|
||||
"cache_read_input_token_cost_above_200k_tokens_priority": 8.9999e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-flash": {
|
||||
"cache_creation_input_token_cost": 6.2503e-07,
|
||||
|
|
@ -21710,7 +21723,13 @@
|
|||
],
|
||||
"supports_vision": true,
|
||||
"supports_audio_input": true,
|
||||
"supports_video_input": true
|
||||
"supports_video_input": true,
|
||||
"input_cost_per_token_priority": 1.12497e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.1249e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.12497e-06,
|
||||
"input_cost_per_audio_token_priority": 2.25001e-06,
|
||||
"cache_read_input_audio_token_cost_priority": 2.2498e-07
|
||||
},
|
||||
"databricks/databricks-gemini-3-pro": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -21757,7 +21776,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-7-flash": {
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
|
|
@ -21785,10 +21808,14 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-6-flash": {
|
||||
"cache_creation_input_token_cost": 1.87502e-06,
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
"cache_read_input_token_cost": 7.497e-08,
|
||||
"input_cost_per_token": 7.4998e-07,
|
||||
"input_dbu_cost_per_token": 1.0714e-05,
|
||||
|
|
@ -21816,7 +21843,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.35002e-06,
|
||||
"output_cost_per_token_priority": 6.75003e-06,
|
||||
"cache_read_input_token_cost_priority": 1.3503e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.35002e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-5-flash": {
|
||||
"cache_creation_input_token_cost": 1.87502e-06,
|
||||
|
|
@ -21846,7 +21877,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 3.37498e-06,
|
||||
"output_cost_per_token_priority": 2.025002e-05,
|
||||
"cache_read_input_token_cost_priority": 3.3747e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.37498e-06
|
||||
},
|
||||
"databricks/databricks-gemini-3-5-flash-lite": {
|
||||
"cache_creation_input_token_cost": 3.7499e-07,
|
||||
|
|
@ -21878,7 +21913,11 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 6.7501e-07,
|
||||
"output_cost_per_token_priority": 5.62499e-06,
|
||||
"cache_read_input_token_cost_priority": 6.748e-08,
|
||||
"cache_creation_input_token_cost_priority": 6.7501e-07
|
||||
},
|
||||
"databricks/databricks-gemma-3-12b": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -21921,7 +21960,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
"supports_vision": false,
|
||||
"input_cost_per_token_priority": 2.45e-06,
|
||||
"output_cost_per_token_priority": 7.7e-06,
|
||||
"cache_read_input_token_cost_priority": 4.55e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.45e-06
|
||||
},
|
||||
"databricks/databricks-glm-5-3": {
|
||||
"cache_creation_input_token_cost": 1.4e-06,
|
||||
|
|
@ -21951,7 +21994,11 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 2.45e-06,
|
||||
"output_cost_per_token_priority": 7.7e-06,
|
||||
"cache_read_input_token_cost_priority": 4.55e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.45e-06
|
||||
},
|
||||
"databricks/databricks-glm-5-3-flash": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -21980,7 +22027,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 2.625e-07,
|
||||
"output_cost_per_token_priority": 8.75e-07,
|
||||
"cache_read_input_token_cost_priority": 5.25e-08,
|
||||
"cache_creation_input_token_cost_priority": 2.625e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -21999,7 +22050,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 2.49998e-06,
|
||||
"output_cost_per_token_priority": 1.999998e-05,
|
||||
"cache_read_input_token_cost_priority": 2.4997e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-1": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -22018,7 +22073,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 2.49998e-06,
|
||||
"output_cost_per_token_priority": 1.999998e-05,
|
||||
"cache_read_input_token_cost_priority": 2.4997e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-2": {
|
||||
"cache_creation_input_token_cost": 1.75e-06,
|
||||
|
|
@ -22037,7 +22096,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 3.5e-06,
|
||||
"output_cost_per_token_priority": 2.8e-05,
|
||||
"cache_read_input_token_cost_priority": 3.5e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.5e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-3-codex": {
|
||||
"cache_creation_input_token_cost": 1.75e-06,
|
||||
|
|
@ -22056,7 +22119,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 3.5e-06,
|
||||
"output_cost_per_token_priority": 2.8e-05,
|
||||
"cache_read_input_token_cost_priority": 3.5e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.5e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -22087,7 +22154,14 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.00003e-06,
|
||||
"output_cost_per_token_priority": 2.999997e-05,
|
||||
"cache_read_input_token_cost_priority": 5.0001e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 5.00003e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.999997e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 5.0001e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-mini": {
|
||||
"cache_creation_input_token_cost": 7.4998e-07,
|
||||
|
|
@ -22115,7 +22189,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.50003e-06,
|
||||
"output_cost_per_token_priority": 8.99997e-06,
|
||||
"cache_read_input_token_cost_priority": 1.5001e-07,
|
||||
"cache_creation_input_token_cost_priority": 1.50003e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-nano": {
|
||||
"cache_creation_input_token_cost": 1.9999e-07,
|
||||
|
|
@ -22175,7 +22253,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 8.00002e-06,
|
||||
"output_cost_per_token_priority": 4.000003e-05,
|
||||
"cache_read_input_token_cost_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_priority": 9.99999e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 1.599997e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 6.000001e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.59999e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.999998e-05
|
||||
},
|
||||
"databricks/databricks-gpt-5-6-terra": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -22207,7 +22293,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 2.399999e-05,
|
||||
"cache_read_input_token_cost_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 3.600002e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-6-luna": {
|
||||
"cache_creation_input_token_cost": 2.4997e-07,
|
||||
|
|
@ -22239,7 +22333,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 3.9998e-07,
|
||||
"output_cost_per_token_priority": 2.40002e-06,
|
||||
"cache_read_input_token_cost_priority": 3.997e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.0001e-07,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.0003e-07,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 3.60003e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.001e-08,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 1.00002e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-5": {
|
||||
"cache_creation_input_token_cost": 5.00003e-06,
|
||||
|
|
@ -22270,7 +22372,14 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.249997e-05,
|
||||
"output_cost_per_token_priority": 7.50001e-05,
|
||||
"cache_read_input_token_cost_priority": 1.24999e-06,
|
||||
"cache_creation_input_token_cost_priority": 1.249997e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 1.249997e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 7.50001e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.24999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-5-pro": {
|
||||
"cache_creation_input_token_cost": 2.999997e-05,
|
||||
|
|
@ -22319,7 +22428,11 @@
|
|||
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
"supports_tool_choice": true,
|
||||
"input_cost_per_token_priority": 4.5003e-07,
|
||||
"output_cost_per_token_priority": 3.60003e-06,
|
||||
"cache_read_input_token_cost_priority": 4.501e-08,
|
||||
"cache_creation_input_token_cost_priority": 4.5003e-07
|
||||
},
|
||||
"databricks/databricks-gpt-5-nano": {
|
||||
"cache_creation_input_token_cost": 4.998e-08,
|
||||
|
|
@ -22469,7 +22582,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.25e-06,
|
||||
"output_cost_per_token_priority": 2.625e-05,
|
||||
"cache_read_input_token_cost_priority": 5.25e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.25e-06
|
||||
},
|
||||
"databricks/system.ai.kimi-k3": {
|
||||
"litellm_provider": "databricks",
|
||||
|
|
@ -22551,7 +22668,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"thinking_always_on": true
|
||||
"thinking_always_on": true,
|
||||
"input_cost_per_token_priority": 3.85e-07,
|
||||
"output_cost_per_token_priority": 3.85e-06,
|
||||
"cache_read_input_token_cost_priority": 3.85e-07,
|
||||
"cache_creation_input_token_cost_priority": 3.85e-07
|
||||
},
|
||||
"databricks/databricks-qwen3-next-80b-a3b-instruct": {
|
||||
"cache_creation_input_token_cost": 1.5001e-07,
|
||||
|
|
@ -77774,7 +77895,11 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 5.25e-07,
|
||||
"output_cost_per_token_priority": 2.1e-06,
|
||||
"cache_read_input_token_cost_priority": 5.25e-08,
|
||||
"cache_creation_input_token_cost_priority": 5.25e-07
|
||||
},
|
||||
"databricks/databricks-deepseek-v4-1-flash-1017": {
|
||||
"cache_creation_input_token_cost": 3.0002e-07,
|
||||
|
|
@ -77897,7 +78022,11 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.50002e-06,
|
||||
"output_cost_per_token_priority": 2.250003e-05,
|
||||
"cache_read_input_token_cost_priority": 2.2498e-07,
|
||||
"cache_creation_input_token_cost_priority": 4.50002e-06
|
||||
},
|
||||
"databricks/databricks-gpt-5-4-pro": {
|
||||
"cache_creation_input_token_cost": 2.999997e-05,
|
||||
|
|
@ -77959,7 +78088,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 1.99997e-05,
|
||||
"cache_read_input_token_cost_priority": 1.9999e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-6-1-terra": {
|
||||
"cache_creation_input_token_cost": 1.24999e-06,
|
||||
|
|
@ -77991,7 +78128,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.99997e-06,
|
||||
"output_cost_per_token_priority": 9.99999e-06,
|
||||
"cache_read_input_token_cost_priority": 1.0003e-07,
|
||||
"cache_creation_input_token_cost_priority": 2.49998e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4.00001e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 1.500002e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06
|
||||
},
|
||||
"databricks/databricks-gpt-6-astra": {
|
||||
"cache_creation_input_token_cost": 1.249997e-05,
|
||||
|
|
@ -78023,7 +78168,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.999998e-05,
|
||||
"output_cost_per_token_priority": 9.99999e-05,
|
||||
"cache_read_input_token_cost_priority": 1.99997e-06,
|
||||
"cache_creation_input_token_cost_priority": 2.500001e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4.000003e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 0.0001500002,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05
|
||||
},
|
||||
"databricks/databricks-gpt-6-luna": {
|
||||
"cache_creation_input_token_cost": 1.2502e-07,
|
||||
|
|
@ -78055,7 +78208,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 1.9999e-07,
|
||||
"output_cost_per_token_priority": 1.0003e-06,
|
||||
"cache_read_input_token_cost_priority": 2.002e-08,
|
||||
"cache_creation_input_token_cost_priority": 2.4997e-07,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 3.9998e-07,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 1.5001e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07
|
||||
},
|
||||
"databricks/databricks-gpt-6-sol": {
|
||||
"cache_creation_input_token_cost": 2.49998e-06,
|
||||
|
|
@ -78087,7 +78248,15 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"input_cost_per_token_priority": 4.00001e-06,
|
||||
"output_cost_per_token_priority": 1.99997e-05,
|
||||
"cache_read_input_token_cost_priority": 3.9998e-07,
|
||||
"cache_creation_input_token_cost_priority": 5.00003e-06,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 8.00002e-06,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 2.99999e-05,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06
|
||||
},
|
||||
"databricks/databricks-gpt-image-1": {
|
||||
"cache_creation_input_token_cost": 5.00003e-06,
|
||||
|
|
|
|||
205
skills/litellm-setup/SKILL.md
Normal file
205
skills/litellm-setup/SKILL.md
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
---
|
||||
name: litellm-setup
|
||||
description: Set up a self-hosted LiteLLM proxy, the OpenAI-compatible LLM gateway, on this machine, from install through provider keys, config.yaml, a master key, optional Postgres and Admin UI, and a verified chat completion. Use when the user asks to set up LiteLLM, install the LiteLLM proxy, self-host an LLM gateway, or put one endpoint in front of OpenAI, Anthropic, Gemini, Azure, Bedrock, Ollama, or other providers. Not for using the LiteLLM Python SDK inside an app
|
||||
---
|
||||
|
||||
# Set up a LiteLLM proxy
|
||||
|
||||
Take the user from nothing to a running LiteLLM proxy that has answered one chat completion. Run every command yourself unless a step says the user does it
|
||||
|
||||
Never print a secret. Provider keys, the master key, the salt key, the database URL, and virtual keys live only in `.env`, and `config.yaml` reads them through `os.environ/NAME`. Check that a value is there with `grep -c '^NAME=.' .env`, never with `cat .env` or `echo "$NAME"`. Load `.env` into a shell with this line, which exports each line as it is, so a `&` or `$` in a value stays literal, and first drops any database settings the shell already had and sets `LITELLM_MODE=PRODUCTION`, which stops `litellm` from also loading a `.env` it finds above its install folder, such as `~/.env`, so only this `.env` picks the database. The steps that need it start with it:
|
||||
|
||||
```bash
|
||||
unset DATABASE_URL DIRECT_URL DATABASE_HOST; export LITELLM_MODE=PRODUCTION; while IFS= read -r line || [ -n "$line" ]; do case "$line" in [A-Za-z_]*=*) export "$line" ;; esac; done < .env
|
||||
```
|
||||
|
||||
## 1. Ask once
|
||||
|
||||
Ask everything in one message, with defaults the user can accept by saying "go":
|
||||
|
||||
- Providers and models. Default is OpenAI with its newest flagship model, using the `OPENAI_API_KEY` in your environment when `[ -n "$OPENAI_API_KEY" ] && echo set` says it is there. Run that check before asking, and say in the question that the default uses that variable, naming it but never its value, since a key in your environment can belong to another account
|
||||
- Whether they want virtual keys, teams, spend tracking, or the Admin UI. Those need Postgres, plain routing does not. Default is no
|
||||
- uv (default) or Docker. With Postgres, Docker runs the repo's quickstart, which brings its own database, and uv needs a Postgres server they already have, with a new, empty database for LiteLLM, so someone with their own Postgres server picks uv
|
||||
- Folder and port. Defaults are `./litellm-proxy` (`./litellm-gateway` for the quickstart) and 4000
|
||||
|
||||
Then follow the steps for their answers. Plain routing skips steps 3 and 6. The quickstart is steps 2, 3, 8, and 9. uv with their own Postgres skips step 3
|
||||
|
||||
In the same message, say that someone who picks uv and would rather answer prompts can run `litellm --setup` in a separate terminal window once LiteLLM is installed, not through Claude Code. Never run it yourself, since it needs a real terminal and without one prints "Setup cancelled." and writes nothing. Give them the command as `LITELLM_MODE=PRODUCTION <bin>/litellm --setup`, with `<bin>` replaced by the folder `uv tool dir --bin` prints, to run from the folder they want the config in. It writes `litellm_config.yaml` there with the provider keys and master key in plain text and offers to start the proxy, which then listens on every network interface and stops with Ctrl+C in that window. When the user is done, run this in that folder. It makes the file readable only by them, keeps it out of git (the wizard does neither), and prints only the model names. Never print that file or run `grep`, `sed`, or `cat` on it:
|
||||
|
||||
```bash
|
||||
chmod 600 litellm_config.yaml; grep -qx 'litellm_config.yaml' .gitignore 2>/dev/null || echo 'litellm_config.yaml' >> .gitignore
|
||||
"$(uv tool dir)/litellm/bin/python" -I -c 'import yaml; print(*(m["model_name"] for m in yaml.safe_load(open("litellm_config.yaml"))["model_list"]), sep="\n")'
|
||||
```
|
||||
|
||||
Then skip to step 8 and run it against their proxy, starting that shell command with this line instead of the loader line, which reads the master key without printing it:
|
||||
|
||||
```bash
|
||||
export LITELLM_MASTER_KEY="$("$(uv tool dir)/litellm/bin/python" -I -c 'import yaml; print(yaml.safe_load(open("litellm_config.yaml"))["general_settings"]["master_key"])')"
|
||||
```
|
||||
|
||||
## 2. Check prerequisites
|
||||
|
||||
```bash
|
||||
uv --version; docker info --format '{{.ServerVersion}}'; docker compose version; curl --version | head -1; openssl version; command -v python3
|
||||
```
|
||||
|
||||
On native Windows, ask the user to run Claude Code inside WSL, since every command here is POSIX shell. Every path needs curl, openssl, and python3, so ask the user to install any that is missing. The uv path needs uv, which downloads a supported Python (LiteLLM supports 3.10 to 3.14) when the `--python` range in step 4 rules out the one on the machine. If uv is missing, offer to install it, and only after the user says yes run `curl -LsSf https://astral.sh/uv/install.sh | env UV_NO_MODIFY_PATH=1 sh`, which leaves their shell profile alone (https://docs.astral.sh/uv/getting-started/installation/). It puts uv in `~/.local/bin`, so start every later command that runs `uv` or `litellm` with `export PATH="$HOME/.local/bin:$PATH";`. The Docker path needs a running Docker daemon, and the quickstart also needs Docker Compose v2. If `docker info` says `permission denied`, the user's account cannot use Docker without sudo. Offer the uv path, or ask them to add themselves to the `docker` group the way https://docs.docker.com/engine/install/linux-postinstall/ shows and start a new login session before you check again
|
||||
|
||||
Check the port is free with `curl -s -o /dev/null --max-time 2 http://127.0.0.1:4000; echo $?`, which prints `7` (could not connect) only when nothing listens there. Don't use `lsof` for this, since it misses ports that root or Docker hold. If it is taken, pick another port and use it everywhere below, because a proxy asked for a busy 4000 silently moves to a random port
|
||||
|
||||
## 3. Postgres with Docker: run the quickstart
|
||||
|
||||
The repo's quickstart starts the gateway, Postgres, and the Admin UI in Docker, listening on this machine only, and puts a `litellm-gateway` command in that folder that stops, starts, updates, or removes them. It asks nothing when it has no terminal. It writes the master key, salt key, and database password to its own `.env` and never prints them. Download it, then run it with the folder and port from step 1. It returns once the gateway is ready:
|
||||
|
||||
```bash
|
||||
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.sh -o quickstart.sh
|
||||
LITELLM_DIR="$PWD/litellm-gateway" LITELLM_PORT=4000 sh quickstart.sh --yes
|
||||
cd litellm-gateway
|
||||
```
|
||||
|
||||
Models live in the database here, so add each one through the API instead of a `config.yaml`, naming it the way step 5 says. Pipe the request body in, so the provider key never shows up in output or on a command line. Unless the user agreed to the key in your environment, as step 5 says, ask them to add the line `OPENAI_API_KEY=<their key>` to this `.env` themselves. Run the loader line first:
|
||||
|
||||
```bash
|
||||
python3 -c 'import json, os, sys; json.dump({"model_name": "gpt-6.1-sol", "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": os.environ["OPENAI_API_KEY"]}}, sys.stdout)' \
|
||||
| curl -sS -o /dev/null -w '%{http_code}\n' http://127.0.0.1:4000/model/new \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json" -d @-
|
||||
```
|
||||
|
||||
Run that in one shell command that starts with the loader line, since the master key is in `.env`. `200` means the model was added. Then go to step 8
|
||||
|
||||
## 4. Install
|
||||
|
||||
uv:
|
||||
|
||||
```bash
|
||||
uv tool install --python '>=3.10,<3.15' 'litellm[proxy]'
|
||||
export PATH="$(uv tool dir --bin):$PATH"; litellm --version
|
||||
```
|
||||
|
||||
With their own Postgres, install `'litellm[proxy,extra_proxy]'` instead, with the same `--python` range, which adds the Prisma client the proxy uses for the database. Running `uv tool install` again with the new extras replaces the old install. Start every later command that runs `litellm` with `export PATH="$(uv tool dir --bin):$PATH";`, so this install runs even when the shell has no `litellm` or an older one from a venv or `pip install litellm`. `uv tool update-shell` puts it on the PATH of new shells by editing the user's shell profile, so run it only after the user says yes
|
||||
|
||||
Docker:
|
||||
|
||||
```bash
|
||||
docker pull docker.litellm.ai/berriai/litellm:main-stable
|
||||
```
|
||||
|
||||
## 5. Write `.env` and `config.yaml`
|
||||
|
||||
Run every command from here on in the folder from step 1. Create `.env` readable only by the user, keep it out of git, and add a master key unless it already has one. The proxy refuses to start without one, and `sk-` is the format of every key it issues:
|
||||
|
||||
```bash
|
||||
mkdir -p litellm-proxy && cd litellm-proxy
|
||||
touch .env && chmod 600 .env; grep -qx '.env' .gitignore 2>/dev/null || echo '.env' >> .gitignore
|
||||
grep -q '^LITELLM_MASTER_KEY=.' .env || printf 'LITELLM_MASTER_KEY=sk-%s\n' "$(openssl rand -hex 32)" >> .env
|
||||
```
|
||||
|
||||
Copy a provider key that is already in your environment only after the user agreed to that variable, in their step 1 answer or when you ask now, and copy it without printing it, as in `printf 'OPENAI_API_KEY=%s\n' "$OPENAI_API_KEY" >> .env`. Otherwise ask the user to add the line `OPENAI_API_KEY=<their key>` to `.env` themselves and tell you when it is there. If they paste a key into the chat anyway, write it to `.env` and never repeat it. Keep one `NAME=value` per line with no quotes, so `docker --env-file` and the loader line read the same value
|
||||
|
||||
Confirm each model's current LiteLLM name in the cost map instead of trusting memory, since model names go stale fast. This prints the chat models of the provider named last (`openai`, `anthropic`, `gemini`, `azure`, `bedrock`, `vertex_ai`, and so on) the way `config.yaml` takes them, with the provider in front. It matches by prefix, so `bedrock` also lists the newer Bedrock models and `vertex_ai` lists every Vertex family. For Ollama, use a model from `ollama list` instead:
|
||||
|
||||
```bash
|
||||
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json | python3 -c '
|
||||
import json, sys
|
||||
for name, row in json.load(sys.stdin).items():
|
||||
provider = str(row.get("litellm_provider", "")) if isinstance(row, dict) else ""
|
||||
if provider.startswith(sys.argv[1]) and row.get("mode") == "chat":
|
||||
prefix = "bedrock" if provider == "bedrock_converse" else "vertex_ai" if provider.startswith("vertex_ai") else provider
|
||||
print(name if name.startswith(prefix + "/") else prefix + "/" + name)
|
||||
' openai
|
||||
```
|
||||
|
||||
In `config.yaml`, `model` is a name from that list, as in `openai/gpt-6.1-sol` or `bedrock/us.anthropic.<name>`, and `model_name` is what clients will send. If the folder already has a `config.yaml`, show the user its `model_name` lines and ask before replacing it. This example uses the newest GPT model in the cost map when this skill was written, so replace it with what the user picked:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-6.1-sol
|
||||
litellm_params:
|
||||
model: openai/gpt-6.1-sol
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
general_settings:
|
||||
master_key: os.environ/LITELLM_MASTER_KEY
|
||||
```
|
||||
|
||||
Add one entry per model. On Azure, ask the user for each deployment's name and the model it runs. Use `azure/<deployment name>` as `model`, since Azure routes by deployment and a name from the lookup works only when the deployment has that same name, and add `model_info:` with `base_model: azure/<model it runs>` to the entry, since LiteLLM prices each call by that name. Providers that need more than a key, like Azure or Bedrock, list their `litellm_params` on their page under https://docs.litellm.ai/docs/providers. Point every secret at `.env` through `os.environ/NAME` the same way
|
||||
|
||||
## 6. Their own Postgres, on the uv path
|
||||
|
||||
Ask the user to create a new, empty database for LiteLLM on their server (`CREATE DATABASE litellm;`) and add `DATABASE_URL=postgresql://...` for it to `.env` themselves. The proxy creates its tables there on first start and stops on a database that already holds other tables. If that URL goes through a transaction pooler such as PgBouncer, they also add `DIRECT_URL` for the same database without the pooler, which the migrations use. Then add a salt key. It encrypts the provider credentials the proxy stores in the database and falls back to the master key when unset, so set it once and never change it:
|
||||
|
||||
```bash
|
||||
grep -q '^LITELLM_SALT_KEY=.' .env || printf 'LITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" >> .env
|
||||
```
|
||||
|
||||
Generate the Prisma client once inside the tool's environment. A pip or uv install doesn't ship it generated (Docker images do):
|
||||
|
||||
```bash
|
||||
tool="$(uv tool dir)/litellm"
|
||||
schema=$("$tool/bin/python" -I -c 'import importlib.util, os; print(os.path.join(importlib.util.find_spec("litellm").submodule_search_locations[0], "proxy", "schema.prisma"))')
|
||||
PATH="$tool/bin:$PATH" prisma generate --schema "$schema"
|
||||
```
|
||||
|
||||
## 7. Start the proxy
|
||||
|
||||
Bind to `127.0.0.1` unless the user wants other machines to reach it
|
||||
|
||||
uv, in one shell command that starts with the loader line, since the installed `litellm` does not read `.env` by itself:
|
||||
|
||||
```bash
|
||||
export PATH="$HOME/.local/bin:$PATH"; export PATH="$(uv tool dir --bin):$PATH"
|
||||
nohup litellm --config config.yaml --host 127.0.0.1 --port 4000 > litellm.log 2>&1 &
|
||||
echo $! > litellm.pid
|
||||
```
|
||||
|
||||
Docker:
|
||||
|
||||
```bash
|
||||
docker run -d --name litellm --env-file .env -v "$PWD/config.yaml:/app/config.yaml:ro" -p 127.0.0.1:4000:4000 \
|
||||
docker.litellm.ai/berriai/litellm:main-stable --config /app/config.yaml --port 4000
|
||||
```
|
||||
|
||||
Wait on the health route, stopping early if the proxy exits. The first start with Postgres applies the database migrations and takes longer. Define `alive` for the path you took, `alive() { kill -0 "$(cat litellm.pid)" 2>/dev/null; }` for uv or `alive() { [ "$(docker inspect -f '{{.State.Running}}' litellm 2>/dev/null)" = true ]; }` for Docker, in the same shell command as the loop:
|
||||
|
||||
```bash
|
||||
for i in $(seq 1 180); do curl -fsS http://127.0.0.1:4000/health/liveliness >/dev/null 2>&1 && break; alive || break; sleep 1; done
|
||||
curl -sS http://127.0.0.1:4000/health/readiness
|
||||
```
|
||||
|
||||
Readiness reports `"db":"connected"` with Postgres and `"db":"Not connected"` without it. If the proxy never comes up, read the end of `litellm.log` (or `docker logs litellm`) and check the troubleshooting list
|
||||
|
||||
## 8. Verify with a chat completion
|
||||
|
||||
Use a `model_name` from `config.yaml` (or step 3), in one shell command that starts with the loader line:
|
||||
|
||||
```bash
|
||||
curl -sS http://127.0.0.1:4000/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json" \
|
||||
-d '{"model": "gpt-6.1-sol", "messages": [{"role": "user", "content": "Say hello in five words."}]}'
|
||||
```
|
||||
|
||||
Show the user the reply in `choices[0].message.content`. An error here means the setup is not done, so fix it before handing off
|
||||
|
||||
## 9. Hand off
|
||||
|
||||
Tell the user, without printing any secret:
|
||||
|
||||
- The base URL is `http://127.0.0.1:4000`. Any OpenAI-compatible client works with it as the base URL and the master key or a virtual key as the API key
|
||||
- The master key is `LITELLM_MASTER_KEY` in the folder's `.env`, or `general_settings.master_key` in `litellm_config.yaml` when they used `litellm --setup`
|
||||
- To stop it, `kill "$(cat litellm.pid)"` (uv), `docker stop litellm` (Docker), `<folder>/litellm-gateway stop` (quickstart), or Ctrl+C in its window (`litellm --setup`). To start it again, `docker start litellm`, `<folder>/litellm-gateway start`, on uv the whole step 7 command as you ran it, loader line and PATH export included, wrapped in `( ... )` so their shell keeps its own settings, shown as a code block and never paraphrased, since a new shell has neither, or for `litellm --setup` `LITELLM_MODE=PRODUCTION <bin>/litellm --config litellm_config.yaml --host 127.0.0.1 --port <port>` in a terminal window in that folder, with `<bin>` and `<port>` filled in, since running the wizard again writes a new config with a new master key. `update` and `uninstall` update or remove a quickstart install. Give the quickstart commands with the folder's full path, since a bare `litellm-gateway` is on PATH only when `~/.local/bin` is, and can belong to an earlier install
|
||||
- With Postgres, the Admin UI is at `http://127.0.0.1:4000/ui` (user `admin`, password the master key), and `POST /key/generate` with the master key creates virtual keys with their own budgets and model lists. Without Postgres, both need step 3 or 6 first
|
||||
- The Admin UI warns that the master-key login is a shared secret, which is fine while only this machine reaches the proxy. Before anyone else can reach it, invite an admin with their own password on the Internal Users page, then set `general_settings.disable_env_credential_login: true` in the proxy config
|
||||
- To manage the running proxy from a coding agent (keys, teams, users, models, MCP servers), install `npx skills add BerriAI/litellm-skills -g`
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- `refused to start: no master key is set` means the loader line did not run in the shell command that started `litellm`, or `--env-file .env` is missing from `docker run`
|
||||
- `a database URL is set but the prisma CLI is neither on PATH nor importable` means the install lacks `extra_proxy`, so rerun step 4 with it, or, on plain routing, that the command that started `litellm` did not start with the loader line
|
||||
- `The Client hasn't been generated yet` or `Unable to find Prisma binaries` means the Prisma generate command in step 6 has not run
|
||||
- `Incorrect API key provided` or another provider `AuthenticationError` means a provider key is wrong. Fix it in `.env` and restart, since the proxy reads the environment only at startup. On uv, stop it and rerun step 7. On Docker, `docker rm -f litellm` and rerun step 7, since `docker restart` keeps the environment the container was created with. On the quickstart the key is saved with the model, so delete the model and add it again: `GET /model/info` lists each model's `model_info.id`, and `POST /model/delete` with `{"id": "<id>"}` deletes it
|
||||
- `Cannot automatically baseline this database` means `DATABASE_URL` points at a database that already holds tables LiteLLM did not create. Point it at a new, empty database
|
||||
- `LLM Provider NOT provided` means a `model` in `config.yaml` (or step 3) lacks its provider in front. Use the name as the step 5 lookup prints it
|
||||
- `Invalid model name passed in model=...` means the request's `model` is not a `model_name` in `config.yaml` or step 3. `GET /v1/models` lists the ones the key can use
|
||||
- `No connected db.` on a proxy without Postgres means the bearer token is not the master key
|
||||
- `Conflict. The container name "/litellm" is already in use` means an earlier run left a container. `docker start litellm` reuses it, or with the user's OK `docker rm -f litellm` and rerun step 7
|
||||
- Inside Docker, `localhost` in `config.yaml` (an Ollama `api_base`, say) is the container itself. Use `host.docker.internal`, and on Linux also pass `--add-host=host.docker.internal:host-gateway`
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
import hashlib
|
||||
import os
|
||||
import queue
|
||||
import secrets
|
||||
|
|
@ -973,6 +974,7 @@ def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_key
|
|||
_insert_spend_log_with_key(conn, "legacy-1", _legacy_bearer_token())
|
||||
_insert_spend_log_with_key(conn, "legacy-2", _legacy_bearer_token())
|
||||
_insert_spend_log_with_key(conn, "hashed-1", "a" * 64)
|
||||
_insert_spend_log_with_key(conn, "repetitive-1", "z" * 3000)
|
||||
ledger_before: Final = _ledger(scratch_database)
|
||||
|
||||
with pytest.raises(RuntimeError) as refused:
|
||||
|
|
@ -1109,6 +1111,68 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does(
|
|||
assert finding.detail == f"with an index entry over 2704 bytes, largest 2712 bytes ({described_as})"
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random_text_of_that_length_is(
|
||||
scratch_database: str,
|
||||
) -> None:
|
||||
"""Postgres compresses a long value before it checks the entry size, so 3000 repeated
|
||||
bytes index fine while 3000 random hex bytes do not. The check must draw the same line."""
|
||||
create_index: Final = sql.SQL('CREATE INDEX "compress_probe_idx" ON "compress_probe" ("k")')
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
conn.execute('CREATE TABLE "compress_probe" ("k" TEXT)')
|
||||
conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s), (%s)', ("a" * 3000, secrets.token_hex(1500)))
|
||||
findings: Final = migration_data_checks.find_blockers(
|
||||
conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
|
||||
)
|
||||
with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"):
|
||||
conn.execute(create_index)
|
||||
conn.execute('DELETE FROM "compress_probe" WHERE "k" <> %s', ("a" * 3000,))
|
||||
conn.execute(create_index)
|
||||
stored: Final = conn.execute('SELECT octet_length("k"), pg_column_size("k") FROM "compress_probe"').fetchone()
|
||||
|
||||
assert stored is not None and stored[0] == 3000 and stored[1] < 2704, stored
|
||||
(finding,) = findings
|
||||
assert isinstance(finding, migration_data_checks.BlockedChange)
|
||||
assert finding.rows == 1, "only the random key is too long for the index once Postgres compresses the other"
|
||||
assert finding.detail == "with an index entry over 2704 bytes, largest 3016 bytes (k up to 3000 bytes)"
|
||||
|
||||
|
||||
def _postgres_rejects_an_index_on(conn: "psycopg.Connection[tuple[object, ...]]", key: str) -> bool:
|
||||
conn.execute('DROP TABLE IF EXISTS "one_key"')
|
||||
conn.execute('CREATE TABLE "one_key" ("k" TEXT)')
|
||||
conn.execute('INSERT INTO "one_key" ("k") VALUES (%s)', (key,))
|
||||
try:
|
||||
conn.execute('CREATE INDEX "one_key_idx" ON "one_key" ("k")')
|
||||
except psycopg.errors.ProgramLimitExceeded:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_entry_limit(scratch_database: str) -> None:
|
||||
"""Keys that compress to a stored size just under and just over the entry limit go through the
|
||||
check and through Postgres one by one. The check must flag a key exactly when Postgres rejects it."""
|
||||
noise: Final = "".join(hashlib.sha256(str(i).encode()).hexdigest() for i in range(40))
|
||||
keys: Final = tuple(noise[:n] + "a" * 1500 for n in range(2400, 2500))
|
||||
create_index: Final = sql.SQL('CREATE INDEX "stored_probe_idx" ON "stored_probe" ("k")')
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
conn.execute('CREATE TABLE "stored_probe" ("k" TEXT)')
|
||||
conn.execute('INSERT INTO "stored_probe" ("k") SELECT unnest(%s::text[])', (list(keys),))
|
||||
compressed: Final = conn.execute(
|
||||
'SELECT count(*) FROM "stored_probe" WHERE pg_column_size("k") < octet_length("k") AND octet_length("k") > 2704'
|
||||
).fetchone()
|
||||
findings: Final = migration_data_checks.find_blockers(
|
||||
conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
|
||||
)
|
||||
rejected: Final = tuple(key for key in keys if _postgres_rejects_an_index_on(conn, key))
|
||||
|
||||
assert compressed == (len(keys),), "every probe key must be longer than the limit yet stored compressed"
|
||||
assert 0 < len(rejected) < len(keys), f"the probe must straddle the limit, Postgres rejected {len(rejected)}"
|
||||
(finding,) = findings
|
||||
assert isinstance(finding, migration_data_checks.BlockedChange)
|
||||
assert finding.rows == len(rejected), f"the check must flag the {len(rejected)} keys Postgres rejects and no other"
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_scan_that_cannot_finish_in_time_is_reported_as_unverified_rather_than_passed(scratch_database: str) -> None:
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ GEMINI: Final = "gemini-3-pro-preview"
|
|||
ZAI_GLM: Final = "glm-5"
|
||||
ECHOED_KEY_MODEL: Final = f"{KIMI}-echoed-key"
|
||||
BEDROCK_CLAUDE: Final = f"us.anthropic.{CLAUDE}"
|
||||
DATABRICKS_CLAUDE: Final = f"databricks-{CLAUDE}"
|
||||
AZURE_API_VERSION: Final = "2025-04-01-preview"
|
||||
SYNTHETIC_KEY: Final = "synthetic-provider-key"
|
||||
_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
|
||||
|
|
@ -215,6 +216,10 @@ def backend(kind: str, wire_url: str) -> Backend:
|
|||
f"bedrock/invoke/{BEDROCK_CLAUDE}",
|
||||
MappingProxyType({"api_base": wire_url, "api_key": SYNTHETIC_KEY, "aws_region_name": "us-east-1"}),
|
||||
)
|
||||
case "databricks_claude":
|
||||
return Backend(
|
||||
f"databricks/{DATABRICKS_CLAUDE}", MappingProxyType({"api_base": wire_url, "api_key": SYNTHETIC_KEY})
|
||||
)
|
||||
case _:
|
||||
raise AssertionError(f"unknown backend kind {kind}")
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
|
|
@ -8,9 +9,11 @@ import anthropic
|
|||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.anthropic_sse import message_json
|
||||
from integration._support.client import Gateway, eventually, list_value, object_value, string_value
|
||||
from integration._support.openai_wire import openai_error
|
||||
from integration._support.wire import Request, wire_server
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from integration.routing._reasoning_effort_levels_support import (
|
||||
AZURE_API_VERSION,
|
||||
CLAUDE,
|
||||
|
|
@ -34,6 +37,7 @@ from integration.routing._reasoning_effort_levels_support import (
|
|||
marker_of,
|
||||
messages_request,
|
||||
new_marker,
|
||||
path_of,
|
||||
prompt,
|
||||
provider_reply,
|
||||
responses_request,
|
||||
|
|
@ -41,6 +45,8 @@ from integration.routing._reasoning_effort_levels_support import (
|
|||
)
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS
|
||||
|
||||
CHAT: Final = "/v1/chat/completions"
|
||||
MESSAGES: Final = "/v1/messages"
|
||||
RESPONSES: Final = "/v1/responses"
|
||||
|
|
@ -49,6 +55,7 @@ _ADAPTIVE: Final[dict[str, JsonValue]] = {"type": "adaptive"}
|
|||
_STRING_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": "2048"}
|
||||
_XHIGH_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 9000}
|
||||
_BOOL_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": True}
|
||||
_FALSE_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": False}
|
||||
_NULL_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": None}
|
||||
_LIST_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": [2048]}
|
||||
_EMPTY_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": ""}
|
||||
|
|
@ -67,6 +74,8 @@ _KIMI_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort
|
|||
_CLAUDE_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": list(CLAUDE_LEVELS)})
|
||||
_GPT_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": list(GPT_LEVELS)})
|
||||
_UNDECLARED: Final[Mapping[str, JsonValue]] = MappingProxyType({})
|
||||
_EMPTY_DECLARATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": []})
|
||||
_NONE_ONLY_DECLARATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": ["none"]})
|
||||
_ECHOED_MAP_KEY: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{"key": "azure/chat-latest", "litellm_provider": "azure"}
|
||||
)
|
||||
|
|
@ -134,7 +143,7 @@ def _invoke_chat(effort: str, thinking: Mapping[str, JsonValue], betas: tuple[st
|
|||
"messages": [{"role": "user", "content": [{"type": "text", "text": prompt(marker)}]}],
|
||||
"thinking": dict(thinking),
|
||||
"output_config": {"effort": effort},
|
||||
"max_tokens": 4096,
|
||||
"max_tokens": DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS,
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"anthropic_beta": list(betas),
|
||||
}
|
||||
|
|
@ -144,6 +153,14 @@ def _with_output_config(expected: Expected, output_config: JsonValue) -> Expecte
|
|||
return lambda marker: {**expected(marker), "output_config": output_config}
|
||||
|
||||
|
||||
def _with_fields(expected: Expected, fields: Mapping[str, JsonValue]) -> Expected:
|
||||
return lambda marker: {**expected(marker), **fields}
|
||||
|
||||
|
||||
def _without(expected: Expected, field: str) -> Expected:
|
||||
return lambda marker: {key: value for key, value in expected(marker).items() if key != field}
|
||||
|
||||
|
||||
def _converse_chat(effort: str, thinking: Mapping[str, JsonValue]) -> Expected:
|
||||
return lambda marker: {
|
||||
"messages": [{"role": "user", "content": [{"text": prompt(marker)}]}],
|
||||
|
|
@ -320,6 +337,7 @@ _OPENAI_RESPONSES: Final = "/v1/responses"
|
|||
_AZURE_RESPONSES: Final = f"/openai/responses?api-version={AZURE_API_VERSION}"
|
||||
_UNSUPPORTED: Final = "does not support parameters"
|
||||
_GEMINI: Final = f"/models/{GEMINI}:generateContent"
|
||||
_DATABRICKS: Final = "/chat/completions"
|
||||
|
||||
CELLS: Final[Mapping[str, Cell]] = MappingProxyType(
|
||||
{
|
||||
|
|
@ -1750,6 +1768,139 @@ CELLS: Final[Mapping[str, Cell]] = MappingProxyType(
|
|||
_anthropic_chat_without_effort(_ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"chat-anthropic-false-budget-400": Cell(
|
||||
CHAT, "anthropic", _CLAUDE_INFO, {"thinking": _FALSE_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
|
||||
),
|
||||
"chat-anthropic-undeclared-null-budget-400": Cell(
|
||||
CHAT, "anthropic", _UNDECLARED, {"thinking": _NULL_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
|
||||
),
|
||||
"chat-anthropic-undeclared-false-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
CHAT,
|
||||
"anthropic",
|
||||
_UNDECLARED,
|
||||
{"thinking": _FALSE_BUDGET, "drop_params": True},
|
||||
_anthropic_chat_without_effort(_ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"chat-anthropic-undeclared-missing-budget-sends-low": Cell(
|
||||
CHAT,
|
||||
"anthropic",
|
||||
_UNDECLARED,
|
||||
{"thinking": _MISSING_BUDGET},
|
||||
_anthropic_chat("low", _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"messages-anthropic-undeclared-false-budget-400": Cell(
|
||||
MESSAGES, "anthropic", _UNDECLARED, {"thinking": _FALSE_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
|
||||
),
|
||||
"messages-anthropic-empty-string-budget-400": Cell(
|
||||
MESSAGES, "anthropic", _CLAUDE_INFO, {"thinking": _EMPTY_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET
|
||||
),
|
||||
"messages-anthropic-undeclared-null-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
MESSAGES,
|
||||
"anthropic",
|
||||
_UNDECLARED,
|
||||
{"thinking": _NULL_BUDGET, "drop_params": True},
|
||||
_native_messages(None, _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"responses-anthropic-false-budget-400": Cell(
|
||||
RESPONSES,
|
||||
"anthropic",
|
||||
_CLAUDE_INFO,
|
||||
{"thinking": _FALSE_BUDGET},
|
||||
None,
|
||||
_ANTHROPIC,
|
||||
400,
|
||||
_NON_NUMERIC_BUDGET,
|
||||
),
|
||||
"chat-converse-undeclared-false-budget-400": Cell(
|
||||
CHAT,
|
||||
"bedrock_converse",
|
||||
_UNDECLARED,
|
||||
{"thinking": _FALSE_BUDGET},
|
||||
None,
|
||||
_CONVERSE,
|
||||
400,
|
||||
_NON_NUMERIC_BUDGET,
|
||||
),
|
||||
"messages-invoke-undeclared-null-budget-400": Cell(
|
||||
MESSAGES, "bedrock_invoke", _UNDECLARED, {"thinking": _NULL_BUDGET}, None, _INVOKE, 400, _NON_NUMERIC_BUDGET
|
||||
),
|
||||
"chat-databricks-claude-undeclared-false-budget-400": Cell(
|
||||
CHAT,
|
||||
"databricks_claude",
|
||||
_UNDECLARED,
|
||||
{"thinking": _FALSE_BUDGET},
|
||||
None,
|
||||
_DATABRICKS,
|
||||
400,
|
||||
_NON_NUMERIC_BUDGET,
|
||||
),
|
||||
"chat-anthropic-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
CHAT,
|
||||
"anthropic",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_anthropic_chat_without_effort(_ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"chat-anthropic-none-only-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
CHAT,
|
||||
"anthropic",
|
||||
_NONE_ONLY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_anthropic_chat_without_effort(_ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"chat-anthropic-empty-declaration-budget-keeps-its-effort": Cell(
|
||||
CHAT,
|
||||
"anthropic",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET},
|
||||
_anthropic_chat("medium", _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"chat-converse-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
CHAT,
|
||||
"bedrock_converse",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_converse_chat_without_effort(_ADAPTIVE),
|
||||
_CONVERSE,
|
||||
),
|
||||
"messages-anthropic-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
MESSAGES,
|
||||
"anthropic",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_native_messages(None, _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"messages-anthropic-none-only-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
MESSAGES,
|
||||
"anthropic",
|
||||
_NONE_ONLY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_native_messages(None, _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"messages-anthropic-empty-declaration-budget-keeps-its-effort": Cell(
|
||||
MESSAGES,
|
||||
"anthropic",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET},
|
||||
_native_messages({"effort": "medium"}, _ADAPTIVE),
|
||||
_ANTHROPIC,
|
||||
),
|
||||
"messages-converse-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell(
|
||||
MESSAGES,
|
||||
"bedrock_converse",
|
||||
_EMPTY_DECLARATION,
|
||||
{"thinking": _BUDGET, "drop_params": True},
|
||||
_converse_messages(None),
|
||||
_CONVERSE,
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -1815,20 +1966,65 @@ def _bedrock_native_schema(body: Mapping[str, JsonValue], *, invoke: bool) -> Js
|
|||
return JSON_OBJECT.validate_json(string_value(json_schema["schema"]))
|
||||
|
||||
|
||||
_CONVERSE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = {
|
||||
"textFormat": {
|
||||
"type": "json_schema",
|
||||
"structure": {"jsonSchema": {"name": "answer", "schema": json.dumps(_NATIVE_OUTPUT_SCHEMA)}},
|
||||
}
|
||||
}
|
||||
_INVOKE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = {
|
||||
"effort": "high",
|
||||
"format": {"type": "json_schema", "schema": _NATIVE_OUTPUT_SCHEMA},
|
||||
}
|
||||
_JSON_TOOL: Final = "json_tool_call"
|
||||
_CONVERSE_TOOL_SCHEMA: Final[dict[str, JsonValue]] = {
|
||||
key: value for key, value in _NATIVE_OUTPUT_SCHEMA.items() if key != "additionalProperties"
|
||||
}
|
||||
_CONVERSE_WITH_NATIVE_OUTPUT: Final = _with_fields(
|
||||
_converse_chat("high", _ADAPTIVE), {"outputConfig": _CONVERSE_NATIVE_OUTPUT}
|
||||
)
|
||||
_INVOKE_WITH_NATIVE_OUTPUT: Final = _with_output_config(
|
||||
_invoke_chat("high", _ADAPTIVE, ("structured-outputs-2025-11-13",)), _INVOKE_NATIVE_OUTPUT
|
||||
)
|
||||
_CONVERSE_WITH_JSON_TOOL: Final = _with_fields(
|
||||
_converse_chat("high", _ADAPTIVE),
|
||||
{
|
||||
"toolConfig": {
|
||||
"tools": [
|
||||
{
|
||||
"toolSpec": {
|
||||
"name": _JSON_TOOL,
|
||||
"description": _JSON_TOOL,
|
||||
"inputSchema": {"json": _CONVERSE_TOOL_SCHEMA},
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
_INVOKE_WITH_JSON_TOOL: Final = _with_fields(
|
||||
_without(_invoke_chat("high", _ADAPTIVE, ()), "anthropic_beta"),
|
||||
{"tools": [{"name": _JSON_TOOL, "input_schema": _NATIVE_OUTPUT_SCHEMA}]},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kind", "target", "fields"),
|
||||
("kind", "target", "fields", "expected"),
|
||||
[
|
||||
pytest.param("bedrock_converse", _CONVERSE, {"thinking": _BUDGET}, id="converse-budget"),
|
||||
pytest.param(
|
||||
"bedrock_converse", _CONVERSE, {"thinking": _BUDGET}, _CONVERSE_WITH_NATIVE_OUTPUT, id="converse-budget"
|
||||
),
|
||||
pytest.param(
|
||||
"bedrock_invoke",
|
||||
_INVOKE,
|
||||
{"thinking": _BUDGET, "output_config": {"effort": "medium"}},
|
||||
_INVOKE_WITH_NATIVE_OUTPUT,
|
||||
id="invoke-output-config-beside-budget",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_a_deployment_opted_into_bedrock_native_output_keeps_it_beside_its_declared_level(
|
||||
gateway: Gateway, kind: str, target: str, fields: Mapping[str, JsonValue]
|
||||
gateway: Gateway, kind: str, target: str, fields: Mapping[str, JsonValue], expected: Expected
|
||||
) -> None:
|
||||
marker: Final = new_marker()
|
||||
response_format: Final[dict[str, JsonValue]] = {
|
||||
|
|
@ -1844,10 +2040,116 @@ def test_a_deployment_opted_into_bedrock_native_output_keeps_it_beside_its_decla
|
|||
assert response.status_code == 200, response.text[:2000]
|
||||
assert answer(marker) in response.text, response.text[:2000]
|
||||
assert [(request.method, request.target) for request in received] == [("POST", target)]
|
||||
body: Final = body_of(received[0])
|
||||
assert _bedrock_native_schema(body, invoke=target == _INVOKE) == _NATIVE_OUTPUT_SCHEMA, body
|
||||
assert {"toolConfig", "tools"}.isdisjoint(body), body
|
||||
assert sent_effort(received[0]) == "high", body
|
||||
assert body_of(received[0]) == expected(marker)
|
||||
|
||||
|
||||
_NATIVE_OUTPUT_REJECTION: Final = (
|
||||
"The model returned the following errors: output_config.format: Extra inputs are not permitted"
|
||||
)
|
||||
|
||||
|
||||
def _rejecting_native_output(request: Request) -> Reply:
|
||||
native: Final = _bedrock_native_schema(body_of(request), invoke=path_of(request).endswith("/invoke"))
|
||||
if native is None:
|
||||
return provider_reply(request)
|
||||
return Reply(status=400, body=json.dumps({"message": _NATIVE_OUTPUT_REJECTION}).encode())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kind", "target", "rejected", "retried"),
|
||||
[
|
||||
pytest.param(
|
||||
"bedrock_converse", _CONVERSE, _CONVERSE_WITH_NATIVE_OUTPUT, _CONVERSE_WITH_JSON_TOOL, id="converse"
|
||||
),
|
||||
pytest.param("bedrock_invoke", _INVOKE, _INVOKE_WITH_NATIVE_OUTPUT, _INVOKE_WITH_JSON_TOOL, id="invoke"),
|
||||
],
|
||||
)
|
||||
def test_the_retry_after_bedrock_rejects_native_output_keeps_the_declared_level(
|
||||
gateway: Gateway, kind: str, target: str, rejected: Expected, retried: Expected
|
||||
) -> None:
|
||||
marker: Final = new_marker()
|
||||
response_format: Final[dict[str, JsonValue]] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "answer", "schema": _NATIVE_OUTPUT_SCHEMA},
|
||||
}
|
||||
with wire_server(_rejecting_native_output) as wire, gateway.scenario() as scenario:
|
||||
model: Final = deployment(scenario, kind, wire, _CLAUDE_INFO_WITH_NATIVE_OUTPUT)
|
||||
response: Final = gateway.request(
|
||||
"POST", CHAT, chat_request(model, marker, response_format=response_format, thinking=_BUDGET)
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert response.status_code == 200, response.text[:2000]
|
||||
assert [(request.method, request.target) for request in received] == [("POST", target), ("POST", target)]
|
||||
assert [body_of(request) for request in received] == [rejected(marker), retried(marker)]
|
||||
|
||||
|
||||
_VERTEX_PROJECT: Final = "scripted-project"
|
||||
_VERTEX_LOCATION: Final = "us-east5"
|
||||
_VERTEX_TOKEN: Final = "/_oauth/token"
|
||||
_VERTEX_RAW_PREDICT: Final = (
|
||||
f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/anthropic/models/{CLAUDE}:rawPredict"
|
||||
)
|
||||
_VERTEX_ACCESS_TOKEN: Final = json.dumps({"access_token": "scripted", "token_type": "Bearer", "expires_in": 3600})
|
||||
|
||||
|
||||
def _vertex(expected: Expected) -> Expected:
|
||||
return _with_fields(_without(expected, "model"), {"anthropic_version": "vertex-2023-10-16"})
|
||||
|
||||
|
||||
def _vertex_reply(request: Request) -> Reply:
|
||||
if request.target == _VERTEX_TOKEN:
|
||||
return Reply(body=_VERTEX_ACCESS_TOKEN.encode())
|
||||
marker: Final = marker_of(request)
|
||||
return Reply(body=message_json(f"msg_{marker}", CLAUDE, answer(marker)))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "fields", "expected", "status"),
|
||||
[
|
||||
pytest.param(
|
||||
CHAT,
|
||||
{"reasoning_effort": "medium"},
|
||||
_vertex(_anthropic_chat("high", {"type": "adaptive", "display": "summarized"})),
|
||||
200,
|
||||
id="chat-medium-to-high",
|
||||
),
|
||||
pytest.param(
|
||||
MESSAGES,
|
||||
{"thinking": _BUDGET},
|
||||
_vertex(_native_messages({"effort": "high"}, _ADAPTIVE)),
|
||||
200,
|
||||
id="messages-budget-to-high",
|
||||
),
|
||||
pytest.param(CHAT, {"thinking": _FALSE_BUDGET}, None, 400, id="chat-false-budget-400"),
|
||||
],
|
||||
)
|
||||
def test_a_vertex_ai_claude_deployment_gets_one_of_its_declared_levels(
|
||||
gateway: Gateway, route: str, fields: dict[str, JsonValue], expected: Expected | None, status: int
|
||||
) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(_vertex_reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"vertex_ai/{CLAUDE}",
|
||||
api_base=wire.url,
|
||||
api_key=None,
|
||||
vertex_project=_VERTEX_PROJECT,
|
||||
vertex_location=_VERTEX_LOCATION,
|
||||
vertex_credentials=service_account_json(_VERTEX_PROJECT, wire.url),
|
||||
model_info=dict(_CLAUDE_INFO),
|
||||
)
|
||||
body: Final = (
|
||||
chat_request(model, marker, **fields) if route == CHAT else messages_request(model, marker, **fields)
|
||||
)
|
||||
response: Final = gateway.request("POST", route, body)
|
||||
received: Final = tuple(request for request in wire.drain() if request.target != _VERTEX_TOKEN)
|
||||
assert response.status_code == status, response.text[:2000]
|
||||
if expected is None:
|
||||
assert _NON_NUMERIC_BUDGET in response.text, response.text[:2000]
|
||||
assert received == (), [request.body[:300] for request in received]
|
||||
return
|
||||
assert answer(marker) in response.text, response.text[:2000]
|
||||
assert [(request.method, request.target) for request in received] == [("POST", _VERTEX_RAW_PREDICT)]
|
||||
assert body_of(received[0]) == expected(marker)
|
||||
|
||||
|
||||
def _single(received: tuple[Request, ...], marker: str) -> Request:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,67 @@
|
|||
"""An SDK caller that reuses one ``metadata=`` dict across ``litellm.completion`` calls to a Bedrock deployment whose
|
||||
``guardrailConfig`` has ``trace: enabled`` is billed, on each call, for that call's guardrail trace only, and the
|
||||
caller's dict is left as the caller wrote it (LIT-10239)"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import eventually
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
|
||||
_MODEL: Final = "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
_PROMPT_TOKENS: Final = 5
|
||||
_COMPLETION_TOKENS: Final = 1
|
||||
_USAGE: Final[Mapping[str, int]] = {"contentPolicyUnits": 1, "topicPolicyUnits": 1}
|
||||
_GUARDRAIL_IDS: Final = ("guardrail-one", "guardrail-two")
|
||||
|
||||
|
||||
def _bedrock(request: Request) -> Reply:
|
||||
"""Bedrock's Converse runtime: a one-token answer and the trace of the guardrail the request named."""
|
||||
guardrail_id: Final = json.loads(request.body)["guardrailConfig"]["guardrailIdentifier"]
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6},
|
||||
"metrics": {"latencyMs": 3},
|
||||
"trace": {"guardrail": {"inputAssessment": {guardrail_id: {"invocationMetrics": {"usage": dict(_USAGE)}}}}},
|
||||
}
|
||||
return Reply(body=json.dumps(body).encode())
|
||||
|
||||
|
||||
def _guardrail_cost() -> float:
|
||||
rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"]
|
||||
return sum(units * rates[counter] for counter, units in _USAGE.items())
|
||||
|
||||
|
||||
def _llm_cost() -> float:
|
||||
prices: Final = litellm.model_cost[_MODEL.removeprefix("bedrock/")]
|
||||
return _PROMPT_TOKENS * prices["input_cost_per_token"] + _COMPLETION_TOKENS * prices["output_cost_per_token"]
|
||||
|
||||
|
||||
def test_each_call_sharing_one_metadata_dict_is_billed_its_own_trace_only() -> None:
|
||||
billed: list[float] = [] # mutable-ok: the success callback appends what each call was billed
|
||||
litellm.success_callback.append(
|
||||
lambda kwargs, response, start, end: billed.append(kwargs["standard_logging_object"]["response_cost"])
|
||||
)
|
||||
shared_metadata: Final = {"tag": "x"}
|
||||
with wire_server(_bedrock) as wire:
|
||||
for guardrail_id in _GUARDRAIL_IDS:
|
||||
litellm.completion(
|
||||
model=_MODEL,
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
api_base=wire.url,
|
||||
aws_access_key_id="AKIASCRIPTED",
|
||||
aws_secret_access_key="scripted",
|
||||
aws_region_name="us-east-1",
|
||||
guardrailConfig={"guardrailIdentifier": guardrail_id, "guardrailVersion": "1", "trace": "enabled"},
|
||||
metadata=shared_metadata,
|
||||
)
|
||||
costs: Final = eventually(lambda: tuple(billed), lambda seen: len(seen) == len(_GUARDRAIL_IDS), seconds=30)
|
||||
|
||||
assert costs == pytest.approx((_llm_cost() + _guardrail_cost(),) * len(_GUARDRAIL_IDS))
|
||||
assert shared_metadata == {"tag": "x"}
|
||||
|
|
@ -12,6 +12,7 @@ from pydantic import JsonValue, TypeAdapter
|
|||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
|
||||
_MODEL: Final = "gpt-5.4"
|
||||
|
|
@ -118,3 +119,30 @@ async def test_cooldown_skips_only_the_failed_duplicate_model_deployment() -> No
|
|||
"/primary/chat/completions",
|
||||
"/backup/chat/completions",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"router_kwargs",
|
||||
[{"fallbacks": [{"primary": "backup"}]}, {"default_fallbacks": "backup"}],
|
||||
ids=["string_fallback_target", "string_default_fallbacks"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(router_kwargs: dict[str, object]) -> None:
|
||||
with wire_server(_peer) as wire:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
_deployment("primary", f"{wire.url}/primary", "primary-deployment"),
|
||||
_deployment("backup", f"{wire.url}/backup", "backup-deployment"),
|
||||
],
|
||||
allowed_fails=0,
|
||||
num_retries=0,
|
||||
**router_kwargs,
|
||||
)
|
||||
with pytest.raises(litellm.AuthenticationError):
|
||||
await router.acompletion(
|
||||
model="primary",
|
||||
messages=[{"role": "user", "content": f"cooldown {uuid.uuid4().hex}"}],
|
||||
)
|
||||
cooled: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None)
|
||||
|
||||
assert tuple(cooled) == ("primary-deployment",)
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -26,7 +26,7 @@
|
|||
},
|
||||
"results": [
|
||||
{
|
||||
"breakdown": {
|
||||
"breakdown": {"guardrails": {},
|
||||
"api_keys": {
|
||||
"key-a": {
|
||||
"metadata": {
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@
|
|||
},
|
||||
"results": [
|
||||
{
|
||||
"breakdown": {
|
||||
"breakdown": {"guardrails": {},
|
||||
"api_keys": {
|
||||
"key-a": {
|
||||
"metadata": {
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@
|
|||
},
|
||||
"results": [
|
||||
{
|
||||
"breakdown": {
|
||||
"breakdown": {"guardrails": {},
|
||||
"api_keys": {
|
||||
"key-a": {
|
||||
"metadata": {
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@
|
|||
},
|
||||
"results": [
|
||||
{
|
||||
"breakdown": {
|
||||
"breakdown": {"guardrails": {},
|
||||
"api_keys": {
|
||||
"key-a": {
|
||||
"metadata": {
|
||||
|
|
|
|||
|
|
@ -243,3 +243,183 @@ async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_mo
|
|||
assert object_value(object_value(models["gpt-5"])["metrics"])["spend"] == 10.0
|
||||
finally:
|
||||
_clean(day)
|
||||
|
||||
|
||||
def _seed_guardrail_units(day: str, rows: Sequence[tuple[str, str, str, float | None, int, int]]) -> None:
|
||||
for guardrail_id, api_key, usage_unit, cost, units, untracked_units in rows:
|
||||
write_rows(
|
||||
'INSERT INTO "LiteLLM_DailyGuardrailUsageUnits" (guardrail_id, date, team_id, api_key, usage_unit,'
|
||||
" units, cost, untracked_units, updated_at) VALUES (%s, %s, '', %s, %s, %s, %s, %s, now())",
|
||||
(
|
||||
guardrail_id,
|
||||
day,
|
||||
api_key,
|
||||
usage_unit,
|
||||
str(units),
|
||||
None if cost is None else str(cost),
|
||||
str(untracked_units),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _guardrail_spend(breakdown: dict[str, JsonValue]) -> dict[str, tuple[float, dict[str, float]]]:
|
||||
return {
|
||||
guardrail_id: (
|
||||
float(object_value(object_value(entry)["metrics"])["spend"]),
|
||||
{
|
||||
api_key: float(object_value(object_value(key_entry)["metrics"])["spend"])
|
||||
for api_key, key_entry in object_value(object_value(entry)["api_key_breakdown"]).items()
|
||||
},
|
||||
)
|
||||
for guardrail_id, entry in object_value(breakdown["guardrails"]).items()
|
||||
}
|
||||
|
||||
|
||||
def _day_breakdown(gateway: Gateway, day: str, **params: str) -> dict[str, JsonValue]:
|
||||
results: Final = _RESULTS.validate_python(_activity(gateway, day, **params)["results"])
|
||||
assert len(results) == 1
|
||||
return object_value(object_value(results[0])["breakdown"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_daily_activity_aggregated_breaks_down_guardrail_spend(gateway: Gateway) -> None:
|
||||
day: Final = _unique_day()
|
||||
_seed(
|
||||
day,
|
||||
[
|
||||
(_row_id(), "user-a", day, "key-a", "gpt-5", "", "bedrock", None, "/v1/chat/completions", 10, 1.0, 1, 1),
|
||||
(_row_id(), "user-b", day, "key-b", "gpt-5", "", "bedrock", None, "/v1/chat/completions", 10, 2.0, 1, 1),
|
||||
],
|
||||
)
|
||||
_seed_guardrail_units(
|
||||
day,
|
||||
[
|
||||
("pii-filter", "key-a", "contentPolicyUnits", 0.10, 5, 2),
|
||||
("pii-filter", "key-a", "sensitiveInformationPolicyUnits", 0.05, 1, 0),
|
||||
("pii-filter", "key-b", "contentPolicyUnits", 0.20, 1, 0),
|
||||
("pii-filter", "key-b", "legacyUnits", None, 7, 0),
|
||||
("topic-filter", "key-a", "topicPolicyUnits", 0.01, 1, 0),
|
||||
("topic-filter", "key-b", "topicPolicyUnits", None, 1, 0),
|
||||
("pii-filter", "orphan-key", "contentPolicyUnits", 9.0, 1, 0),
|
||||
],
|
||||
)
|
||||
try:
|
||||
everyone: Final = _day_breakdown(gateway, day)
|
||||
assert _guardrail_spend(everyone) == {
|
||||
"pii-filter": (pytest.approx(0.35), {"key-a": pytest.approx(0.15), "key-b": pytest.approx(0.20)}),
|
||||
"topic-filter": (pytest.approx(0.01), {"key-a": pytest.approx(0.01), "key-b": 0.0}),
|
||||
}
|
||||
api_keys: Final = object_value(everyone["api_keys"])
|
||||
pii_filter: Final = object_value(object_value(everyone["guardrails"])["pii-filter"])
|
||||
assert object_value(pii_filter["metrics"]) == {"spend": pytest.approx(0.35), "untracked_units": 9}
|
||||
pii_api_keys: Final = object_value(pii_filter["api_key_breakdown"])
|
||||
assert object_value(object_value(pii_api_keys["key-a"])["metrics"]) == {
|
||||
"spend": pytest.approx(0.15),
|
||||
"untracked_units": 2,
|
||||
}
|
||||
assert object_value(object_value(pii_api_keys["key-b"])["metrics"]) == {
|
||||
"spend": pytest.approx(0.20),
|
||||
"untracked_units": 7,
|
||||
}
|
||||
|
||||
assert object_value(everyone["guardrails"])["topic-filter"] == {
|
||||
"metrics": {"spend": pytest.approx(0.01), "untracked_units": 1},
|
||||
"metadata": {"guardrail_name": "topic-filter"},
|
||||
"api_key_breakdown": {
|
||||
"key-a": {
|
||||
"metrics": {"spend": pytest.approx(0.01), "untracked_units": 0},
|
||||
"metadata": object_value(api_keys["key-a"])["metadata"],
|
||||
},
|
||||
"key-b": {
|
||||
"metrics": {"spend": 0.0, "untracked_units": 1},
|
||||
"metadata": object_value(api_keys["key-b"])["metadata"],
|
||||
},
|
||||
},
|
||||
}
|
||||
assert object_value(object_value(object_value(everyone["models"])["gpt-5"])["metrics"])["spend"] == 3.0
|
||||
|
||||
assert _guardrail_spend(_day_breakdown(gateway, day, user_id="user-a")) == {
|
||||
"pii-filter": (pytest.approx(0.15), {"key-a": pytest.approx(0.15)}),
|
||||
"topic-filter": (pytest.approx(0.01), {"key-a": pytest.approx(0.01)}),
|
||||
}
|
||||
assert _guardrail_spend(_day_breakdown(gateway, day, api_key_limit="1")) == {
|
||||
"pii-filter": (pytest.approx(0.35), {"key-b": pytest.approx(0.20)}),
|
||||
"topic-filter": (pytest.approx(0.01), {"key-b": 0.0}),
|
||||
}
|
||||
limited_pii_filter: Final = object_value(
|
||||
object_value(_day_breakdown(gateway, day, api_key_limit="1")["guardrails"])["pii-filter"]
|
||||
)
|
||||
assert object_value(limited_pii_filter["metrics"])["untracked_units"] == 9
|
||||
assert set(object_value(limited_pii_filter["api_key_breakdown"])) == {"key-b"}
|
||||
assert _guardrail_spend(_day_breakdown(gateway, day, model="gpt-5")) == {}
|
||||
assert _guardrail_spend(_day_breakdown(gateway, day, model="")) == _guardrail_spend(everyone)
|
||||
finally:
|
||||
_clean(day)
|
||||
write_rows('DELETE FROM "LiteLLM_DailyGuardrailUsageUnits" WHERE date = %s', (day,))
|
||||
|
||||
|
||||
def _entity_breakdown(gateway: Gateway, url: str, day: str, **params: str) -> dict[str, JsonValue]:
|
||||
response: Final = gateway.request("GET", url, params={"start_date": day, "end_date": day, **params})
|
||||
assert response.status_code == 200, response.text
|
||||
results: Final = _RESULTS.validate_python(object_value(response.json())["results"])
|
||||
assert len(results) == 1
|
||||
return object_value(object_value(results[0])["breakdown"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("url", "table", "entity_column", "filter_param"),
|
||||
[
|
||||
("/team/daily/activity/aggregated", "LiteLLM_DailyTeamSpend", "team_id", "team_ids"),
|
||||
(
|
||||
"/organization/daily/activity/aggregated",
|
||||
"LiteLLM_DailyOrganizationSpend",
|
||||
"organization_id",
|
||||
"organization_ids",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_team_and_organization_activity_break_down_guardrail_spend_by_name(
|
||||
gateway: Gateway, url: str, table: str, entity_column: str, filter_param: str
|
||||
) -> None:
|
||||
day: Final = _unique_day()
|
||||
db_guardrail_id: Final = str(uuid.uuid4())
|
||||
db_guardrail_name: Final = f"pii-filter-{db_guardrail_id[:8]}"
|
||||
for entity_id, api_key, spend in (("entity-a", "key-a", 1.0), ("entity-b", "key-b", 2.0)):
|
||||
write_rows(
|
||||
f'INSERT INTO "{table}" (id, {entity_column}, date, api_key, model, spend, api_requests,'
|
||||
" successful_requests, updated_at) VALUES (%s, %s, %s, %s, 'gpt-5', %s, 1, 1, now())",
|
||||
(_row_id(), entity_id, day, api_key, str(spend)),
|
||||
)
|
||||
write_rows(
|
||||
'INSERT INTO "LiteLLM_GuardrailsTable" (guardrail_id, guardrail_name, litellm_params, updated_at)'
|
||||
" VALUES (%s, %s, '{}', now())",
|
||||
(db_guardrail_id, db_guardrail_name),
|
||||
)
|
||||
_seed_guardrail_units(
|
||||
day,
|
||||
[
|
||||
(db_guardrail_id, "key-a", "contentPolicyUnits", 0.10, 1, 0),
|
||||
(db_guardrail_id, "key-b", "contentPolicyUnits", 0.20, 1, 0),
|
||||
("config-topic-filter", "key-a", "topicPolicyUnits", 0.05, 1, 0),
|
||||
(db_guardrail_id, "orphan-key", "contentPolicyUnits", 9.0, 1, 0),
|
||||
],
|
||||
)
|
||||
try:
|
||||
everyone: Final = _entity_breakdown(gateway, url, day)
|
||||
assert _guardrail_spend(everyone) == {
|
||||
db_guardrail_id: (pytest.approx(0.30), {"key-a": pytest.approx(0.10), "key-b": pytest.approx(0.20)}),
|
||||
"config-topic-filter": (pytest.approx(0.05), {"key-a": pytest.approx(0.05)}),
|
||||
}
|
||||
guardrails: Final = object_value(everyone["guardrails"])
|
||||
assert object_value(guardrails[db_guardrail_id])["metadata"] == {"guardrail_name": db_guardrail_name}
|
||||
assert object_value(guardrails["config-topic-filter"])["metadata"] == {"guardrail_name": "config-topic-filter"}
|
||||
|
||||
assert _guardrail_spend(_entity_breakdown(gateway, url, day, **{filter_param: "entity-a"})) == {
|
||||
db_guardrail_id: (pytest.approx(0.10), {"key-a": pytest.approx(0.10)}),
|
||||
"config-topic-filter": (pytest.approx(0.05), {"key-a": pytest.approx(0.05)}),
|
||||
}
|
||||
finally:
|
||||
write_rows(f'DELETE FROM "{table}" WHERE date = %s', (day,))
|
||||
write_rows('DELETE FROM "LiteLLM_DailyGuardrailUsageUnits" WHERE date = %s', (day,))
|
||||
write_rows('DELETE FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id = %s', (db_guardrail_id,))
|
||||
|
|
|
|||
|
|
@ -13775,3 +13775,35 @@ def test_record_native_bedrock_guardrail_bills_nothing_for_a_response_served_fro
|
|||
logging_obj.record_native_bedrock_guardrail({"trace": trace})
|
||||
|
||||
assert "standard_logging_guardrail_information" not in logging_obj.model_call_details["litellm_params"]["metadata"]
|
||||
|
||||
|
||||
def test_record_native_bedrock_guardrail_bills_each_call_its_own_trace_when_calls_share_one_metadata_dict() -> None:
|
||||
import datetime
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
|
||||
|
||||
shared_metadata = {"tag": "x"}
|
||||
traces = tuple(
|
||||
{"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": units}}}}}}
|
||||
for units in (1, 2)
|
||||
)
|
||||
costs = []
|
||||
for call_index, trace in enumerate(traces):
|
||||
logging_obj = Logging(
|
||||
model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time=datetime.datetime.now(),
|
||||
litellm_call_id=f"call-shared-metadata-{call_index}",
|
||||
function_id="fn",
|
||||
)
|
||||
logging_obj.model_call_details["litellm_params"] = {"metadata": shared_metadata}
|
||||
logging_obj.record_native_bedrock_guardrail({"trace": trace})
|
||||
recorded = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
|
||||
costs.append(guardrail_information_cost(recorded))
|
||||
|
||||
(first_cost, second_cost) = costs
|
||||
assert second_cost == pytest.approx(2 * first_cost)
|
||||
assert shared_metadata == {"tag": "x"}
|
||||
|
|
|
|||
|
|
@ -7,7 +7,9 @@ from litellm_proxy_extras.migration_data_checks import (
|
|||
MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR,
|
||||
BlockedChange,
|
||||
Check,
|
||||
ColumnType,
|
||||
CouldNotVerify,
|
||||
_count_query,
|
||||
btree_max_entry_bytes,
|
||||
checks_for_index,
|
||||
checks_for_migration,
|
||||
|
|
@ -113,6 +115,19 @@ def test_the_btree_entry_limit_on_the_default_page_is_the_number_in_postgres_err
|
|||
assert btree_max_entry_bytes(8192) == 2704
|
||||
|
||||
|
||||
def test_the_entry_size_query_sizes_a_text_column_by_its_stored_bytes_when_postgres_compressed_it() -> None:
|
||||
"""Postgres compresses a long value before it checks the entry size, so a repetitive
|
||||
3000-byte key indexes fine. The query has to size it as stored, never by its length alone."""
|
||||
(check,) = checks_for_index(
|
||||
"LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime")
|
||||
)
|
||||
types: Final = (ColumnType("text", -1, "i", -1), ColumnType("timestamp", 8, "d", 3))
|
||||
query: Final = _count_query(check, "public", types, 2704).as_string(None)
|
||||
stored: Final = 'least(octet_length("api_key"), pg_column_size("api_key"))'
|
||||
assert stored in query
|
||||
assert 'octet_length("api_key")' not in query.replace(stored, ""), query
|
||||
|
||||
|
||||
def test_the_report_names_the_rows_the_statement_and_the_fix_and_says_nothing_was_applied() -> None:
|
||||
(check,) = checks_for_index(
|
||||
"LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime")
|
||||
|
|
|
|||
|
|
@ -1523,6 +1523,7 @@ async def test_summary_model_denied_when_user_over_model_budget():
|
|||
user_id="user-over-budget",
|
||||
user_model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}},
|
||||
model="claude-haiku-4-5",
|
||||
team_id=None,
|
||||
)
|
||||
import inspect
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import socket
|
|||
import sys
|
||||
import struct
|
||||
import threading
|
||||
import weakref
|
||||
from typing import Final
|
||||
|
||||
import aiohttp
|
||||
|
|
@ -1383,3 +1384,57 @@ def test_a_transport_collected_while_its_loop_runs_elsewhere_closes_its_session_
|
|||
loop.call_soon_threadsafe(loop.stop)
|
||||
thread.join(timeout=10)
|
||||
loop.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_stream_keeps_reading_after_the_client_that_opened_it_is_collected():
|
||||
from aiohttp import web
|
||||
|
||||
release: Final = asyncio.Event()
|
||||
body: Final = tuple(f"chunk{i}\n".encode() for i in range(4))
|
||||
|
||||
async def handler(request: web.Request) -> web.StreamResponse:
|
||||
response: Final = web.StreamResponse()
|
||||
await response.prepare(request)
|
||||
await response.write(body[0])
|
||||
await release.wait()
|
||||
for chunk in body[1:]:
|
||||
await response.write(chunk)
|
||||
await response.write_eof()
|
||||
return response
|
||||
|
||||
app: Final = web.Application()
|
||||
app.router.add_get("/stream", handler)
|
||||
runner: Final = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
sock: Final = socket.socket()
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
port: Final[int] = sock.getsockname()[1]
|
||||
await web.SockSite(runner, sock).start()
|
||||
|
||||
transport: Final = LiteLLMAiohttpTransport(client=aiohttp.ClientSession)
|
||||
client: Final = httpx.AsyncClient(transport=transport)
|
||||
transport_ref: Final = weakref.ref(transport)
|
||||
try:
|
||||
response: Final = await client.send(client.build_request("GET", f"http://127.0.0.1:{port}/stream"), stream=True)
|
||||
session: Final = transport.client
|
||||
assert isinstance(session, aiohttp.ClientSession)
|
||||
chunks: Final = response.aiter_bytes()
|
||||
first: Final = await chunks.__anext__()
|
||||
|
||||
del client, transport
|
||||
gc.collect()
|
||||
await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
|
||||
release.set()
|
||||
|
||||
rest: Final = [chunk async for chunk in chunks]
|
||||
assert first + b"".join(rest) == b"".join(body), (first, rest)
|
||||
assert not session.closed, "the session was closed under a stream that was still being read"
|
||||
|
||||
del response, chunks
|
||||
gc.collect()
|
||||
await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
|
||||
assert transport_ref() is None, "the finished stream kept the transport alive"
|
||||
assert session.closed, "the owned session stayed open after its last stream finished"
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
|
|
|||
|
|
@ -253,3 +253,29 @@ def test_cost_per_token_bills_the_served_priority_tier(
|
|||
prompt_cost, completion_cost = cost_per_token(model="databricks/dbrx-tiered-test", usage=usage)
|
||||
assert prompt_cost == pytest.approx(30 * 0.001)
|
||||
assert completion_cost == pytest.approx(40 * 0.002)
|
||||
|
||||
|
||||
RATES_WITHOUT_A_PRIORITY_FIELD: Final = frozenset({"input_cost_per_image_token"})
|
||||
|
||||
|
||||
def _is_standard_rate_key(key: str) -> bool:
|
||||
return (
|
||||
("_cost_per_" in key or key.endswith("_token_cost"))
|
||||
and "dbu" not in key
|
||||
and key not in RATES_WITHOUT_A_PRIORITY_FIELD
|
||||
and not key.endswith(("_priority", "_flex", "_batches"))
|
||||
)
|
||||
|
||||
|
||||
def test_every_databricks_priority_row_prices_every_standard_rate_at_the_priority_tier() -> None:
|
||||
cost_map: Final = json.loads(MAIN_PRICES.read_text())
|
||||
priority_rows: Final = {
|
||||
key: row
|
||||
for key, row in cost_map.items()
|
||||
if key.startswith("databricks/") and "input_cost_per_token_priority" in row
|
||||
}
|
||||
|
||||
assert priority_rows
|
||||
for key, row in priority_rows.items():
|
||||
missing = sorted(f for f in row if _is_standard_rate_key(f) and f"{f}_priority" not in row)
|
||||
assert not missing, f"{key} bills these at the standard rate on priority requests: {missing}"
|
||||
|
|
|
|||
|
|
@ -10655,3 +10655,34 @@ async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_n
|
|||
sent: Final = await _sent_thinking_config("gemini-3.8-flash", {"reasoning_effort": "high"})
|
||||
|
||||
assert sent["includeThoughts"] is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("chunk_indices", [[0, 1], [2, 1, 0]])
|
||||
def test_grounding_support_preserves_every_web_source(chunk_indices: list[int]) -> None:
|
||||
metadata: Final = {
|
||||
"groundingChunks": [
|
||||
{"web": {"uri": "https://example.com/a", "title": "A"}},
|
||||
{"web": {"uri": "https://example.com/b", "title": "B"}},
|
||||
{"retrievedContext": {"uri": "gs://document"}},
|
||||
],
|
||||
"groundingSupports": [
|
||||
{"segment": {"startIndex": 0, "endIndex": 5}, "groundingChunkIndices": chunk_indices},
|
||||
{"segment": {"startIndex": 6, "endIndex": 11}, "groundingChunkIndices": [1]},
|
||||
],
|
||||
}
|
||||
|
||||
annotations: Final = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata], "First claim")
|
||||
|
||||
expected_sources: Final = [(index, 0, 5) for index in chunk_indices if index in (0, 1)] + [(1, 6, 11)]
|
||||
assert annotations == [
|
||||
{
|
||||
"type": "url_citation",
|
||||
"url_citation": {
|
||||
"start_index": start,
|
||||
"end_index": end,
|
||||
"url": metadata["groundingChunks"][index]["web"]["uri"],
|
||||
"title": metadata["groundingChunks"][index]["web"]["title"],
|
||||
},
|
||||
}
|
||||
for index, start, end in expected_sources
|
||||
]
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@
|
|||
import litellm.proxy
|
||||
import litellm.proxy.proxy_server
|
||||
|
||||
from typing import Dict, Final, List, Optional
|
||||
from typing import Dict, Final, List, Mapping, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -1367,7 +1367,9 @@ async def test_check_user_model_budget(user_id, user_model_max_budget, expected_
|
|||
calls = []
|
||||
|
||||
class _Limiter:
|
||||
async def is_user_within_model_budget(self, user_id, user_model_max_budget, model):
|
||||
async def is_user_within_model_budget(
|
||||
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None
|
||||
) -> bool:
|
||||
calls.append((user_id, user_model_max_budget, model))
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import pytest
|
|||
from litellm.proxy.db.budget_window_spend_writer import (
|
||||
WindowSeedTotals,
|
||||
commit_window_spend_updates,
|
||||
open_window_spend_row,
|
||||
roll_window_spend_row,
|
||||
spend_logs_seed_totals,
|
||||
)
|
||||
|
|
@ -405,6 +406,29 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o
|
|||
assert params[:4] == ("team", "t1", "30d", datetime(2026, 8, 31))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("replaces_window_start", "replaced_param"), [(None, None), (WINDOW_A, datetime(2026, 8, 1))]
|
||||
)
|
||||
async def test_open_window_spend_row_moves_only_the_window_the_caller_saw_end(replaces_window_start, replaced_param):
|
||||
db = _FakeDB()
|
||||
|
||||
await open_window_spend_row(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
entity_type="model_max_budget",
|
||||
entity_id="k1",
|
||||
window_duration="1d",
|
||||
window_start=WINDOW_B,
|
||||
replaces_window_start=replaces_window_start,
|
||||
)
|
||||
|
||||
((query, params),) = db.execute_raw_calls
|
||||
normalized = " ".join(query.split())
|
||||
assert "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE" in normalized
|
||||
assert normalized.endswith("WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')")
|
||||
assert params == ("model_max_budget", "k1", "1d", datetime(2026, 8, 31), replaced_param)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_receives_the_batch_earliest_start_as_its_cutoff():
|
||||
db = _FakeDB(existing_rows=[])
|
||||
|
|
|
|||
199
tests/unit/proxy/hooks/test_model_budget_shared_windows.py
Normal file
199
tests/unit/proxy/hooks/test_model_budget_shared_windows.py
Normal file
|
|
@ -0,0 +1,199 @@
|
|||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction
|
||||
from litellm.proxy.hooks.model_budget_shared_windows import DatabaseSharedWindows, SharedWindow
|
||||
|
||||
SPEND_KEY: Final = "virtual_key_spend:key-hash:claude-opus*:1d"
|
||||
WINDOW_START: Final = datetime(2026, 10, 20, 9, 30, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
class _Table:
|
||||
def __init__(self, row: object | None = None, error: Exception | None = None) -> None:
|
||||
self.row = row
|
||||
self.error = error
|
||||
self.reads = 0
|
||||
|
||||
async def find_unique(self, where: object) -> object | None:
|
||||
self.reads = self.reads + 1
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
return self.row
|
||||
|
||||
|
||||
class _Database:
|
||||
def __init__(self, table: _Table) -> None:
|
||||
self.litellm_budgetwindowspend = table
|
||||
self.opened_with: tuple[object, ...] = ()
|
||||
|
||||
async def execute_raw(self, query: str, *params: object) -> int:
|
||||
self.opened_with = params
|
||||
self.litellm_budgetwindowspend.row = SimpleNamespace(window_start=params[3], spend=0.0)
|
||||
return 1
|
||||
|
||||
|
||||
class _OutracedDatabase:
|
||||
def __init__(self, table: _Table, winner: object) -> None:
|
||||
self.litellm_budgetwindowspend = table
|
||||
self._winner = winner
|
||||
|
||||
async def execute_raw(self, query: str, *params: object) -> int:
|
||||
self.litellm_budgetwindowspend.row = self._winner
|
||||
return 0
|
||||
|
||||
|
||||
class _Sink:
|
||||
def __init__(self) -> None:
|
||||
self.transactions: tuple[WindowSpendTransaction, ...] = ()
|
||||
|
||||
async def __call__(self, transaction: WindowSpendTransaction) -> None:
|
||||
self.transactions = (*self.transactions, transaction)
|
||||
|
||||
|
||||
class _LaggingReplica:
|
||||
def __init__(self, writer: _Database) -> None:
|
||||
self.litellm_budgetwindowspend = _Table()
|
||||
self._writer = writer
|
||||
|
||||
async def execute_raw(self, query: str, *params: object) -> int:
|
||||
return await self._writer.execute_raw(query, *params)
|
||||
|
||||
|
||||
def _database_windows(table: _Table | None, sink: _Sink, lagging_replica: bool = False) -> DatabaseSharedWindows:
|
||||
writer: Final = None if table is None else _Database(table)
|
||||
prisma: Final = (
|
||||
None
|
||||
if writer is None
|
||||
else SimpleNamespace(db=_LaggingReplica(writer) if lagging_replica else writer, writer_db=writer)
|
||||
)
|
||||
return DatabaseSharedWindows(
|
||||
prisma_client_source=lambda: prisma, # pyright: ignore[reportArgumentType] # duck-typed prisma double
|
||||
sink=sink,
|
||||
)
|
||||
|
||||
|
||||
def test_a_window_is_open_until_its_period_has_passed():
|
||||
window: Final = SharedWindow(started_at=WINDOW_START.timestamp(), spend=1.0)
|
||||
day: Final = 24 * 60 * 60
|
||||
|
||||
assert window.is_open(budget_duration="1d", now=WINDOW_START.timestamp() + day - 1)
|
||||
assert not window.is_open(budget_duration="1d", now=WINDOW_START.timestamp() + day)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_returns_the_stored_window_and_rereads_only_after_the_cache_expires():
|
||||
table: Final = _Table(row=SimpleNamespace(window_start=WINDOW_START.replace(tzinfo=None), spend=2.5))
|
||||
windows: Final = _database_windows(table, _Sink())
|
||||
|
||||
first: Final = await windows.read(spend_key=SPEND_KEY, budget_duration="1d")
|
||||
second: Final = await windows.read(spend_key=SPEND_KEY, budget_duration="1d")
|
||||
|
||||
assert first == SharedWindow(started_at=WINDOW_START.timestamp(), spend=2.5)
|
||||
assert second == first
|
||||
assert table.reads == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_unreadable_table_reads_as_no_shared_window():
|
||||
table: Final = _Table(error=RuntimeError("database is down"))
|
||||
windows: Final = _database_windows(table, _Sink())
|
||||
|
||||
assert await windows.read(spend_key=SPEND_KEY, budget_duration="1d") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_queues_the_charge_for_the_window_it_belongs_to():
|
||||
sink: Final = _Sink()
|
||||
windows: Final = _database_windows(_Table(), sink)
|
||||
|
||||
await windows.add(spend_key=SPEND_KEY, budget_duration="1d", started_at=WINDOW_START.timestamp(), spend=0.25)
|
||||
|
||||
assert len(sink.transactions) == 1
|
||||
queued: Final = sink.transactions[0]
|
||||
assert (queued["entity_id"], queued["window_duration"], queued["spend"]) == (SPEND_KEY, "1d", 0.25)
|
||||
assert datetime.fromisoformat(queued["window_start"]).replace(tzinfo=timezone.utc) == WINDOW_START
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nothing_is_queued_without_a_database():
|
||||
sink: Final = _Sink()
|
||||
windows: Final = _database_windows(None, sink)
|
||||
|
||||
await windows.add(spend_key=SPEND_KEY, budget_duration="1d", started_at=WINDOW_START.timestamp(), spend=0.25)
|
||||
|
||||
assert sink.transactions == ()
|
||||
assert await windows.read(spend_key=SPEND_KEY, budget_duration="1d") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_joins_the_window_another_replica_already_started():
|
||||
table: Final = _Table(row=SimpleNamespace(window_start=WINDOW_START.replace(tzinfo=None), spend=2.5))
|
||||
windows: Final = _database_windows(table, _Sink())
|
||||
|
||||
started_at: Final = await windows.open(
|
||||
spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp() + 3600
|
||||
)
|
||||
|
||||
assert started_at == WINDOW_START.timestamp()
|
||||
assert table.row.spend == 2.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_reads_the_window_from_the_writer_while_the_read_replica_lags():
|
||||
table: Final = _Table(row=SimpleNamespace(window_start=WINDOW_START.replace(tzinfo=None), spend=2.5))
|
||||
windows: Final = _database_windows(table, _Sink(), lagging_replica=True)
|
||||
|
||||
started_at: Final = await windows.open(
|
||||
spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp() + 3600
|
||||
)
|
||||
|
||||
assert started_at == WINDOW_START.timestamp()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_joins_the_window_a_racing_replica_opened_first():
|
||||
winner: Final = SimpleNamespace(window_start=WINDOW_START.replace(tzinfo=None), spend=2.5)
|
||||
database: Final = _OutracedDatabase(_Table(), winner)
|
||||
windows: Final = DatabaseSharedWindows(
|
||||
prisma_client_source=lambda: SimpleNamespace(db=database, writer_db=database), # pyright: ignore[reportArgumentType] # duck-typed prisma double
|
||||
sink=_Sink(),
|
||||
)
|
||||
|
||||
started_at: Final = await windows.open(spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp() + 60)
|
||||
|
||||
assert started_at == WINDOW_START.timestamp()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stored_start", [None, WINDOW_START - timedelta(days=2)])
|
||||
async def test_open_starts_a_window_at_now_when_none_is_open(stored_start: datetime | None):
|
||||
table: Final = _Table(
|
||||
row=None if stored_start is None else SimpleNamespace(window_start=stored_start.replace(tzinfo=None), spend=9.0)
|
||||
)
|
||||
windows: Final = _database_windows(table, _Sink())
|
||||
|
||||
started_at: Final = await windows.open(spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp())
|
||||
|
||||
assert started_at == WINDOW_START.timestamp()
|
||||
assert await windows.read(spend_key=SPEND_KEY, budget_duration="1d") == SharedWindow(
|
||||
started_at=WINDOW_START.timestamp(), spend=0.0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_starts_this_replicas_own_window_without_a_database():
|
||||
windows: Final = _database_windows(None, _Sink())
|
||||
|
||||
assert await windows.open(spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp()) == (
|
||||
WINDOW_START.timestamp()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_reports_no_shared_window_while_the_database_is_unreachable():
|
||||
windows: Final = _database_windows(_Table(error=RuntimeError("database is down")), _Sink())
|
||||
|
||||
assert await windows.open(spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp()) is None
|
||||
|
|
@ -1,19 +1,30 @@
|
|||
import asyncio
|
||||
import functools
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import Litellm_EntityType
|
||||
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction
|
||||
from litellm.proxy.hooks.model_budget_shared_windows import DatabaseSharedWindows, SharedWindows
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
build_model_max_budget_usage,
|
||||
model_access_groups_of,
|
||||
model_budget_spend_cache_key,
|
||||
resolve_model_budget,
|
||||
routed_models_of,
|
||||
team_member_budget_entity_id,
|
||||
)
|
||||
from litellm.router import Router
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.utils import LiteLLMBatch, Usage
|
||||
|
||||
|
|
@ -74,10 +85,14 @@ def _event(
|
|||
}
|
||||
|
||||
|
||||
async def _log_success(
|
||||
limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, event: dict[str, object], response_obj: object = None
|
||||
) -> None:
|
||||
await limiter.async_log_success_event(event, response_obj=response_obj, start_time=None, end_time=None)
|
||||
|
||||
|
||||
async def _poll(limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, batch: LiteLLMBatch, response_cost: float) -> None:
|
||||
await limiter.async_log_success_event(
|
||||
_event("aretrieve_batch", response_cost), response_obj=batch, start_time=None, end_time=None
|
||||
)
|
||||
await _log_success(limiter, _event("aretrieve_batch", response_cost), response_obj=batch)
|
||||
|
||||
|
||||
async def _chat(
|
||||
|
|
@ -86,7 +101,8 @@ async def _chat(
|
|||
user_id: str = USER_ID,
|
||||
team_id: str = TEAM_ID,
|
||||
) -> None:
|
||||
await limiter.async_log_success_event(
|
||||
await _log_success(
|
||||
limiter,
|
||||
_event(
|
||||
"acompletion",
|
||||
CHAT_COST,
|
||||
|
|
@ -94,9 +110,6 @@ async def _chat(
|
|||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
),
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -311,3 +324,568 @@ async def test_a_batch_polled_within_every_budget_window_is_never_charged_again(
|
|||
await _poll(limiter, finished, BATCH_COST)
|
||||
|
||||
assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST)
|
||||
|
||||
|
||||
OPUS_GROUP: Final = "claude-opus"
|
||||
OPUS_GROUP_MEMBERS: Final = MappingProxyType(
|
||||
{"claude-opus-4-8": frozenset({OPUS_GROUP}), "claude-opus-5-5": frozenset({OPUS_GROUP})}
|
||||
)
|
||||
FAMILY_CAP: Final = {"budget_limit": CHAT_COST, "time_period": "1d"} # mutable-ok: BudgetConfig validates dicts only
|
||||
ROOMY_CAP: Final = {
|
||||
"budget_limit": CHAT_COST * 100,
|
||||
"time_period": "1d",
|
||||
} # mutable-ok: BudgetConfig validates dicts only
|
||||
WINDOW_EPOCH: Final = datetime(2026, 10, 20, 9, 30, tzinfo=timezone.utc).timestamp()
|
||||
|
||||
|
||||
def _model_access_groups(
|
||||
groups_by_model: Mapping[str, frozenset[str]],
|
||||
) -> Callable[[str, str | None], frozenset[str]]:
|
||||
return lambda model, _team_id: groups_by_model.get(model, frozenset())
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _WindowRow:
|
||||
window_start: datetime
|
||||
spend: float
|
||||
|
||||
|
||||
class _WindowSpendDatabase:
|
||||
def __init__(self) -> None:
|
||||
self.rows: Mapping[str, _WindowRow] = MappingProxyType({})
|
||||
self.queued: tuple[WindowSpendTransaction, ...] = ()
|
||||
self.reachable = True
|
||||
|
||||
@property
|
||||
def litellm_budgetwindowspend(self) -> "_WindowSpendDatabase":
|
||||
return self
|
||||
|
||||
async def find_unique(self, where: Mapping[str, Mapping[str, str]]) -> _WindowRow | None:
|
||||
await asyncio.sleep(0)
|
||||
self._answer()
|
||||
return self.rows.get(where["entity_type_entity_id_window_duration"]["entity_id"])
|
||||
|
||||
async def execute_raw(
|
||||
self,
|
||||
query: str,
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_duration: str,
|
||||
window_start: datetime,
|
||||
replaces_window_start: datetime | None,
|
||||
) -> int:
|
||||
await asyncio.sleep(0)
|
||||
self._answer()
|
||||
stored: Final = self.rows.get(entity_id)
|
||||
if stored is not None and (replaces_window_start is None or stored.window_start > replaces_window_start):
|
||||
return 0
|
||||
self.rows = MappingProxyType({**self.rows, entity_id: _WindowRow(window_start=window_start, spend=0.0)})
|
||||
return 1
|
||||
|
||||
async def queue(self, transaction: WindowSpendTransaction) -> None:
|
||||
self.queued = (*self.queued, transaction)
|
||||
|
||||
def flush(self) -> None:
|
||||
self.rows = functools.reduce(_upsert_window_spend, self.queued, self.rows)
|
||||
self.queued = ()
|
||||
|
||||
def replica(self) -> DatabaseSharedWindows:
|
||||
return DatabaseSharedWindows(
|
||||
prisma_client_source=lambda: SimpleNamespace(db=self, writer_db=self), # pyright: ignore[reportArgumentType] # duck-typed prisma double
|
||||
sink=self.queue,
|
||||
)
|
||||
|
||||
def _answer(self) -> None:
|
||||
if not self.reachable:
|
||||
raise ConnectionError("database is unreachable")
|
||||
|
||||
|
||||
def _upsert_window_spend(
|
||||
rows: Mapping[str, _WindowRow], transaction: WindowSpendTransaction
|
||||
) -> Mapping[str, _WindowRow]:
|
||||
window_start: Final = datetime.fromisoformat(transaction["window_start"])
|
||||
stored: Final = rows.get(transaction["entity_id"])
|
||||
row: Final = (
|
||||
_WindowRow(window_start=stored.window_start, spend=stored.spend + transaction["spend"])
|
||||
if stored is not None and stored.window_start >= window_start
|
||||
else _WindowRow(window_start=window_start, spend=transaction["spend"])
|
||||
)
|
||||
return MappingProxyType({**rows, transaction["entity_id"]: row})
|
||||
|
||||
|
||||
def _family_limiter(
|
||||
groups_by_model: Mapping[str, frozenset[str]] = OPUS_GROUP_MEMBERS,
|
||||
dual_cache: DualCache | None = None,
|
||||
shared_windows: SharedWindows | None = None,
|
||||
clock: Callable[[], float] = time.time,
|
||||
) -> PROXY_VirtualKeyModelMaxBudgetLimiter:
|
||||
return PROXY_VirtualKeyModelMaxBudgetLimiter(
|
||||
dual_cache=dual_cache or DualCache(),
|
||||
clock=clock,
|
||||
model_access_groups=_model_access_groups(groups_by_model),
|
||||
shared_windows=shared_windows or _WindowSpendDatabase().replica(),
|
||||
)
|
||||
|
||||
|
||||
def _key(model_max_budget: Mapping[str, object]) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(token=KEY_HASH, key_alias="opus-family", model_max_budget=dict(model_max_budget))
|
||||
|
||||
|
||||
async def _spend_on(
|
||||
limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, model: str, model_max_budget: Mapping[str, object]
|
||||
) -> None:
|
||||
await _log_success(
|
||||
limiter,
|
||||
{
|
||||
"call_type": "acompletion",
|
||||
"standard_logging_object": {
|
||||
"call_type": "acompletion",
|
||||
"response_cost": CHAT_COST,
|
||||
"model": f"anthropic/{model}",
|
||||
"model_group": model,
|
||||
"metadata": {"user_api_key_hash": KEY_HASH},
|
||||
},
|
||||
"litellm_params": {"metadata": {"user_api_key_model_max_budget": dict(model_max_budget)}},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("family_entry", ["claude-opus*", OPUS_GROUP])
|
||||
async def test_every_opus_version_draws_from_one_family_entry(family_entry: str):
|
||||
budget: Final = {family_entry: FAMILY_CAP}
|
||||
limiter: Final = _family_limiter()
|
||||
assert await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-5-5")
|
||||
|
||||
await _spend_on(limiter, model="claude-opus-4-8", model_max_budget=budget)
|
||||
|
||||
with pytest.raises(
|
||||
litellm.BudgetExceededError, match=rf"model=claude-opus-5-5, budget_model={re.escape(family_entry)}$"
|
||||
):
|
||||
await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-5-5")
|
||||
assert await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-sonnet-5")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_model_added_to_the_access_group_later_is_covered_with_no_key_change():
|
||||
budget: Final = {OPUS_GROUP: FAMILY_CAP}
|
||||
shared_cache: Final = DualCache()
|
||||
before_group_change: Final = _family_limiter(dual_cache=shared_cache)
|
||||
after_group_change: Final = _family_limiter(
|
||||
groups_by_model={**OPUS_GROUP_MEMBERS, "claude-opus-6": frozenset({OPUS_GROUP})}, dual_cache=shared_cache
|
||||
)
|
||||
|
||||
await _spend_on(before_group_change, model="claude-opus-4-8", model_max_budget=budget)
|
||||
|
||||
assert await before_group_change.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-6")
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await after_group_change.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-6")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("budget", "checked_model", "blocked"),
|
||||
[
|
||||
({"claude-opus-4-8": ROOMY_CAP, "claude-opus*": FAMILY_CAP}, "claude-opus-4-8", False),
|
||||
({"claude-opus-4-8": ROOMY_CAP, "claude-opus*": FAMILY_CAP}, "claude-opus-5-5", False),
|
||||
({OPUS_GROUP: FAMILY_CAP, "claude-opus*": ROOMY_CAP}, "claude-opus-5-5", True),
|
||||
({"claude-opus-4*": FAMILY_CAP, "claude-opus*": ROOMY_CAP}, "claude-opus-4-7", True),
|
||||
({"claude-opus-4*": FAMILY_CAP, "claude-opus*": ROOMY_CAP}, "claude-opus-5-5", False),
|
||||
({"claude-opus-4*": {"budget_limit": CHAT_COST}, "claude-opus*": FAMILY_CAP}, "claude-opus-5-5", True),
|
||||
({"claude-opus-4*": FAMILY_CAP, "c" + "*" * 15: ROOMY_CAP}, "claude-opus-4-7", True),
|
||||
],
|
||||
)
|
||||
async def test_the_most_specific_entry_takes_the_spend_and_gates_the_request(
|
||||
budget: Mapping[str, object], checked_model: str, blocked: bool
|
||||
):
|
||||
limiter: Final = _family_limiter()
|
||||
|
||||
await _spend_on(limiter, model="claude-opus-4-8", model_max_budget=budget)
|
||||
|
||||
if blocked:
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model=checked_model)
|
||||
else:
|
||||
assert await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model=checked_model)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_exact_entry_refusal_reads_as_before_for_a_provider_prefixed_model():
|
||||
budget: Final = {"claude-opus-4-8": FAMILY_CAP}
|
||||
limiter: Final = _family_limiter()
|
||||
|
||||
await _spend_on(limiter, model="claude-opus-4-8", model_max_budget=budget)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError, match=r"model=anthropic/claude-opus-4-8$"):
|
||||
await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model="anthropic/claude-opus-4-8")
|
||||
|
||||
|
||||
async def _spend_on_team_key(
|
||||
limiter: PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
model: str,
|
||||
key_model_max_budget: Mapping[str, object],
|
||||
team_model_max_budget: Mapping[str, object],
|
||||
) -> None:
|
||||
await _log_success(
|
||||
limiter,
|
||||
{
|
||||
"call_type": "acompletion",
|
||||
"standard_logging_object": {
|
||||
"call_type": "acompletion",
|
||||
"response_cost": CHAT_COST,
|
||||
"model": f"anthropic/{model}",
|
||||
"model_group": model,
|
||||
"metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_team_id": TEAM_ID},
|
||||
},
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"user_api_key_model_max_budget": dict(key_model_max_budget),
|
||||
"user_api_key_team_model_max_budget": dict(team_model_max_budget),
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("team_model_max_budget", "team_blocks"),
|
||||
[({"claude-opus*": FAMILY_CAP}, True), ({OPUS_GROUP: FAMILY_CAP}, True), ({"claude-opus-4-8": FAMILY_CAP}, False)],
|
||||
)
|
||||
async def test_a_key_entry_named_after_the_model_replaces_only_the_team_entry_of_that_name(
|
||||
team_model_max_budget: Mapping[str, object], team_blocks: bool
|
||||
):
|
||||
key_model_max_budget: Final = {"claude-opus-4-8": ROOMY_CAP}
|
||||
limiter: Final = _family_limiter()
|
||||
|
||||
await _spend_on_team_key(
|
||||
limiter,
|
||||
model="claude-opus-4-8",
|
||||
key_model_max_budget=key_model_max_budget,
|
||||
team_model_max_budget=team_model_max_budget,
|
||||
)
|
||||
team_check: Final = limiter.is_team_within_model_budget(
|
||||
team_id=TEAM_ID,
|
||||
team_model_max_budget=team_model_max_budget,
|
||||
key_model_max_budget=key_model_max_budget,
|
||||
model="claude-opus-4-8",
|
||||
)
|
||||
|
||||
if team_blocks:
|
||||
with pytest.raises(litellm.BudgetExceededError, match=r"budget_model="):
|
||||
await team_check
|
||||
else:
|
||||
assert await team_check
|
||||
|
||||
|
||||
TEAM_OPUS_GROUP: Final = "team-opus"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("budget_entry", [TEAM_OPUS_GROUP, "claude-opus*"])
|
||||
async def test_a_team_model_is_gated_and_charged_under_its_public_name(budget_entry: str):
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": f"model_name_{TEAM_ID}_opus",
|
||||
"litellm_params": {"model": "anthropic/claude-opus-4-8", "api_key": "sk-test"},
|
||||
"model_info": {
|
||||
"id": "team-opus-deployment",
|
||||
"team_id": TEAM_ID,
|
||||
"team_public_model_name": "claude-opus-4-8",
|
||||
"access_groups": [TEAM_OPUS_GROUP],
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(
|
||||
dual_cache=DualCache(),
|
||||
model_access_groups=lambda model, team_id: model_access_groups_of(router=router, model=model, team_id=team_id),
|
||||
shared_windows=_WindowSpendDatabase().replica(),
|
||||
)
|
||||
budget: Final = {budget_entry: FAMILY_CAP}
|
||||
team_key: Final = UserAPIKeyAuth(token=KEY_HASH, team_id=TEAM_ID, model_max_budget=budget)
|
||||
assert await limiter.is_key_within_model_budget(user_api_key_dict=team_key, model="claude-opus-4-8")
|
||||
|
||||
await _log_success(
|
||||
limiter,
|
||||
{
|
||||
"call_type": "acompletion",
|
||||
"standard_logging_object": {
|
||||
"call_type": "acompletion",
|
||||
"response_cost": CHAT_COST,
|
||||
"model": "anthropic/claude-opus-4-8",
|
||||
"model_group": f"model_name_{TEAM_ID}_opus",
|
||||
"metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_team_id": TEAM_ID},
|
||||
},
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"user_api_key_model_max_budget": budget,
|
||||
"model_info": router.get_deployment(model_id="team-opus-deployment").model_info.model_dump(),
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError, match=rf"budget_model={re.escape(budget_entry)}$"):
|
||||
await limiter.is_key_within_model_budget(user_api_key_dict=team_key, model="claude-opus-4-8")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_charged_while_the_database_is_unreachable_reaches_the_shared_window_without_wiping_it():
|
||||
budget: Final = {"claude-opus*": ROOMY_CAP}
|
||||
database: Final = _WindowSpendDatabase()
|
||||
clock: Final = _Clock()
|
||||
clock.advance(WINDOW_EPOCH)
|
||||
replica_a: Final = _family_limiter(shared_windows=database.replica(), clock=clock.now)
|
||||
replica_b: Final = _family_limiter(shared_windows=database.replica(), clock=clock.now)
|
||||
|
||||
await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget)
|
||||
database.flush()
|
||||
clock.advance(60)
|
||||
database.reachable = False
|
||||
await _spend_on(replica_b, model="claude-opus-4-8", model_max_budget=budget)
|
||||
database.reachable = True
|
||||
database.flush()
|
||||
clock.advance(60)
|
||||
await _spend_on(replica_b, model="claude-opus-5-5", model_max_budget=budget)
|
||||
database.flush()
|
||||
|
||||
usage: Final = await build_model_max_budget_usage(
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
entity_id=KEY_HASH,
|
||||
model_max_budget=budget,
|
||||
cache=DualCache(),
|
||||
now=datetime.fromtimestamp(clock.now(), tz=timezone.utc),
|
||||
shared_windows=database.replica(),
|
||||
)
|
||||
assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 3, 4)
|
||||
assert usage["claude-opus*"]["budget_reset_at"] == datetime.fromtimestamp(
|
||||
WINDOW_EPOCH + 24 * 3600, tz=timezone.utc
|
||||
).isoformat()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_user_windows_stay_on_the_replica_without_redis():
|
||||
budget: Final = {"claude-opus*": FAMILY_CAP}
|
||||
table: Final = _WindowSpendDatabase()
|
||||
limiter: Final = _family_limiter(shared_windows=table.replica())
|
||||
|
||||
await _log_success(
|
||||
limiter,
|
||||
{
|
||||
"call_type": "acompletion",
|
||||
"standard_logging_object": {
|
||||
"call_type": "acompletion",
|
||||
"response_cost": CHAT_COST,
|
||||
"model": "anthropic/claude-opus-4-8",
|
||||
"model_group": "claude-opus-4-8",
|
||||
"end_user": "end-user-1",
|
||||
"metadata": {},
|
||||
},
|
||||
"litellm_params": {"metadata": {"user_api_key_end_user_model_max_budget": budget}},
|
||||
},
|
||||
)
|
||||
table.flush()
|
||||
|
||||
assert dict(table.rows) == {}
|
||||
with pytest.raises(litellm.BudgetExceededError, match="model=claude-opus-5-5"):
|
||||
await limiter.is_end_user_within_model_budget(
|
||||
end_user_id="end-user-1", end_user_model_max_budget=budget, model="claude-opus-5-5"
|
||||
)
|
||||
|
||||
|
||||
def test_an_access_group_named_after_a_served_model_never_widens_that_models_entry():
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-4o",
|
||||
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test"},
|
||||
"model_info": {"access_groups": ["gpt-4o", "openai-models"]},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-4o-mini",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
|
||||
"model_info": {"access_groups": ["gpt-4o", "openai-models"]},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
def governing_entry(model_max_budget: Mapping[str, object]) -> str | None:
|
||||
resolved: Final = resolve_model_budget(
|
||||
model="gpt-4o-mini",
|
||||
model_max_budget=model_max_budget,
|
||||
model_access_groups=lambda model, team_id: model_access_groups_of(router=router, model=model, team_id=team_id),
|
||||
)
|
||||
return None if resolved is None else resolved.budget_model
|
||||
|
||||
assert governing_entry({"gpt-4o": FAMILY_CAP}) is None
|
||||
assert governing_entry({"openai-models": FAMILY_CAP}) == "openai-models"
|
||||
assert routed_models_of(router).members("gpt-4o") == frozenset()
|
||||
assert routed_models_of(router).members("openai-models") == frozenset({"gpt-4o", "gpt-4o-mini"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_info_usage_reports_the_family_total():
|
||||
budget: Final = {"claude-opus*": ROOMY_CAP}
|
||||
limiter: Final = _family_limiter()
|
||||
|
||||
await _spend_on(limiter, model="claude-opus-4-8", model_max_budget=budget)
|
||||
await _spend_on(limiter, model="claude-opus-5-5", model_max_budget=budget)
|
||||
|
||||
usage: Final = await build_model_max_budget_usage(
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
entity_id=KEY_HASH,
|
||||
model_max_budget=budget,
|
||||
cache=limiter.dual_cache,
|
||||
)
|
||||
assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 2, 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("pattern", "model", "governs"),
|
||||
[
|
||||
("claude-opus*", "claude-opus-4-8", True),
|
||||
("claude-opus*", "anthropic/claude-opus-4-8", True),
|
||||
("anthropic/claude-opus*", "anthropic/claude-opus-4-8", True),
|
||||
("anthropic/claude-opus*", "claude-opus-4-8", False),
|
||||
("claude-opus*", "my-claude-opus-4-8", False),
|
||||
("claude-opus*", "claude-sonnet-5", False),
|
||||
("gpt-4.1*", "gpt-401-mini", False),
|
||||
("*", "claude-sonnet-5", True),
|
||||
],
|
||||
)
|
||||
def test_glob_entries_match_the_whole_model_name_with_only_star_as_a_wildcard(pattern: str, model: str, governs: bool):
|
||||
resolved: Final = resolve_model_budget(
|
||||
model=model, model_max_budget={pattern: FAMILY_CAP}, model_access_groups=_model_access_groups({})
|
||||
)
|
||||
|
||||
assert (resolved is not None and resolved.budget_model == pattern) is governs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replicas_without_redis_enforce_one_family_cap_together():
|
||||
budget: Final = {"claude-opus*": FAMILY_CAP}
|
||||
table: Final = _WindowSpendDatabase()
|
||||
replica_a: Final = _family_limiter(shared_windows=table.replica())
|
||||
replica_b: Final = _family_limiter(shared_windows=table.replica())
|
||||
|
||||
await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget)
|
||||
table.flush()
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError, match="model=claude-opus-5-5"):
|
||||
await replica_b.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-5-5")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replicas_sharing_redis_enforce_the_family_cap_without_the_database():
|
||||
budget: Final = {"claude-opus*": FAMILY_CAP}
|
||||
redis: Final = _SharedRedisDouble()
|
||||
table: Final = _WindowSpendDatabase()
|
||||
replica_a: Final = _family_limiter(
|
||||
dual_cache=DualCache(redis_cache=redis), # pyright: ignore[reportArgumentType] # duck-typed Redis double
|
||||
shared_windows=table.replica(),
|
||||
)
|
||||
replica_b: Final = _family_limiter(
|
||||
dual_cache=DualCache(redis_cache=redis), # pyright: ignore[reportArgumentType] # duck-typed Redis double
|
||||
shared_windows=table.replica(),
|
||||
)
|
||||
|
||||
await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget)
|
||||
await _drain_redis_pushes()
|
||||
table.flush()
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError, match="model=claude-opus-5-5"):
|
||||
await replica_b.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-5-5")
|
||||
assert dict(table.rows) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_replica_charges_the_open_shared_window_instead_of_restarting_it():
|
||||
budget: Final = {"claude-opus*": {"budget_limit": CHAT_COST * 1.5, "time_period": "1d"}}
|
||||
table: Final = _WindowSpendDatabase()
|
||||
clock: Final = _Clock()
|
||||
clock.advance(WINDOW_EPOCH)
|
||||
replica_a: Final = _family_limiter(shared_windows=table.replica(), clock=clock.now)
|
||||
replica_b: Final = _family_limiter(shared_windows=table.replica(), clock=clock.now)
|
||||
|
||||
await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget)
|
||||
table.flush()
|
||||
clock.advance(3600)
|
||||
await _spend_on(replica_b, model="claude-opus-5-5", model_max_budget=budget)
|
||||
table.flush()
|
||||
|
||||
for replica in (replica_a, replica_b):
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await replica.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-5-5")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_info_on_one_replica_reports_spend_from_every_replica():
|
||||
budget: Final = {"claude-opus*": ROOMY_CAP}
|
||||
table: Final = _WindowSpendDatabase()
|
||||
replica_a: Final = _family_limiter(shared_windows=table.replica())
|
||||
replica_b: Final = _family_limiter(shared_windows=table.replica())
|
||||
|
||||
await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget)
|
||||
await _spend_on(replica_b, model="claude-opus-5-5", model_max_budget=budget)
|
||||
table.flush()
|
||||
|
||||
usage: Final = await build_model_max_budget_usage(
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
entity_id=KEY_HASH,
|
||||
model_max_budget=budget,
|
||||
cache=replica_a.dual_cache,
|
||||
shared_windows=table.replica(),
|
||||
)
|
||||
assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 2, 4)
|
||||
assert usage["claude-opus*"]["budget_reset_at"] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("budget_model", ["claude-opus*", "claude-opus-4-8"])
|
||||
async def test_replicas_starting_a_window_moments_apart_keep_each_others_spend(budget_model: str):
|
||||
budget: Final = {budget_model: ROOMY_CAP}
|
||||
table: Final = _WindowSpendDatabase()
|
||||
clock: Final = _Clock()
|
||||
clock.advance(WINDOW_EPOCH)
|
||||
replica_a: Final = _family_limiter(shared_windows=table.replica(), clock=clock.now)
|
||||
replica_b: Final = _family_limiter(shared_windows=table.replica(), clock=clock.now)
|
||||
|
||||
await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget)
|
||||
clock.advance(1)
|
||||
await _spend_on(replica_b, model="claude-opus-4-8", model_max_budget=budget)
|
||||
table.flush()
|
||||
|
||||
usage: Final = await build_model_max_budget_usage(
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
entity_id=KEY_HASH,
|
||||
model_max_budget=budget,
|
||||
cache=DualCache(),
|
||||
now=datetime.fromtimestamp(clock.now(), tz=timezone.utc),
|
||||
shared_windows=table.replica(),
|
||||
)
|
||||
assert usage[budget_model]["current_spend"] == round(CHAT_COST * 2, 4)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_first_charges_on_one_replica_both_count():
|
||||
budget: Final = {"claude-opus*": ROOMY_CAP}
|
||||
clock: Final = _Clock()
|
||||
clock.advance(WINDOW_EPOCH)
|
||||
limiter: Final = _family_limiter(shared_windows=_WindowSpendDatabase().replica(), clock=clock.now)
|
||||
|
||||
await asyncio.gather(
|
||||
_spend_on(limiter, model="claude-opus-4-8", model_max_budget=budget),
|
||||
_spend_on(limiter, model="claude-opus-5-5", model_max_budget=budget),
|
||||
)
|
||||
|
||||
usage: Final = await build_model_max_budget_usage(
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
entity_id=KEY_HASH,
|
||||
model_max_budget=budget,
|
||||
cache=limiter.dual_cache,
|
||||
now=datetime.fromtimestamp(clock.now(), tz=timezone.utc),
|
||||
shared_windows=_WindowSpendDatabase().replica(),
|
||||
)
|
||||
assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 2, 4)
|
||||
|
|
|
|||
|
|
@ -248,6 +248,7 @@ class _FakeRepository:
|
|||
def __init__(self, rows: tuple[_Activity, ...]) -> None:
|
||||
self._rows: Final = rows
|
||||
self.aggregated = AsyncMock(side_effect=self._aggregated)
|
||||
self.guardrail_spend = AsyncMock(return_value=())
|
||||
self.key_page = AsyncMock(side_effect=self._key_page)
|
||||
self.key_page_call: tuple[DailyActivityScope, int, int] | None = None
|
||||
self.search_keys = AsyncMock(side_effect=self._search_keys)
|
||||
|
|
|
|||
|
|
@ -16580,6 +16580,7 @@ async def test_info_key_fn_includes_model_max_budget_usage(monkeypatch):
|
|||
}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -16703,6 +16704,7 @@ async def test_info_key_fn_v2_includes_model_max_budget_usage(monkeypatch):
|
|||
model_max_budget = {"gpt-4o": {"budget_limit": 1.00, "time_period": "7d"}}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -16767,6 +16769,7 @@ async def test_info_key_fn_budget_table_fallback(monkeypatch):
|
|||
}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -16840,6 +16843,7 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch):
|
|||
}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -17215,6 +17219,7 @@ async def test_info_key_fn_reads_the_configured_budget_model_key(monkeypatch):
|
|||
}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ from litellm.proxy._types import (
|
|||
UpdateTeamRequest,
|
||||
UserAPIKeyAuth, # Import UserAPIKeyAuth
|
||||
)
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import RoutedModels, routed_models_of
|
||||
from litellm.proxy.management.teams.authz import TeamAccess
|
||||
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
|
||||
PROXY_CONTROLLED_TEAM_METADATA_KEYS,
|
||||
|
|
@ -17570,6 +17571,26 @@ def test_team_admin_cannot_loosen_team_model_caps(requested) -> None:
|
|||
assert "proxy admin" in exc.value.detail["error"]
|
||||
|
||||
|
||||
def test_team_cap_refusal_names_the_entry_holding_the_cap() -> None:
|
||||
from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_check_team_model_budget_update_authority(
|
||||
data=UpdateTeamRequest(
|
||||
team_id="t1",
|
||||
model_max_budget={
|
||||
**_EXISTING_TEAM_MODEL_CAPS,
|
||||
"openai/gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"},
|
||||
},
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"),
|
||||
existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS,
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
assert "for 'openai/gpt-4o'" in exc.value.detail["error"]
|
||||
assert "(entry 'gpt-4o')" in exc.value.detail["error"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"requested",
|
||||
[
|
||||
|
|
@ -17593,6 +17614,253 @@ def test_team_admin_can_tighten_or_keep_team_model_caps(requested) -> None:
|
|||
)
|
||||
|
||||
|
||||
_TEAM_OPUS_FAMILY_CAP: Final = MappingProxyType({"claude-opus*": {"max_budget": 10.0, "budget_duration": "1d"}})
|
||||
_TEAM_OPUS_GROUP_CAP: Final = MappingProxyType({"z-opus-group": {"max_budget": 10.0, "budget_duration": "1d"}})
|
||||
_TEAM_ROUTED_MODELS: Final = RoutedModels(
|
||||
served_names=frozenset(
|
||||
{"claude-opus-4-8", "claude-opus-5-5", "gemini-3-pro", "gpt-4o", "my-team/gpt-4o", "team-route/*"}
|
||||
),
|
||||
group_members=MappingProxyType(
|
||||
{
|
||||
"opus-group": frozenset({"claude-opus-4-8"}),
|
||||
"z-opus-group": frozenset({"claude-opus-4-8"}),
|
||||
"a-opus-group": frozenset({"claude-opus-4-8", "claude-opus-5-5"}),
|
||||
"zz-opus-group": frozenset({"claude-opus-4-8"}),
|
||||
"a-gemini-group": frozenset({"gemini-3-pro"}),
|
||||
"z-anthropic-group": frozenset({"anthropic/*"}),
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _first_bedrock_anthropic_model(model_cost: Mapping[str, Mapping[str, object]]) -> str:
|
||||
return next(
|
||||
model.removeprefix("anthropic.")
|
||||
for model, cost in model_cost.items()
|
||||
if model.startswith("anthropic.") and str(cost.get("litellm_provider", "")).startswith("bedrock")
|
||||
)
|
||||
|
||||
|
||||
_BEDROCK_ANTHROPIC_MODEL: Final = _first_bedrock_anthropic_model(litellm.model_cost)
|
||||
_TEAM_BEDROCK_VENDOR_CAP: Final = MappingProxyType({"anthropic.*": {"max_budget": 10.0, "budget_duration": "1d"}})
|
||||
_BEDROCK_ROUTE_ADDED_ENTRIES: Final = (_BEDROCK_ANTHROPIC_MODEL, f"{_BEDROCK_ANTHROPIC_MODEL[:12]}*")
|
||||
|
||||
|
||||
def _check_team_admin_model_caps(
|
||||
existing: Mapping[str, object],
|
||||
requested: Mapping[str, object],
|
||||
routed_models: RoutedModels = _TEAM_ROUTED_MODELS,
|
||||
) -> None:
|
||||
from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority
|
||||
|
||||
_check_team_model_budget_update_authority(
|
||||
data=UpdateTeamRequest(team_id="t1", model_max_budget=dict(requested)),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"),
|
||||
existing_model_max_budget=existing,
|
||||
routed_models_source=lambda _team_id: routed_models,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("existing", "added"),
|
||||
[
|
||||
(_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 1000.0, "budget_duration": "1d"}}),
|
||||
(_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 10.0, "budget_duration": "1h"}}),
|
||||
(_TEAM_OPUS_FAMILY_CAP, {"*laude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}),
|
||||
(_TEAM_OPUS_GROUP_CAP, {"a-opus-group": {"max_budget": 1000.0, "budget_duration": "1d"}}),
|
||||
(_TEAM_OPUS_GROUP_CAP, {"claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}),
|
||||
(
|
||||
{"z-anthropic-group": {"max_budget": 10.0, "budget_duration": "1d"}},
|
||||
{"claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}},
|
||||
),
|
||||
(_TEAM_OPUS_FAMILY_CAP, {"claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}),
|
||||
(
|
||||
{"anthropic/*": {"max_budget": 10.0, "budget_duration": "1d"}},
|
||||
{"claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}},
|
||||
),
|
||||
(
|
||||
{"my-team/*": {"max_budget": 10.0, "budget_duration": "1d"}},
|
||||
{"gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}},
|
||||
),
|
||||
(
|
||||
{"team-route/*": {"max_budget": 10.0, "budget_duration": "1d"}},
|
||||
{"gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}},
|
||||
),
|
||||
(_TEAM_OPUS_FAMILY_CAP, {"claude-opus-4-8": {"max_budget": 10.0, "budget_duration": "1d"}}),
|
||||
(_TEAM_OPUS_FAMILY_CAP, {"claude-opus-4*": {"max_budget": 10.0, "budget_duration": "1d"}}),
|
||||
(_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 5.0, "budget_duration": "1d"}}),
|
||||
(_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 10.0, "budget_duration": "2d"}}),
|
||||
(_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 0.01, "budget_duration": "1h"}}),
|
||||
(_TEAM_OPUS_FAMILY_CAP, {"claude-opus-4-8": {"rpm_limit": 10}}),
|
||||
],
|
||||
ids=[
|
||||
"group_over_pattern",
|
||||
"group_with_a_faster_window",
|
||||
"more_specific_pattern",
|
||||
"earlier_group_sharing_a_model",
|
||||
"model_inside_group",
|
||||
"model_under_wildcard_group",
|
||||
"model_matching_pattern",
|
||||
"model_under_provider_pattern",
|
||||
"model_under_served_prefix",
|
||||
"model_under_wildcard_route",
|
||||
"model_carved_out_at_the_same_cap",
|
||||
"pattern_carved_out_at_the_same_cap",
|
||||
"tighter_group",
|
||||
"group_with_a_slower_window",
|
||||
"group_with_a_faster_window_within_the_cap",
|
||||
"uncapped_model",
|
||||
],
|
||||
)
|
||||
def test_team_admin_cannot_add_an_entry_that_takes_precedence_over_a_cap(
|
||||
existing: Mapping[str, object], added: Mapping[str, object]
|
||||
) -> None:
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_check_team_admin_model_caps(existing=existing, requested={**existing, **added})
|
||||
assert exc.value.status_code == 403
|
||||
assert "proxy admin" in exc.value.detail["error"]
|
||||
assert repr(next(iter(existing))) in exc.value.detail["error"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("existing", "requested"),
|
||||
[
|
||||
(_TEAM_OPUS_FAMILY_CAP, {**_TEAM_OPUS_FAMILY_CAP, "claude*": {"max_budget": 1000.0, "budget_duration": "1d"}}),
|
||||
(_TEAM_OPUS_GROUP_CAP, {**_TEAM_OPUS_GROUP_CAP, "a-gemini-group": {"max_budget": 1000.0, "budget_duration": "1d"}}),
|
||||
(_TEAM_OPUS_GROUP_CAP, {**_TEAM_OPUS_GROUP_CAP, "zz-opus-group": {"max_budget": 1000.0, "budget_duration": "1d"}}),
|
||||
(_TEAM_OPUS_GROUP_CAP, {**_TEAM_OPUS_GROUP_CAP, "gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}}),
|
||||
(_TEAM_OPUS_FAMILY_CAP, {**_TEAM_OPUS_FAMILY_CAP, "gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}}),
|
||||
(
|
||||
_TEAM_OPUS_FAMILY_CAP,
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "a-gemini-group": {"max_budget": 1000.0, "budget_duration": "1d"}},
|
||||
),
|
||||
(
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 100.0, "budget_duration": "1d"}},
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 100.0, "budget_duration": "1d"}},
|
||||
),
|
||||
(
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 100.0, "budget_duration": "1d"}},
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 50.0, "budget_duration": "1d"}},
|
||||
),
|
||||
(
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"rpm_limit": 10}},
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"rpm_limit": 20}},
|
||||
),
|
||||
(
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "claude-opus-4-8": {"rpm_limit": 10, "budget_duration": "1d"}},
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}},
|
||||
),
|
||||
(
|
||||
_TEAM_OPUS_FAMILY_CAP,
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "gemini-3-pro*": {"max_budget": 1000.0, "budget_duration": "1d"}},
|
||||
),
|
||||
(
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "gpt-4o": {"max_budget": 10.0}},
|
||||
{**_TEAM_OPUS_FAMILY_CAP, "gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}},
|
||||
),
|
||||
(
|
||||
{"z-anthropic-group": {"max_budget": 10.0}},
|
||||
{"z-anthropic-group": {"max_budget": 10.0, "budget_duration": "1d"}},
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"less_specific_pattern",
|
||||
"group_without_shared_models",
|
||||
"later_group_sharing_a_model",
|
||||
"model_outside_group",
|
||||
"model_outside_pattern",
|
||||
"group_outside_pattern",
|
||||
"unchanged_looser_entry",
|
||||
"tightened_looser_entry",
|
||||
"uncapped_entry_edited",
|
||||
"enforced_uncapped_entry_capped",
|
||||
"pattern_sharing_only_a_wildcard_route",
|
||||
"unenforced_capped_model_given_a_duration",
|
||||
"unenforced_capped_wildcard_group_given_a_duration",
|
||||
],
|
||||
)
|
||||
def test_team_admin_can_keep_or_tighten_entries_and_add_ones_that_outrank_no_cap(
|
||||
existing: Mapping[str, object], requested: Mapping[str, object]
|
||||
) -> None:
|
||||
assert _check_team_admin_model_caps(existing=existing, requested=requested) is None
|
||||
|
||||
|
||||
def test_team_admin_cannot_make_an_unenforced_entry_take_precedence_over_a_cap() -> None:
|
||||
existing: Final = {**_TEAM_OPUS_FAMILY_CAP, "claude-opus-4-8": {"rpm_limit": 10}}
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_check_team_admin_model_caps(
|
||||
existing=existing,
|
||||
requested={**existing, "claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}},
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
assert "'claude-opus*'" in exc.value.detail["error"]
|
||||
|
||||
|
||||
def test_team_admin_cannot_raise_an_unenforced_cap_while_enforcing_it() -> None:
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_check_team_admin_model_caps(
|
||||
existing={"gpt-4o": {"max_budget": 10.0}},
|
||||
requested={"gpt-4o": {"max_budget": 20.0, "budget_duration": "1d"}},
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
assert "(entry 'gpt-4o')" in exc.value.detail["error"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("added_entry", _BEDROCK_ROUTE_ADDED_ENTRIES, ids=["model", "pattern"])
|
||||
def test_team_admin_cannot_lift_a_bedrock_vendor_cap_through_a_bedrock_wildcard_route(added_entry: str) -> None:
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_check_team_admin_model_caps(
|
||||
existing=_TEAM_BEDROCK_VENDOR_CAP,
|
||||
requested={**_TEAM_BEDROCK_VENDOR_CAP, added_entry: {"max_budget": 1000.0, "budget_duration": "1d"}},
|
||||
routed_models=RoutedModels(
|
||||
served_names=frozenset({"bedrock/*", "gpt-4o"}), group_members=MappingProxyType({})
|
||||
),
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
assert "'anthropic.*'" in exc.value.detail["error"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("added_entry", _BEDROCK_ROUTE_ADDED_ENTRIES, ids=["model", "pattern"])
|
||||
def test_team_admin_can_add_entries_beside_a_bedrock_vendor_cap_no_route_serves(added_entry: str) -> None:
|
||||
assert (
|
||||
_check_team_admin_model_caps(
|
||||
existing=_TEAM_BEDROCK_VENDOR_CAP,
|
||||
requested={**_TEAM_BEDROCK_VENDOR_CAP, added_entry: {"max_budget": 1000.0, "budget_duration": "1d"}},
|
||||
routed_models=RoutedModels(served_names=frozenset({"gpt-4o"}), group_members=MappingProxyType({})),
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_team_admin_cannot_outrank_a_group_cap_through_a_team_model_both_groups_serve() -> None:
|
||||
from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority
|
||||
|
||||
team_router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": f"model_name_t1_{group}",
|
||||
"litellm_params": {"model": "anthropic/claude-opus-4-8", "api_key": "sk-test"},
|
||||
"model_info": {"team_id": "t1", "team_public_model_name": "claude-opus-4-8", "access_groups": [group]},
|
||||
}
|
||||
for group in ("a-opus", "z-opus")
|
||||
]
|
||||
)
|
||||
existing: Final = {"z-opus": {"max_budget": 10.0, "budget_duration": "1d"}}
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_check_team_model_budget_update_authority(
|
||||
data=UpdateTeamRequest(
|
||||
team_id="t1", model_max_budget={**existing, "a-opus": {"max_budget": 100.0, "budget_duration": "1d"}}
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"),
|
||||
existing_model_max_budget=existing,
|
||||
routed_models_source=lambda team_id: routed_models_of(router=team_router, team_id=team_id),
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
assert "'z-opus'" in exc.value.detail["error"]
|
||||
|
||||
|
||||
def test_team_model_cap_authority_skips_omitted_field_malformed_rows_and_proxy_admins() -> None:
|
||||
from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority
|
||||
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callba
|
|||
from litellm.router_utils.cooldown_handlers import (
|
||||
_get_deployment_cooldown_policy,
|
||||
_has_explicit_allowed_fails_policy_for_exception,
|
||||
_has_fallback_group,
|
||||
_increment_allowed_fails,
|
||||
_is_cooldown_required,
|
||||
_resolve_allowed_fails_from_policy,
|
||||
|
|
@ -3149,6 +3150,8 @@ BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, Router
|
|||
{},
|
||||
BENCHED,
|
||||
),
|
||||
({"allowed_fails": 1, "fallbacks": [{"test-model": "backup-model"}]}, {}, BENCHED),
|
||||
({"allowed_fails": 1, "default_fallbacks": "backup-model"}, {}, BENCHED),
|
||||
],
|
||||
ids=[
|
||||
"allowed_fails",
|
||||
|
|
@ -3157,6 +3160,8 @@ BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, Router
|
|||
"policy_names_timeout",
|
||||
"allowed_fails_with_a_fallback_group",
|
||||
"allowed_fails_with_a_dict_fallback_target",
|
||||
"allowed_fails_with_a_string_fallback_target",
|
||||
"allowed_fails_with_a_string_default_fallbacks",
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -3180,6 +3185,33 @@ async def test_router_allowed_fails_benches_a_lone_deployment_only_when_the_poli
|
|||
assert len(route.calls) == outcomes.count(litellm.Timeout)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("fallbacks", "expected"),
|
||||
[
|
||||
([{"test-model": "backup-model"}], True),
|
||||
([{"*": "backup-model"}], True),
|
||||
([{"test-model": ["backup-model"]}], True),
|
||||
([{"test-model": 5}, {"*": "backup-model"}], True),
|
||||
([{"test-model": 5}], False),
|
||||
([{"other-model": ["backup-model"]}], False),
|
||||
],
|
||||
ids=[
|
||||
"string_target",
|
||||
"string_catch_all_from_default_fallbacks",
|
||||
"list_target",
|
||||
"unparseable_entry_does_not_hide_the_catch_all",
|
||||
"unparseable_entry_alone",
|
||||
"another_groups_chain",
|
||||
],
|
||||
)
|
||||
def test_has_fallback_group_reads_string_targets_and_skips_unparseable_entries(
|
||||
fallbacks: Sequence[Mapping[str, object]], expected: bool
|
||||
) -> None:
|
||||
router: Final = Router(model_list=[TIMED_OUT_DEPLOYMENT], fallbacks=list(fallbacks))
|
||||
|
||||
assert _has_fallback_group(router, "test-model") is expected
|
||||
|
||||
|
||||
async def _group_call_outcome(router: Router) -> str:
|
||||
try:
|
||||
response: Final = await router.acompletion(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: slack-gif-creator
|
||||
description: Knowledge and utilities for creating animated GIFs optimized for Slack. Provides constraints, validation tools, and animation concepts. Use when users request animated GIFs for Slack like "make me a GIF of X doing Y for Slack."
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Slack GIF Creator
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: test-delete-skill
|
||||
description: A test skill created specifically for deletion testing
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Test Delete Skill
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: test-skill-litellm
|
||||
description: A test skill created by LiteLLM automated tests
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Test Skill
|
||||
|
|
|
|||
|
|
@ -74,6 +74,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset(
|
|||
"litellm/proxy/container_endpoints/ownership.py",
|
||||
"litellm/proxy/discovery_endpoints/agent_skills_endpoints.py",
|
||||
"litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py",
|
||||
"litellm/proxy/hooks/model_budget_shared_windows.py",
|
||||
"litellm/proxy/spend_tracking/key_metadata_recovery.py",
|
||||
"litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py",
|
||||
"litellm/responses/litellm_completion_transformation/transformation.py",
|
||||
|
|
|
|||
68
tests/unit/test_repo_agent_skills.py
Normal file
68
tests/unit/test_repo_agent_skills.py
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
from collections.abc import Iterator
|
||||
from fnmatch import fnmatch
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
_REPO_ROOT: Final = Path(__file__).resolve().parents[2]
|
||||
_SKILLS_DIR: Final = _REPO_ROOT / "skills"
|
||||
_SKILL_DIRECTORIES: Final = tuple(sorted(skill_file.parent for skill_file in _SKILLS_DIR.glob("*/SKILL.md")))
|
||||
_FULL_DEPTH_MAX_DEPTH: Final = 5
|
||||
_FULL_DEPTH_SKIPPED_NAMES: Final = frozenset(("node_modules", ".git", "dist", "build", "__pycache__"))
|
||||
_GITIGNORED_NAME_PATTERNS: Final = tuple(
|
||||
line.strip().rstrip("/")
|
||||
for line in (_REPO_ROOT / ".gitignore").read_text(encoding="utf-8").splitlines()
|
||||
if line.strip() and not line.startswith(("#", "!")) and "/" not in line.strip().rstrip("/")
|
||||
)
|
||||
|
||||
|
||||
def _absent_from_a_clone_or_skipped_by_the_cli(name: str) -> bool:
|
||||
return name in _FULL_DEPTH_SKIPPED_NAMES or any(fnmatch(name, pattern) for pattern in _GITIGNORED_NAME_PATTERNS)
|
||||
|
||||
|
||||
def _full_depth_skill_directories(directory: Path, depth: int) -> Iterator[Path]:
|
||||
if (directory / "SKILL.md").is_file():
|
||||
yield directory
|
||||
if depth == _FULL_DEPTH_MAX_DEPTH:
|
||||
return
|
||||
for child in sorted(directory.iterdir()):
|
||||
if child.is_dir() and not _absent_from_a_clone_or_skipped_by_the_cli(child.name):
|
||||
yield from _full_depth_skill_directories(child, depth + 1)
|
||||
|
||||
|
||||
_NON_END_USER_SKILL_DIRECTORIES: Final = tuple(
|
||||
directory for directory in _full_depth_skill_directories(_REPO_ROOT, 0) if directory.parent != _SKILLS_DIR
|
||||
)
|
||||
|
||||
|
||||
def _frontmatter(skill_directory: Path) -> object:
|
||||
opening, frontmatter, _body = (skill_directory / "SKILL.md").read_text(encoding="utf-8").split("---\n", 2)
|
||||
assert opening == "", f"{skill_directory.name}/SKILL.md must open with a --- frontmatter block"
|
||||
return yaml.safe_load(frontmatter)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("skill_directory", _SKILL_DIRECTORIES, ids=lambda path: path.name)
|
||||
def test_skill_frontmatter_has_string_name_matching_its_directory_and_a_description(skill_directory: Path) -> None:
|
||||
frontmatter: Final = _frontmatter(skill_directory)
|
||||
assert isinstance(frontmatter, dict), frontmatter
|
||||
assert frontmatter.get("name") == skill_directory.name, frontmatter
|
||||
description: Final = frontmatter.get("description")
|
||||
assert isinstance(description, str) and description.strip(), frontmatter
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"skill_directory",
|
||||
_NON_END_USER_SKILL_DIRECTORIES,
|
||||
ids=lambda path: str(path.relative_to(_REPO_ROOT)),
|
||||
)
|
||||
def test_skill_outside_skills_folder_is_hidden_from_end_user_skill_installs(skill_directory: Path) -> None:
|
||||
frontmatter: Final = _frontmatter(skill_directory)
|
||||
assert isinstance(frontmatter, dict), frontmatter
|
||||
metadata: Final = frontmatter.get("metadata")
|
||||
assert isinstance(metadata, dict) and metadata.get("internal") is True, (
|
||||
f"{skill_directory} needs `metadata: {{internal: true}}`, or `npx skills add BerriAI/litellm` installs it "
|
||||
"for end users, from a root dot-folder always and from anywhere with --full-depth (skills CLI 1.7.4, "
|
||||
"checked 2026-10-28)"
|
||||
)
|
||||
|
|
@ -1126,6 +1126,36 @@ def test_arouter_test_team_model():
|
|||
assert result is not None
|
||||
|
||||
|
||||
def test_team_public_model_names_of_lists_only_that_teams_live_public_names():
|
||||
def team_deployment(deployment_id: str, team_id: str, public_model_name: str) -> DeploymentTypedDict:
|
||||
return {
|
||||
"model_name": f"model_name_{team_id}_{deployment_id}",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
|
||||
"model_info": {"id": deployment_id, "team_id": team_id, "team_public_model_name": public_model_name},
|
||||
}
|
||||
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
team_deployment("team-a-1", "team-a", "shared-model"),
|
||||
team_deployment("team-a-2", "team-a", "shared-model"),
|
||||
team_deployment("team-a-solo", "team-a", "solo-model"),
|
||||
team_deployment("team-b-1", "team-b", "team-b-model"),
|
||||
{
|
||||
"model_name": "plain-model",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
|
||||
"model_info": {"id": "plain-1"},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
assert router.team_public_model_names_of("team-a") == frozenset({"shared-model", "solo-model"})
|
||||
assert router.team_public_model_names_of("team-c") == frozenset()
|
||||
|
||||
router.delete_deployment(id="team-a-solo")
|
||||
|
||||
assert router.team_public_model_names_of("team-a") == frozenset({"shared-model"})
|
||||
|
||||
|
||||
def test_team_model_has_alternatives():
|
||||
def team_deployment(
|
||||
deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@ import sys
|
|||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
_SCRIPT: Final = Path(__file__).resolve().parents[2] / "ci_cd" / "spend_logs_guard.py"
|
||||
_SPEC: Final = importlib.util.spec_from_file_location("spend_logs_guard", _SCRIPT)
|
||||
assert _SPEC is not None and _SPEC.loader is not None
|
||||
|
|
@ -35,6 +37,13 @@ class TestSpendLogHits:
|
|||
elsewhere: Final = _diff("litellm/proxy/utils.py", 'table = "LiteLLM_SpendLogs"', 5)
|
||||
assert guard.spend_log_hits(removed + other + elsewhere) == ()
|
||||
|
||||
def test_a_table_whose_name_only_starts_with_the_watched_one_is_not_a_hit(self):
|
||||
archive: Final = _diff("schema.prisma", "model LiteLLM_SpendLogsArchive {", 900)
|
||||
index: Final = _diff("schema.prisma", " @@index([team_id], map: \"LiteLLM_SpendLogs_team_id_idx\")", 760)
|
||||
assert guard.spend_log_hits(archive + index) == (
|
||||
'schema.prisma: @@index([team_id], map: "LiteLLM_SpendLogs_team_id_idx")',
|
||||
)
|
||||
|
||||
def test_two_files_in_one_diff_are_attributed_to_their_own_path(self):
|
||||
diff: Final = _diff("litellm/proxy/schema.prisma", " a LiteLLM_SpendLogs", 1) + _diff(
|
||||
"schema.prisma", " b String", 1
|
||||
|
|
@ -59,6 +68,20 @@ class TestApprovedBy:
|
|||
assert guard.approved_by(reviews, "kerry-berri") is True
|
||||
assert guard.approved_by(reviews[::-1], "kerry-berri") is False
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("states", "approved"),
|
||||
(
|
||||
(("APPROVED", "COMMENTED"), True),
|
||||
(("APPROVED", "PENDING"), True),
|
||||
(("CHANGES_REQUESTED", "COMMENTED"), False),
|
||||
(("APPROVED", "DISMISSED"), False),
|
||||
(("COMMENTED",), False),
|
||||
),
|
||||
)
|
||||
def test_a_comment_or_pending_review_does_not_change_the_answer(self, states, approved):
|
||||
reviews: Final = tuple(guard.Review(state=state, user=guard.Reviewer(login="kerry-berri")) for state in states)
|
||||
assert guard.approved_by(reviews, "kerry-berri") is approved
|
||||
|
||||
def test_no_review_from_the_approver_is_not_an_approval(self):
|
||||
other: Final = guard.Review(state="APPROVED", user=guard.Reviewer(login="someone-else"))
|
||||
assert guard.approved_by((other,), "kerry-berri") is False
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: url-state
|
||||
description: Read or write dashboard URL query state (tabs, filters, pagination, deep links, demo flags) with nuqs instead of raw history or URLSearchParams
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# URL state
|
||||
|
|
|
|||
|
|
@ -498,9 +498,6 @@
|
|||
}
|
||||
},
|
||||
"src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx": {
|
||||
"no-nested-ternary": {
|
||||
"count": 1
|
||||
},
|
||||
"react-hooks/set-state-in-effect": {
|
||||
"count": 1
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,118 @@
|
|||
/* @vitest-environment jsdom */
|
||||
import React from "react";
|
||||
import { act, renderHook, waitFor } from "@testing-library/react";
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { mcpServerToolKeys, useMCPServerTools } from "./useMCPServerTools";
|
||||
import * as networking from "@/components/networking";
|
||||
|
||||
vi.mock("@/components/networking", () => ({
|
||||
listMCPTools: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
|
||||
default: vi.fn(() => ({ accessToken: "test-token-123" })),
|
||||
}));
|
||||
|
||||
const wrapper = ({ children }: { children: React.ReactNode }) => {
|
||||
const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } });
|
||||
return React.createElement(QueryClientProvider, { client: queryClient }, children);
|
||||
};
|
||||
|
||||
const toolsFor = (serverId: string) => ({
|
||||
tools: [{ name: `${serverId}_tool`, description: `from ${serverId}` }],
|
||||
error: null,
|
||||
});
|
||||
|
||||
describe("useMCPServerTools", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("fetches only the enabled servers and leaves the rest idle", async () => {
|
||||
vi.mocked(networking.listMCPTools).mockImplementation(async (_token, serverId) => toolsFor(serverId));
|
||||
|
||||
const { result } = renderHook(() => useMCPServerTools(["a", "b"], new Set(["a"])), { wrapper });
|
||||
|
||||
await waitFor(() => expect(result.current.get("a")?.tools).toEqual([{ name: "a_tool", description: "from a" }]));
|
||||
expect(result.current.get("b")).toEqual({ tools: [], status: "idle" });
|
||||
expect(networking.listMCPTools).toHaveBeenCalledTimes(1);
|
||||
expect(networking.listMCPTools).toHaveBeenCalledWith("test-token-123", "a");
|
||||
});
|
||||
|
||||
it("isolates a failing server from the others", async () => {
|
||||
vi.mocked(networking.listMCPTools).mockImplementation(async (_token, serverId) =>
|
||||
serverId === "bad" ? { tools: [], error: "http_500", message: "upstream down" } : toolsFor(serverId),
|
||||
);
|
||||
|
||||
const { result } = renderHook(() => useMCPServerTools(["good", "bad"], new Set(["good", "bad"])), { wrapper });
|
||||
|
||||
await waitFor(() => expect(result.current.get("bad")).toEqual({ tools: [], status: "error" }));
|
||||
expect(result.current.get("good")).toEqual({
|
||||
tools: [{ name: "good_tool", description: "from good" }],
|
||||
status: "success",
|
||||
});
|
||||
});
|
||||
|
||||
it("reports a server that is being fetched as loading", () => {
|
||||
vi.mocked(networking.listMCPTools).mockReturnValue(new Promise(() => {}));
|
||||
|
||||
const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper });
|
||||
|
||||
expect(result.current.get("a")).toEqual({ tools: [], status: "loading" });
|
||||
});
|
||||
|
||||
it("normalizes a bare array response, tools without descriptions, and drops tools without a name", async () => {
|
||||
vi.mocked(networking.listMCPTools).mockResolvedValue([
|
||||
{ tool_name: "legacy" },
|
||||
"plain",
|
||||
{ description: "nameless" },
|
||||
]);
|
||||
|
||||
const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper });
|
||||
|
||||
await waitFor(() =>
|
||||
expect(result.current.get("a")?.tools).toEqual([
|
||||
{ name: "legacy", description: "" },
|
||||
{ name: "plain", description: "" },
|
||||
]),
|
||||
);
|
||||
});
|
||||
|
||||
it("reports an error when a refetch fails for a server that had loaded no tools", async () => {
|
||||
const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
|
||||
const sharedWrapper = ({ children }: { children: React.ReactNode }) =>
|
||||
React.createElement(QueryClientProvider, { client: queryClient }, children);
|
||||
vi.mocked(networking.listMCPTools)
|
||||
.mockResolvedValueOnce({ tools: [], error: null })
|
||||
.mockResolvedValueOnce({ tools: [], error: "http_502", message: "upstream down" });
|
||||
|
||||
const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper: sharedWrapper });
|
||||
await waitFor(() => expect(result.current.get("a")?.status).toBe("success"));
|
||||
|
||||
await act(() => queryClient.refetchQueries({ queryKey: mcpServerToolKeys.detail("a") }));
|
||||
|
||||
await waitFor(() => expect(result.current.get("a")).toEqual({ tools: [], status: "error" }));
|
||||
});
|
||||
|
||||
it("keeps showing the loaded tools when a later refetch fails", async () => {
|
||||
const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
|
||||
const sharedWrapper = ({ children }: { children: React.ReactNode }) =>
|
||||
React.createElement(QueryClientProvider, { client: queryClient }, children);
|
||||
vi.mocked(networking.listMCPTools)
|
||||
.mockResolvedValueOnce(toolsFor("a"))
|
||||
.mockResolvedValueOnce({ tools: [], error: "http_502", message: "upstream down" });
|
||||
|
||||
const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper: sharedWrapper });
|
||||
await waitFor(() => expect(result.current.get("a")?.status).toBe("success"));
|
||||
|
||||
await act(() => queryClient.refetchQueries({ queryKey: mcpServerToolKeys.detail("a") }));
|
||||
|
||||
await waitFor(() => expect(queryClient.getQueryState(mcpServerToolKeys.detail("a"))?.status).toBe("error"));
|
||||
expect(networking.listMCPTools).toHaveBeenCalledTimes(2);
|
||||
expect(result.current.get("a")).toEqual({
|
||||
tools: [{ name: "a_tool", description: "from a" }],
|
||||
status: "success",
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,70 @@
|
|||
import { useQueries, type QueryObserverResult } from "@tanstack/react-query";
|
||||
import { createQueryKeys } from "../common/queryKeysFactory";
|
||||
import { listMCPTools } from "@/components/networking";
|
||||
import useAuthorized from "../useAuthorized";
|
||||
|
||||
export const mcpServerToolKeys = createQueryKeys("mcpServerTools");
|
||||
|
||||
export interface MCPServerToolEntry {
|
||||
name: string;
|
||||
description: string;
|
||||
}
|
||||
|
||||
export type MCPServerToolsStatus = "idle" | "loading" | "error" | "success";
|
||||
|
||||
export interface MCPServerToolsState {
|
||||
tools: MCPServerToolEntry[];
|
||||
status: MCPServerToolsStatus;
|
||||
}
|
||||
|
||||
type RawTool = string | { name?: string; tool_name?: string; description?: string };
|
||||
|
||||
const toEntry = (tool: RawTool): MCPServerToolEntry =>
|
||||
typeof tool === "string"
|
||||
? { name: tool, description: "" }
|
||||
: { name: tool.name ?? tool.tool_name ?? "", description: tool.description ?? "" };
|
||||
|
||||
// listMCPTools never throws: it reports failures as `{ tools: [], error }`.
|
||||
// Throw here so a broken server surfaces as a query error instead of an empty tool list.
|
||||
export const fetchServerTools = async (accessToken: string, serverId: string): Promise<MCPServerToolEntry[]> => {
|
||||
const result = await listMCPTools(accessToken, serverId);
|
||||
if (!Array.isArray(result) && result?.error) {
|
||||
throw new Error(result.message || result.error);
|
||||
}
|
||||
const toolList: RawTool[] = Array.isArray(result) ? result : result?.tools ?? [];
|
||||
return toolList.map(toEntry).filter((tool) => tool.name !== "");
|
||||
};
|
||||
|
||||
const toStatus = ({
|
||||
status,
|
||||
fetchStatus,
|
||||
data,
|
||||
}: Pick<QueryObserverResult<MCPServerToolEntry[]>, "status" | "fetchStatus" | "data">): MCPServerToolsStatus => {
|
||||
if (status === "error" && data !== undefined && data.length > 0) return "success";
|
||||
if (status !== "pending") return status;
|
||||
return fetchStatus === "idle" ? "idle" : "loading";
|
||||
};
|
||||
|
||||
export const useMCPServerTools = (
|
||||
serverIds: readonly string[],
|
||||
enabledServerIds: ReadonlySet<string>,
|
||||
): ReadonlyMap<string, MCPServerToolsState> => {
|
||||
const { accessToken } = useAuthorized();
|
||||
return useQueries({
|
||||
queries: serverIds.map((serverId) => ({
|
||||
queryKey: mcpServerToolKeys.detail(serverId),
|
||||
queryFn: () => fetchServerTools(accessToken!, serverId),
|
||||
enabled: !!accessToken && enabledServerIds.has(serverId),
|
||||
staleTime: 5 * 60 * 1000,
|
||||
retry: false,
|
||||
refetchOnWindowFocus: false,
|
||||
})),
|
||||
combine: (results) =>
|
||||
new Map(
|
||||
results.map((result, idx): [string, MCPServerToolsState] => [
|
||||
serverIds[idx],
|
||||
{ tools: result.data ?? [], status: toStatus(result) },
|
||||
]),
|
||||
),
|
||||
});
|
||||
};
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
import React from "react";
|
||||
import { fireEvent, render, screen, waitFor, within } from "@testing-library/react";
|
||||
import { act, fireEvent, render, screen, waitFor, within } from "@testing-library/react";
|
||||
import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event";
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
|
|
@ -7,6 +7,7 @@ import { MCPToolsetsTab } from "./MCPToolsetsTab";
|
|||
import * as networking from "@/components/networking";
|
||||
import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets";
|
||||
import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
|
||||
import { mcpServerToolKeys } from "@/app/(dashboard)/hooks/mcpServers/useMCPServerTools";
|
||||
import { MCPToolset } from "@/components/mcp_tools/types";
|
||||
|
||||
vi.mock("@/components/networking", () => ({
|
||||
|
|
@ -20,6 +21,7 @@ vi.mock("@/components/networking", () => ({
|
|||
|
||||
vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ useMCPToolsets: vi.fn() }));
|
||||
vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ useMCPServers: vi.fn() }));
|
||||
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "sk-test" }) }));
|
||||
|
||||
const setup = () => userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
|
||||
|
||||
|
|
@ -197,6 +199,221 @@ describe("MCPToolsetsTab create/edit toolset form", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("MCPToolsetsTab tool search and bulk selection", () => {
|
||||
const servers = [
|
||||
{ server_id: "gh", alias: "github", server_name: "GitHub" },
|
||||
{ server_id: "lin", alias: "linear", server_name: "Linear" },
|
||||
{ server_id: "broken", alias: "flaky", server_name: "Flaky" },
|
||||
];
|
||||
const catalogs: Record<string, { name: string; description: string }[]> = {
|
||||
gh: [
|
||||
{ name: "create_issue", description: "Open a GitHub issue" },
|
||||
{ name: "list_pulls", description: "List pull requests" },
|
||||
],
|
||||
lin: [
|
||||
{ name: "create_ticket", description: "Create a Linear issue" },
|
||||
{ name: "list_projects", description: "List projects" },
|
||||
],
|
||||
};
|
||||
|
||||
const renderWithServers = (toolsets: MCPToolset[] = []) => {
|
||||
vi.mocked(useMCPToolsets).mockReturnValue({
|
||||
data: toolsets,
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useMCPToolsets>);
|
||||
vi.mocked(useMCPServers).mockReturnValue({ data: servers, isSuccess: true } as unknown as ReturnType<
|
||||
typeof useMCPServers
|
||||
>);
|
||||
const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } });
|
||||
render(
|
||||
<QueryClientProvider client={queryClient}>
|
||||
<MCPToolsetsTab accessToken="sk-test" userRole="Admin" />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
return queryClient;
|
||||
};
|
||||
|
||||
const editToolset = (tools: MCPToolset["tools"]) =>
|
||||
({ toolset_id: "ts-1", toolset_name: "triage", description: "", tools }) as MCPToolset;
|
||||
|
||||
const expectedUpdate = (tools: MCPToolset["tools"]) => ({
|
||||
toolset_id: "ts-1",
|
||||
toolset_name: "triage",
|
||||
description: "",
|
||||
tools,
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(networking.listMCPTools).mockImplementation(async (_token, serverId) =>
|
||||
serverId in catalogs
|
||||
? { tools: catalogs[serverId], error: null }
|
||||
: { tools: [], error: "http_502", message: "upstream down" },
|
||||
);
|
||||
vi.mocked(networking.createMCPToolset).mockResolvedValue(
|
||||
{} as Awaited<ReturnType<typeof networking.createMCPToolset>>,
|
||||
);
|
||||
vi.mocked(networking.updateMCPToolset).mockResolvedValue(
|
||||
{} as Awaited<ReturnType<typeof networking.updateMCPToolset>>,
|
||||
);
|
||||
});
|
||||
|
||||
it("searches tool names and descriptions across collapsed servers and selects every match", async () => {
|
||||
const user = setup();
|
||||
renderWithServers();
|
||||
|
||||
await openCreate(user);
|
||||
fireEvent.change(screen.getByPlaceholderText("e.g. github-linear-tools"), { target: { value: "issues" } });
|
||||
await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "issue");
|
||||
|
||||
expect(await screen.findByRole("checkbox", { name: /^create_ticket/ })).toBeInTheDocument();
|
||||
expect(screen.getByRole("checkbox", { name: /^create_issue/ })).toBeInTheDocument();
|
||||
expect(screen.queryByRole("checkbox", { name: /^list_pulls/ })).not.toBeInTheDocument();
|
||||
expect(await screen.findByText("2 matching tools")).toBeInTheDocument();
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Select all matching" }));
|
||||
await user.click(screen.getByRole("button", { name: "Create Toolset" }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPToolset).toHaveBeenCalledWith("sk-test", {
|
||||
toolset_name: "issues",
|
||||
description: "",
|
||||
tools: [
|
||||
{ server_id: "gh", tool_name: "create_issue" },
|
||||
{ server_id: "lin", tool_name: "create_ticket" },
|
||||
],
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
it("deselecting all matches keeps selected tools the search hides", async () => {
|
||||
const user = setup();
|
||||
renderWithServers([
|
||||
editToolset([
|
||||
{ server_id: "gh", tool_name: "list_pulls" },
|
||||
{ server_id: "gh", tool_name: "create_issue" },
|
||||
{ server_id: "lin", tool_name: "create_ticket" },
|
||||
]),
|
||||
]);
|
||||
|
||||
await openEditFor(user);
|
||||
await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "create");
|
||||
expect(await screen.findByRole("checkbox", { name: /^create_ticket/ })).toBeChecked();
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Deselect all matching" }));
|
||||
await user.click(screen.getByRole("button", { name: "Save Changes" }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.updateMCPToolset).toHaveBeenCalledWith(
|
||||
"sk-test",
|
||||
expectedUpdate([{ server_id: "gh", tool_name: "list_pulls" }]),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it("selects all of one server's tools from its header and toggles single rows, fetching only that server", async () => {
|
||||
const user = setup();
|
||||
renderWithServers();
|
||||
|
||||
await openCreate(user);
|
||||
fireEvent.change(screen.getByPlaceholderText("e.g. github-linear-tools"), { target: { value: "gh-only" } });
|
||||
await user.click(screen.getByRole("button", { name: /^github/ }));
|
||||
expect(await screen.findByRole("checkbox", { name: /^list_pulls/ })).toBeInTheDocument();
|
||||
|
||||
await user.click(screen.getByRole("checkbox", { name: "Select all tools on github" }));
|
||||
expect(screen.getByRole("checkbox", { name: /^list_pulls/ })).toBeChecked();
|
||||
// Clicking anywhere on a tool row toggles it exactly once.
|
||||
await user.click(screen.getByText("List pull requests"));
|
||||
expect(screen.getByRole("checkbox", { name: /^list_pulls/ })).not.toBeChecked();
|
||||
await user.click(screen.getByRole("button", { name: "Create Toolset" }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPToolset).toHaveBeenCalledWith("sk-test", {
|
||||
toolset_name: "gh-only",
|
||||
description: "",
|
||||
tools: [{ server_id: "gh", tool_name: "create_issue" }],
|
||||
});
|
||||
});
|
||||
expect(networking.listMCPTools).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("keeps and flags tools whose server is unreachable or gone when saving", async () => {
|
||||
const user = setup();
|
||||
const tools = [
|
||||
{ server_id: "broken", tool_name: "flaky_tool" },
|
||||
{ server_id: "deleted", tool_name: "old_tool" },
|
||||
{ server_id: "gh", tool_name: "create_issue" },
|
||||
];
|
||||
renderWithServers([editToolset(tools)]);
|
||||
|
||||
await openEditFor(user);
|
||||
await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "create");
|
||||
|
||||
expect(await screen.findByText("Couldn't load tools for this server.")).toBeInTheDocument();
|
||||
expect(await screen.findByText("server unreachable")).toBeInTheDocument();
|
||||
expect(screen.getByText("server not found")).toBeInTheDocument();
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Save Changes" }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.updateMCPToolset).toHaveBeenCalledWith("sk-test", expectedUpdate(tools));
|
||||
});
|
||||
});
|
||||
|
||||
it("stops fetching server tools once the editor is closed after a search", async () => {
|
||||
const user = setup();
|
||||
const queryClient = renderWithServers();
|
||||
|
||||
await openCreate(user);
|
||||
await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "issue");
|
||||
expect(await screen.findByRole("checkbox", { name: /^create_ticket/ })).toBeInTheDocument();
|
||||
await user.click(screen.getByRole("button", { name: "Cancel" }));
|
||||
await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument());
|
||||
const callsWhileOpen = vi.mocked(networking.listMCPTools).mock.calls.length;
|
||||
|
||||
await act(() => queryClient.invalidateQueries({ queryKey: mcpServerToolKeys.all }));
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
expect(networking.listMCPTools).toHaveBeenCalledTimes(callsWhileOpen);
|
||||
|
||||
await openCreate(user);
|
||||
expect(screen.getByRole("textbox", { name: "Search tools or servers" })).toHaveValue("");
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
expect(networking.listMCPTools).toHaveBeenCalledTimes(callsWhileOpen);
|
||||
});
|
||||
|
||||
it("does not refetch a previously expanded server when the editor is reopened", async () => {
|
||||
const user = setup();
|
||||
const queryClient = renderWithServers();
|
||||
|
||||
await openCreate(user);
|
||||
await user.click(screen.getByRole("button", { name: /^github/ }));
|
||||
expect(await screen.findByRole("checkbox", { name: /^list_pulls/ })).toBeInTheDocument();
|
||||
await user.click(screen.getByRole("button", { name: "Cancel" }));
|
||||
await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument());
|
||||
await act(() => queryClient.invalidateQueries({ queryKey: mcpServerToolKeys.all }));
|
||||
|
||||
await openCreate(user);
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
|
||||
expect(networking.listMCPTools).toHaveBeenCalledTimes(1);
|
||||
expect(screen.queryByRole("checkbox", { name: /^list_pulls/ })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("clears the whole toolset with Clear all", async () => {
|
||||
const user = setup();
|
||||
renderWithServers([editToolset([{ server_id: "gh", tool_name: "create_issue" }])]);
|
||||
|
||||
await openEditFor(user);
|
||||
await user.click(screen.getByRole("button", { name: "Clear all" }));
|
||||
expect(screen.getByText("No tools added yet")).toBeInTheDocument();
|
||||
await user.click(screen.getByRole("button", { name: "Save Changes" }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.updateMCPToolset).toHaveBeenCalledWith("sk-test", expectedUpdate([]));
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("MCPToolsetsTab usage guide", () => {
|
||||
it("shows the toolset URL on the gateway host", () => {
|
||||
renderTab();
|
||||
|
|
|
|||
|
|
@ -1,28 +1,39 @@
|
|||
import React, { useState, useCallback } from "react";
|
||||
import React, { useState } from "react";
|
||||
import { z } from "zod";
|
||||
import { SortingState } from "@tanstack/react-table";
|
||||
import { Inbox, Plus, X } from "lucide-react";
|
||||
import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets";
|
||||
import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
import { DataTable } from "@/components/shared/DataTable";
|
||||
import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
|
||||
import {
|
||||
createMCPToolset,
|
||||
updateMCPToolset,
|
||||
deleteMCPToolset,
|
||||
listMCPTools,
|
||||
getGatewayBaseUrl,
|
||||
} from "@/components/networking";
|
||||
import { MCPToolset, MCPToolsetTool } from "@/components/mcp_tools/types";
|
||||
useMCPServerTools,
|
||||
type MCPServerToolEntry,
|
||||
type MCPServerToolsState,
|
||||
} from "@/app/(dashboard)/hooks/mcpServers/useMCPServerTools";
|
||||
import { DataTable } from "@/components/shared/DataTable";
|
||||
import { createMCPToolset, updateMCPToolset, deleteMCPToolset, getGatewayBaseUrl } from "@/components/networking";
|
||||
import { MCPServer, MCPToolset, MCPToolsetTool } from "@/components/mcp_tools/types";
|
||||
import { FieldGroup } from "@/components/ui/field";
|
||||
import { FormField } from "@/components/shared/form/FormField";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Checkbox } from "@/components/ui/checkbox";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group";
|
||||
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
|
||||
import { useZodForm } from "@/lib/forms/useZodForm";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { displayToolName, getMCPToolsetTableColumns } from "./MCPToolsetTableColumns";
|
||||
import {
|
||||
deselectTools,
|
||||
filterServerTools,
|
||||
isToolSelected,
|
||||
matchesServerName,
|
||||
selectTools,
|
||||
selectionState,
|
||||
toggleTool,
|
||||
type SelectionState,
|
||||
} from "./toolsetSelection";
|
||||
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
|
||||
|
||||
interface MCPToolsetsTabProps {
|
||||
|
|
@ -37,110 +48,191 @@ const toolsetSchema = z.object({
|
|||
|
||||
type ToolsetFormValues = z.infer<typeof toolsetSchema>;
|
||||
|
||||
interface MCPToolListProps {
|
||||
const SEARCH_DEBOUNCE_WAIT_MS = 300;
|
||||
|
||||
const IDLE_SERVER_TOOLS: MCPServerToolsState = { tools: [], status: "idle" };
|
||||
|
||||
const NO_SERVERS: ReadonlySet<string> = new Set();
|
||||
|
||||
// The modal stays mounted while closed, so fetch only while it is open; otherwise a finished
|
||||
// search keeps every server's query live and refetches all of them in the background.
|
||||
const serversToFetch = (
|
||||
open: boolean,
|
||||
searching: boolean,
|
||||
serverIds: readonly string[],
|
||||
expandedServerIds: ReadonlySet<string>,
|
||||
): ReadonlySet<string> => {
|
||||
if (!open) return NO_SERVERS;
|
||||
return searching ? new Set(serverIds) : expandedServerIds;
|
||||
};
|
||||
|
||||
const checkboxStateProps = (state: SelectionState) => ({
|
||||
checked: state === "all",
|
||||
indeterminate: state === "some",
|
||||
});
|
||||
|
||||
interface ServerView {
|
||||
server: MCPServer;
|
||||
serverName: string;
|
||||
state: MCPServerToolsState;
|
||||
visibleTools: readonly MCPServerToolEntry[];
|
||||
}
|
||||
|
||||
const isPending = (state: MCPServerToolsState) => state.status === "idle" || state.status === "loading";
|
||||
|
||||
// While searching, a server stays listed until its tools load and turn out not to match.
|
||||
const isHiddenBySearch = (
|
||||
search: string,
|
||||
serverNames: readonly string[],
|
||||
state: MCPServerToolsState,
|
||||
visibleTools: readonly MCPServerToolEntry[],
|
||||
) => {
|
||||
if (!search.trim() || state.status !== "success") return false;
|
||||
return visibleTools.length === 0 && !matchesServerName(search, serverNames);
|
||||
};
|
||||
|
||||
const buildServerViews = (
|
||||
servers: readonly MCPServer[],
|
||||
toolsByServer: ReadonlyMap<string, MCPServerToolsState>,
|
||||
search: string,
|
||||
): ServerView[] =>
|
||||
servers.flatMap((server) => {
|
||||
const state = toolsByServer.get(server.server_id) ?? IDLE_SERVER_TOOLS;
|
||||
const names = [server.alias ?? "", server.server_name ?? ""];
|
||||
const visibleTools = filterServerTools(search, names, state.tools);
|
||||
if (isHiddenBySearch(search, names, state, visibleTools)) return [];
|
||||
return [{ server, serverName: server.alias || server.server_name || server.server_id, state, visibleTools }];
|
||||
});
|
||||
|
||||
const summarizeSearch = (serverViews: readonly ServerView[]) => ({
|
||||
matchingTools: serverViews.flatMap(({ server, visibleTools }) =>
|
||||
visibleTools.map((tool) => ({ server_id: server.server_id, tool_name: tool.name })),
|
||||
),
|
||||
pendingServerCount: serverViews.filter(({ state }) => isPending(state)).length,
|
||||
});
|
||||
|
||||
// Why a selected tool cannot be confirmed against its server. The tool stays selected and is
|
||||
// still saved; the note only keeps it from looking lost.
|
||||
const selectedToolIssue = (
|
||||
tool: MCPToolsetTool,
|
||||
serverNameById: ReadonlyMap<string, string>,
|
||||
toolsByServer: ReadonlyMap<string, MCPServerToolsState>,
|
||||
): string | undefined => {
|
||||
if (!serverNameById.has(tool.server_id)) return "server not found";
|
||||
const state = toolsByServer.get(tool.server_id);
|
||||
if (state?.status === "error") return "server unreachable";
|
||||
if (state?.status === "success" && !state.tools.some((t) => t.name === tool.tool_name)) return "tool not found";
|
||||
return undefined;
|
||||
};
|
||||
|
||||
interface ServerToolGroupProps {
|
||||
serverId: string;
|
||||
serverName: string;
|
||||
accessToken: string | null;
|
||||
state: MCPServerToolsState;
|
||||
visibleTools: readonly MCPServerToolEntry[];
|
||||
expanded: boolean;
|
||||
searchActive: boolean;
|
||||
selectedTools: MCPToolsetTool[];
|
||||
onToggle: (tool: MCPToolsetTool) => void;
|
||||
onToggleExpanded: () => void;
|
||||
onToggleTool: (tool: MCPToolsetTool) => void;
|
||||
onSetToolsSelected: (tools: MCPToolsetTool[], selected: boolean) => void;
|
||||
}
|
||||
|
||||
interface ToolEntry {
|
||||
name: string;
|
||||
description?: string;
|
||||
}
|
||||
function ServerToolGroup({
|
||||
serverId,
|
||||
serverName,
|
||||
state,
|
||||
visibleTools,
|
||||
expanded,
|
||||
searchActive,
|
||||
selectedTools,
|
||||
onToggleExpanded,
|
||||
onToggleTool,
|
||||
onSetToolsSelected,
|
||||
}: ServerToolGroupProps) {
|
||||
const selectedCount = selectedTools.filter((t) => t.server_id === serverId).length;
|
||||
const visibleRefs = visibleTools.map((tool) => ({ server_id: serverId, tool_name: tool.name }));
|
||||
const groupState = selectionState(selectedTools, visibleRefs);
|
||||
|
||||
function MCPToolList({ serverId, serverName, accessToken, selectedTools, onToggle }: MCPToolListProps) {
|
||||
const [tools, setTools] = useState<ToolEntry[]>([]);
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [expanded, setExpanded] = useState(false);
|
||||
|
||||
const selectedSet = new Set(selectedTools.filter((t) => t.server_id === serverId).map((t) => t.tool_name));
|
||||
|
||||
const fetchTools = useCallback(async () => {
|
||||
if (!accessToken || tools.length > 0) return;
|
||||
setLoading(true);
|
||||
try {
|
||||
const result = await listMCPTools(accessToken, serverId);
|
||||
const toolList = Array.isArray(result) ? result : result?.tools ?? [];
|
||||
setTools(toolList.map((t: any) => ({ name: t.name ?? t.tool_name ?? t, description: t.description ?? "" })));
|
||||
} catch {
|
||||
setTools([]);
|
||||
} finally {
|
||||
setLoading(false);
|
||||
const renderBody = () => {
|
||||
if (isPending(state)) {
|
||||
return (
|
||||
<div className="flex justify-center py-3">
|
||||
<UiLoadingSpinner className="size-4" />
|
||||
</div>
|
||||
);
|
||||
}
|
||||
}, [accessToken, serverId, tools.length]);
|
||||
|
||||
const handleToggle = () => {
|
||||
if (!expanded) fetchTools();
|
||||
setExpanded(!expanded);
|
||||
if (state.status === "error") {
|
||||
return <p className="text-xs text-destructive px-2 py-2">Couldn't load tools for this server.</p>;
|
||||
}
|
||||
if (visibleTools.length === 0) {
|
||||
return (
|
||||
<p className="text-xs text-muted-foreground px-2 py-2">
|
||||
{searchActive ? "No matching tools on this server." : "No tools found for this server."}
|
||||
</p>
|
||||
);
|
||||
}
|
||||
return (
|
||||
<div className="flex flex-col gap-1">
|
||||
{visibleRefs.map((ref, idx) => {
|
||||
const tool = visibleTools[idx];
|
||||
const selected = isToolSelected(selectedTools, ref);
|
||||
return (
|
||||
<label
|
||||
key={tool.name}
|
||||
className={`flex cursor-pointer items-start gap-3 px-3 py-2 rounded-lg transition-colors ${
|
||||
selected
|
||||
? "bg-purple-50 border border-purple-300 dark:bg-purple-950 dark:border-purple-700"
|
||||
: "bg-card border border-border hover:bg-muted"
|
||||
}`}
|
||||
>
|
||||
<Checkbox checked={selected} onCheckedChange={() => onToggleTool(ref)} className="mt-0.5" />
|
||||
<div className="min-w-0 flex-1">
|
||||
<p
|
||||
className={`text-sm font-medium leading-tight ${selected ? "text-purple-800 dark:text-purple-200" : "text-foreground"}`}
|
||||
>
|
||||
{tool.name}
|
||||
</p>
|
||||
{tool.description && (
|
||||
<p className="text-xs text-muted-foreground mt-0.5 leading-tight line-clamp-2">{tool.description}</p>
|
||||
)}
|
||||
</div>
|
||||
</label>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="border border-border rounded-lg overflow-hidden">
|
||||
<button
|
||||
type="button"
|
||||
className="w-full flex items-center justify-between px-4 py-3 bg-muted hover:bg-accent transition-colors"
|
||||
onClick={handleToggle}
|
||||
>
|
||||
<span className="text-sm font-medium text-foreground flex items-center gap-2">
|
||||
<span className="inline-block w-2 h-2 rounded-full bg-info shrink-0" />
|
||||
{serverName}
|
||||
{selectedSet.size > 0 && (
|
||||
<span className="ml-1 text-xs text-purple-600 font-semibold dark:text-purple-400">
|
||||
{selectedSet.size} selected
|
||||
</span>
|
||||
)}
|
||||
</span>
|
||||
<span className="text-muted-foreground text-xs">{expanded ? "▲" : "▼"}</span>
|
||||
</button>
|
||||
{expanded && (
|
||||
<div className="p-2">
|
||||
{loading ? (
|
||||
<div className="flex justify-center py-3">
|
||||
<UiLoadingSpinner className="size-4" />
|
||||
</div>
|
||||
) : tools.length === 0 ? (
|
||||
<p className="text-xs text-muted-foreground px-2 py-2">No tools found for this server.</p>
|
||||
) : (
|
||||
<div className="flex flex-col gap-1">
|
||||
{tools.map((tool) => {
|
||||
const selected = selectedSet.has(tool.name);
|
||||
return (
|
||||
<button
|
||||
key={tool.name}
|
||||
type="button"
|
||||
onClick={() => onToggle({ server_id: serverId, tool_name: tool.name })}
|
||||
className={`flex items-start justify-between px-3 py-2 rounded-lg text-left transition-colors ${
|
||||
selected
|
||||
? "bg-purple-50 border border-purple-300 dark:bg-purple-950 dark:border-purple-700"
|
||||
: "bg-card border border-border hover:bg-muted"
|
||||
}`}
|
||||
>
|
||||
<div className="min-w-0 flex-1">
|
||||
<p
|
||||
className={`text-sm font-medium leading-tight ${selected ? "text-purple-800 dark:text-purple-200" : "text-foreground"}`}
|
||||
>
|
||||
{tool.name}
|
||||
</p>
|
||||
{tool.description && (
|
||||
<p className="text-xs text-muted-foreground mt-0.5 leading-tight line-clamp-2">
|
||||
{tool.description}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
{selected && (
|
||||
<span className="text-purple-500 text-xs font-semibold ml-2 shrink-0 mt-0.5 dark:text-purple-400">
|
||||
✓
|
||||
</span>
|
||||
)}
|
||||
</button>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
<div className="flex items-center gap-3 pl-4 bg-muted">
|
||||
<Checkbox
|
||||
{...checkboxStateProps(groupState)}
|
||||
disabled={visibleRefs.length === 0}
|
||||
// Base UI renders a span, which the global button cursor rule does not cover.
|
||||
className="cursor-pointer data-disabled:cursor-not-allowed"
|
||||
aria-label={`${searchActive ? "Select matching tools" : "Select all tools"} on ${serverName}`}
|
||||
onCheckedChange={(checked) => onSetToolsSelected(visibleRefs, checked)}
|
||||
/>
|
||||
<button
|
||||
type="button"
|
||||
aria-expanded={expanded}
|
||||
className="flex-1 flex items-center justify-between pr-4 py-3 hover:bg-accent transition-colors"
|
||||
onClick={onToggleExpanded}
|
||||
>
|
||||
<span className="text-sm font-medium text-foreground flex items-center gap-2">
|
||||
{serverName}
|
||||
{selectedCount > 0 && (
|
||||
<span className="ml-1 text-xs text-purple-600 font-semibold dark:text-purple-400">
|
||||
{selectedCount} selected
|
||||
</span>
|
||||
)}
|
||||
</span>
|
||||
<span className="text-muted-foreground text-xs">{expanded ? "▲" : "▼"}</span>
|
||||
</button>
|
||||
</div>
|
||||
{expanded && <div className="p-2">{renderBody()}</div>}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -149,21 +241,29 @@ interface CreateToolsetModalProps {
|
|||
open: boolean;
|
||||
onClose: () => void;
|
||||
onSave: (name: string, description: string | undefined, tools: MCPToolsetTool[]) => Promise<void>;
|
||||
accessToken: string | null;
|
||||
initialToolset?: MCPToolset;
|
||||
}
|
||||
|
||||
function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset }: CreateToolsetModalProps) {
|
||||
const form = useZodForm(toolsetSchema, {
|
||||
defaultValues: {
|
||||
toolset_name: initialToolset?.toolset_name || "",
|
||||
description: initialToolset?.description || "",
|
||||
},
|
||||
});
|
||||
const toolsetFormDefaults = (toolset?: MCPToolset): ToolsetFormValues => ({
|
||||
toolset_name: toolset?.toolset_name || "",
|
||||
description: toolset?.description || "",
|
||||
});
|
||||
|
||||
function CreateToolsetModal({ open, onClose, onSave, initialToolset }: CreateToolsetModalProps) {
|
||||
const form = useZodForm(toolsetSchema, { defaultValues: toolsetFormDefaults(initialToolset) });
|
||||
const [selectedTools, setSelectedTools] = useState<MCPToolsetTool[]>(initialToolset?.tools || []);
|
||||
const [saving, setSaving] = useState(false);
|
||||
const [serverSearch, setServerSearch] = useState("");
|
||||
const { data: mcpServers = [] } = useMCPServers();
|
||||
const [search, setSearch] = useState("");
|
||||
const [debouncedSearch] = useDebouncedValue(search, { wait: SEARCH_DEBOUNCE_WAIT_MS });
|
||||
const [expandedServerIds, setExpandedServerIds] = useState<ReadonlySet<string>>(() => new Set());
|
||||
const { data: mcpServers = [], isSuccess: serversLoaded } = useMCPServers();
|
||||
const serverIds = React.useMemo(() => mcpServers.map((s) => s.server_id), [mcpServers]);
|
||||
const searching = search.trim() !== "" && debouncedSearch.trim() !== "";
|
||||
const fetchServerIds = React.useMemo(
|
||||
() => serversToFetch(open, searching, serverIds, expandedServerIds),
|
||||
[open, searching, serverIds, expandedServerIds],
|
||||
);
|
||||
const toolsByServer = useMCPServerTools(serverIds, fetchServerIds);
|
||||
const serverPrefixById = React.useMemo(
|
||||
() => new Map(mcpServers.map((s) => [s.server_id, s.alias || s.server_name || s.server_id])),
|
||||
[mcpServers],
|
||||
|
|
@ -171,41 +271,54 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset
|
|||
|
||||
React.useEffect(() => {
|
||||
if (open) {
|
||||
form.reset({
|
||||
toolset_name: initialToolset?.toolset_name || "",
|
||||
description: initialToolset?.description || "",
|
||||
});
|
||||
form.reset(toolsetFormDefaults(initialToolset));
|
||||
setSelectedTools(initialToolset?.tools || []);
|
||||
setServerSearch("");
|
||||
setSearch("");
|
||||
setExpandedServerIds(new Set());
|
||||
}
|
||||
}, [open, initialToolset, form]);
|
||||
|
||||
const handleToggleTool = (tool: MCPToolsetTool) => {
|
||||
setSelectedTools((prev) => {
|
||||
const exists = prev.some((t) => t.server_id === tool.server_id && t.tool_name === tool.tool_name);
|
||||
return exists
|
||||
? prev.filter((t) => !(t.server_id === tool.server_id && t.tool_name === tool.tool_name))
|
||||
: [...prev, tool];
|
||||
});
|
||||
setSelectedTools((prev) => toggleTool(prev, tool));
|
||||
};
|
||||
|
||||
const handleSetToolsSelected = (tools: MCPToolsetTool[], selected: boolean) => {
|
||||
setSelectedTools((prev) => (selected ? selectTools(prev, tools) : deselectTools(prev, tools)));
|
||||
};
|
||||
|
||||
const handleToggleExpanded = (serverId: string) => {
|
||||
setExpandedServerIds((prev) =>
|
||||
prev.has(serverId) ? new Set([...prev].filter((id) => id !== serverId)) : new Set([...prev, serverId]),
|
||||
);
|
||||
};
|
||||
|
||||
// The create modal stays mounted, so state left from the last session would fetch servers
|
||||
// on the first render after reopening, before the open effect resets it.
|
||||
const handleClose = () => {
|
||||
setSearch("");
|
||||
setExpandedServerIds(new Set());
|
||||
onClose();
|
||||
};
|
||||
|
||||
const handleSubmit = async (values: ToolsetFormValues) => {
|
||||
setSaving(true);
|
||||
try {
|
||||
await onSave(values.toolset_name, values.description, selectedTools);
|
||||
onClose();
|
||||
handleClose();
|
||||
} finally {
|
||||
setSaving(false);
|
||||
}
|
||||
};
|
||||
|
||||
const filteredServers = mcpServers.filter((s) => {
|
||||
const q = serverSearch.toLowerCase();
|
||||
return !q || (s.alias || "").toLowerCase().includes(q) || (s.server_name || "").toLowerCase().includes(q);
|
||||
});
|
||||
const searchActive = search.trim() !== "";
|
||||
const serverViews = buildServerViews(mcpServers, toolsByServer, search);
|
||||
const { matchingTools, pendingServerCount } = searchActive
|
||||
? summarizeSearch(serverViews)
|
||||
: { matchingTools: [], pendingServerCount: 0 };
|
||||
const matchingState = selectionState(selectedTools, matchingTools);
|
||||
|
||||
return (
|
||||
<Dialog open={open} onOpenChange={(open) => !open && onClose()}>
|
||||
<Dialog open={open} onOpenChange={(open) => !open && handleClose()}>
|
||||
<DialogContent className="max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-[960px]">
|
||||
<DialogHeader>
|
||||
<DialogTitle>{initialToolset ? "Edit Toolset" : "New Toolset"}</DialogTitle>
|
||||
|
|
@ -229,32 +342,65 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset
|
|||
</div>
|
||||
<InputGroup className="mb-2">
|
||||
<InputGroupInput
|
||||
placeholder="Search MCP servers..."
|
||||
value={serverSearch}
|
||||
onChange={(e) => setServerSearch(e.target.value)}
|
||||
placeholder="Search tools or servers..."
|
||||
aria-label="Search tools or servers"
|
||||
value={search}
|
||||
onChange={(e) => setSearch(e.target.value)}
|
||||
/>
|
||||
{serverSearch && (
|
||||
{search && (
|
||||
<InputGroupAddon align="inline-end">
|
||||
<InputGroupButton size="icon-xs" aria-label="Clear search" onClick={() => setServerSearch("")}>
|
||||
<InputGroupButton size="icon-xs" aria-label="Clear search" onClick={() => setSearch("")}>
|
||||
<X />
|
||||
</InputGroupButton>
|
||||
</InputGroupAddon>
|
||||
)}
|
||||
</InputGroup>
|
||||
{searchActive && (
|
||||
<div className="mb-2 flex items-center justify-between gap-2">
|
||||
<p className="text-xs text-muted-foreground" aria-live="polite">
|
||||
{matchingTools.length} matching {matchingTools.length === 1 ? "tool" : "tools"}
|
||||
{pendingServerCount > 0 &&
|
||||
` · loading ${pendingServerCount} ${pendingServerCount === 1 ? "server" : "servers"}…`}
|
||||
</p>
|
||||
<div className="flex gap-1 shrink-0">
|
||||
<Button
|
||||
size="xs"
|
||||
variant="outline"
|
||||
disabled={matchingTools.length === 0 || matchingState === "all"}
|
||||
onClick={() => handleSetToolsSelected(matchingTools, true)}
|
||||
>
|
||||
Select all matching
|
||||
</Button>
|
||||
<Button
|
||||
size="xs"
|
||||
variant="outline"
|
||||
disabled={matchingState === "none"}
|
||||
onClick={() => handleSetToolsSelected(matchingTools, false)}
|
||||
>
|
||||
Deselect all matching
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
<div className="space-y-2 overflow-y-auto" style={{ maxHeight: 300 }}>
|
||||
{filteredServers.length === 0 ? (
|
||||
{serverViews.length === 0 ? (
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{mcpServers.length === 0 ? "No MCP servers configured" : "No servers match your search"}
|
||||
{mcpServers.length === 0 ? "No MCP servers configured" : "No tools or servers match your search"}
|
||||
</p>
|
||||
) : (
|
||||
filteredServers.map((server) => (
|
||||
<MCPToolList
|
||||
serverViews.map(({ server, serverName, state, visibleTools }) => (
|
||||
<ServerToolGroup
|
||||
key={server.server_id}
|
||||
serverId={server.server_id}
|
||||
serverName={server.alias || server.server_name || server.server_id}
|
||||
accessToken={accessToken}
|
||||
serverName={serverName}
|
||||
state={state}
|
||||
visibleTools={visibleTools}
|
||||
expanded={searchActive || expandedServerIds.has(server.server_id)}
|
||||
searchActive={searchActive}
|
||||
selectedTools={selectedTools}
|
||||
onToggle={handleToggleTool}
|
||||
onToggleExpanded={() => handleToggleExpanded(server.server_id)}
|
||||
onToggleTool={handleToggleTool}
|
||||
onSetToolsSelected={handleSetToolsSelected}
|
||||
/>
|
||||
))
|
||||
)}
|
||||
|
|
@ -266,41 +412,60 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset
|
|||
|
||||
{/* Right panel: Your Toolset */}
|
||||
<div className="w-72 shrink-0">
|
||||
<p className="text-sm font-semibold text-foreground mb-2 block">
|
||||
Your Toolset{" "}
|
||||
<span className="text-xs font-normal text-muted-foreground">({selectedTools.length} tools)</span>
|
||||
</p>
|
||||
<div className="flex items-center justify-between mb-2">
|
||||
<p className="text-sm font-semibold text-foreground">
|
||||
Your Toolset{" "}
|
||||
<span className="text-xs font-normal text-muted-foreground">({selectedTools.length} tools)</span>
|
||||
</p>
|
||||
<Button
|
||||
size="xs"
|
||||
variant="ghost"
|
||||
disabled={selectedTools.length === 0}
|
||||
onClick={() => setSelectedTools([])}
|
||||
>
|
||||
Clear all
|
||||
</Button>
|
||||
</div>
|
||||
<div className="space-y-1 overflow-y-auto" style={{ maxHeight: 340 }}>
|
||||
{selectedTools.length === 0 ? (
|
||||
<p className="text-muted-foreground text-sm">No tools added yet</p>
|
||||
) : (
|
||||
selectedTools.map((tool, idx) => (
|
||||
<button
|
||||
key={idx}
|
||||
type="button"
|
||||
onClick={() => handleToggleTool(tool)}
|
||||
className="w-full flex items-center justify-between px-3 py-1.5 rounded-lg border border-purple-200 bg-purple-50 hover:bg-destructive/10 hover:border-destructive/20 group transition-colors dark:border-purple-800 dark:bg-purple-950"
|
||||
>
|
||||
<div className="min-w-0 text-left">
|
||||
<span className="text-xs font-medium text-purple-800 group-hover:text-destructive truncate block dark:text-purple-200">
|
||||
{displayToolName(serverPrefixById.get(tool.server_id), tool.tool_name)}
|
||||
selectedTools.map((tool) => {
|
||||
const toolName = displayToolName(serverPrefixById.get(tool.server_id), tool.tool_name);
|
||||
const issue = serversLoaded ? selectedToolIssue(tool, serverPrefixById, toolsByServer) : undefined;
|
||||
return (
|
||||
<button
|
||||
key={`${tool.server_id}:${tool.tool_name}`}
|
||||
type="button"
|
||||
aria-label={`Remove ${toolName}`}
|
||||
onClick={() => handleToggleTool(tool)}
|
||||
className="w-full flex items-center justify-between px-3 py-1.5 rounded-lg border border-purple-200 bg-purple-50 hover:bg-destructive/10 hover:border-destructive/20 group transition-colors dark:border-purple-800 dark:bg-purple-950"
|
||||
>
|
||||
<div className="min-w-0 text-left">
|
||||
<span className="text-xs font-medium text-purple-800 group-hover:text-destructive truncate block dark:text-purple-200">
|
||||
{toolName}
|
||||
</span>
|
||||
{issue ? (
|
||||
<span className="text-[10px] text-warning truncate block">{issue}</span>
|
||||
) : (
|
||||
<span className="text-[10px] text-purple-400 truncate block dark:text-purple-500">
|
||||
{tool.server_id.slice(0, 8)}…
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
<span className="ml-2 text-purple-300 group-hover:text-destructive text-xs shrink-0 dark:text-purple-600">
|
||||
✕
|
||||
</span>
|
||||
<span className="text-[10px] text-purple-400 truncate block dark:text-purple-500">
|
||||
{tool.server_id.slice(0, 8)}…
|
||||
</span>
|
||||
</div>
|
||||
<span className="ml-2 text-purple-300 group-hover:text-destructive text-xs shrink-0 dark:text-purple-600">
|
||||
✕
|
||||
</span>
|
||||
</button>
|
||||
))
|
||||
</button>
|
||||
);
|
||||
})
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="flex justify-end gap-2 mt-4 pt-4 border-t border-border">
|
||||
<Button variant="outline" onClick={onClose}>
|
||||
<Button variant="outline" onClick={handleClose}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button onClick={() => void form.handleSubmit(handleSubmit)()} disabled={saving} aria-busy={saving}>
|
||||
|
|
@ -463,19 +628,13 @@ export function MCPToolsetsTab({ accessToken, userRole }: MCPToolsetsTabProps) {
|
|||
size="compact"
|
||||
/>
|
||||
|
||||
<CreateToolsetModal
|
||||
open={createOpen}
|
||||
onClose={() => setCreateOpen(false)}
|
||||
onSave={handleCreate}
|
||||
accessToken={accessToken}
|
||||
/>
|
||||
<CreateToolsetModal open={createOpen} onClose={() => setCreateOpen(false)} onSave={handleCreate} />
|
||||
|
||||
{editToolset && (
|
||||
<CreateToolsetModal
|
||||
open={!!editToolset}
|
||||
onClose={() => setEditToolset(null)}
|
||||
onSave={handleUpdate}
|
||||
accessToken={accessToken}
|
||||
initialToolset={editToolset}
|
||||
/>
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,90 @@
|
|||
import { describe, it, expect } from "vitest";
|
||||
import {
|
||||
deselectTools,
|
||||
filterServerTools,
|
||||
isToolSelected,
|
||||
matchesServerName,
|
||||
selectTools,
|
||||
selectionState,
|
||||
toggleTool,
|
||||
} from "./toolsetSelection";
|
||||
|
||||
const t = (server_id: string, tool_name: string) => ({ server_id, tool_name });
|
||||
|
||||
const githubTools = [
|
||||
{ name: "create_issue", description: "Open a new issue" },
|
||||
{ name: "list_pulls", description: "List pull requests" },
|
||||
{ name: "search_code", description: "Search code across repos" },
|
||||
];
|
||||
|
||||
describe("filterServerTools", () => {
|
||||
it("returns every tool for an empty or whitespace query", () => {
|
||||
expect(filterServerTools("", ["github"], githubTools)).toBe(githubTools);
|
||||
expect(filterServerTools(" ", ["github"], githubTools)).toBe(githubTools);
|
||||
});
|
||||
|
||||
it("returns every tool when the query matches a server name, case-insensitively", () => {
|
||||
expect(filterServerTools("GitH", ["github_prod", "GitHub"], githubTools)).toBe(githubTools);
|
||||
});
|
||||
|
||||
it("matches tool names and descriptions when the server name does not match", () => {
|
||||
expect(filterServerTools("issue", ["github"], githubTools).map((tool) => tool.name)).toEqual(["create_issue"]);
|
||||
expect(filterServerTools("PULL REQ", ["github"], githubTools).map((tool) => tool.name)).toEqual(["list_pulls"]);
|
||||
});
|
||||
|
||||
it("returns nothing when neither the server nor any tool matches", () => {
|
||||
expect(filterServerTools("slack", ["github"], githubTools)).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("matchesServerName", () => {
|
||||
it("matches any of the server names, case-insensitively, and never for an empty query", () => {
|
||||
expect(matchesServerName("PROD", ["github_prod", "GitHub"])).toBe(true);
|
||||
expect(matchesServerName("slack", ["github_prod", "GitHub"])).toBe(false);
|
||||
expect(matchesServerName(" ", ["github"])).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("selectTools", () => {
|
||||
it("appends new tools after the existing selection in order", () => {
|
||||
expect(selectTools([t("a", "x")], [t("b", "y"), t("b", "z")])).toEqual([t("a", "x"), t("b", "y"), t("b", "z")]);
|
||||
});
|
||||
|
||||
it("skips tools already selected and duplicates within the input", () => {
|
||||
expect(selectTools([t("a", "x")], [t("a", "x"), t("a", "y"), t("a", "y")])).toEqual([t("a", "x"), t("a", "y")]);
|
||||
});
|
||||
|
||||
it("treats the same tool name on different servers as different tools", () => {
|
||||
expect(selectTools([t("a", "search")], [t("b", "search")])).toEqual([t("a", "search"), t("b", "search")]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("deselectTools", () => {
|
||||
it("removes only the given tools and keeps the rest of the selection", () => {
|
||||
const selected = [t("a", "x"), t("b", "y"), t("gone", "z")];
|
||||
expect(deselectTools(selected, [t("b", "y"), t("b", "not-selected")])).toEqual([t("a", "x"), t("gone", "z")]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("toggleTool and isToolSelected", () => {
|
||||
it("adds an unselected tool and removes a selected one", () => {
|
||||
const added = toggleTool([t("a", "x")], t("a", "y"));
|
||||
expect(added).toEqual([t("a", "x"), t("a", "y")]);
|
||||
expect(isToolSelected(added, t("a", "y"))).toBe(true);
|
||||
expect(toggleTool(added, t("a", "x"))).toEqual([t("a", "y")]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("selectionState", () => {
|
||||
const tools = [t("a", "x"), t("a", "y")];
|
||||
|
||||
it("is none for an empty tool list", () => {
|
||||
expect(selectionState([t("a", "x")], [])).toBe("none");
|
||||
});
|
||||
|
||||
it("reports none, some and all against the given tools only", () => {
|
||||
expect(selectionState([t("b", "x")], tools)).toBe("none");
|
||||
expect(selectionState([t("a", "x"), t("b", "z")], tools)).toBe("some");
|
||||
expect(selectionState([t("a", "y"), t("a", "x")], tools)).toBe("all");
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,63 @@
|
|||
import type { MCPToolsetTool } from "@/components/mcp_tools/types";
|
||||
|
||||
export interface SearchableTool {
|
||||
name: string;
|
||||
description: string;
|
||||
}
|
||||
|
||||
export type SelectionState = "all" | "some" | "none";
|
||||
|
||||
const toolKey = (tool: MCPToolsetTool): string => `${tool.server_id}\u0000${tool.tool_name}`;
|
||||
|
||||
const normalize = (query: string): string => query.trim().toLowerCase();
|
||||
|
||||
export const matchesServerName = (query: string, serverNames: readonly string[]): boolean => {
|
||||
const q = normalize(query);
|
||||
return q !== "" && serverNames.some((name) => name.toLowerCase().includes(q));
|
||||
};
|
||||
|
||||
export const filterServerTools = <T extends SearchableTool>(
|
||||
query: string,
|
||||
serverNames: readonly string[],
|
||||
tools: readonly T[],
|
||||
): readonly T[] => {
|
||||
const q = normalize(query);
|
||||
if (!q || matchesServerName(query, serverNames)) {
|
||||
return tools;
|
||||
}
|
||||
return tools.filter((tool) => tool.name.toLowerCase().includes(q) || tool.description.toLowerCase().includes(q));
|
||||
};
|
||||
|
||||
export const selectTools = (
|
||||
selected: readonly MCPToolsetTool[],
|
||||
tools: readonly MCPToolsetTool[],
|
||||
): MCPToolsetTool[] => {
|
||||
const selectedKeys = new Set(selected.map(toolKey));
|
||||
const uniqueTools = [...new Map(tools.map((tool) => [toolKey(tool), tool])).values()];
|
||||
return [...selected, ...uniqueTools.filter((tool) => !selectedKeys.has(toolKey(tool)))];
|
||||
};
|
||||
|
||||
export const deselectTools = (
|
||||
selected: readonly MCPToolsetTool[],
|
||||
tools: readonly MCPToolsetTool[],
|
||||
): MCPToolsetTool[] => {
|
||||
const removed = new Set(tools.map(toolKey));
|
||||
return selected.filter((tool) => !removed.has(toolKey(tool)));
|
||||
};
|
||||
|
||||
export const toggleTool = (selected: readonly MCPToolsetTool[], tool: MCPToolsetTool): MCPToolsetTool[] =>
|
||||
isToolSelected(selected, tool) ? deselectTools(selected, [tool]) : selectTools(selected, [tool]);
|
||||
|
||||
export const isToolSelected = (selected: readonly MCPToolsetTool[], tool: MCPToolsetTool): boolean =>
|
||||
selected.some((t) => toolKey(t) === toolKey(tool));
|
||||
|
||||
export const selectionState = (
|
||||
selected: readonly MCPToolsetTool[],
|
||||
tools: readonly MCPToolsetTool[],
|
||||
): SelectionState => {
|
||||
if (tools.length === 0) return "none";
|
||||
const selectedKeys = new Set(selected.map(toolKey));
|
||||
const count = tools.filter((tool) => selectedKeys.has(toolKey(tool))).length;
|
||||
if (count === 0) return "none";
|
||||
return count === tools.length ? "all" : "some";
|
||||
};
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
import userEvent from "@testing-library/user-event";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { renderWithProviders, screen } from "../../../tests/test-utils";
|
||||
import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils";
|
||||
import { MODEL_MAX_BUDGET_PREMIUM_HINT, ModelMaxBudgetEditor, type ModelMaxBudget } from "./ModelMaxBudgetEditor";
|
||||
|
||||
const STORED: ModelMaxBudget = { "gpt-4o": { budget_limit: 5, time_period: "30d" } };
|
||||
|
|
@ -67,3 +68,65 @@ describe("ModelMaxBudgetEditor with an enterprise license", () => {
|
|||
expect(addButton()).toBeEnabled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("ModelMaxBudgetEditor model families", () => {
|
||||
it("stores a typed pattern that matches no listed model as one budget entry", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onChange = vi.fn();
|
||||
renderWithProviders(
|
||||
<ModelMaxBudgetEditor
|
||||
value={{}}
|
||||
onChange={onChange}
|
||||
availableModels={["gpt-4o", "claude-opus-4-8"]}
|
||||
premiumUser={true}
|
||||
/>,
|
||||
);
|
||||
|
||||
await user.click(addButton());
|
||||
const modelField = screen.getByPlaceholderText("Select model");
|
||||
await user.type(modelField, "claude-opus*");
|
||||
await user.click(await screen.findByRole("option", { name: /claude-opus\*/ }));
|
||||
fireEvent.change(screen.getByPlaceholderText("Max spend ($)"), { target: { value: "10" } });
|
||||
|
||||
expect(onChange).toHaveBeenLastCalledWith({ "claude-opus*": { budget_limit: 10, time_period: "30d" } });
|
||||
});
|
||||
|
||||
it("does not offer a typed pattern that another row already budgets", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(
|
||||
<ModelMaxBudgetEditor
|
||||
value={{ "claude-opus*": { budget_limit: 10, time_period: "1d" } }}
|
||||
onChange={vi.fn()}
|
||||
availableModels={["gpt-4o", "claude-opus-4-8"]}
|
||||
premiumUser={true}
|
||||
/>,
|
||||
);
|
||||
|
||||
await user.click(addButton());
|
||||
const [, newRowModelField] = screen.getAllByPlaceholderText("Select model");
|
||||
await user.type(newRowModelField, "claude-opus*");
|
||||
|
||||
expect(screen.queryByRole("option", { name: /claude-opus\*/ })).not.toBeInTheDocument();
|
||||
await user.clear(newRowModelField);
|
||||
await user.type(newRowModelField, "claude-sonnet*");
|
||||
expect(await screen.findByRole("option", { name: /claude-sonnet\*/ })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not offer a typed name that is neither listed nor a pattern", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(
|
||||
<ModelMaxBudgetEditor
|
||||
value={{}}
|
||||
onChange={vi.fn()}
|
||||
availableModels={["gpt-4o", "claude-opus-4-8"]}
|
||||
premiumUser={true}
|
||||
/>,
|
||||
);
|
||||
|
||||
await user.click(addButton());
|
||||
await user.type(screen.getByPlaceholderText("Select model"), "claude-opsu-4-8");
|
||||
|
||||
expect(screen.queryByRole("option", { name: /claude-opsu-4-8/ })).not.toBeInTheDocument();
|
||||
expect(screen.getByText("No models found")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -50,6 +50,8 @@ export const MODEL_BUDGET_PERIOD_OPTIONS = [
|
|||
|
||||
const DEFAULT_PERIOD = "30d";
|
||||
|
||||
const MODEL_PATTERN_VALUE = { hint: "Model pattern", accepts: (typed: string) => typed.includes("*") };
|
||||
|
||||
export const entriesToModelMaxBudget = (entries: readonly ModelBudgetEntry[]): ModelMaxBudget =>
|
||||
Object.fromEntries(
|
||||
entries
|
||||
|
|
@ -121,7 +123,7 @@ export function ModelMaxBudgetEditor({
|
|||
const blurb = (
|
||||
<div className="text-xs text-muted-foreground">
|
||||
{premiumUser
|
||||
? "Cap spend per model over its own window. A budget set on the bare model name also covers the provider-prefixed spelling of that model."
|
||||
? "Cap spend per model over its own window. A budget set on the bare model name also covers the provider-prefixed spelling of that model. Pick an access group, or type a pattern like claude-opus*, to share one budget across every model it covers."
|
||||
: MODEL_MAX_BUDGET_PREMIUM_HINT}
|
||||
</div>
|
||||
);
|
||||
|
|
@ -165,6 +167,8 @@ export function ModelMaxBudgetEditor({
|
|||
onValueChange={(model) => updateEntry(entry.id, { model })}
|
||||
placeholder="Select model"
|
||||
emptyText="No models found"
|
||||
customValue={MODEL_PATTERN_VALUE}
|
||||
reservedValues={entries.flatMap((other) => (other.id !== entry.id && other.model ? [other.model] : []))}
|
||||
disabled={!premiumUser}
|
||||
/>
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -63,6 +63,24 @@ describe("SearchSelect", () => {
|
|||
expect(onValueChange).toHaveBeenLastCalledWith("team-3");
|
||||
});
|
||||
|
||||
it("offers the typed text as a value only when the caller accepts it", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onValueChange = vi.fn();
|
||||
const customValue = { hint: "Custom value", accepts: (typed: string) => typed.endsWith("*") };
|
||||
const { rerender } = render(<SearchSelect options={OPTIONS} onValueChange={onValueChange} />);
|
||||
const input = screen.getByRole("combobox");
|
||||
await user.click(input);
|
||||
fireEvent.change(input, { target: { value: "team-*" } });
|
||||
expect(screen.queryByRole("option", { name: /team-\*/ })).not.toBeInTheDocument();
|
||||
|
||||
rerender(<SearchSelect options={OPTIONS} onValueChange={onValueChange} customValue={customValue} />);
|
||||
fireEvent.change(input, { target: { value: "team-4" } });
|
||||
expect(screen.queryByRole("option", { name: /team-4/ })).not.toBeInTheDocument();
|
||||
fireEvent.change(input, { target: { value: "team-*" } });
|
||||
await user.click(await screen.findByRole("option", { name: /team-\*/ }));
|
||||
expect(onValueChange).toHaveBeenLastCalledWith("team-*");
|
||||
});
|
||||
|
||||
it("filters the options client-side as you type", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<SearchSelect options={OPTIONS} onValueChange={vi.fn()} />);
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
"use client";
|
||||
|
||||
import React from "react";
|
||||
import React, { useState } from "react";
|
||||
import {
|
||||
Combobox,
|
||||
ComboboxContent,
|
||||
|
|
@ -19,6 +19,11 @@ export interface SearchSelectOption {
|
|||
icon?: React.ReactNode;
|
||||
}
|
||||
|
||||
interface CustomValue {
|
||||
hint: string;
|
||||
accepts: (typed: string) => boolean;
|
||||
}
|
||||
|
||||
interface SearchSelectProps {
|
||||
options: SearchSelectOption[];
|
||||
value?: string | null;
|
||||
|
|
@ -29,6 +34,8 @@ interface SearchSelectProps {
|
|||
className?: string;
|
||||
inputId?: string;
|
||||
allowClear?: boolean;
|
||||
customValue?: CustomValue;
|
||||
reservedValues?: readonly string[];
|
||||
"aria-label"?: string;
|
||||
id?: string;
|
||||
}
|
||||
|
|
@ -49,13 +56,22 @@ export function SearchSelect({
|
|||
className,
|
||||
inputId,
|
||||
allowClear = true,
|
||||
customValue,
|
||||
reservedValues = [],
|
||||
"aria-label": ariaLabel,
|
||||
id,
|
||||
}: SearchSelectProps) {
|
||||
const [query, setQuery] = useState("");
|
||||
const selected =
|
||||
value == null || value === "" ? null : options.find((option) => option.value === value) ?? { label: value, value };
|
||||
const items =
|
||||
const known =
|
||||
selected !== null && !options.some((option) => option.value === selected.value) ? [selected, ...options] : options;
|
||||
const typed = query.trim();
|
||||
const isNewValue = typed !== "" && !reservedValues.includes(typed) && !known.some((option) => option.value === typed);
|
||||
const items =
|
||||
customValue !== undefined && isNewValue && customValue.accepts(typed)
|
||||
? [...known, { label: typed, value: typed, sublabel: customValue.hint }]
|
||||
: known;
|
||||
|
||||
return (
|
||||
<Combobox
|
||||
|
|
@ -65,6 +81,7 @@ export function SearchSelect({
|
|||
isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value}
|
||||
itemToStringLabel={(item: SearchSelectOption) => item.label}
|
||||
filter={matchesQuery}
|
||||
onInputValueChange={setQuery}
|
||||
disabled={disabled}
|
||||
>
|
||||
<ComboboxInput
|
||||
|
|
|
|||
54
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
54
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -8475,7 +8475,7 @@ export interface paths {
|
|||
* - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
|
||||
* - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only.
|
||||
* - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||||
* - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||||
|
|
@ -8721,7 +8721,7 @@ export interface paths {
|
|||
* - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
|
||||
* - spend: Optional[float] - Amount spent by key
|
||||
* - max_budget: Optional[float] - Max budget for key
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
* - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||||
|
|
@ -8798,7 +8798,7 @@ export interface paths {
|
|||
* - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
* - guardrails: Optional[List[str]] - List of active guardrails for the key
|
||||
* - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||||
* - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||||
|
|
@ -8931,7 +8931,7 @@ export interface paths {
|
|||
* - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
|
||||
* - spend: Optional[float] - Amount spent by key
|
||||
* - max_budget: Optional[float] - Max budget for key
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
* - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget.
|
||||
|
|
@ -9020,7 +9020,7 @@ export interface paths {
|
|||
* - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
|
||||
* - spend: Optional[float] - Amount spent by key
|
||||
* - max_budget: Optional[float] - Max budget for key
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
* - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||||
|
|
@ -17100,7 +17100,7 @@ export interface paths {
|
|||
* - prompts: Optional[List[str]] - List of prompts that the team is allowed to use.
|
||||
* - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
|
||||
* - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
|
||||
* - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
* - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
* - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
|
||||
* - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
|
||||
* - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
|
||||
|
|
@ -17330,7 +17330,7 @@ export interface paths {
|
|||
* - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
|
||||
* - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
|
||||
* - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
|
||||
* - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
* - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
|
||||
* - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
|
||||
* - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
|
||||
* - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only.
|
||||
|
|
@ -18780,7 +18780,7 @@ export interface paths {
|
|||
* - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
|
||||
* - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
* - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
* - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
* - model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
* - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
|
||||
|
|
@ -18919,7 +18919,7 @@ export interface paths {
|
|||
* - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
|
||||
* - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
* - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
* - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
* - model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
* - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
|
||||
|
|
@ -27725,6 +27725,10 @@ export interface components {
|
|||
entities?: {
|
||||
[key: string]: components["schemas"]["MetricWithMetadata"];
|
||||
};
|
||||
/** Guardrails */
|
||||
guardrails?: {
|
||||
[key: string]: components["schemas"]["GuardrailMetric"];
|
||||
};
|
||||
/** Mcp Servers */
|
||||
mcp_servers?: {
|
||||
[key: string]: components["schemas"]["MetricWithMetadata"];
|
||||
|
|
@ -33240,6 +33244,38 @@ export interface components {
|
|||
/** Updated At */
|
||||
updated_at?: string | null;
|
||||
};
|
||||
/** GuardrailKeyMetric */
|
||||
GuardrailKeyMetric: {
|
||||
metadata: components["schemas"]["KeyMetadata"];
|
||||
metrics: components["schemas"]["GuardrailSpendMetrics"];
|
||||
};
|
||||
/** GuardrailMetadata */
|
||||
GuardrailMetadata: {
|
||||
/** Guardrail Name */
|
||||
guardrail_name: string;
|
||||
};
|
||||
/** GuardrailMetric */
|
||||
GuardrailMetric: {
|
||||
/** Api Key Breakdown */
|
||||
api_key_breakdown?: {
|
||||
[key: string]: components["schemas"]["GuardrailKeyMetric"];
|
||||
};
|
||||
metadata: components["schemas"]["GuardrailMetadata"];
|
||||
metrics: components["schemas"]["GuardrailSpendMetrics"];
|
||||
};
|
||||
/** GuardrailSpendMetrics */
|
||||
GuardrailSpendMetrics: {
|
||||
/**
|
||||
* Spend
|
||||
* @default 0
|
||||
*/
|
||||
spend: number;
|
||||
/**
|
||||
* Untracked Units
|
||||
* @default 0
|
||||
*/
|
||||
untracked_units: number;
|
||||
};
|
||||
/** GuardrailSubmissionItem */
|
||||
GuardrailSubmissionItem: {
|
||||
/** Created At */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue