mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge branch 'litellm_internal_staging' of https://github.com/BerriAI/litellm into litellm_fix_v1_messages_midstream_timeout_failure_logging
# Conflicts: # litellm/litellm_core_utils/litellm_logging.py
This commit is contained in:
commit
2f3d1ca575
60 changed files with 3338 additions and 568 deletions
5
.github/ci-coverage-allowlist.yml
vendored
5
.github/ci-coverage-allowlist.yml
vendored
|
|
@ -79,6 +79,11 @@ test_paths:
|
|||
- tests/load_tests/test_otel_load_test.py
|
||||
- tests/load_tests/test_vertex_embeddings_load_test.py
|
||||
- tests/load_tests/test_vertex_load_tests.py
|
||||
- reason: >-
|
||||
Env-gated saturation benchmark requires a live proxy and provider credentials, so it is run
|
||||
locally rather than in pull-request jobs
|
||||
paths:
|
||||
- tests/load_tests/test_granian_admission_saturation.py
|
||||
- reason: >-
|
||||
A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
|
||||
localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
"limit": 14074
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2214
|
||||
"limit": 2206
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 319
|
||||
|
|
@ -57,7 +57,7 @@
|
|||
"limit": 5601
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15287
|
||||
"limit": 15285
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -99,19 +99,19 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44362
|
||||
"limit": 44360
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 38323
|
||||
"limit": 38311
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19624
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 29861
|
||||
"limit": 29847
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 111
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
"""Anthropic error format type definitions."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Literal
|
||||
|
||||
from typing_extensions import Required, TypedDict
|
||||
from typing_extensions import NotRequired, ReadOnly, Required, TypedDict
|
||||
|
||||
# Known Anthropic error types
|
||||
# Source: https://docs.anthropic.com/en/api/errors
|
||||
|
|
@ -23,6 +24,7 @@ class AnthropicErrorDetail(TypedDict):
|
|||
|
||||
type: AnthropicErrorType
|
||||
message: str
|
||||
provider_specific_fields: NotRequired[ReadOnly[Mapping[str, object]]]
|
||||
|
||||
|
||||
class AnthropicErrorResponse(TypedDict, total=False):
|
||||
|
|
|
|||
|
|
@ -19,13 +19,13 @@ connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-cov
|
|||
retry-exhaustion) is unchanged from upstream, since those already carry real evidence the
|
||||
topology changed.
|
||||
|
||||
redis-py 8.x fixed this upstream with gentler machinery than this override's
|
||||
``node.disconnect()`` (which also kills connections other coroutines are mid-operation
|
||||
on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per
|
||||
killed connection): it marks in-use connections for reconnect only after their current
|
||||
operation completes, disconnects only the idle pooled ones, and defers reinitialization
|
||||
to the outer retry loop. When the installed ``ClusterNode`` has that per-connection
|
||||
recovery API, the factory returns the base ``RedisCluster`` unmodified.
|
||||
redis-py 8.x recovers connections per-connection, so the copied override is not used. Upstream
|
||||
still flips the shared ``_initialize`` flag on any node's timeout, funneling every concurrent
|
||||
caller through the reinit lock and, if ``CLUSTER SLOTS`` lands on the slow node, into a full
|
||||
teardown. For those versions the factory returns a thin wrapper around upstream's
|
||||
``_execute_command`` that clears the flag again after an isolated timeout (a ConnectionError,
|
||||
a third consecutive timeout on the same node, or a concurrent request from any other command
|
||||
or ``aclose()`` still reinits).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
|
@ -44,6 +44,8 @@ class _ClusterNodeAttrs(Protocol):
|
|||
mode; typing ``target_node`` as this Protocol at the one boundary keeps the override's
|
||||
own logic fully typed without a banned ``typing.cast``."""
|
||||
|
||||
name: str
|
||||
|
||||
async def execute_command(
|
||||
self,
|
||||
*args: object,
|
||||
|
|
@ -78,18 +80,20 @@ class _ClusterAttrs(Protocol):
|
|||
#: this override can't see (Python won't error -- it'll just run our now-stale copy), so
|
||||
#: construction logs a loud warning rather than silently trusting an unverified copy.
|
||||
_VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"})
|
||||
_CONSECUTIVE_TIMEOUTS_BEFORE_REINIT: Final = 3
|
||||
|
||||
|
||||
def get_litellm_async_redis_cluster_class(
|
||||
def get_litellm_async_redis_cluster_class( # noqa: C901 # supports redis-py version-specific cluster implementations
|
||||
cluster_node_class: type | None = None,
|
||||
base_cluster_class: type | None = None,
|
||||
) -> type["_AsyncRedisClusterType"]:
|
||||
"""Returns the base ``RedisCluster`` when the installed redis-py already recovers a
|
||||
node-level connection error per-connection (8.x+), else builds the ``RedisCluster``
|
||||
subclass with the per-node isolation fix for older versions whose upstream branch
|
||||
tears down the whole cluster client.
|
||||
"""Returns a timeout-tolerant ``RedisCluster`` subclass when installed redis-py already
|
||||
recovers node-level connections per-connection (8.x+), else builds the ``RedisCluster``
|
||||
subclass with the per-node isolation fix for older versions whose upstream branch tears
|
||||
down the whole cluster client.
|
||||
|
||||
``cluster_node_class`` exists for dependency injection in tests; production callers
|
||||
leave it unset and the installed ``ClusterNode`` is used.
|
||||
``cluster_node_class`` and ``base_cluster_class`` exist for dependency injection in tests;
|
||||
production callers leave them unset and the installed redis-py classes are used.
|
||||
|
||||
Imported lazily because this module is reachable from a base ``import litellm`` while
|
||||
redis is not a base dependency. Cheap to call repeatedly: the underlying redis
|
||||
|
|
@ -118,13 +122,68 @@ def get_litellm_async_redis_cluster_class(
|
|||
from redis.exceptions import TimeoutError as _RedisTimeoutError
|
||||
|
||||
node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode
|
||||
base_class: Final = base_cluster_class if base_cluster_class is not None else _BaseAsyncRedisCluster
|
||||
if hasattr(node_class, "update_active_connections_for_reconnect"):
|
||||
verbose_logger.debug(
|
||||
"redis-py %s recovers a node-level connection error per-connection upstream; "
|
||||
"using the base RedisCluster without litellm's node-isolation override.",
|
||||
"redis-py %s recovers node connections per-connection upstream; using "
|
||||
"LiteLLM's timeout-tolerant RedisCluster wrapper.",
|
||||
redis.__version__,
|
||||
)
|
||||
return _BaseAsyncRedisCluster
|
||||
|
||||
class LiteLLMAsyncRedisClusterTimeoutTolerant(
|
||||
base_class # pyright: ignore[reportGeneralTypeIssues, reportUntypedBaseClass] # the injected base class is selected at runtime
|
||||
):
|
||||
def __init__(
|
||||
self,
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: passes redis-py's constructor kwargs through untouched
|
||||
) -> None:
|
||||
self._litellm_initialize = False
|
||||
self._litellm_reinit_requests = 0
|
||||
self._litellm_tolerated_timeouts = 0
|
||||
super().__init__(*args, **kwargs)
|
||||
self._litellm_consecutive_timeouts: dict[ # mutable-ok: per-node counter updated on the command hot path
|
||||
str, int
|
||||
] = {}
|
||||
|
||||
@property
|
||||
def _initialize(self) -> bool:
|
||||
return self._litellm_initialize
|
||||
|
||||
@_initialize.setter
|
||||
def _initialize(self, value: bool) -> None:
|
||||
if value:
|
||||
self._litellm_reinit_requests += 1
|
||||
self._litellm_initialize = value
|
||||
|
||||
async def _execute_command(
|
||||
self,
|
||||
target_node: _ClusterNodeAttrs,
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: matches redis-py's own command dispatch signature
|
||||
) -> object:
|
||||
outstanding_before: Final = self._litellm_reinit_requests - self._litellm_tolerated_timeouts
|
||||
pending_before: Final = self._litellm_initialize
|
||||
try:
|
||||
result: Final = await super()._execute_command(target_node, *args, **kwargs)
|
||||
except _RedisTimeoutError:
|
||||
timeouts: Final = self._litellm_consecutive_timeouts.get(target_node.name, 0) + 1
|
||||
if timeouts >= _CONSECUTIVE_TIMEOUTS_BEFORE_REINIT:
|
||||
self._litellm_consecutive_timeouts.pop(target_node.name, None)
|
||||
raise
|
||||
self._litellm_consecutive_timeouts[target_node.name] = timeouts
|
||||
self._litellm_tolerated_timeouts += 1
|
||||
if (
|
||||
not pending_before
|
||||
and self._litellm_reinit_requests - self._litellm_tolerated_timeouts == outstanding_before
|
||||
):
|
||||
self._initialize = False
|
||||
raise
|
||||
if self._litellm_consecutive_timeouts:
|
||||
self._litellm_consecutive_timeouts.pop(target_node.name, None)
|
||||
return result
|
||||
|
||||
return LiteLLMAsyncRedisClusterTimeoutTolerant
|
||||
|
||||
if redis.__version__ not in _VERIFIED_REDIS_VERSIONS:
|
||||
verbose_logger.warning(
|
||||
|
|
|
|||
|
|
@ -151,6 +151,7 @@ DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_
|
|||
MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60"))
|
||||
MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200"))
|
||||
MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600"))
|
||||
MCP_SSO_ASSERTION_CACHE_TTL_SECONDS: Final = int(os.getenv("MCP_SSO_ASSERTION_CACHE_TTL_SECONDS", "60"))
|
||||
|
||||
# Default npm cache directory for STDIO MCP servers.
|
||||
# npm/npx needs a writable cache dir; in containers the default (~/.npm)
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ import subprocess
|
|||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from datetime import datetime as dt_object
|
||||
from functools import lru_cache
|
||||
from types import MappingProxyType, TracebackType
|
||||
|
|
@ -5884,14 +5884,28 @@ def _get_status_fields(
|
|||
#########################################################
|
||||
# Map - guardrail_information.guardrail_status to guardrail_status
|
||||
#########################################################
|
||||
guardrail_status: GuardrailStatus = "not_run"
|
||||
if guardrail_information and isinstance(guardrail_information, list):
|
||||
for information in guardrail_information:
|
||||
if isinstance(information, dict):
|
||||
raw_status = information.get("guardrail_status", "not_run")
|
||||
if raw_status != "not_run":
|
||||
guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run")
|
||||
break
|
||||
# Severity order, least severe first. The status aggregates across ALL
|
||||
# guardrail entries rather than taking the first non-"not_run" one: a
|
||||
# pre_call guardrail that passed (e.g. a mask) records its entry before a
|
||||
# later guardrail's block, and first-wins would report a blocked request
|
||||
# as "success".
|
||||
GUARDRAIL_STATUS_SEVERITY: Final[tuple[GuardrailStatus, ...]] = (
|
||||
"not_run",
|
||||
"success",
|
||||
"guardrail_failed_to_respond",
|
||||
"guardrail_intervened",
|
||||
)
|
||||
entries: Final[Sequence[object]] = guardrail_information if isinstance(guardrail_information, list) else ()
|
||||
raw_statuses: Final[Iterator[object]] = (
|
||||
entry.get("guardrail_status", "not_run") for entry in entries if isinstance(entry, dict)
|
||||
)
|
||||
# A guardrail is free to write any value here, and an unhashable one would
|
||||
# raise TypeError on the mapping lookup and drop the whole payload.
|
||||
guardrail_status: Final[GuardrailStatus] = max(
|
||||
(GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") for raw_status in raw_statuses if isinstance(raw_status, str)),
|
||||
key=GUARDRAIL_STATUS_SEVERITY.index,
|
||||
default="not_run",
|
||||
)
|
||||
|
||||
return StandardLoggingPayloadStatusFields(llm_api_status=llm_api_status, guardrail_status=guardrail_status)
|
||||
|
||||
|
|
|
|||
|
|
@ -8,23 +8,31 @@ Routes to native Cortex REST API endpoints based on model:
|
|||
Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api
|
||||
"""
|
||||
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
anthropic_process_openai_file_message,
|
||||
convert_to_anthropic_tool_result,
|
||||
create_anthropic_image_param,
|
||||
select_anthropic_content_block_type_for_file,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolMessage
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionMessageToolCall,
|
||||
ChatCompletionUsageBlock,
|
||||
Choices,
|
||||
Function,
|
||||
GenericStreamingChunk,
|
||||
Message,
|
||||
ModelResponse,
|
||||
Usage,
|
||||
ModelResponseStream,
|
||||
)
|
||||
|
||||
from ...base_llm.base_model_iterator import BaseModelResponseIterator
|
||||
|
|
@ -93,6 +101,103 @@ def _is_claude_model(model: str) -> bool:
|
|||
return any(name.startswith(p) for p in _CLAUDE_MODEL_PREFIXES)
|
||||
|
||||
|
||||
def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object:
|
||||
"""One OpenAI ``image_url`` block in the native shape Cortex accepts.
|
||||
|
||||
Cortex documents base64 sources only, so remote URLs are inlined the way every
|
||||
other base64-only Anthropic dialect (Bedrock invoke, Vertex) inlines them, and
|
||||
pdf/text data URIs become document blocks rather than malformed image blocks.
|
||||
"""
|
||||
image_url: Final = block.get("image_url")
|
||||
url: Final = image_url if isinstance(image_url, str) else _image_url_field(image_url, "url")
|
||||
if not url:
|
||||
return block
|
||||
|
||||
converted: Final = (
|
||||
anthropic_process_openai_file_message({"type": "file", "file": {"file_data": url}})
|
||||
if select_anthropic_content_block_type_for_file(_data_uri_media_type(url)) == "document"
|
||||
else create_anthropic_image_param(
|
||||
image_url if isinstance(image_url, dict) else url, # mutable-ok: caller's JSON block
|
||||
format=_image_url_field(image_url, "format"),
|
||||
is_bedrock_invoke=True,
|
||||
)
|
||||
)
|
||||
cache_control: Final = block.get("cache_control")
|
||||
if cache_control is None:
|
||||
return converted
|
||||
return {**converted, "cache_control": cache_control} # mutable-ok: JSON wire block
|
||||
|
||||
|
||||
def _image_url_field(image_url: object, key: str) -> str | None:
|
||||
value: Final = image_url.get(key) if isinstance(image_url, dict) else None
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
def _data_uri_media_type(url: str) -> str:
|
||||
match: Final = re.match(r"data:([^;,]+)", url)
|
||||
return match.group(1) if match else ""
|
||||
|
||||
|
||||
def _convert_image_url_blocks_to_anthropic(content: object) -> object:
|
||||
if not isinstance(content, list):
|
||||
return content
|
||||
return [ # mutable-ok: JSON wire blocks
|
||||
_convert_image_url_to_anthropic(block)
|
||||
if isinstance(block, Mapping) and block.get("type") == "image_url"
|
||||
else block
|
||||
for block in content
|
||||
]
|
||||
|
||||
|
||||
def _convert_tool_result_to_anthropic(
|
||||
content: object, tool_call_id: str, cache_control: object
|
||||
) -> Mapping[str, object]:
|
||||
"""The Anthropic ``tool_result`` block for one OpenAI tool message.
|
||||
|
||||
Delegating to the shared converter keeps image, document and per-block cache
|
||||
breakpoints identical to every other Anthropic dialect; only the plain-string
|
||||
and non-list shapes it does not model are handled here.
|
||||
"""
|
||||
if not isinstance(content, list):
|
||||
plain: Final[dict[str, object]] = { # mutable-ok: JSON wire block
|
||||
"type": "tool_result",
|
||||
"tool_use_id": tool_call_id,
|
||||
"content": content if isinstance(content, str) else json.dumps(content),
|
||||
}
|
||||
return {**plain, "cache_control": cache_control} if cache_control is not None else plain
|
||||
converted: Final = convert_to_anthropic_tool_result(
|
||||
ChatCompletionToolMessage(role="tool", tool_call_id=tool_call_id, content=content),
|
||||
force_base64=True,
|
||||
)
|
||||
if cache_control is None:
|
||||
return converted
|
||||
return {**converted, "cache_control": cache_control} # mutable-ok: JSON wire block
|
||||
|
||||
|
||||
def _signed_thinking_blocks(msg: object) -> list[dict[str, object]]: # mutable-ok: JSON wire blocks
|
||||
"""The assistant turn's thinking blocks that can legally be echoed back.
|
||||
|
||||
Only signed blocks round-trip: Cortex rejects a thinking block whose signature is
|
||||
missing, which is what an unsigned block from a non-thinking turn would produce.
|
||||
"""
|
||||
blocks: Final = msg.get("thinking_blocks") if isinstance(msg, dict) else getattr(msg, "thinking_blocks", None)
|
||||
if not isinstance(blocks, list):
|
||||
return [] # mutable-ok: JSON wire blocks
|
||||
return [ # mutable-ok: JSON wire blocks
|
||||
dict(block)
|
||||
for block in blocks
|
||||
if isinstance(block, Mapping) and (block.get("signature") or block.get("type") == "redacted_thinking")
|
||||
]
|
||||
|
||||
|
||||
def _clean_input_schema(schema: object) -> object: # mutable-ok: JSON schema copy
|
||||
return (
|
||||
{key: value for key, value in schema.items() if key != "$schema"}
|
||||
if isinstance(schema, Mapping)
|
||||
else schema # mutable-ok: JSON schema copy
|
||||
) # mutable-ok: JSON schema copy
|
||||
|
||||
|
||||
class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
||||
"""
|
||||
Snowflake Cortex REST API — unified provider.
|
||||
|
|
@ -178,7 +283,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
if "description" in func:
|
||||
anthropic_tool["description"] = func["description"]
|
||||
if "parameters" in func:
|
||||
anthropic_tool["input_schema"] = func["parameters"]
|
||||
anthropic_tool["input_schema"] = _clean_input_schema(func["parameters"])
|
||||
else:
|
||||
anthropic_tool["input_schema"] = {
|
||||
"type": "object",
|
||||
|
|
@ -186,10 +291,16 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
}
|
||||
anthropic_tools.append(anthropic_tool)
|
||||
else:
|
||||
anthropic_tools.append(tool)
|
||||
anthropic_tools.append(
|
||||
{**tool, "input_schema": _clean_input_schema(tool["input_schema"])} # mutable-ok: JSON wire tool
|
||||
if "input_schema" in tool
|
||||
else tool
|
||||
)
|
||||
return anthropic_tools
|
||||
|
||||
def _extract_system_and_messages(self, messages: list[AllMessageValues]) -> tuple[str | None, list[dict]]:
|
||||
def _extract_system_and_messages( # mutable-ok: JSON wire messages
|
||||
self, messages: list[AllMessageValues]
|
||||
) -> tuple[list[dict] | None, list[dict]]:
|
||||
"""
|
||||
Split messages into system prompt and conversation turns for Anthropic format.
|
||||
|
||||
|
|
@ -197,26 +308,39 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
- assistant messages with tool_calls → tool_use content blocks
|
||||
- tool role messages → user role with tool_result content blocks
|
||||
"""
|
||||
system_parts: Final[list[str]] = []
|
||||
conversation: Final[list[dict]] = []
|
||||
system_parts: Final[list[dict]] = [] # mutable-ok: JSON wire messages
|
||||
conversation: Final[list[dict]] = [] # mutable-ok: JSON wire messages
|
||||
|
||||
for msg in messages:
|
||||
if isinstance(msg, dict):
|
||||
role = msg.get("role", "")
|
||||
content: Any = msg.get("content", "")
|
||||
msg_cache_control: object = msg.get("cache_control")
|
||||
else:
|
||||
role = getattr(msg, "role", "")
|
||||
content = getattr(msg, "content", "")
|
||||
msg_cache_control = getattr(msg, "cache_control", None)
|
||||
|
||||
if role == "system":
|
||||
if isinstance(content, str) and content:
|
||||
system_parts.append(content)
|
||||
system_parts.append({"type": "text", "text": content}) # mutable-ok: JSON wire system block
|
||||
elif isinstance(content, list):
|
||||
system_parts.append("\n".join(b.get("text", "") for b in content if b.get("type") == "text"))
|
||||
system_parts.extend(
|
||||
{ # mutable-ok: JSON wire system block
|
||||
"type": "text",
|
||||
"text": block.get("text", ""),
|
||||
**(
|
||||
{"cache_control": block["cache_control"]} if "cache_control" in block else {}
|
||||
), # mutable-ok: JSON wire block
|
||||
}
|
||||
for block in content
|
||||
if isinstance(block, Mapping) and block.get("type") == "text"
|
||||
)
|
||||
elif role == "assistant":
|
||||
tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None)
|
||||
thinking_blocks = _signed_thinking_blocks(msg)
|
||||
if tool_calls:
|
||||
content_blocks: list[dict[str, object]] = []
|
||||
content_blocks: list[dict[str, object]] = list(thinking_blocks) # mutable-ok: JSON wire blocks
|
||||
if content:
|
||||
content_blocks.append({"type": "text", "text": content})
|
||||
for tc in tool_calls:
|
||||
|
|
@ -239,18 +363,26 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
}
|
||||
)
|
||||
conversation.append({"role": "assistant", "content": content_blocks})
|
||||
elif thinking_blocks:
|
||||
thinking_content = (
|
||||
[
|
||||
*thinking_blocks,
|
||||
*copy.deepcopy(content),
|
||||
]
|
||||
if isinstance(content, list)
|
||||
else [*thinking_blocks, *([{"type": "text", "text": content}] if content else [])]
|
||||
) # rebind-ok: loop-local normalized content
|
||||
conversation.append({"role": "assistant", "content": thinking_content})
|
||||
else:
|
||||
conversation.append({"role": "assistant", "content": content})
|
||||
elif role == "tool":
|
||||
tool_call_id = (
|
||||
tool_call_id_value = (
|
||||
msg.get("tool_call_id", "") if isinstance(msg, dict) else getattr(msg, "tool_call_id", "")
|
||||
)
|
||||
tool_content = content if isinstance(content, str) else json.dumps(content)
|
||||
tool_result_block = {
|
||||
"type": "tool_result",
|
||||
"tool_use_id": tool_call_id,
|
||||
"content": tool_content,
|
||||
}
|
||||
tool_call_id = (
|
||||
tool_call_id_value if isinstance(tool_call_id_value, str) else ""
|
||||
) # rebind-ok: normalized loop value
|
||||
tool_result_block = _convert_tool_result_to_anthropic(content, tool_call_id, msg_cache_control)
|
||||
if (
|
||||
conversation
|
||||
and conversation[-1]["role"] == "user"
|
||||
|
|
@ -260,11 +392,18 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
):
|
||||
conversation[-1]["content"].append(tool_result_block)
|
||||
else:
|
||||
conversation.append({"role": "user", "content": [tool_result_block]})
|
||||
conversation.append(
|
||||
{"role": "user", "content": [tool_result_block]} # mutable-ok: JSON wire message
|
||||
) # mutable-ok: JSON wire message
|
||||
else:
|
||||
conversation.append({"role": role, "content": content})
|
||||
conversation.append( # mutable-ok: JSON wire message
|
||||
{ # mutable-ok: JSON wire message
|
||||
"role": role,
|
||||
"content": _convert_image_url_blocks_to_anthropic(content),
|
||||
} # mutable-ok: JSON wire message
|
||||
)
|
||||
|
||||
system: Final[str | None] = "\n\n".join(system_parts) if system_parts else None
|
||||
system: Final[list[dict] | None] = system_parts if system_parts else None # mutable-ok: JSON wire messages
|
||||
return system, conversation
|
||||
|
||||
def transform_request(
|
||||
|
|
@ -339,7 +478,9 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
extra_body: dict,
|
||||
) -> dict:
|
||||
"""Anthropic Messages format for /messages endpoint."""
|
||||
system, conversation = self._extract_system_and_messages(messages)
|
||||
passthrough_system: Final = optional_params.pop("system", None)
|
||||
extracted_system, conversation = self._extract_system_and_messages(messages)
|
||||
system: Final = passthrough_system if passthrough_system is not None else extracted_system
|
||||
|
||||
if "tools" in optional_params:
|
||||
optional_params["tools"] = self._transform_tools_to_anthropic(optional_params["tools"])
|
||||
|
|
@ -353,16 +494,19 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
|
||||
model_name: Final = model.removeprefix("snowflake/")
|
||||
|
||||
body: Final[dict[str, object]] = {
|
||||
"model": model_name,
|
||||
"messages": conversation,
|
||||
"stream": stream,
|
||||
**optional_params,
|
||||
**extra_body,
|
||||
}
|
||||
|
||||
body: Final[dict[str, object]] = normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire body
|
||||
{ # mutable-ok: JSON wire body
|
||||
"model": model_name,
|
||||
"messages": conversation,
|
||||
"stream": stream,
|
||||
**optional_params,
|
||||
**extra_body, # mutable-ok: JSON wire body
|
||||
}
|
||||
)
|
||||
if system is not None:
|
||||
body["system"] = system
|
||||
body["system"] = normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire payload
|
||||
{"system": system} # mutable-ok: JSON wire payload
|
||||
)["system"]
|
||||
|
||||
if "max_tokens" not in body:
|
||||
body["max_tokens"] = 4096 # reasonable default; Anthropic API max varies by model
|
||||
|
|
@ -435,23 +579,10 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
additional_args={"complete_input_dict": request_data},
|
||||
)
|
||||
|
||||
text_content = ""
|
||||
tool_calls: Final = []
|
||||
|
||||
for block in response_json.get("content", []):
|
||||
if block.get("type") == "text":
|
||||
text_content += block.get("text", "")
|
||||
elif block.get("type") == "tool_use":
|
||||
tool_calls.append(
|
||||
ChatCompletionMessageToolCall(
|
||||
id=block.get("id", ""),
|
||||
type="function",
|
||||
function=Function(
|
||||
name=block.get("name", ""),
|
||||
arguments=json.dumps(block.get("input", {})),
|
||||
),
|
||||
)
|
||||
)
|
||||
anthropic_config: Final = AnthropicConfig()
|
||||
text_content, _, thinking_blocks, reasoning_content, tool_calls, _, _, _ = (
|
||||
anthropic_config.extract_response_content(completion_response=dict(response_json))
|
||||
)
|
||||
|
||||
_stop_reason_map: Final = {
|
||||
"end_turn": "stop",
|
||||
|
|
@ -461,9 +592,13 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
}
|
||||
finish_reason: Final = _stop_reason_map.get(response_json.get("stop_reason", "end_turn"), "stop")
|
||||
|
||||
message: Final = Message(content=text_content or None, role="assistant")
|
||||
if tool_calls:
|
||||
message.tool_calls = tool_calls
|
||||
message: Final = Message(
|
||||
content=text_content or None,
|
||||
role="assistant",
|
||||
tool_calls=tool_calls or None,
|
||||
thinking_blocks=thinking_blocks,
|
||||
reasoning_content=reasoning_content,
|
||||
)
|
||||
|
||||
choice: Final = Choices(
|
||||
finish_reason=finish_reason,
|
||||
|
|
@ -471,11 +606,13 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
message=message,
|
||||
)
|
||||
|
||||
usage_data: Final = response_json.get("usage", {})
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=usage_data.get("input_tokens", 0),
|
||||
completion_tokens=usage_data.get("output_tokens", 0),
|
||||
total_tokens=usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0),
|
||||
# Cortex reports prompt-cache creation/read counts alongside input_tokens; the
|
||||
# shared calculator folds them into prompt_tokens_details so cached input is
|
||||
# visible and billed at its own rate.
|
||||
usage: Final = anthropic_config.calculate_usage(
|
||||
usage_object=response_json.get("usage", {}),
|
||||
reasoning_content=reasoning_content,
|
||||
completion_response=dict(response_json),
|
||||
)
|
||||
|
||||
model_response.choices = [choice]
|
||||
|
|
@ -516,15 +653,19 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator):
|
|||
json_mode: bool | None = False,
|
||||
):
|
||||
super().__init__(streaming_response=streaming_response, sync_stream=sync_stream)
|
||||
self._tool_index = 0
|
||||
self._tool_id = ""
|
||||
self._tool_name = ""
|
||||
self._input_tokens = 0
|
||||
# Cortex streams the Anthropic SSE dialect on /messages, so its events are parsed
|
||||
# by Anthropic's own parser: thinking deltas, signatures and prompt-cache usage
|
||||
# all arrive the way they do on every other Anthropic-dialect provider.
|
||||
self._anthropic_parser: Final = AnthropicStreamParser(
|
||||
streaming_response=streaming_response,
|
||||
sync_stream=sync_stream,
|
||||
json_mode=json_mode,
|
||||
)
|
||||
|
||||
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk:
|
||||
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream:
|
||||
if "choices" in chunk:
|
||||
return self._parse_openai_chunk(chunk)
|
||||
return self._parse_anthropic_chunk(chunk)
|
||||
return self._anthropic_parser.chunk_parser(chunk)
|
||||
|
||||
def _parse_openai_chunk(self, chunk: dict) -> GenericStreamingChunk:
|
||||
choices: Final = chunk.get("choices", [])
|
||||
|
|
@ -566,117 +707,3 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator):
|
|||
index=choice.get("index", 0),
|
||||
tool_use=tool_use,
|
||||
)
|
||||
|
||||
def _parse_anthropic_chunk(self, chunk: dict) -> GenericStreamingChunk:
|
||||
event_type: Final = chunk.get("type", "")
|
||||
|
||||
if event_type == "message_start":
|
||||
message: Final = chunk.get("message", {})
|
||||
usage_data = message.get("usage", {})
|
||||
self._input_tokens = usage_data.get("input_tokens", 0)
|
||||
return GenericStreamingChunk(
|
||||
text="",
|
||||
is_finished=False,
|
||||
finish_reason="",
|
||||
usage=None,
|
||||
index=0,
|
||||
tool_use=None,
|
||||
)
|
||||
|
||||
elif event_type == "content_block_delta":
|
||||
delta = chunk.get("delta", {})
|
||||
delta_type: Final = delta.get("type", "")
|
||||
|
||||
if delta_type == "text_delta":
|
||||
return GenericStreamingChunk(
|
||||
text=delta.get("text", ""),
|
||||
is_finished=False,
|
||||
finish_reason="",
|
||||
usage=None,
|
||||
index=chunk.get("index", 0),
|
||||
tool_use=None,
|
||||
)
|
||||
elif delta_type == "input_json_delta":
|
||||
return GenericStreamingChunk(
|
||||
text="",
|
||||
is_finished=False,
|
||||
finish_reason="",
|
||||
usage=None,
|
||||
index=chunk.get("index", 0),
|
||||
tool_use=ChatCompletionToolCallChunk(
|
||||
id=self._tool_id,
|
||||
type="function",
|
||||
function={
|
||||
"name": self._tool_name,
|
||||
"arguments": delta.get("partial_json", ""),
|
||||
},
|
||||
index=self._tool_index,
|
||||
),
|
||||
)
|
||||
|
||||
elif event_type == "content_block_start":
|
||||
content_block: Final = chunk.get("content_block", {})
|
||||
if content_block.get("type") == "tool_use":
|
||||
self._tool_id = content_block.get("id", "")
|
||||
self._tool_name = content_block.get("name", "")
|
||||
self._tool_index = chunk.get("index", 0)
|
||||
return GenericStreamingChunk(
|
||||
text="",
|
||||
is_finished=False,
|
||||
finish_reason="",
|
||||
usage=None,
|
||||
index=chunk.get("index", 0),
|
||||
tool_use=ChatCompletionToolCallChunk(
|
||||
id=self._tool_id,
|
||||
type="function",
|
||||
function={"name": self._tool_name, "arguments": ""},
|
||||
index=self._tool_index,
|
||||
),
|
||||
)
|
||||
|
||||
elif event_type == "message_delta":
|
||||
delta = chunk.get("delta", {})
|
||||
stop_reason: Final = delta.get("stop_reason", "")
|
||||
usage_data = chunk.get("usage", {})
|
||||
_stop_map: Final = {
|
||||
"end_turn": "stop",
|
||||
"max_tokens": "length",
|
||||
"tool_use": "tool_calls",
|
||||
"stop_sequence": "stop",
|
||||
}
|
||||
usage = None
|
||||
if usage_data or self._input_tokens:
|
||||
output_t: Final = usage_data.get("output_tokens", 0)
|
||||
input_t: Final = self._input_tokens or usage_data.get("input_tokens", 0)
|
||||
usage = ChatCompletionUsageBlock(
|
||||
prompt_tokens=input_t,
|
||||
completion_tokens=output_t,
|
||||
total_tokens=input_t + output_t,
|
||||
)
|
||||
return GenericStreamingChunk(
|
||||
text="",
|
||||
is_finished=True,
|
||||
finish_reason=_stop_map.get(stop_reason, "stop"),
|
||||
usage=usage,
|
||||
index=0,
|
||||
tool_use=None,
|
||||
)
|
||||
|
||||
elif event_type == "message_stop":
|
||||
return GenericStreamingChunk(
|
||||
text="",
|
||||
is_finished=True,
|
||||
finish_reason="stop",
|
||||
usage=None,
|
||||
index=0,
|
||||
tool_use=None,
|
||||
)
|
||||
|
||||
return GenericStreamingChunk(
|
||||
text="",
|
||||
is_finished=False,
|
||||
finish_reason="",
|
||||
usage=None,
|
||||
index=0,
|
||||
tool_use=None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -11,7 +11,8 @@ being registered, so a gateway with no EMA upstream never stores bearer material
|
|||
The row is one encrypted payload per user, latest login wins. ``expires_at`` mirrors the
|
||||
id_token ``exp`` claim and is judged by the reader, never enforced by deletion here: an
|
||||
expired assertion with a refresh token is still renewable, and the DB row is the source of
|
||||
truth, the same contract as the per-user OAuth credential store.
|
||||
truth, the same contract as the per-user OAuth credential store. Reads use a per-process cache with
|
||||
TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against stale in-flight reads.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -24,6 +25,8 @@ import jwt
|
|||
from pydantic import BaseModel, ConfigDict, SecretStr, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_SSO_ASSERTION_CACHE_TTL_SECONDS
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
|
@ -45,6 +48,46 @@ class SSOIdentityAssertion(BaseModel):
|
|||
expires_at: datetime | None = None
|
||||
|
||||
|
||||
class SSOAssertionCache:
|
||||
"""Process-local read cache. ``invalidate`` bumps a process-wide epoch so a fetch that started
|
||||
before a login cannot repopulate the old assertion after it."""
|
||||
|
||||
def __init__(self, ttl_seconds: int = MCP_SSO_ASSERTION_CACHE_TTL_SECONDS) -> None:
|
||||
self._entries = InMemoryCache(
|
||||
max_size_in_memory=MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE,
|
||||
default_ttl=ttl_seconds,
|
||||
)
|
||||
self._epoch: int = 0
|
||||
|
||||
def epoch(self) -> int:
|
||||
return self._epoch
|
||||
|
||||
def get(self, user_id: str) -> SSOIdentityAssertion | None:
|
||||
cached: Final = self._entries.get_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
|
||||
user_id
|
||||
)
|
||||
return cached if isinstance(cached, SSOIdentityAssertion) else None
|
||||
|
||||
def set_if_unchanged(self, user_id: str, assertion: SSOIdentityAssertion, seen_epoch: int) -> None:
|
||||
if self._epoch != seen_epoch:
|
||||
return
|
||||
self._entries.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
|
||||
user_id, assertion
|
||||
)
|
||||
|
||||
def invalidate(self, user_id: str) -> None:
|
||||
self._epoch += 1
|
||||
self._entries.delete_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
|
||||
user_id
|
||||
)
|
||||
|
||||
def flush(self) -> None:
|
||||
self._entries.flush_cache() # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
|
||||
|
||||
|
||||
_ASSERTION_CACHE: Final = SSOAssertionCache()
|
||||
|
||||
|
||||
class _IdTokenClaims(BaseModel):
|
||||
exp: float | None = None
|
||||
iss: str | None = None
|
||||
|
|
@ -107,7 +150,9 @@ async def ema_assertion_retention_enabled() -> bool:
|
|||
return row is not None
|
||||
|
||||
|
||||
async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAssertion) -> None:
|
||||
async def persist_sso_identity_assertion(
|
||||
user_id: str, assertion: SSOIdentityAssertion, cache: SSOAssertionCache = _ASSERTION_CACHE
|
||||
) -> None:
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper # noqa: PLC0415 # runtime global
|
||||
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global
|
||||
|
||||
|
|
@ -127,11 +172,10 @@ async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAss
|
|||
"update": {"assertion_b64": encoded},
|
||||
},
|
||||
)
|
||||
cache.invalidate(user_id)
|
||||
|
||||
|
||||
async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | None:
|
||||
"""The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key
|
||||
rotation), or unparseable. Expiry is not judged here; the reader owns that policy."""
|
||||
async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None:
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper # noqa: PLC0415 # runtime global
|
||||
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global
|
||||
|
||||
|
|
@ -160,6 +204,21 @@ async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | N
|
|||
)
|
||||
|
||||
|
||||
async def fetch_sso_identity_assertion(
|
||||
user_id: str, cache: SSOAssertionCache = _ASSERTION_CACHE
|
||||
) -> SSOIdentityAssertion | None:
|
||||
"""The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key
|
||||
rotation), or unparseable. Expiry is not judged here; the reader owns that policy."""
|
||||
cached: Final = cache.get(user_id)
|
||||
if cached is not None:
|
||||
return cached
|
||||
seen_epoch: Final = cache.epoch()
|
||||
assertion: Final = await _read_assertion_from_db(user_id)
|
||||
if assertion is not None:
|
||||
cache.set_if_unchanged(user_id, assertion, seen_epoch)
|
||||
return assertion
|
||||
|
||||
|
||||
class AssertionStoreUnavailable(Exception):
|
||||
"""Raised by ``fetch`` when the backing store is unreachable (e.g. the DB is down).
|
||||
|
||||
|
|
@ -189,9 +248,12 @@ class DbSSOAssertionStore:
|
|||
from credential resolution and from the upstream-401 retry.
|
||||
"""
|
||||
|
||||
def __init__(self, cache: SSOAssertionCache = _ASSERTION_CACHE) -> None:
|
||||
self._cache = cache
|
||||
|
||||
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
|
||||
try:
|
||||
return await fetch_sso_identity_assertion(user_id)
|
||||
return await fetch_sso_identity_assertion(user_id, cache=self._cache)
|
||||
except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence
|
||||
raise AssertionStoreUnavailable(str(exc)) from exc
|
||||
|
||||
|
|
|
|||
|
|
@ -2404,6 +2404,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
"""
|
||||
|
||||
completion_model: str | None = Field(None, description="proxy level default model for all chat completion calls")
|
||||
max_in_flight_requests_per_worker: int | None = Field(
|
||||
None, gt=0, description="maximum concurrent requests handled by each worker"
|
||||
)
|
||||
max_queued_requests_per_worker: int | None = Field(
|
||||
None, ge=0, description="maximum requests waiting for a worker slot"
|
||||
)
|
||||
admission_queue_timeout_seconds: float = Field(
|
||||
1.0, gt=0, description="maximum time a request waits for a worker slot"
|
||||
)
|
||||
plugins: list[PluginConfig] | None = Field(
|
||||
None, description="external services registered as embeddable UI plugins"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from fastapi.responses import JSONResponse
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.anthropic_interface.exceptions import AnthropicExceptionMapping
|
||||
from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping
|
||||
from litellm.integrations.custom_guardrail import ModifyResponseException
|
||||
from litellm.llms.anthropic.experimental_pass_through.context_management import (
|
||||
AnthropicContextManagementError,
|
||||
|
|
@ -30,6 +30,27 @@ from litellm.types.utils import TokenCountResponse
|
|||
router: Final = APIRouter()
|
||||
|
||||
|
||||
def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSONResponse:
|
||||
from litellm.proxy.proxy_server import (
|
||||
_close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does
|
||||
)
|
||||
|
||||
status_code: Final = int(exc.code) if exc.code is not None and exc.code.isdigit() else 500
|
||||
_close_dangling_otel_server_span(request, status_code, exc=exc)
|
||||
envelope: Final = AnthropicExceptionMapping.transform_to_anthropic_error(
|
||||
status_code=status_code,
|
||||
raw_message=exc.message,
|
||||
request_id=request.headers.get("x-request-id"),
|
||||
)
|
||||
if not exc.provider_specific_fields:
|
||||
return JSONResponse(status_code=status_code, content=envelope, headers=exc.headers)
|
||||
content: Final[AnthropicErrorResponse] = {
|
||||
**envelope,
|
||||
"error": {**envelope["error"], "provider_specific_fields": exc.provider_specific_fields},
|
||||
}
|
||||
return JSONResponse(status_code=status_code, content=content, headers=exc.headers)
|
||||
|
||||
|
||||
def _strip_total_tokens_from_anthropic_response(response: Any) -> None:
|
||||
"""Remove the OpenAI-flavored `usage.total_tokens` field that LiteLLM
|
||||
injects into Anthropic /v1/messages responses.
|
||||
|
|
@ -195,7 +216,7 @@ async def anthropic_response(
|
|||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e)
|
||||
|
||||
if isinstance(e, ProxyException):
|
||||
raise
|
||||
return _anthropic_error_json_response(e, request)
|
||||
|
||||
# Extract model_id from request metadata (same as success path)
|
||||
litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
|
||||
|
|
@ -216,15 +237,18 @@ async def anthropic_response(
|
|||
)
|
||||
|
||||
if isinstance(e, HTTPException):
|
||||
raise proxy_exception_from_http_exception(e, headers)
|
||||
return _anthropic_error_json_response(proxy_exception_from_http_exception(e, headers), request)
|
||||
|
||||
error_msg: Final = f"{e}"
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
headers=headers,
|
||||
return _anthropic_error_json_response(
|
||||
ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
headers=headers,
|
||||
),
|
||||
request,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -50,6 +50,9 @@ from litellm.proxy.health_check import (
|
|||
perform_health_check,
|
||||
run_with_timeout,
|
||||
)
|
||||
from litellm.proxy.middleware.admission_control_middleware import (
|
||||
get_admission_control_stats,
|
||||
)
|
||||
from litellm.proxy.middleware.in_flight_requests_middleware import (
|
||||
get_in_flight_requests,
|
||||
)
|
||||
|
|
@ -63,6 +66,13 @@ from litellm.secret_managers.main import get_secret_bool
|
|||
#### Health ENDPOINTS ####
|
||||
|
||||
|
||||
class _HealthBacklogResponse(TypedDict):
|
||||
in_flight_requests: ReadOnly[int]
|
||||
admitted_requests: ReadOnly[int]
|
||||
queued_requests: ReadOnly[int]
|
||||
rejected_requests: ReadOnly[int]
|
||||
|
||||
|
||||
def _reject_os_environ_references(params: dict) -> None:
|
||||
"""
|
||||
Validate that the provided params do not contain any ``os.environ/``
|
||||
|
|
@ -1759,7 +1769,14 @@ async def health_backlog():
|
|||
for the event loop to get to them, adding latency before LiteLLM even starts
|
||||
its own timer.
|
||||
"""
|
||||
return {"in_flight_requests": get_in_flight_requests()}
|
||||
stats: Final = get_admission_control_stats()
|
||||
response: Final[_HealthBacklogResponse] = {
|
||||
"in_flight_requests": get_in_flight_requests(),
|
||||
"admitted_requests": stats.admitted,
|
||||
"queued_requests": stats.queued,
|
||||
"rejected_requests": stats.rejected_total,
|
||||
}
|
||||
return response
|
||||
|
||||
|
||||
@router.get(
|
||||
|
|
|
|||
315
litellm/proxy/middleware/admission_control_middleware.py
Normal file
315
litellm/proxy/middleware/admission_control_middleware.py
Normal file
|
|
@ -0,0 +1,315 @@
|
|||
import asyncio
|
||||
import os
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from typing import Annotated, Final, Protocol, TypeAlias, runtime_checkable
|
||||
|
||||
from pydantic import Field, TypeAdapter, ValidationError
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
_EXEMPT_PATHS: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
"/health/liveliness",
|
||||
"/health/liveness",
|
||||
"/health/readiness",
|
||||
"/health/readiness/details",
|
||||
"/health/backlog",
|
||||
"/health/drain",
|
||||
"/metrics",
|
||||
"/metrics/",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AdmissionControlSettings:
|
||||
max_in_flight_requests: int
|
||||
max_queued_requests: int
|
||||
queue_timeout_seconds: float
|
||||
|
||||
|
||||
AdmissionControlSettingsGetter: TypeAlias = Callable[[], AdmissionControlSettings | None] # mutable-ok: Callable params
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AdmissionControlStats:
|
||||
admitted: int
|
||||
queued: int
|
||||
rejected_total: int
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _Gauge(Protocol):
|
||||
def inc(self, amount: float = 1) -> None: ...
|
||||
|
||||
def dec(self, amount: float = 1) -> None: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _CounterChild(Protocol):
|
||||
def inc(self, amount: float = 1) -> None: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _Counter(Protocol):
|
||||
def labels(self, reason: str) -> _CounterChild: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AdmissionControlMetrics:
|
||||
admitted_gauge: _Gauge
|
||||
queued_gauge: _Gauge
|
||||
rejected_counter: _Counter
|
||||
|
||||
|
||||
AdmissionControlMetricsFactory: TypeAlias = Callable[[], AdmissionControlMetrics | None] # mutable-ok: Callable params
|
||||
|
||||
|
||||
class AdmissionControlState:
|
||||
"""Per-process admission counters and the in-flight semaphore shared by one worker's requests."""
|
||||
|
||||
def __init__(self, metrics_factory: AdmissionControlMetricsFactory) -> None:
|
||||
self._metrics_factory = metrics_factory
|
||||
self._metrics: AdmissionControlMetrics | None = None
|
||||
self._metrics_init_attempted = False
|
||||
self._admitted = 0
|
||||
self._queued = 0
|
||||
self._rejected_total = 0
|
||||
self._semaphore: asyncio.Semaphore | None = None
|
||||
self._semaphore_loop: asyncio.AbstractEventLoop | None = None
|
||||
|
||||
def get_stats(self) -> AdmissionControlStats:
|
||||
return AdmissionControlStats(
|
||||
admitted=self._admitted,
|
||||
queued=self._queued,
|
||||
rejected_total=self._rejected_total,
|
||||
)
|
||||
|
||||
def get_semaphore(self, max_in_flight_requests: int) -> asyncio.Semaphore:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
if self._semaphore_loop is not loop:
|
||||
self._semaphore = asyncio.Semaphore(max_in_flight_requests)
|
||||
self._semaphore_loop = loop
|
||||
semaphore: Final = self._semaphore
|
||||
if semaphore is None:
|
||||
raise RuntimeError("Admission control semaphore was not initialized")
|
||||
return semaphore
|
||||
|
||||
def record_admission(self) -> None:
|
||||
self._admitted += 1
|
||||
metrics: Final = self._get_metrics()
|
||||
if metrics is not None:
|
||||
metrics.admitted_gauge.inc()
|
||||
|
||||
def record_release(self) -> None:
|
||||
self._admitted -= 1
|
||||
metrics: Final = self._get_metrics()
|
||||
if metrics is not None:
|
||||
metrics.admitted_gauge.dec()
|
||||
|
||||
def record_queue(self) -> None:
|
||||
self._queued += 1
|
||||
metrics: Final = self._get_metrics()
|
||||
if metrics is not None:
|
||||
metrics.queued_gauge.inc()
|
||||
|
||||
def record_dequeue(self) -> None:
|
||||
self._queued -= 1
|
||||
metrics: Final = self._get_metrics()
|
||||
if metrics is not None:
|
||||
metrics.queued_gauge.dec()
|
||||
|
||||
def record_rejection(self, reason: str) -> None:
|
||||
self._rejected_total += 1
|
||||
metrics: Final = self._get_metrics()
|
||||
if metrics is not None:
|
||||
metrics.rejected_counter.labels(reason=reason).inc()
|
||||
|
||||
def _get_metrics(self) -> AdmissionControlMetrics | None:
|
||||
if not self._metrics_init_attempted:
|
||||
self._metrics_init_attempted = True
|
||||
self._metrics = self._metrics_factory()
|
||||
return self._metrics
|
||||
|
||||
|
||||
class AdmissionControlMiddleware:
|
||||
def __init__(
|
||||
self,
|
||||
app: ASGIApp,
|
||||
get_settings: AdmissionControlSettingsGetter,
|
||||
state: AdmissionControlState,
|
||||
) -> None:
|
||||
self.app = app
|
||||
self.get_settings = get_settings
|
||||
self.state = state
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
settings: Final = self.get_settings()
|
||||
if settings is None or _get_route_path(scope) in _EXEMPT_PATHS:
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
state: Final = self.state
|
||||
semaphore: Final = state.get_semaphore(settings.max_in_flight_requests)
|
||||
if not semaphore.locked():
|
||||
await semaphore.acquire()
|
||||
state.record_admission()
|
||||
elif state.get_stats().queued >= settings.max_queued_requests:
|
||||
state.record_rejection("queue_full")
|
||||
await _overloaded_response(state)(scope, receive, send)
|
||||
return
|
||||
else:
|
||||
state.record_queue()
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
semaphore.acquire(),
|
||||
timeout=settings.queue_timeout_seconds,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
state.record_dequeue()
|
||||
state.record_rejection("queue_timeout")
|
||||
await _overloaded_response(state)(scope, receive, send)
|
||||
return
|
||||
except asyncio.CancelledError:
|
||||
state.record_dequeue()
|
||||
raise
|
||||
state.record_dequeue()
|
||||
state.record_admission()
|
||||
|
||||
try:
|
||||
await self.app(scope, receive, send)
|
||||
finally:
|
||||
semaphore.release()
|
||||
state.record_release()
|
||||
|
||||
|
||||
def _get_route_path(scope: Scope) -> str:
|
||||
"""Strip the ASGI root_path (SERVER_ROOT_PATH) the same way Starlette does before route matching."""
|
||||
path: Final[str] = scope["path"]
|
||||
root_path: Final[str] = scope.get("root_path", "")
|
||||
if not root_path or not path.startswith(root_path):
|
||||
return path
|
||||
if path == root_path:
|
||||
return ""
|
||||
if path[len(root_path)] == "/":
|
||||
return path[len(root_path) :]
|
||||
return path
|
||||
|
||||
|
||||
def _create_gauge(gauge_type: Callable[..., object], name: str, description: str) -> _Gauge:
|
||||
metric: Final = (
|
||||
gauge_type(name, description, multiprocess_mode="livesum")
|
||||
if "PROMETHEUS_MULTIPROC_DIR" in os.environ
|
||||
else gauge_type(name, description)
|
||||
)
|
||||
if not isinstance(metric, _Gauge):
|
||||
raise TypeError("Admission gauge has an unexpected type")
|
||||
return metric
|
||||
|
||||
|
||||
def create_prometheus_admission_metrics() -> AdmissionControlMetrics | None:
|
||||
try:
|
||||
from prometheus_client import Counter, Gauge
|
||||
|
||||
return AdmissionControlMetrics(
|
||||
admitted_gauge=_create_gauge(
|
||||
Gauge,
|
||||
"litellm_admission_admitted_requests",
|
||||
"Number of requests admitted by this worker",
|
||||
),
|
||||
queued_gauge=_create_gauge(
|
||||
Gauge,
|
||||
"litellm_admission_queued_requests",
|
||||
"Number of requests queued by this worker",
|
||||
),
|
||||
rejected_counter=Counter( # mutable-ok: Prometheus requires runtime Counter construction
|
||||
"litellm_admission_rejected_requests_total",
|
||||
"Number of requests rejected by this worker",
|
||||
labelnames=("reason",),
|
||||
),
|
||||
)
|
||||
except (ImportError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
admission_control_state: Final = AdmissionControlState(create_prometheus_admission_metrics)
|
||||
|
||||
|
||||
def get_admission_control_stats() -> AdmissionControlStats:
|
||||
return admission_control_state.get_stats()
|
||||
|
||||
|
||||
_PositiveInt: TypeAlias = Annotated[int, Field(gt=0)]
|
||||
_NonNegativeInt: TypeAlias = Annotated[int, Field(ge=0)]
|
||||
_PositiveFloat: TypeAlias = Annotated[float, Field(gt=0)]
|
||||
_AdmissionControlRaw: TypeAlias = int | float | str | None
|
||||
|
||||
|
||||
def _hashable(value: object) -> _AdmissionControlRaw:
|
||||
return value if value is None or isinstance(value, (int, float, str)) else repr(value)
|
||||
|
||||
|
||||
_POSITIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_PositiveInt)
|
||||
_NON_NEGATIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_NonNegativeInt)
|
||||
_POSITIVE_FLOAT_ADAPTER: Final[TypeAdapter[float]] = TypeAdapter(_PositiveFloat)
|
||||
|
||||
|
||||
@lru_cache(maxsize=16)
|
||||
def _parse_admission_control_settings(
|
||||
max_in_flight_raw: _AdmissionControlRaw,
|
||||
max_queued_raw: _AdmissionControlRaw,
|
||||
queue_timeout_raw: _AdmissionControlRaw,
|
||||
) -> AdmissionControlSettings | None:
|
||||
try:
|
||||
max_in_flight: Final = _POSITIVE_INT_ADAPTER.validate_python(max_in_flight_raw)
|
||||
max_queued: Final = (
|
||||
max_in_flight if max_queued_raw is None else _NON_NEGATIVE_INT_ADAPTER.validate_python(max_queued_raw)
|
||||
)
|
||||
queue_timeout: Final = _POSITIVE_FLOAT_ADAPTER.validate_python(queue_timeout_raw)
|
||||
except ValidationError as exc:
|
||||
verbose_proxy_logger.error(
|
||||
"Ignoring invalid admission control settings, per-worker admission control is disabled: %s",
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
return AdmissionControlSettings(
|
||||
max_in_flight_requests=max_in_flight,
|
||||
max_queued_requests=max_queued,
|
||||
queue_timeout_seconds=queue_timeout,
|
||||
)
|
||||
|
||||
|
||||
def get_admission_control_settings(settings: Mapping[str, object]) -> AdmissionControlSettings | None:
|
||||
max_in_flight_raw: Final = settings.get("max_in_flight_requests_per_worker")
|
||||
if max_in_flight_raw is None:
|
||||
return None
|
||||
return _parse_admission_control_settings(
|
||||
_hashable(max_in_flight_raw),
|
||||
_hashable(settings.get("max_queued_requests_per_worker")),
|
||||
_hashable(settings.get("admission_queue_timeout_seconds", 1.0)),
|
||||
)
|
||||
|
||||
|
||||
def _overloaded_response(state: AdmissionControlState) -> JSONResponse:
|
||||
stats: Final = state.get_stats()
|
||||
return JSONResponse(
|
||||
status_code=503,
|
||||
headers={"retry-after": "1"}, # mutable-ok: Starlette expects a plain headers mapping
|
||||
content={ # mutable-ok: Starlette serializes a plain response mapping
|
||||
"error": { # mutable-ok: nested response mapping
|
||||
"message": (
|
||||
f"Worker at capacity: {stats.admitted} in-flight, {stats.queued} queued requests. Retry later."
|
||||
),
|
||||
"type": "overloaded_error",
|
||||
"code": "503",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
|
@ -1659,6 +1659,12 @@ async def azure_proxy_route(
|
|||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
_VERTEX_LOCATION_REQUIRED_DETAIL: Final = (
|
||||
"No Vertex AI location for this request. Include /projects/<project>/locations/<location>/ in the "
|
||||
"route, set vertex_location in default_vertex_config (or DEFAULT_VERTEXAI_LOCATION), or add the "
|
||||
"model to model_list with use_in_pass_through: true."
|
||||
)
|
||||
|
||||
|
||||
class BaseVertexAIPassThroughHandler(ABC):
|
||||
@staticmethod
|
||||
|
|
@ -1666,29 +1672,18 @@ class BaseVertexAIPassThroughHandler(ABC):
|
|||
def get_default_base_target_url(vertex_location: str | None) -> str:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
|
||||
pass
|
||||
|
||||
|
||||
class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler):
|
||||
@staticmethod
|
||||
def get_default_base_target_url(vertex_location: str | None) -> str:
|
||||
return "https://discoveryengine.googleapis.com/"
|
||||
|
||||
@staticmethod
|
||||
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
|
||||
return base_target_url
|
||||
|
||||
|
||||
class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler):
|
||||
@staticmethod
|
||||
def get_default_base_target_url(vertex_location: str | None) -> str:
|
||||
return get_vertex_base_url(vertex_location)
|
||||
|
||||
@staticmethod
|
||||
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
|
||||
if vertex_location is None:
|
||||
raise HTTPException(status_code=400, detail=_VERTEX_LOCATION_REQUIRED_DETAIL)
|
||||
return get_vertex_base_url(vertex_location)
|
||||
|
||||
|
||||
|
|
@ -1911,10 +1906,8 @@ async def _prepare_vertex_auth_headers(
|
|||
router_credentials: LiteLLM_ManagedVectorStore | None,
|
||||
vertex_project: str | None,
|
||||
vertex_location: str | None,
|
||||
base_target_url: str | None,
|
||||
get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> tuple[Mapping[str, str], str | None, bool, str | None, str | None]:
|
||||
) -> tuple[Mapping[str, str], bool, str | None, str | None]:
|
||||
"""
|
||||
Prepare authentication headers for Vertex AI pass-through requests.
|
||||
|
||||
|
|
@ -1924,15 +1917,12 @@ async def _prepare_vertex_auth_headers(
|
|||
router_credentials: Optional vector store credentials from registry
|
||||
vertex_project: Vertex project ID
|
||||
vertex_location: Vertex location
|
||||
base_target_url: Base URL for the Vertex AI service
|
||||
get_vertex_pass_through_handler: Handler for the specific Vertex AI service
|
||||
user_api_key_dict: The caller's resolved authentication, so only the secret that
|
||||
authenticated them is stripped on the credential-less branch
|
||||
|
||||
Returns:
|
||||
tuple containing:
|
||||
- headers: dict - Authentication headers to use
|
||||
- base_target_url: str | None - Updated base target URL
|
||||
- headers_passed_through: bool - Whether headers were passed through from request
|
||||
- vertex_project: str | None - Updated vertex project ID
|
||||
- vertex_location: str | None - Updated vertex location
|
||||
|
|
@ -1985,14 +1975,8 @@ async def _prepare_vertex_auth_headers(
|
|||
# Add the Authorization header with vendor credentials
|
||||
headers["Authorization"] = f"Bearer {auth_header}"
|
||||
|
||||
if base_target_url is not None:
|
||||
base_target_url = get_vertex_pass_through_handler.update_base_target_url_with_credential_location(
|
||||
base_target_url, vertex_location
|
||||
)
|
||||
|
||||
return (
|
||||
headers,
|
||||
base_target_url,
|
||||
headers_passed_through,
|
||||
vertex_project,
|
||||
vertex_location,
|
||||
|
|
@ -2085,12 +2069,9 @@ async def _base_vertex_proxy_route(
|
|||
location=vertex_location,
|
||||
)
|
||||
|
||||
base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
|
||||
|
||||
# Prepare authentication headers
|
||||
(
|
||||
headers,
|
||||
base_target_url,
|
||||
headers_passed_through,
|
||||
vertex_project,
|
||||
vertex_location,
|
||||
|
|
@ -2100,13 +2081,10 @@ async def _base_vertex_proxy_route(
|
|||
router_credentials=router_credentials,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
base_target_url=base_target_url,
|
||||
get_vertex_pass_through_handler=get_vertex_pass_through_handler,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
if base_target_url is None:
|
||||
base_target_url = get_vertex_base_url(vertex_location)
|
||||
base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
|
||||
|
||||
request_route: Final = encoded_endpoint
|
||||
verbose_proxy_logger.debug("request_route %s", request_route)
|
||||
|
|
|
|||
|
|
@ -583,6 +583,11 @@ try:
|
|||
except ImportError:
|
||||
build_billing_metrics_recorder = None
|
||||
shutdown_billing_metrics_recorder = None
|
||||
from litellm.proxy.middleware.admission_control_middleware import (
|
||||
AdmissionControlMiddleware,
|
||||
admission_control_state,
|
||||
get_admission_control_settings,
|
||||
)
|
||||
from litellm.proxy.middleware.in_flight_requests_middleware import (
|
||||
InFlightRequestsMiddleware,
|
||||
)
|
||||
|
|
@ -16502,6 +16507,9 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
|
|||
{
|
||||
"max_parallel_requests": "Integer",
|
||||
"global_max_parallel_requests": "Integer",
|
||||
"max_in_flight_requests_per_worker": "Integer",
|
||||
"max_queued_requests_per_worker": "Integer",
|
||||
"admission_queue_timeout_seconds": "Float",
|
||||
"max_request_size_mb": "Integer",
|
||||
"max_batch_file_size_mb": "Integer",
|
||||
"max_file_size_mb": "Integer",
|
||||
|
|
@ -18177,6 +18185,11 @@ app.add_middleware(
|
|||
get_max_request_size_mb=lambda: general_settings.get("max_request_size_mb"),
|
||||
is_request_size_limit_enabled=lambda: premium_user is True,
|
||||
)
|
||||
app.add_middleware(
|
||||
AdmissionControlMiddleware,
|
||||
get_settings=lambda: get_admission_control_settings(general_settings),
|
||||
state=admission_control_state,
|
||||
)
|
||||
|
||||
|
||||
async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse":
|
||||
|
|
|
|||
|
|
@ -3,7 +3,8 @@ import collections
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
|
|
@ -16,7 +17,6 @@ from typing import (
|
|||
TypeAlias,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
cast, # noqa: TID251 # prisma group_by returns untyped aggregate mappings
|
||||
)
|
||||
|
||||
import fastapi
|
||||
|
|
@ -201,16 +201,12 @@ class _SessionSpendStats(NamedTuple):
|
|||
_SessionSpendMap: TypeAlias = Mapping[tuple[str, str], _SessionSpendStats]
|
||||
|
||||
|
||||
class _SpendSumAggregate(TypedDict, total=False):
|
||||
spend: ReadOnly[float]
|
||||
|
||||
|
||||
class _SpendGroupByRow(TypedDict):
|
||||
class _SpendDailySummaryRow(TypedDict):
|
||||
day: ReadOnly[str]
|
||||
api_key: ReadOnly[str]
|
||||
user: ReadOnly[str | None]
|
||||
model: ReadOnly[str]
|
||||
startTime: ReadOnly[object]
|
||||
_sum: ReadOnly[_SpendSumAggregate]
|
||||
spend: ReadOnly[float]
|
||||
|
||||
|
||||
async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]:
|
||||
|
|
@ -251,6 +247,66 @@ def _verification_token_table(prisma_client: PrismaClient) -> _VerificationToken
|
|||
return VerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _spend_logs_daily_summary_sql(
|
||||
*,
|
||||
start_date_iso: str,
|
||||
end_date_iso: str,
|
||||
api_key: str | None,
|
||||
request_id: str | None,
|
||||
user_id: str | None,
|
||||
) -> tuple[str, tuple[object, ...]]:
|
||||
filter_params: Final[tuple[tuple[str, object], ...]] = tuple(
|
||||
(column, value)
|
||||
for column, value in (
|
||||
("api_key", api_key),
|
||||
("request_id", request_id),
|
||||
('"user"', user_id),
|
||||
)
|
||||
if value is not None
|
||||
)
|
||||
filter_clauses: Final[tuple[str, ...]] = tuple(
|
||||
f"AND {column} = ${index}" for index, (column, _) in enumerate(filter_params, start=3)
|
||||
)
|
||||
filter_sql: Final = "\n".join(filter_clauses)
|
||||
sql_query: Final = f"""
|
||||
SELECT
|
||||
to_char(date_trunc('day', "startTime"), 'YYYY-MM-DD') AS day,
|
||||
api_key,
|
||||
"user",
|
||||
model,
|
||||
SUM(spend) AS spend
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') AND "startTime" <= ($2::timestamptz AT TIME ZONE 'UTC')
|
||||
{filter_sql}
|
||||
GROUP BY 1, 2, 3, 4
|
||||
ORDER BY 1
|
||||
"""
|
||||
params: Final[tuple[object, ...]] = (
|
||||
start_date_iso,
|
||||
end_date_iso,
|
||||
*(value for _, value in filter_params),
|
||||
)
|
||||
return sql_query, params
|
||||
|
||||
|
||||
def _sum_spend_by(
|
||||
rows: Sequence[_SpendDailySummaryRow], column: Literal["api_key", "user", "model"]
|
||||
) -> Mapping[str | None, float]:
|
||||
keys: Final = frozenset(row[column] for row in rows)
|
||||
return {key: sum(float(row["spend"]) for row in rows if row[column] == key) for key in keys}
|
||||
|
||||
|
||||
def _daily_summary_item(summary_date: date, rows: Sequence[_SpendDailySummaryRow]) -> Mapping[str, object]:
|
||||
api_key_spend: Final = {key: value for key, value in _sum_spend_by(rows, "api_key").items() if key is not None}
|
||||
return {
|
||||
**api_key_spend,
|
||||
"startTime": summary_date,
|
||||
"spend": sum(float(row["spend"]) for row in rows),
|
||||
"users": _sum_spend_by(rows, "user"),
|
||||
"models": _sum_spend_by(rows, "model"),
|
||||
}
|
||||
|
||||
|
||||
async def _find_spend_logs(
|
||||
prisma_client: PrismaClient,
|
||||
where: Mapping[str, object],
|
||||
|
|
@ -3266,18 +3322,22 @@ async def view_spend_logs(
|
|||
start_date_iso: Final = start_date_obj.isoformat()
|
||||
end_date_iso: Final = end_date_obj.isoformat()
|
||||
|
||||
filter_query: Final = {
|
||||
filter_query: Final[
|
||||
dict[str, object]
|
||||
] = { # mutable-ok: legacy filters are extended for optional parameters
|
||||
"startTime": {
|
||||
"gte": start_date_iso, # Greater than or equal to Start Date
|
||||
"lte": end_date_iso, # Less than or equal to End Date
|
||||
}
|
||||
}
|
||||
|
||||
summary_api_key: Final[str | None] = (
|
||||
prisma_client.hash_token(token=api_key)
|
||||
if api_key is not None and api_key.startswith("sk-")
|
||||
else api_key
|
||||
)
|
||||
if api_key is not None and isinstance(api_key, str):
|
||||
if api_key.startswith("sk-"):
|
||||
filter_query["api_key"] = prisma_client.hash_token(token=api_key)
|
||||
else:
|
||||
filter_query["api_key"] = api_key
|
||||
filter_query["api_key"] = summary_api_key
|
||||
if request_id is not None and isinstance(request_id, str):
|
||||
filter_query["request_id"] = request_id
|
||||
if user_id is not None and isinstance(user_id, str):
|
||||
|
|
@ -3296,58 +3356,34 @@ async def view_spend_logs(
|
|||
return data
|
||||
|
||||
# Legacy behavior: return summarized data (when summarize=true)
|
||||
# SQL query
|
||||
response: Final = await SpendLogsRepository(prisma_client).table.group_by(
|
||||
by=["api_key", "user", "model", "startTime"],
|
||||
where=filter_query,
|
||||
sum={
|
||||
"spend": True,
|
||||
},
|
||||
summary_sql_and_params: Final = _spend_logs_daily_summary_sql(
|
||||
start_date_iso=start_date_iso,
|
||||
end_date_iso=end_date_iso,
|
||||
api_key=summary_api_key,
|
||||
request_id=request_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
sql_query, params = summary_sql_and_params
|
||||
rows: Final[Sequence[_SpendDailySummaryRow]] = await _query_raw(prisma_client, sql_query, *params)
|
||||
if len(rows) == 0:
|
||||
return [] # pyright: ignore[reportUnknownVariableType] # empty summary has no element type
|
||||
|
||||
if isinstance(response, list) and len(response) > 0 and isinstance(response[0], dict):
|
||||
spend_rows: Final = cast(Sequence[_SpendGroupByRow], response) # cast-ok: by/sum fix the shape
|
||||
result: Final[dict] = {}
|
||||
for record in spend_rows:
|
||||
dt_object = datetime.strptime(str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ")
|
||||
date = dt_object.date()
|
||||
if date not in result:
|
||||
result[date] = {"users": {}, "models": {}}
|
||||
api_key = record["api_key"]
|
||||
user_id = record["user"]
|
||||
model = record["model"]
|
||||
result[date]["spend"] = result[date].get("spend", 0) + record.get("_sum", {}).get("spend", 0)
|
||||
result[date][api_key] = result[date].get(api_key, 0) + record.get("_sum", {}).get("spend", 0)
|
||||
result[date]["users"][user_id] = result[date]["users"].get(user_id, 0) + record.get("_sum", {}).get(
|
||||
"spend", 0
|
||||
)
|
||||
result[date]["models"][model] = result[date]["models"].get(model, 0) + record.get("_sum", {}).get(
|
||||
"spend", 0
|
||||
)
|
||||
return_list: Final = []
|
||||
final_date = None
|
||||
for k, v in sorted(result.items()):
|
||||
return_list.append({**v, "startTime": k})
|
||||
final_date = k
|
||||
|
||||
end_date_date: Final = end_date_obj.date()
|
||||
if final_date is not None and final_date < end_date_date:
|
||||
current_date = final_date + timedelta(days=1)
|
||||
while current_date <= end_date_date:
|
||||
# Represent current_date as string because original response has it this way
|
||||
return_list.append(
|
||||
{
|
||||
"startTime": current_date,
|
||||
"spend": 0,
|
||||
"users": {},
|
||||
"models": {},
|
||||
}
|
||||
) # If no data, will stay as zero
|
||||
current_date += timedelta(days=1) # Move on to the next day
|
||||
|
||||
return return_list
|
||||
|
||||
return response
|
||||
summary_items: Final = tuple(
|
||||
_daily_summary_item(date.fromisoformat(day), tuple(day_rows))
|
||||
for day, day_rows in groupby(rows, key=lambda row: row["day"])
|
||||
)
|
||||
final_date: Final = date.fromisoformat(rows[-1]["day"])
|
||||
end_date_date: Final = end_date_obj.date()
|
||||
padding: Final[tuple[Mapping[str, object], ...]] = tuple(
|
||||
{
|
||||
"startTime": final_date + timedelta(days=offset),
|
||||
"spend": 0,
|
||||
"users": {},
|
||||
"models": {},
|
||||
}
|
||||
for offset in range(1, (end_date_date - final_date).days + 1)
|
||||
)
|
||||
return [*summary_items, *padding]
|
||||
|
||||
else:
|
||||
scoped_filter: Final[dict[str, str]] = {}
|
||||
|
|
|
|||
|
|
@ -9546,6 +9546,7 @@ class Router:
|
|||
public_model_name for _, public_model_name in self.team_model_to_deployment_indices
|
||||
)
|
||||
|
||||
self.pattern_router.remove_deployment(model_id)
|
||||
for team_id in list(self.team_pattern_routers.keys()):
|
||||
team_pattern_router = self.team_pattern_routers[team_id]
|
||||
team_pattern_router.remove_deployment(model_id)
|
||||
|
|
|
|||
|
|
@ -190,6 +190,9 @@ model_list:
|
|||
|
||||
# Let that replacement also override a kept session pin, for image turns only (default: false)
|
||||
modality_pin_override: true
|
||||
|
||||
# Refreshes on every pin reuse, so this is idle time rather than total session length (default: 3600)
|
||||
session_affinity_ttl_seconds: 300
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
|
@ -240,6 +243,10 @@ affinity write happens upstream of the gate and stores the session's own model,
|
|||
turn replays the original pin and the override is never pinned in its place. It does nothing
|
||||
unless `modality_routing` is also on.
|
||||
|
||||
### Session pin retention
|
||||
|
||||
`session_affinity_ttl_seconds` is the idle window for both the model pin selected by session affinity and the deployment pin. Every request that reuses a pin refreshes its TTL, so a session actively sending requests stays pinned. After the window passes with no pin reuse, the next request classifies again and creates a fresh pin. Omit the setting to track the default of 3600 seconds.
|
||||
|
||||
### Heuristic-first chaining
|
||||
|
||||
`classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM
|
||||
|
|
|
|||
|
|
@ -96,7 +96,7 @@
|
|||
"limit": 10
|
||||
},
|
||||
"DTZ007": {
|
||||
"limit": 17
|
||||
"limit": 6
|
||||
},
|
||||
"DTZ011": {
|
||||
"limit": 3
|
||||
|
|
|
|||
218
tests/e2e/ui/tests/tables/tableScrolling.spec.ts
Normal file
218
tests/e2e/ui/tests/tables/tableScrolling.spec.ts
Normal file
|
|
@ -0,0 +1,218 @@
|
|||
import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test";
|
||||
import { ADMIN_STORAGE_PATH } from "../../constants";
|
||||
import { Page } from "../../fixtures/pages";
|
||||
import { navigateToPage } from "../../helpers/navigation";
|
||||
import { CHAT_MODEL_A, masterKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic";
|
||||
|
||||
const VIEWPORT = { width: 1280, height: 720 };
|
||||
const SEED_ROWS = 40;
|
||||
const LOG_ROWS = 20;
|
||||
const BODY_SCROLL_PX = 500;
|
||||
const MAX_FOOTER_GAP_PX = 40;
|
||||
|
||||
interface GeneratedKey {
|
||||
key: string;
|
||||
}
|
||||
|
||||
interface CreatedTeam {
|
||||
team_id: string;
|
||||
}
|
||||
|
||||
interface CreatedModel {
|
||||
model_info: { id: string };
|
||||
}
|
||||
|
||||
interface BoxMetrics {
|
||||
top: number;
|
||||
bottom: number;
|
||||
scrollHeight: number;
|
||||
clientHeight: number;
|
||||
scrollWidth: number;
|
||||
clientWidth: number;
|
||||
}
|
||||
|
||||
const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
|
||||
|
||||
const adminHeaders = (): Record<string, string> => ({ Authorization: `Bearer ${masterKey()}` });
|
||||
|
||||
const appShellMain = (page: PlaywrightPage): Locator => page.locator("main").first();
|
||||
|
||||
const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true });
|
||||
|
||||
const visibleDataTable = (page: PlaywrightPage): Locator => visibleTestId(page, "data-table-root").first();
|
||||
|
||||
const visibleRows = (page: PlaywrightPage): Locator => visibleDataTable(page).locator("tbody tr");
|
||||
|
||||
const metrics = (locator: Locator): Promise<BoxMetrics> =>
|
||||
locator.evaluate((el) => {
|
||||
const rect = el.getBoundingClientRect();
|
||||
return {
|
||||
top: rect.top,
|
||||
bottom: rect.bottom,
|
||||
scrollHeight: el.scrollHeight,
|
||||
clientHeight: el.clientHeight,
|
||||
scrollWidth: el.scrollWidth,
|
||||
clientWidth: el.clientWidth,
|
||||
};
|
||||
});
|
||||
|
||||
async function postOk<T>(request: APIRequestContext, path: string, data: Record<string, unknown>): Promise<T> {
|
||||
const res = await request.post(path, { headers: adminHeaders(), data });
|
||||
expect(res.ok(), `POST ${path} failed (${res.status()}): ${await res.text()}`).toBe(true);
|
||||
return (await res.json()) as T;
|
||||
}
|
||||
|
||||
const oneAtATime = <T>(count: number, call: (index: number) => Promise<T>): Promise<readonly T[]> =>
|
||||
Array.from({ length: count }, (_, i) => i).reduce<Promise<readonly T[]>>(
|
||||
async (previous, i) => [...(await previous), await call(i)],
|
||||
Promise.resolve([]),
|
||||
);
|
||||
|
||||
async function expectRowsAtLeast(page: PlaywrightPage, count: number): Promise<void> {
|
||||
await expect.poll(() => visibleRows(page).count(), { timeout: 30_000 }).toBeGreaterThanOrEqual(count);
|
||||
}
|
||||
|
||||
async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): Promise<void> {
|
||||
await visibleTestId(page, "pagination-page-size").click();
|
||||
await page.getByRole("option", { name: size, exact: true }).click();
|
||||
}
|
||||
|
||||
async function expectBodyIsTheOnlyScroller(page: PlaywrightPage): Promise<void> {
|
||||
const scroller = await metrics(appShellMain(page));
|
||||
const body = visibleTestId(page, "data-table-scroller");
|
||||
const bodyBefore = await metrics(body);
|
||||
const headBefore = await metrics(visibleTestId(page, "data-table-head"));
|
||||
const footer = await metrics(visibleDataTable(page));
|
||||
|
||||
expect(scroller.scrollHeight, "page scroller must not overflow vertically").toBe(scroller.clientHeight);
|
||||
expect(scroller.scrollWidth, "page scroller must not overflow horizontally").toBe(scroller.clientWidth);
|
||||
expect(bodyBefore.scrollHeight, "table body must be the element that scrolls").toBeGreaterThan(
|
||||
bodyBefore.clientHeight,
|
||||
);
|
||||
expect(footer.bottom, "pagination footer must be inside the page").toBeLessThanOrEqual(scroller.bottom);
|
||||
expect(scroller.bottom - footer.bottom, "pagination footer must sit at the bottom of the page").toBeLessThanOrEqual(
|
||||
MAX_FOOTER_GAP_PX,
|
||||
);
|
||||
|
||||
await body.evaluate((el, px) => {
|
||||
el.scrollTop = px;
|
||||
}, BODY_SCROLL_PX);
|
||||
await expect.poll(() => body.evaluate((el) => el.scrollTop)).toBeGreaterThan(0);
|
||||
const headAfter = await metrics(visibleTestId(page, "data-table-head"));
|
||||
expect(Math.round(headAfter.top), "header must stay put while the body scrolls").toBe(Math.round(headBefore.top));
|
||||
}
|
||||
|
||||
const rowsPaintingPastAnAncestor = (page: PlaywrightPage): Promise<string[]> =>
|
||||
visibleDataTable(page)
|
||||
.locator("table")
|
||||
.evaluate((table) => {
|
||||
const scrollsVertically = (el: Element): boolean =>
|
||||
/auto|scroll/.test(getComputedStyle(el).overflowY) && el.scrollHeight > el.clientHeight + 1;
|
||||
const boxesUpToTheScroller = (el: Element | null): Element[] =>
|
||||
el === null || el === document.body || scrollsVertically(el)
|
||||
? []
|
||||
: [el, ...boxesUpToTheScroller(el.parentElement)];
|
||||
const describe = (el: Element): string =>
|
||||
`<${el.tagName.toLowerCase()} class="${el.getAttribute("class") ?? ""}">`;
|
||||
return Array.from(table.querySelectorAll("tbody tr")).flatMap((row, index) => {
|
||||
const rowBottom = row.getBoundingClientRect().bottom;
|
||||
return boxesUpToTheScroller(row.parentElement)
|
||||
.filter((box) => rowBottom > box.getBoundingClientRect().bottom + 1)
|
||||
.map(
|
||||
(box) =>
|
||||
`row ${index} bottom ${Math.round(rowBottom)} past ${describe(box)} bottom ${Math.round(box.getBoundingClientRect().bottom)}`,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
test.describe("Admin tables scroll inside the page", () => {
|
||||
test.use({ storageState: ADMIN_STORAGE_PATH, viewport: VIEWPORT });
|
||||
|
||||
test("Virtual Keys: rows scroll under a sticky header and the page itself never scrolls", async ({
|
||||
page,
|
||||
request,
|
||||
}) => {
|
||||
const suffix = uniqueSuffix();
|
||||
const keys = await oneAtATime(SEED_ROWS, (i) =>
|
||||
postOk<GeneratedKey>(request, "/key/generate", { key_alias: `e2e-scroll-key-${suffix}-${i}` }),
|
||||
);
|
||||
try {
|
||||
await navigateToPage(page, Page.ApiKeys);
|
||||
await expectRowsAtLeast(page, SEED_ROWS);
|
||||
await expectBodyIsTheOnlyScroller(page);
|
||||
} finally {
|
||||
await request.post("/key/delete", { headers: adminHeaders(), data: { keys: keys.map((k) => k.key) } });
|
||||
}
|
||||
});
|
||||
|
||||
test("Teams: rows scroll under a sticky header and the page itself never scrolls", async ({ page, request }) => {
|
||||
const suffix = uniqueSuffix();
|
||||
const teams = await oneAtATime(SEED_ROWS, (i) =>
|
||||
postOk<CreatedTeam>(request, "/team/new", { team_alias: `e2e-scroll-team-${suffix}-${i}` }),
|
||||
);
|
||||
try {
|
||||
await navigateToPage(page, Page.Teams);
|
||||
await expectRowsAtLeast(page, SEED_ROWS);
|
||||
await expectBodyIsTheOnlyScroller(page);
|
||||
} finally {
|
||||
await request.post("/team/delete", { headers: adminHeaders(), data: { team_ids: teams.map((t) => t.team_id) } });
|
||||
}
|
||||
});
|
||||
|
||||
test("Request Logs: rows scroll under a sticky header and the page itself never scrolls", async ({
|
||||
page,
|
||||
request,
|
||||
}) => {
|
||||
const suffix = uniqueSuffix();
|
||||
const ids = await oneAtATime(LOG_ROWS, (i) =>
|
||||
sendChatCompletion(request, { model: CHAT_MODEL_A, prompt: `scroll ${suffix} ${i}` }),
|
||||
);
|
||||
await waitForSpendLog(request, ids[ids.length - 1]);
|
||||
|
||||
await navigateToPage(page, Page.Logs);
|
||||
await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 });
|
||||
await setRowsPerPage(page, "25");
|
||||
await expectRowsAtLeast(page, LOG_ROWS);
|
||||
await expectBodyIsTheOnlyScroller(page);
|
||||
});
|
||||
|
||||
test("Tags: no row paints past the box it lives in", async ({ page, request }) => {
|
||||
const suffix = uniqueSuffix();
|
||||
const names = Array.from({ length: SEED_ROWS }, (_, i) => `e2e-scroll-tag-${suffix}-${i}`);
|
||||
await oneAtATime(SEED_ROWS, (i) =>
|
||||
postOk<unknown>(request, "/tag/new", { name: names[i], description: "LIT-4738 scroll" }),
|
||||
);
|
||||
try {
|
||||
await navigateToPage(page, Page.TagManagement);
|
||||
await expectRowsAtLeast(page, SEED_ROWS);
|
||||
expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
|
||||
} finally {
|
||||
await oneAtATime(SEED_ROWS, (i) =>
|
||||
request.post("/tag/delete", { headers: adminHeaders(), data: { name: names[i] } }),
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
test("Model Hub: no row paints past the box it lives in", async ({ page, request }) => {
|
||||
const suffix = uniqueSuffix();
|
||||
const models = await oneAtATime(SEED_ROWS, (i) =>
|
||||
postOk<CreatedModel>(request, "/model/new", {
|
||||
model_name: `e2e-scroll-model-${suffix}-${i}`,
|
||||
litellm_params: {
|
||||
model: "openai/fake-gpt-4",
|
||||
api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`,
|
||||
api_key: "fake-key",
|
||||
},
|
||||
}),
|
||||
);
|
||||
try {
|
||||
await navigateToPage(page, Page.ModelHubTable);
|
||||
await expectRowsAtLeast(page, SEED_ROWS);
|
||||
expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
|
||||
} finally {
|
||||
await oneAtATime(SEED_ROWS, (i) =>
|
||||
request.post("/model/delete", { headers: adminHeaders(), data: { id: models[i].model_info.id } }),
|
||||
);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
|
@ -17,7 +17,7 @@ test.describe("Internal Users Search", () => {
|
|||
test("narrows the table to the matching email, and restores it when cleared", async ({ page }) => {
|
||||
await goToInternalUsers(page);
|
||||
|
||||
const search = page.getByPlaceholder("Search by email…");
|
||||
const search = page.getByPlaceholder("Search by email or ID…");
|
||||
await expect(search).toBeVisible();
|
||||
|
||||
await search.fill("noteam@");
|
||||
|
|
|
|||
|
|
@ -806,3 +806,91 @@ def test_guardrail_status_fields_computation():
|
|||
)
|
||||
assert status_fields_no_guardrail.get("llm_api_status") == "success"
|
||||
assert status_fields_no_guardrail.get("guardrail_status") == "not_run"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status, guardrail_information, expected_guardrail_status",
|
||||
[
|
||||
pytest.param(
|
||||
"failure",
|
||||
[
|
||||
{"guardrail_status": "success"},
|
||||
{"guardrail_status": "guardrail_intervened"},
|
||||
],
|
||||
"guardrail_intervened",
|
||||
id="pre_call_success_before_blocker",
|
||||
),
|
||||
pytest.param(
|
||||
"failure",
|
||||
[
|
||||
{"guardrail_status": "guardrail_intervened"},
|
||||
{"guardrail_status": "success"},
|
||||
],
|
||||
"guardrail_intervened",
|
||||
id="blocker_before_success",
|
||||
),
|
||||
pytest.param(
|
||||
"failure",
|
||||
[
|
||||
{"guardrail_status": "success"},
|
||||
{"guardrail_status": "guardrail_failed_to_respond"},
|
||||
],
|
||||
"guardrail_failed_to_respond",
|
||||
id="failure_outranks_success",
|
||||
),
|
||||
pytest.param(
|
||||
"failure",
|
||||
[
|
||||
{"guardrail_status": "guardrail_failed_to_respond"},
|
||||
{"guardrail_status": "guardrail_intervened"},
|
||||
],
|
||||
"guardrail_intervened",
|
||||
id="intervention_outranks_failure",
|
||||
),
|
||||
pytest.param(
|
||||
"success",
|
||||
[
|
||||
{"guardrail_status": "success"},
|
||||
{"guardrail_status": "success"},
|
||||
],
|
||||
"success",
|
||||
id="all_success_stays_success",
|
||||
),
|
||||
pytest.param(
|
||||
"failure",
|
||||
[
|
||||
{"guardrail_status": "some_new_status"},
|
||||
{"guardrail_status": "blocked"},
|
||||
],
|
||||
"guardrail_intervened",
|
||||
id="unknown_status_does_not_mask_blocker",
|
||||
),
|
||||
pytest.param(
|
||||
"failure",
|
||||
[
|
||||
{"guardrail_status": {"unhashable": True}},
|
||||
{"guardrail_status": "guardrail_intervened"},
|
||||
],
|
||||
"guardrail_intervened",
|
||||
id="unhashable_status_is_skipped",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_guardrail_status_fields_severity_across_entries(
|
||||
status, guardrail_information, expected_guardrail_status
|
||||
):
|
||||
"""
|
||||
A blocked request must never be reported as a guardrail success.
|
||||
|
||||
With multiple guardrails on one request (e.g. a pre_call mask that passes,
|
||||
then a post_call guardrail that blocks), entries are recorded in execution
|
||||
order, so the earlier "success" entry must not shadow the later
|
||||
"guardrail_intervened" entry: the aggregate takes the most severe status,
|
||||
regardless of entry order.
|
||||
"""
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_status_fields
|
||||
|
||||
fields = _get_status_fields(
|
||||
status=status, guardrail_information=guardrail_information, error_str=None
|
||||
)
|
||||
assert fields.get("guardrail_status") == expected_guardrail_status
|
||||
|
|
|
|||
150
tests/load_tests/test_granian_admission_saturation.py
Normal file
150
tests/load_tests/test_granian_admission_saturation.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
import asyncio
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
os.environ.get("LITELLM_RUN_SATURATION_BENCHMARK") != "1",
|
||||
reason="set LITELLM_RUN_SATURATION_BENCHMARK=1 to run the saturation benchmark",
|
||||
)
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
return int(listener.getsockname()[1])
|
||||
|
||||
|
||||
def _percentile(values: list[float], percentile: float) -> float:
|
||||
return sorted(values)[min(int(len(values) * percentile), len(values) - 1)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_granian_admission_control_saturation(tmp_path: Path) -> None:
|
||||
fake_port: Final = _free_port()
|
||||
proxy_port: Final = _free_port()
|
||||
fake_script: Final = Path(__file__).parents[1] / "_fake_openai_endpoint_server.py"
|
||||
config_path: Final = tmp_path / "saturation_config.yaml"
|
||||
config_path.write_text(
|
||||
f"""model_list:
|
||||
- model_name: slow-endpoint
|
||||
litellm_params:
|
||||
model: openai/slow-endpoint
|
||||
api_base: http://127.0.0.1:{fake_port}/v1
|
||||
general_settings:
|
||||
master_key: sk-saturation
|
||||
max_in_flight_requests_per_worker: 8
|
||||
max_queued_requests_per_worker: 8
|
||||
admission_queue_timeout_seconds: 0.5
|
||||
"""
|
||||
)
|
||||
fake_process: Final = subprocess.Popen(
|
||||
[sys.executable, str(fake_script), "--host", "127.0.0.1", "--port", str(fake_port)],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
try:
|
||||
proxy_process: Final = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"litellm.proxy.proxy_cli",
|
||||
"--config",
|
||||
str(config_path),
|
||||
"--run_granian",
|
||||
"--num_workers",
|
||||
"1",
|
||||
"--port",
|
||||
str(proxy_port),
|
||||
],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
try:
|
||||
async with httpx.AsyncClient(base_url=f"http://127.0.0.1:{proxy_port}") as client:
|
||||
deadline: Final = time.monotonic() + 60
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
response: Final = await client.get("/health/liveliness", timeout=2)
|
||||
if response.status_code == 200:
|
||||
break
|
||||
except httpx.HTTPError:
|
||||
pass
|
||||
await asyncio.sleep(0.25)
|
||||
else:
|
||||
raise AssertionError("Granian proxy did not become healthy")
|
||||
|
||||
liveness_latencies: Final[list[float]] = []
|
||||
stop_sampling: Final = asyncio.Event()
|
||||
|
||||
async def sample_liveness() -> None:
|
||||
while not stop_sampling.is_set():
|
||||
start: Final = time.perf_counter()
|
||||
try:
|
||||
response = await client.get("/health/liveliness", timeout=2)
|
||||
response.raise_for_status()
|
||||
liveness_latencies.append(time.perf_counter() - start)
|
||||
except httpx.HTTPError:
|
||||
pass
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
async def send_completion() -> tuple[int, float, bool]:
|
||||
start: Final = time.perf_counter()
|
||||
response = await client.post(
|
||||
"/chat/completions",
|
||||
headers={"Authorization": "Bearer sk-saturation"},
|
||||
json={
|
||||
"model": "slow-endpoint",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
},
|
||||
timeout=10,
|
||||
)
|
||||
return response.status_code, time.perf_counter() - start, "retry-after" in response.headers
|
||||
|
||||
sampler: Final = asyncio.create_task(sample_liveness())
|
||||
results: Final = await asyncio.gather(*(send_completion() for _ in range(200)))
|
||||
stop_sampling.set()
|
||||
await sampler
|
||||
|
||||
statuses: Final = [result[0] for result in results]
|
||||
latencies: Final = [result[1] for result in results]
|
||||
rejected: Final = [result for result in results if result[0] == 503]
|
||||
assert set(statuses) <= {200, 503}
|
||||
assert rejected
|
||||
assert all(result[2] for result in rejected)
|
||||
assert _percentile(latencies, 0.99) < 5
|
||||
assert liveness_latencies
|
||||
assert _percentile(liveness_latencies, 0.95) < 0.5
|
||||
|
||||
duration: Final = max(latencies)
|
||||
print(
|
||||
"\nmetric value\n"
|
||||
f"rps {len(results) / duration:.2f}\n"
|
||||
f"200 count {statuses.count(200)}\n"
|
||||
f"503 count {statuses.count(503)}\n"
|
||||
f"p50 {_percentile(latencies, 0.50):.3f}s\n"
|
||||
f"p95 {_percentile(latencies, 0.95):.3f}s\n"
|
||||
f"p99 {_percentile(latencies, 0.99):.3f}s\n"
|
||||
f"liveness p95 {_percentile(liveness_latencies, 0.95):.3f}s"
|
||||
)
|
||||
finally:
|
||||
proxy_process.terminate()
|
||||
try:
|
||||
proxy_process.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
proxy_process.kill()
|
||||
proxy_process.wait()
|
||||
finally:
|
||||
fake_process.terminate()
|
||||
try:
|
||||
fake_process.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
fake_process.kill()
|
||||
fake_process.wait()
|
||||
|
|
@ -5,20 +5,26 @@ CLIENT PAUSE) showed 100% of concurrent commands to the other two, untouched nod
|
|||
stalling for the full pause duration before this fix, and zero after -- these tests pin
|
||||
the same behavior at the unit level so it can run without a live Redis Cluster."""
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from redis.exceptions import (
|
||||
AskError,
|
||||
BusyLoadingError,
|
||||
ClusterDownError,
|
||||
ClusterError,
|
||||
MaxConnectionsError,
|
||||
MovedError,
|
||||
TryAgainError,
|
||||
)
|
||||
from redis.exceptions import (
|
||||
ConnectionError as RedisConnectionError,
|
||||
)
|
||||
from redis.exceptions import TimeoutError as RedisTimeoutError
|
||||
from redis.exceptions import (
|
||||
TimeoutError as RedisTimeoutError,
|
||||
)
|
||||
|
||||
from litellm.caching.redis_cluster_node_isolation import (
|
||||
get_litellm_async_redis_cluster_class,
|
||||
|
|
@ -39,10 +45,35 @@ class _NodeClassWithoutPerConnectionRecovery:
|
|||
class _FakeClusterNode:
|
||||
def __init__(self, name: str, raises: Exception | None = None, response: object = None) -> None:
|
||||
self.name = name
|
||||
self.execute_command = AsyncMock(side_effect=raises, return_value=response)
|
||||
|
||||
async def execute_command(*args: object, **kwargs: object) -> object:
|
||||
await asyncio.sleep(0)
|
||||
if raises is not None:
|
||||
raise raises
|
||||
return response
|
||||
|
||||
self.execute_command = AsyncMock(side_effect=execute_command)
|
||||
self.disconnect = AsyncMock()
|
||||
|
||||
|
||||
class _Fake8xRedisCluster:
|
||||
def __init__(self) -> None:
|
||||
self._initialize = False
|
||||
|
||||
async def _execute_command(
|
||||
self, target_node: _FakeClusterNode, *args: object, **kwargs: object
|
||||
) -> object:
|
||||
try:
|
||||
return await target_node.execute_command(*args, **kwargs)
|
||||
except (RedisConnectionError, RedisTimeoutError):
|
||||
self._initialize = True
|
||||
await asyncio.sleep(0)
|
||||
raise
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self._initialize = True
|
||||
|
||||
|
||||
class _FakeNodesManager:
|
||||
def __init__(self, node_to_return: _FakeClusterNode) -> None:
|
||||
self._moved_exception: object = None
|
||||
|
|
@ -68,31 +99,198 @@ def _build_cluster_instance() -> "_AsyncRedisClusterType":
|
|||
return instance
|
||||
|
||||
|
||||
def test_per_connection_recovery_redis_py_gets_the_unmodified_upstream_class() -> None:
|
||||
"""Regression (redis-py 8.x): when upstream ClusterNode already recovers a node-level
|
||||
connection error per-connection, the factory must NOT install the copied override,
|
||||
whose node.disconnect() also kills connections other coroutines are mid-operation on."""
|
||||
from redis.asyncio.cluster import RedisCluster
|
||||
|
||||
def _build_8x_cluster_instance() -> _Fake8xRedisCluster:
|
||||
cluster_cls = get_litellm_async_redis_cluster_class(
|
||||
cluster_node_class=_NodeClassWithPerConnectionRecovery
|
||||
cluster_node_class=_NodeClassWithPerConnectionRecovery,
|
||||
base_cluster_class=_Fake8xRedisCluster,
|
||||
)
|
||||
return cluster_cls()
|
||||
|
||||
|
||||
def test_unverified_redis_version_logs_warning(caplog: pytest.LogCaptureFixture) -> None:
|
||||
import redis
|
||||
|
||||
with patch.object(redis, "__version__", "8.0.1"):
|
||||
get_litellm_async_redis_cluster_class(cluster_node_class=_NodeClassWithoutPerConnectionRecovery)
|
||||
|
||||
assert "not in the set this cluster-teardown-storm fix was verified against" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_timeout_does_not_request_topology_reinit() -> None:
|
||||
error = RedisTimeoutError("timeout")
|
||||
target_node = _FakeClusterNode("node-a")
|
||||
target_node.execute_command.side_effect = error
|
||||
instance = _build_8x_cluster_instance()
|
||||
|
||||
with pytest.raises(RedisTimeoutError) as exc_info:
|
||||
await instance._execute_command(target_node, "GET", "k")
|
||||
|
||||
assert exc_info.value is error
|
||||
assert instance._initialize is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_error_preserves_upstream_topology_reinit() -> None:
|
||||
error = RedisConnectionError("connection error")
|
||||
target_node = _FakeClusterNode("node-a")
|
||||
target_node.execute_command.side_effect = error
|
||||
instance = _build_8x_cluster_instance()
|
||||
|
||||
with pytest.raises(RedisConnectionError) as exc_info:
|
||||
await instance._execute_command(target_node, "GET", "k")
|
||||
|
||||
assert exc_info.value is error
|
||||
assert instance._initialize is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_three_consecutive_timeouts_request_topology_reinit_and_reset_counter() -> None:
|
||||
errors = [
|
||||
RedisTimeoutError("timeout-1"),
|
||||
RedisTimeoutError("timeout-2"),
|
||||
RedisTimeoutError("timeout-3"),
|
||||
]
|
||||
fourth_error = RedisTimeoutError("timeout-4")
|
||||
target_node = _FakeClusterNode("node-a")
|
||||
target_node.execute_command.side_effect = [*errors, fourth_error]
|
||||
instance = _build_8x_cluster_instance()
|
||||
|
||||
for error in errors:
|
||||
with pytest.raises(RedisTimeoutError) as exc_info:
|
||||
await instance._execute_command(target_node, "GET", "k")
|
||||
assert exc_info.value is error
|
||||
|
||||
assert instance._initialize is True
|
||||
instance._initialize = False
|
||||
|
||||
with pytest.raises(RedisTimeoutError) as exc_info:
|
||||
await instance._execute_command(target_node, "GET", "k")
|
||||
|
||||
assert exc_info.value is fourth_error
|
||||
assert instance._initialize is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_resets_consecutive_timeout_counter() -> None:
|
||||
errors = [RedisTimeoutError("timeout-1"), RedisTimeoutError("timeout-2")]
|
||||
final_error = RedisTimeoutError("timeout-3")
|
||||
target_node = _FakeClusterNode("node-a")
|
||||
target_node.execute_command.side_effect = [*errors, b"value", final_error]
|
||||
instance = _build_8x_cluster_instance()
|
||||
|
||||
for error in errors:
|
||||
with pytest.raises(RedisTimeoutError) as exc_info:
|
||||
await instance._execute_command(target_node, "GET", "k")
|
||||
assert exc_info.value is error
|
||||
assert instance._initialize is False
|
||||
|
||||
result = await instance._execute_command(target_node, "GET", "k")
|
||||
assert result == b"value"
|
||||
assert instance._initialize is False
|
||||
|
||||
with pytest.raises(RedisTimeoutError) as exc_info:
|
||||
await instance._execute_command(target_node, "GET", "k")
|
||||
|
||||
assert exc_info.value is final_error
|
||||
assert instance._initialize is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_counters_are_per_node() -> None:
|
||||
node_a_errors = [RedisTimeoutError("node-a-1"), RedisTimeoutError("node-a-2")]
|
||||
node_b_error = RedisTimeoutError("node-b-1")
|
||||
node_a = _FakeClusterNode("node-a")
|
||||
node_b = _FakeClusterNode("node-b")
|
||||
node_a.execute_command.side_effect = node_a_errors
|
||||
node_b.execute_command.side_effect = node_b_error
|
||||
instance = _build_8x_cluster_instance()
|
||||
|
||||
for target_node, error in (
|
||||
(node_a, node_a_errors[0]),
|
||||
(node_b, node_b_error),
|
||||
(node_a, node_a_errors[1]),
|
||||
):
|
||||
with pytest.raises(RedisTimeoutError) as exc_info:
|
||||
await instance._execute_command(target_node, "GET", "k")
|
||||
assert exc_info.value is error
|
||||
|
||||
assert instance._initialize is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_does_not_clear_concurrent_topology_reinit_request() -> None:
|
||||
error = RedisTimeoutError("timeout")
|
||||
instance = _build_8x_cluster_instance()
|
||||
|
||||
async def request_reinit(*args: object, **kwargs: object) -> object:
|
||||
await instance.aclose()
|
||||
raise error
|
||||
|
||||
target_node = _FakeClusterNode("node-a")
|
||||
target_node.execute_command.side_effect = request_reinit
|
||||
|
||||
with pytest.raises(RedisTimeoutError) as exc_info:
|
||||
await instance._execute_command(target_node, "GET", "k")
|
||||
|
||||
assert exc_info.value is error
|
||||
assert instance._initialize is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tolerated_timeout_does_not_erase_concurrent_connection_error_reinit() -> None:
|
||||
instance = _build_8x_cluster_instance()
|
||||
failing_node = _FakeClusterNode("node-a", raises=RedisConnectionError("gone"))
|
||||
slow_node = _FakeClusterNode("node-b", raises=RedisTimeoutError("slow"))
|
||||
|
||||
results = await asyncio.gather(
|
||||
instance._execute_command(failing_node, "GET", "a"),
|
||||
instance._execute_command(slow_node, "GET", "b"),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
assert cluster_cls is RedisCluster
|
||||
assert isinstance(results[0], RedisConnectionError)
|
||||
assert isinstance(results[1], RedisTimeoutError)
|
||||
assert instance._initialize is True
|
||||
|
||||
|
||||
def test_pre_recovery_redis_py_still_gets_the_node_isolation_override() -> None:
|
||||
"""Old redis-py (5.x) responds to a node-level error with a full-cluster aclose(),
|
||||
so those versions must keep litellm's per-node isolation override."""
|
||||
from redis.asyncio.cluster import RedisCluster
|
||||
@pytest.mark.asyncio
|
||||
async def test_overlapping_tolerated_timeouts_do_not_request_topology_reinit() -> None:
|
||||
instance = _build_8x_cluster_instance()
|
||||
node_a = _FakeClusterNode("node-a", raises=RedisTimeoutError("slow-a"))
|
||||
node_b = _FakeClusterNode("node-b", raises=RedisTimeoutError("slow-b"))
|
||||
|
||||
cluster_cls = get_litellm_async_redis_cluster_class(
|
||||
cluster_node_class=_NodeClassWithoutPerConnectionRecovery
|
||||
results = await asyncio.gather(
|
||||
instance._execute_command(node_a, "GET", "a"),
|
||||
instance._execute_command(node_b, "GET", "b"),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
assert cluster_cls is not RedisCluster
|
||||
assert issubclass(cluster_cls, RedisCluster)
|
||||
assert "_execute_command" in cluster_cls.__dict__
|
||||
assert all(isinstance(result, RedisTimeoutError) for result in results)
|
||||
assert instance._initialize is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tolerated_timeout_does_not_clear_pending_reinit() -> None:
|
||||
instance = _build_8x_cluster_instance()
|
||||
instance._initialize = True
|
||||
target_node = _FakeClusterNode("node-a", raises=RedisTimeoutError("slow"))
|
||||
|
||||
with pytest.raises(RedisTimeoutError):
|
||||
await instance._execute_command(target_node, "GET", "k")
|
||||
|
||||
assert instance._initialize is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_returns_value_without_topology_reinit() -> None:
|
||||
target_node = _FakeClusterNode("node-a", response=b"value")
|
||||
instance = _build_8x_cluster_instance()
|
||||
|
||||
result = await instance._execute_command(target_node, "GET", "k")
|
||||
|
||||
assert result == b"value"
|
||||
assert instance._initialize is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -110,6 +308,48 @@ async def test_node_level_error_resets_only_that_node_not_the_whole_client(error
|
|||
instance.aclose.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_moved_error_retries_without_full_reinit_before_threshold() -> None:
|
||||
moved_error = MovedError("1 127.0.0.1:7001")
|
||||
target_node = _FakeClusterNode("node-a")
|
||||
target_node.execute_command = AsyncMock(side_effect=[moved_error, b"value"])
|
||||
instance = _build_cluster_instance()
|
||||
instance.RedisClusterRequestTTL = 2
|
||||
instance.nodes_manager = _FakeNodesManager(node_to_return=target_node)
|
||||
instance._determine_slot = AsyncMock(return_value=0)
|
||||
|
||||
result = await instance._execute_command(target_node, "GET", "k")
|
||||
|
||||
assert result == b"value"
|
||||
assert instance.nodes_manager._moved_exception is moved_error
|
||||
instance.aclose.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ask_error_sends_asking_and_retries_on_redirected_node() -> None:
|
||||
ask_error = AskError("0 127.0.0.1:7001")
|
||||
target_node = _FakeClusterNode("node-a")
|
||||
target_node.execute_command = AsyncMock(side_effect=[ask_error, None, b"value"])
|
||||
instance = _build_cluster_instance()
|
||||
instance.RedisClusterRequestTTL = 2
|
||||
instance.get_node = Mock(return_value=target_node)
|
||||
|
||||
result = await instance._execute_command(target_node, "GET", "k")
|
||||
|
||||
assert result == b"value"
|
||||
instance.get_node.assert_called_once_with(node_name="127.0.0.1:7001")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_try_again_error_exhausts_ttl() -> None:
|
||||
target_node = _FakeClusterNode("node-a", raises=TryAgainError("try again"))
|
||||
instance = _build_cluster_instance()
|
||||
instance.RedisClusterRequestTTL = 2
|
||||
|
||||
with pytest.raises(ClusterError):
|
||||
await instance._execute_command(target_node, "GET", "k")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_command_touches_neither_disconnect_nor_aclose() -> None:
|
||||
target_node = _FakeClusterNode("node-a", response=b"v")
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ import pytest
|
|||
import litellm
|
||||
from litellm import completion, acompletion
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.snowflake.chat.transformation import SnowflakeConfig
|
||||
from litellm.llms.snowflake.chat.transformation import SnowflakeConfig, SnowflakeStreamingHandler
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
|
|
@ -114,8 +114,7 @@ class TestSnowflakeToolTransformation:
|
|||
)
|
||||
|
||||
assert transformed_request["tool_choice"] == value, (
|
||||
f"tool_choice='{value}' should pass through unchanged, "
|
||||
f"got {transformed_request['tool_choice']}"
|
||||
f"tool_choice='{value}' should pass through unchanged, got {transformed_request['tool_choice']}"
|
||||
)
|
||||
|
||||
def test_transform_response_with_tool_calls(self):
|
||||
|
|
@ -159,9 +158,7 @@ class TestSnowflakeToolTransformation:
|
|||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
model_response = ModelResponse(
|
||||
choices=[litellm.Choices(index=0, message=litellm.Message())]
|
||||
)
|
||||
model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())])
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
||||
|
|
@ -232,9 +229,7 @@ class TestSnowflakeToolTransformation:
|
|||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
model_response = ModelResponse(
|
||||
choices=[litellm.Choices(index=0, message=litellm.Message())]
|
||||
)
|
||||
model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())])
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
||||
|
|
@ -280,9 +275,7 @@ class TestSnowflakeToolTransformation:
|
|||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
model_response = ModelResponse(
|
||||
choices=[litellm.Choices(index=0, message=litellm.Message())]
|
||||
)
|
||||
model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())])
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
||||
|
|
@ -300,10 +293,7 @@ class TestSnowflakeToolTransformation:
|
|||
|
||||
# Verify standard response works
|
||||
assert isinstance(result, ModelResponse)
|
||||
assert (
|
||||
result.choices[0].message.content
|
||||
== "Hello! I'm doing well, thank you for asking."
|
||||
)
|
||||
assert result.choices[0].message.content == "Hello! I'm doing well, thank you for asking."
|
||||
|
||||
def test_get_supported_openai_params_includes_tools(self):
|
||||
"""
|
||||
|
|
@ -318,6 +308,385 @@ class TestSnowflakeToolTransformation:
|
|||
assert "max_tokens" in supported_params
|
||||
|
||||
|
||||
class TestSnowflakeCortexClaudeFixes:
|
||||
def setup_method(self):
|
||||
self.config = SnowflakeConfig()
|
||||
|
||||
@staticmethod
|
||||
def _transform(messages, optional_params=None):
|
||||
return SnowflakeConfig().transform_request(
|
||||
model="snowflake/claude-sonnet-4-6",
|
||||
messages=messages,
|
||||
optional_params=optional_params or {},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
def test_thinking_is_offered_on_every_claude_model(self):
|
||||
"""Cortex documents extended thinking (budget_tokens) for Claude generally, so a
|
||||
4.6-only gate would silently drop it on the models that do support it."""
|
||||
for model in (
|
||||
"snowflake/claude-sonnet-4-6",
|
||||
"snowflake/claude-sonnet-4-5",
|
||||
"snowflake/claude-3-7-sonnet",
|
||||
"snowflake/claude-4-opus",
|
||||
):
|
||||
assert "thinking" in self.config.get_supported_openai_params(model), model
|
||||
assert "thinking" not in self.config.get_supported_openai_params("snowflake/llama3.1-70b")
|
||||
|
||||
def test_system_blocks_preserve_cache_control_and_strip_ttl(self):
|
||||
body = self._transform(
|
||||
[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are helpful",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "hi"},
|
||||
]
|
||||
)
|
||||
assert body["system"] == [{"type": "text", "text": "You are helpful", "cache_control": {"type": "ephemeral"}}]
|
||||
|
||||
def test_direct_system_param_is_normalized(self):
|
||||
body = self._transform(
|
||||
[{"role": "user", "content": "hi"}],
|
||||
{"system": [{"type": "text", "text": "direct", "cache_control": {"type": "ephemeral", "ttl": "1h"}}]},
|
||||
)
|
||||
assert body["system"] == [{"type": "text", "text": "direct", "cache_control": {"type": "ephemeral"}}]
|
||||
|
||||
def test_message_and_tool_cache_control_are_normalized(self):
|
||||
body = self._transform(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}}],
|
||||
}
|
||||
],
|
||||
{
|
||||
"tools": [
|
||||
{
|
||||
"name": "f",
|
||||
"input_schema": {"type": "object", "properties": {}},
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
assert body["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"}
|
||||
assert body["tools"][0]["cache_control"] == {"type": "ephemeral"}
|
||||
|
||||
def test_extra_body_message_override_is_normalized(self):
|
||||
body = self._transform(
|
||||
[{"role": "user", "content": "original"}],
|
||||
{
|
||||
"extra_body": {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "override",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert body["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"}
|
||||
|
||||
def test_image_blocks_are_converted_to_anthropic_source(self):
|
||||
body = self._transform(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "data:image/png;base64,ZmFrZQ==", "format": "image/jpeg"},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
assert body["messages"][0]["content"] == [
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "ZmFrZQ=="}}
|
||||
]
|
||||
|
||||
def test_tool_result_image_list_is_converted(self):
|
||||
body = self._transform(
|
||||
[
|
||||
{"role": "user", "content": "look"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,ZmFrZQ=="}}],
|
||||
},
|
||||
]
|
||||
)
|
||||
assert body["messages"][2]["content"][0]["content"] == [
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "ZmFrZQ=="}}
|
||||
]
|
||||
|
||||
def test_tool_result_preserves_cache_control(self):
|
||||
"""A cache breakpoint the bridge puts on a tool message must survive onto the tool_result."""
|
||||
for tool_content in ("done", [{"type": "text", "text": "done"}]):
|
||||
body = self._transform(
|
||||
[
|
||||
{"role": "user", "content": "look"},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": tool_content,
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
},
|
||||
]
|
||||
)
|
||||
tool_result = body["messages"][1]["content"][0]
|
||||
assert tool_result["cache_control"] == {"type": "ephemeral"}, tool_content
|
||||
|
||||
def test_pdf_data_uri_becomes_a_document_block(self):
|
||||
"""A bridged pdf data URI is a document block; forwarding it as an image is malformed."""
|
||||
body = self._transform(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": "data:application/pdf;base64,ZmFrZQ=="}},
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
assert body["messages"][0]["content"] == [
|
||||
{
|
||||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": "ZmFrZQ=="},
|
||||
}
|
||||
]
|
||||
|
||||
def test_multipart_tool_result_preserves_text_and_converts_image(self):
|
||||
body = self._transform(
|
||||
[
|
||||
{"role": "user", "content": "look"},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": [
|
||||
{"type": "text", "text": "first"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,ZmFrZQ=="}},
|
||||
{"type": "text", "text": "last"},
|
||||
],
|
||||
},
|
||||
]
|
||||
)
|
||||
assert body["messages"][1]["content"][0]["content"] == [
|
||||
{"type": "text", "text": "first"},
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "ZmFrZQ=="}},
|
||||
{"type": "text", "text": "last"},
|
||||
]
|
||||
|
||||
def test_plain_text_tool_result_remains_string(self):
|
||||
body = self._transform(
|
||||
[{"role": "user", "content": "look"}, {"role": "tool", "tool_call_id": "call_1", "content": "done"}]
|
||||
)
|
||||
assert body["messages"][1]["content"][0]["content"] == "done"
|
||||
|
||||
def test_anthropic_tool_schema_strips_only_top_level_schema_key(self):
|
||||
tools = [
|
||||
{
|
||||
"name": "f",
|
||||
"input_schema": {"$schema": "schema", "type": "object", "properties": {"$schema": {"type": "string"}}},
|
||||
}
|
||||
]
|
||||
body = self._transform([{"role": "user", "content": "hi"}], {"tools": tools})
|
||||
schema = body["tools"][0]["input_schema"]
|
||||
assert "$schema" not in schema
|
||||
assert "$schema" in schema["properties"]
|
||||
|
||||
def test_tool_schema_strips_only_top_level_schema_key(self):
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "f",
|
||||
"parameters": {
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"type": "object",
|
||||
"properties": {"$schema": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
body = self._transform([{"role": "user", "content": "hi"}], {"tools": tools})
|
||||
schema = body["tools"][0]["input_schema"]
|
||||
assert "$schema" not in schema
|
||||
assert "$schema" in schema["properties"]
|
||||
|
||||
def test_streaming_tool_identity_is_emitted_only_on_start(self):
|
||||
handler = SnowflakeStreamingHandler(streaming_response=[], sync_stream=True)
|
||||
start = handler.chunk_parser(
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "tool_use", "id": "tool_1", "name": "read"},
|
||||
}
|
||||
)
|
||||
first_delta = handler.chunk_parser(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"path":'},
|
||||
}
|
||||
)
|
||||
second_delta = handler.chunk_parser(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '"/tmp"}'},
|
||||
}
|
||||
)
|
||||
|
||||
def _tool_call(chunk):
|
||||
return chunk.choices[0].delta.tool_calls[0]
|
||||
|
||||
assert _tool_call(start).id == "tool_1"
|
||||
assert _tool_call(start).function.name == "read"
|
||||
assert _tool_call(first_delta).id is None
|
||||
assert _tool_call(first_delta).function.name is None
|
||||
assert _tool_call(second_delta).id is None
|
||||
assert _tool_call(second_delta).function.name is None
|
||||
assert _tool_call(first_delta).function.arguments == '{"path":'
|
||||
assert _tool_call(second_delta).function.arguments == '"/tmp"}'
|
||||
|
||||
def test_signed_thinking_blocks_lead_the_assistant_turn(self):
|
||||
"""Multi-turn tool use with thinking only works if the signed block is echoed back first."""
|
||||
body = self._transform(
|
||||
[
|
||||
{"role": "user", "content": "hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"thinking_blocks": [
|
||||
{"type": "thinking", "thinking": "391", "signature": "Eto"},
|
||||
{"type": "thinking", "thinking": "unsigned"},
|
||||
],
|
||||
"tool_calls": [
|
||||
{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}
|
||||
],
|
||||
},
|
||||
]
|
||||
)
|
||||
blocks = body["messages"][1]["content"]
|
||||
assert blocks[0] == {"type": "thinking", "thinking": "391", "signature": "Eto"}
|
||||
assert [b["type"] for b in blocks] == ["thinking", "tool_use"]
|
||||
|
||||
def test_signed_thinking_blocks_lead_a_plain_text_assistant_turn(self):
|
||||
"""A thinking response without a tool call must also round-trip on the next request."""
|
||||
body = self._transform(
|
||||
[
|
||||
{"role": "user", "content": "hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "391",
|
||||
"thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}],
|
||||
},
|
||||
{"role": "user", "content": "continue"},
|
||||
]
|
||||
)
|
||||
assert body["messages"][1] == {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "391", "signature": "Eto"},
|
||||
{"type": "text", "text": "391"},
|
||||
],
|
||||
}
|
||||
|
||||
def test_signed_thinking_blocks_preserve_list_content(self):
|
||||
"""Cached assistant text reaches this transform as a content list, not a string."""
|
||||
body = self._transform(
|
||||
[
|
||||
{"role": "user", "content": "hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "391", "cache_control": {"type": "ephemeral"}}],
|
||||
"thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}],
|
||||
},
|
||||
{"role": "user", "content": "continue"},
|
||||
]
|
||||
)
|
||||
assert body["messages"][1] == {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "391", "signature": "Eto"},
|
||||
{"type": "text", "text": "391", "cache_control": {"type": "ephemeral"}},
|
||||
],
|
||||
}
|
||||
|
||||
def test_thinking_only_assistant_turn_sends_no_empty_text_block(self):
|
||||
"""Anthropic-shaped APIs reject empty text blocks, so a content-less thinking turn is thinking only."""
|
||||
body = self._transform(
|
||||
[
|
||||
{"role": "user", "content": "hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}],
|
||||
},
|
||||
{"role": "user", "content": "continue"},
|
||||
]
|
||||
)
|
||||
assert body["messages"][1]["content"] == [{"type": "thinking", "thinking": "391", "signature": "Eto"}]
|
||||
|
||||
def test_streaming_surfaces_thinking_and_prompt_cache_usage(self):
|
||||
"""Cortex streams thinking deltas, signatures and cache counts; all must reach the caller."""
|
||||
handler = SnowflakeStreamingHandler(streaming_response=[], sync_stream=True)
|
||||
handler.chunk_parser(
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {"usage": {"input_tokens": 18, "cache_creation_input_tokens": 1323}},
|
||||
}
|
||||
)
|
||||
thinking = handler.chunk_parser(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "thinking_delta", "thinking": "391"},
|
||||
}
|
||||
)
|
||||
signature = handler.chunk_parser(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "signature_delta", "signature": "Eto"},
|
||||
}
|
||||
)
|
||||
final = handler.chunk_parser(
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn"},
|
||||
"usage": {"output_tokens": 8, "cache_read_input_tokens": 1323},
|
||||
}
|
||||
)
|
||||
|
||||
assert thinking.choices[0].delta.reasoning_content == "391"
|
||||
assert signature.choices[0].delta.thinking_blocks[0]["signature"] == "Eto"
|
||||
assert final.usage.prompt_tokens_details.cached_tokens == 1323
|
||||
|
||||
|
||||
class TestSnowFlakeCompletion:
|
||||
model_name = "mistral"
|
||||
|
||||
|
|
@ -380,10 +749,7 @@ class TestSnowFlakeCompletion:
|
|||
# PAT key was used
|
||||
post_kwargs = mock_post.call_args_list[-1][1]
|
||||
assert "xxxxx" in post_kwargs["headers"]["Authorization"]
|
||||
assert (
|
||||
post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"]
|
||||
== "PROGRAMMATIC_ACCESS_TOKEN"
|
||||
)
|
||||
assert post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"] == "PROGRAMMATIC_ACCESS_TOKEN"
|
||||
|
||||
# account id was used
|
||||
assert "AAAA-BBBB" in post_kwargs["url"]
|
||||
|
|
@ -495,9 +861,7 @@ class TestSnowflakeChatCompletion:
|
|||
)
|
||||
mock_post.assert_called_once()
|
||||
else:
|
||||
with patch.object(
|
||||
AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp
|
||||
) as mock_post:
|
||||
with patch.object(AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp) as mock_post:
|
||||
response = asyncio.run(
|
||||
acompletion(
|
||||
model="snowflake/mistral-7b",
|
||||
|
|
@ -580,8 +944,4 @@ class TestSnowflakeChatCompletion:
|
|||
chunks_received = asyncio.run(_run())
|
||||
|
||||
assert len(chunks_received) > 0
|
||||
content = "".join(
|
||||
c.choices[0].delta.content
|
||||
for c in chunks_received
|
||||
if c.choices[0].delta.content
|
||||
)
|
||||
content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content)
|
||||
|
|
|
|||
|
|
@ -338,7 +338,7 @@ class TestAnthropicConfigRequest:
|
|||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert body["system"] == "You are helpful."
|
||||
assert body["system"] == [{"type": "text", "text": "You are helpful."}]
|
||||
assert all(m["role"] != "system" for m in body["messages"])
|
||||
assert body["messages"][0] == {"role": "user", "content": "Hello"}
|
||||
|
||||
|
|
@ -422,6 +422,64 @@ class TestAnthropicConfigResponse:
|
|||
assert result.usage.completion_tokens == 5
|
||||
assert result.usage.total_tokens == 15
|
||||
|
||||
def test_prompt_cache_usage_is_surfaced(self):
|
||||
"""Cortex reports cache creation/read counts; dropping them hides caching and bills cached input at full price."""
|
||||
raw = httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_1",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 18, "cache_creation_input_tokens": 1323, "cache_read_input_tokens": 0},
|
||||
},
|
||||
)
|
||||
result = self.cfg.transform_response(
|
||||
model="snowflake/claude-sonnet-4-6",
|
||||
raw_response=raw,
|
||||
model_response=ModelResponse(),
|
||||
logging_obj=_mock_logging(),
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
assert result.usage.prompt_tokens == 1341
|
||||
assert result.usage.prompt_tokens_details.cache_creation_tokens == 1323
|
||||
assert result.usage.prompt_tokens_details.cached_tokens == 0
|
||||
|
||||
def test_thinking_block_and_signature_are_preserved(self):
|
||||
"""The signature must survive so a client can echo the thinking block on the next turn."""
|
||||
raw = httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_1",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "391", "signature": "Eto"},
|
||||
{"type": "text", "text": "391"},
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
},
|
||||
)
|
||||
result = self.cfg.transform_response(
|
||||
model="snowflake/claude-sonnet-4-6",
|
||||
raw_response=raw,
|
||||
model_response=ModelResponse(),
|
||||
logging_obj=_mock_logging(),
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
message = result.choices[0].message
|
||||
assert message.content == "391"
|
||||
assert message.reasoning_content == "391"
|
||||
assert message.thinking_blocks[0]["signature"] == "Eto"
|
||||
|
||||
def test_stop_reason_end_turn_maps_to_stop(self):
|
||||
raw = _make_anthropic_response()
|
||||
result = self.cfg.transform_response(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ round-trips exactly, a store failure never escapes into the login path, and a sa
|
|||
rotation re-encrypts stored rows like the sibling per-user credential tables.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
|
@ -16,8 +17,10 @@ import jwt as pyjwt
|
|||
import pytest
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
|
||||
_ASSERTION_CACHE,
|
||||
AssertionStoreUnavailable,
|
||||
DbSSOAssertionStore,
|
||||
SSOAssertionCache,
|
||||
assertion_from_sso_login,
|
||||
ema_assertion_retention_enabled,
|
||||
fetch_sso_identity_assertion,
|
||||
|
|
@ -25,7 +28,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s
|
|||
retain_sso_identity_assertion_for_ema,
|
||||
rotate_sso_identity_assertions_master_key,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
SALT_KEY = "test-salt-key-for-sso-assertion-tests-1234"
|
||||
|
|
@ -38,6 +41,11 @@ def _set_salt_key(monkeypatch):
|
|||
monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _flush_assertion_cache():
|
||||
_ASSERTION_CACHE.flush()
|
||||
|
||||
|
||||
def _make_id_token(exp_offset: int = 3600, iss: str = ISSUER) -> str:
|
||||
return pyjwt.encode(
|
||||
{"iss": iss, "sub": "u1", "exp": int(time.time()) + exp_offset},
|
||||
|
|
@ -52,9 +60,7 @@ def _make_prisma(stored: dict, db_has_id_jag_server: bool = False):
|
|||
``db_has_id_jag_server`` drives the retention gate's authoritative DB fallback;
|
||||
it is wired explicitly so the gate never reads a truthy bare MagicMock."""
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_mcpservertable.find_first = AsyncMock(
|
||||
return_value=MagicMock() if db_has_id_jag_server else None
|
||||
)
|
||||
prisma.db.litellm_mcpservertable.find_first = AsyncMock(return_value=MagicMock() if db_has_id_jag_server else None)
|
||||
|
||||
async def _upsert(where, data):
|
||||
stored[where["user_id"]] = data["update"]["assertion_b64"]
|
||||
|
|
@ -235,6 +241,78 @@ async def test_persist_overwrites_previous_login():
|
|||
assert fetched.refresh_token is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_serves_second_read_from_cache_without_db_read():
|
||||
stored = {}
|
||||
prisma = _make_prisma(stored)
|
||||
cache = SSOAssertionCache()
|
||||
token = _make_id_token()
|
||||
assertion = assertion_from_sso_login(token, "rt_1")
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
|
||||
await persist_sso_identity_assertion("user-a", assertion, cache=cache)
|
||||
first = await fetch_sso_identity_assertion("user-a", cache=cache)
|
||||
second = await fetch_sso_identity_assertion("user-a", cache=cache)
|
||||
assert first is not None
|
||||
assert second is not None
|
||||
assert first.id_token.get_secret_value() == token
|
||||
assert second.id_token.get_secret_value() == token
|
||||
prisma.db.litellm_ssoidentityassertion.find_unique.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persist_busts_cache_so_relogin_is_visible_immediately():
|
||||
stored = {}
|
||||
prisma = _make_prisma(stored)
|
||||
cache = SSOAssertionCache()
|
||||
first_token = _make_id_token(exp_offset=100)
|
||||
second_token = _make_id_token(exp_offset=7200)
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
|
||||
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
|
||||
first = await fetch_sso_identity_assertion("user-a", cache=cache)
|
||||
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(second_token, "rt_new"), cache=cache)
|
||||
second = await fetch_sso_identity_assertion("user-a", cache=cache)
|
||||
assert first is not None
|
||||
assert second is not None
|
||||
assert first.id_token.get_secret_value() == first_token
|
||||
assert second.id_token.get_secret_value() == second_token
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_racing_a_relogin_does_not_cache_the_previous_assertion():
|
||||
stored = {}
|
||||
prisma = _make_prisma(stored)
|
||||
cache = SSOAssertionCache()
|
||||
first_token = _make_id_token(exp_offset=100)
|
||||
second_token = _make_id_token(exp_offset=7200)
|
||||
db_read_started = asyncio.Event()
|
||||
relogin_done = asyncio.Event()
|
||||
unpaused_find_unique = prisma.db.litellm_ssoidentityassertion.find_unique
|
||||
|
||||
async def _paused_find_unique(where):
|
||||
row = await unpaused_find_unique(where=where)
|
||||
db_read_started.set()
|
||||
await relogin_done.wait()
|
||||
return row
|
||||
|
||||
prisma.db.litellm_ssoidentityassertion.find_unique = _paused_find_unique
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
|
||||
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
|
||||
racing_fetch = asyncio.create_task(fetch_sso_identity_assertion("user-a", cache=cache))
|
||||
await db_read_started.wait()
|
||||
await persist_sso_identity_assertion(
|
||||
"user-a",
|
||||
assertion_from_sso_login(second_token, "rt_new"),
|
||||
cache=cache,
|
||||
)
|
||||
relogin_done.set()
|
||||
raced = await racing_fetch
|
||||
after = await fetch_sso_identity_assertion("user-a", cache=cache)
|
||||
assert raced is not None
|
||||
assert after is not None
|
||||
assert raced.id_token.get_secret_value() == first_token
|
||||
assert after.id_token.get_secret_value() == second_token
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_missing_row_returns_none():
|
||||
prisma = _make_prisma({})
|
||||
|
|
@ -242,6 +320,18 @@ async def test_fetch_missing_row_returns_none():
|
|||
assert await fetch_sso_identity_assertion("nobody") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_does_not_cache_a_missing_row():
|
||||
prisma = _make_prisma({})
|
||||
cache = SSOAssertionCache()
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
|
||||
first = await fetch_sso_identity_assertion("nobody", cache=cache)
|
||||
second = await fetch_sso_identity_assertion("nobody", cache=cache)
|
||||
assert first is None
|
||||
assert second is None
|
||||
assert prisma.db.litellm_ssoidentityassertion.find_unique.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_undecryptable_row_returns_none():
|
||||
prisma = _make_prisma({"user-a": "not-an-encrypted-blob"})
|
||||
|
|
@ -251,13 +341,37 @@ async def test_fetch_undecryptable_row_returns_none():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_unparseable_payload_returns_none():
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
|
||||
prisma = _make_prisma({"user-a": encrypt_value_helper("]]not json")})
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
|
||||
assert await fetch_sso_identity_assertion("user-a") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_assertion_expires_after_ttl():
|
||||
stored = {}
|
||||
prisma = _make_prisma(stored)
|
||||
cache = SSOAssertionCache(ttl_seconds=1)
|
||||
first_token = _make_id_token(exp_offset=100)
|
||||
second_token = _make_id_token(exp_offset=7200)
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
|
||||
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
|
||||
first = await fetch_sso_identity_assertion("user-a", cache=cache)
|
||||
await persist_sso_identity_assertion(
|
||||
"user-a",
|
||||
assertion_from_sso_login(second_token, "rt_new"),
|
||||
cache=SSOAssertionCache(),
|
||||
)
|
||||
cached = await fetch_sso_identity_assertion("user-a", cache=cache)
|
||||
time.sleep(1.1)
|
||||
expired = await fetch_sso_identity_assertion("user-a", cache=cache)
|
||||
assert first is not None
|
||||
assert cached is not None
|
||||
assert expired is not None
|
||||
assert first.id_token.get_secret_value() == first_token
|
||||
assert cached.id_token.get_secret_value() == first_token
|
||||
assert expired.id_token.get_secret_value() == second_token
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_noop_when_no_id_jag_server():
|
||||
stored = {}
|
||||
|
|
@ -357,6 +471,24 @@ async def test_db_store_converts_a_driver_failure_into_assertion_store_unavailab
|
|||
await DbSSOAssertionStore().fetch("alice")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_db_store_uses_injected_cache():
|
||||
stored = {}
|
||||
prisma = _make_prisma(stored)
|
||||
cache = SSOAssertionCache()
|
||||
token = _make_id_token()
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
|
||||
await persist_sso_identity_assertion("alice", assertion_from_sso_login(token, None), cache=cache)
|
||||
store = DbSSOAssertionStore(cache=cache)
|
||||
first = await store.fetch("alice")
|
||||
second = await store.fetch("alice")
|
||||
assert first is not None
|
||||
assert second is not None
|
||||
assert first.id_token.get_secret_value() == token
|
||||
assert second.id_token.get_secret_value() == token
|
||||
prisma.db.litellm_ssoidentityassertion.find_unique.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_db_store_returns_none_for_a_user_with_no_stored_assertion():
|
||||
"""An absent row stays an absence, not an outage, so a user who never signed in still gets the
|
||||
|
|
|
|||
|
|
@ -125,11 +125,12 @@ class TestBlockedResponseUsage:
|
|||
mock_logging.post_call_failure_hook.assert_awaited_once()
|
||||
|
||||
|
||||
class TestProxyExceptionPassthrough:
|
||||
class TestProxyExceptionAnthropicEnvelope:
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_response_reraises_proxy_exception_unwrapped(self):
|
||||
"""A 400 ProxyException from request validation must surface as-is,
|
||||
not be re-wrapped into a code-500 ProxyException."""
|
||||
async def test_anthropic_response_maps_proxy_exception_to_anthropic_envelope(self):
|
||||
"""LIT-6468: a 400 ProxyException from request validation must surface as
|
||||
Anthropic's documented {"type": "error", "error": {...}} envelope with the
|
||||
original status and message, not the OpenAI {"error": {...}} envelope."""
|
||||
import litellm.proxy.anthropic_endpoints.endpoints as ep
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
|
|
@ -140,6 +141,8 @@ class TestProxyExceptionPassthrough:
|
|||
param="metadata",
|
||||
code=400,
|
||||
)
|
||||
request = MagicMock()
|
||||
request.headers = {"x-request-id": "req_test_6468"}
|
||||
|
||||
with (
|
||||
patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})),
|
||||
|
|
@ -151,30 +154,61 @@ class TestProxyExceptionPassthrough:
|
|||
patch.object(proxy_server, "proxy_logging_obj") as mock_logging,
|
||||
):
|
||||
mock_logging.post_call_failure_hook = AsyncMock()
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await ep.anthropic_response(
|
||||
fastapi_response=MagicMock(),
|
||||
request=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
)
|
||||
response = await ep.anthropic_response(
|
||||
fastapi_response=MagicMock(),
|
||||
request=request,
|
||||
user_api_key_dict=MagicMock(),
|
||||
)
|
||||
|
||||
assert exc_info.value is exc
|
||||
assert exc_info.value.code == "400"
|
||||
assert exc_info.value.param == "metadata"
|
||||
assert response.status_code == 400
|
||||
body = json.loads(response.body)
|
||||
assert body == {
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"message": "Invalid type for 'metadata': expected an object, but got a string instead.",
|
||||
},
|
||||
"request_id": "req_test_6468",
|
||||
}
|
||||
mock_logging.post_call_failure_hook.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_response_maps_429_to_rate_limit_error(self):
|
||||
"""The Anthropic error type follows the status code (429 -> rate_limit_error),
|
||||
and a code-less exception falls back to 500 api_error."""
|
||||
import litellm.proxy.anthropic_endpoints.endpoints as ep
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
request = MagicMock()
|
||||
request.headers = {}
|
||||
|
||||
response = ep._anthropic_error_json_response(
|
||||
ProxyException(message="Rate limit exceeded", type="rate_limit_error", param=None, code=429),
|
||||
request,
|
||||
)
|
||||
assert response.status_code == 429
|
||||
assert json.loads(response.body)["error"]["type"] == "rate_limit_error"
|
||||
|
||||
fallback = ep._anthropic_error_json_response(
|
||||
ProxyException(message="boom", type="None", param=None, code=None),
|
||||
request,
|
||||
)
|
||||
assert fallback.status_code == 500
|
||||
assert json.loads(fallback.body)["error"]["type"] == "api_error"
|
||||
|
||||
|
||||
class TestHttpExceptionDictDetail:
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_response_serializes_dict_detail_http_exception(self):
|
||||
"""LIT-6466: a post_call guardrail's HTTPException(detail=<dict>) must
|
||||
surface with a clean message plus provider_specific_fields, matching
|
||||
/v1/chat/completions and /v1/responses, not the str() of the exception."""
|
||||
"""LIT-6466 + LIT-6468: a post_call guardrail's HTTPException(detail=<dict>)
|
||||
must surface as Anthropic's {"type": "error", "error": {...}} envelope with
|
||||
the guardrail's clean message plus provider_specific_fields, not the str()
|
||||
of the exception and not the OpenAI envelope."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm.proxy.anthropic_endpoints.endpoints as ep
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
detail = {
|
||||
"error": "Content blocked: keyword 'kumquat' detected",
|
||||
|
|
@ -182,6 +216,8 @@ class TestHttpExceptionDictDetail:
|
|||
"guardrail": "keyword-block",
|
||||
}
|
||||
exc = HTTPException(status_code=400, detail=detail)
|
||||
request = MagicMock()
|
||||
request.headers = {}
|
||||
|
||||
with (
|
||||
patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})), # test-quality-ok: endpoint reads the body via a module function; no injection seam
|
||||
|
|
@ -193,17 +229,19 @@ class TestHttpExceptionDictDetail:
|
|||
patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
|
||||
):
|
||||
mock_logging.post_call_failure_hook = AsyncMock()
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await ep.anthropic_response(
|
||||
fastapi_response=MagicMock(),
|
||||
request=MagicMock(),
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
response = await ep.anthropic_response(
|
||||
fastapi_response=MagicMock(),
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
assert exc_info.value.message == "Content blocked: keyword 'kumquat' detected"
|
||||
assert "{'error'" not in exc_info.value.message
|
||||
assert exc_info.value.provider_specific_fields == detail
|
||||
assert exc_info.value.code == "400"
|
||||
assert response.status_code == 400
|
||||
body = json.loads(response.body)
|
||||
assert body["type"] == "error"
|
||||
assert body["error"]["type"] == "invalid_request_error"
|
||||
assert body["error"]["message"] == "Content blocked: keyword 'kumquat' detected"
|
||||
assert "{'error'" not in body["error"]["message"]
|
||||
assert body["error"]["provider_specific_fields"] == detail
|
||||
mock_logging.post_call_failure_hook.assert_awaited_once()
|
||||
|
||||
|
||||
|
|
@ -215,7 +253,7 @@ class TestFailureHookRequestData:
|
|||
handler must pass that replaced dict, not the raw request body dict."""
|
||||
import litellm.proxy.anthropic_endpoints.endpoints as ep
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
captured = {}
|
||||
|
||||
|
|
@ -224,18 +262,23 @@ class TestFailureHookRequestData:
|
|||
captured["processor_data"] = self.data
|
||||
raise RuntimeError("provider timeout")
|
||||
|
||||
request = MagicMock()
|
||||
request.headers = {}
|
||||
|
||||
with (
|
||||
patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})),
|
||||
patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process),
|
||||
patch.object(proxy_server, "proxy_logging_obj") as mock_logging,
|
||||
):
|
||||
mock_logging.post_call_failure_hook = AsyncMock()
|
||||
with pytest.raises(ProxyException):
|
||||
await ep.anthropic_response(
|
||||
fastapi_response=MagicMock(),
|
||||
request=MagicMock(),
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
response = await ep.anthropic_response(
|
||||
fastapi_response=MagicMock(),
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
assert response.status_code == 500
|
||||
assert json.loads(response.body)["error"]["message"] == "provider timeout"
|
||||
|
||||
hook_request_data = mock_logging.post_call_failure_hook.await_args.kwargs["request_data"]
|
||||
assert hook_request_data is captured["processor_data"]
|
||||
|
|
|
|||
|
|
@ -1238,6 +1238,18 @@ def test_health_liveness_endpoint(proxy_client):
|
|||
print(f"\n/health/liveness response time: {duration_ms:.2f}ms")
|
||||
|
||||
|
||||
def test_health_backlog_includes_admission_control_stats(proxy_client):
|
||||
response = proxy_client.get("/health/backlog")
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert set(response.json()) == {
|
||||
"in_flight_requests",
|
||||
"admitted_requests",
|
||||
"queued_requests",
|
||||
"rejected_requests",
|
||||
}
|
||||
|
||||
|
||||
def test_health_readiness(proxy_client):
|
||||
"""
|
||||
Test /health/readiness endpoint.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,402 @@
|
|||
import asyncio
|
||||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
from litellm.proxy.middleware.admission_control_middleware import (
|
||||
AdmissionControlMetrics,
|
||||
AdmissionControlMiddleware,
|
||||
AdmissionControlSettings,
|
||||
AdmissionControlState,
|
||||
AdmissionControlStats,
|
||||
_parse_admission_control_settings,
|
||||
create_prometheus_admission_metrics,
|
||||
get_admission_control_settings,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def state() -> AdmissionControlState:
|
||||
return AdmissionControlState(lambda: None)
|
||||
|
||||
|
||||
async def _call(
|
||||
middleware: AdmissionControlMiddleware,
|
||||
path: str = "/",
|
||||
root_path: str = "",
|
||||
) -> tuple[Message, ...]:
|
||||
messages: Final[list[Message]] = []
|
||||
|
||||
async def receive() -> Message:
|
||||
return {"type": "http.request", "body": b"", "more_body": False}
|
||||
|
||||
async def send(message: Message) -> None:
|
||||
messages.append(message)
|
||||
|
||||
scope: Final[Scope] = {
|
||||
"type": "http",
|
||||
"path": path,
|
||||
"root_path": root_path,
|
||||
"method": "GET",
|
||||
"headers": [],
|
||||
}
|
||||
await middleware(scope, receive, send)
|
||||
return tuple(messages)
|
||||
|
||||
|
||||
def _handler_with_release(
|
||||
started: asyncio.Event,
|
||||
release: asyncio.Event,
|
||||
) -> ASGIApp:
|
||||
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
started.set()
|
||||
await release.wait()
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b"ok", "more_body": False})
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def test_is_not_base_http_middleware() -> None:
|
||||
assert not issubclass(AdmissionControlMiddleware, BaseHTTPMiddleware)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_capacity_rejects_excess_and_releases_queued_request(state: AdmissionControlState) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
_handler_with_release(started, release),
|
||||
lambda: AdmissionControlSettings(1, 1, 1.0),
|
||||
state,
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
second: Final = asyncio.create_task(_call(middleware))
|
||||
await asyncio.sleep(0)
|
||||
assert state.get_stats().queued == 1
|
||||
|
||||
third: Final = await _call(middleware)
|
||||
assert third[0]["status"] == 503
|
||||
headers: Final = dict(third[0]["headers"])
|
||||
assert headers[b"retry-after"] == b"1"
|
||||
assert headers[b"content-type"] == b"application/json"
|
||||
assert json.loads(third[1]["body"])["error"] == {
|
||||
"message": "Worker at capacity: 1 in-flight, 1 queued requests. Retry later.",
|
||||
"type": "overloaded_error",
|
||||
"code": "503",
|
||||
}
|
||||
assert state.get_stats().rejected_total == 1
|
||||
|
||||
release.set()
|
||||
assert (await first)[0]["status"] == 200
|
||||
assert (await second)[0]["status"] == 200
|
||||
assert state.get_stats() == AdmissionControlStats(0, 0, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_waiter_is_not_skipped_after_admission_is_released(state: AdmissionControlState) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
third_trigger: Final = asyncio.Event()
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
_handler_with_release(started, release),
|
||||
lambda: AdmissionControlSettings(1, 2, 1.0),
|
||||
state,
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
second: Final = asyncio.create_task(_call(middleware))
|
||||
await asyncio.sleep(0)
|
||||
|
||||
async def call_third() -> tuple[Message, ...]:
|
||||
await third_trigger.wait()
|
||||
return await _call(middleware)
|
||||
|
||||
third: Final = asyncio.create_task(call_third())
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
third_trigger.set()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert state.get_stats().queued == 2
|
||||
await asyncio.gather(first, second, third)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_timeout_rejects_and_decrements_queue(state: AdmissionControlState) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
_handler_with_release(started, release),
|
||||
lambda: AdmissionControlSettings(1, 1, 0.05),
|
||||
state,
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
start_time: Final = asyncio.get_running_loop().time()
|
||||
second: Final = await _call(middleware)
|
||||
elapsed: Final = asyncio.get_running_loop().time() - start_time
|
||||
|
||||
assert second[0]["status"] == 503
|
||||
assert elapsed < 0.5
|
||||
assert state.get_stats().queued == 0
|
||||
assert state.get_stats().rejected_total == 1
|
||||
release.set()
|
||||
await first
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("root_path", "probe_path"),
|
||||
(
|
||||
("", "/health/liveliness"),
|
||||
("/proxy", "/proxy/health/liveliness"),
|
||||
("/proxy", "/proxy/metrics"),
|
||||
),
|
||||
)
|
||||
async def test_exempt_path_passes_through_when_saturated(
|
||||
state: AdmissionControlState,
|
||||
root_path: str,
|
||||
probe_path: str,
|
||||
) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
|
||||
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["path"] == "/":
|
||||
started.set()
|
||||
await release.wait()
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b"ok", "more_body": False})
|
||||
|
||||
middleware: Final = AdmissionControlMiddleware(handler, lambda: AdmissionControlSettings(1, 0, 1.0), state)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
health: Final = await _call(middleware, probe_path, root_path)
|
||||
assert health[0]["status"] == 200
|
||||
blocked: Final = await _call(middleware, "/proxy/v1/chat/completions", root_path)
|
||||
assert blocked[0]["status"] == 503
|
||||
lookalike: Final = await _call(middleware, "/proxyhealth/liveliness", "/proxy")
|
||||
assert lookalike[0]["status"] == 503
|
||||
release.set()
|
||||
await first
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_http_scope_passes_through_when_saturated(state: AdmissionControlState) -> None:
|
||||
seen: Final[list[str]] = []
|
||||
|
||||
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
seen.append(scope["type"])
|
||||
|
||||
middleware: Final = AdmissionControlMiddleware(handler, lambda: AdmissionControlSettings(1, 0, 1.0), state)
|
||||
state.record_admission()
|
||||
|
||||
async def receive() -> Message:
|
||||
return {"type": "lifespan.startup"}
|
||||
|
||||
async def send(message: Message) -> None:
|
||||
return None
|
||||
|
||||
await middleware({"type": "lifespan"}, receive, send)
|
||||
assert seen == ["lifespan"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_settings_does_not_limit_concurrency() -> None:
|
||||
active: Final = [0]
|
||||
peak: Final = [0]
|
||||
all_started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
|
||||
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
active[0] += 1
|
||||
peak[0] = max(peak[0], active[0])
|
||||
if active[0] == 3:
|
||||
all_started.set()
|
||||
await release.wait()
|
||||
active[0] -= 1
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b"ok", "more_body": False})
|
||||
|
||||
middleware: Final = AdmissionControlMiddleware(handler, lambda: None, AdmissionControlState(lambda: None))
|
||||
requests: Final = tuple(asyncio.create_task(_call(middleware)) for _ in range(3))
|
||||
await all_started.wait()
|
||||
assert peak[0] == 3
|
||||
release.set()
|
||||
results: Final = await asyncio.gather(*requests)
|
||||
assert tuple(result[0]["status"] for result in results) == (200, 200, 200)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelling_queued_request_does_not_leak_counter(state: AdmissionControlState) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
_handler_with_release(started, release),
|
||||
lambda: AdmissionControlSettings(1, 1, 1.0),
|
||||
state,
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
queued: Final = asyncio.create_task(_call(middleware))
|
||||
await asyncio.sleep(0)
|
||||
queued.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await queued
|
||||
assert state.get_stats().queued == 0
|
||||
release.set()
|
||||
await first
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_response_holds_admission_until_final_body(state: AdmissionControlState) -> None:
|
||||
first_chunk_sent: Final = asyncio.Event()
|
||||
finish_stream: Final = asyncio.Event()
|
||||
|
||||
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b"first", "more_body": True})
|
||||
first_chunk_sent.set()
|
||||
await finish_stream.wait()
|
||||
await send({"type": "http.response.body", "body": b"last", "more_body": False})
|
||||
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
handler,
|
||||
lambda: AdmissionControlSettings(1, 1, 1.0),
|
||||
state,
|
||||
)
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await first_chunk_sent.wait()
|
||||
second: Final = asyncio.create_task(_call(middleware))
|
||||
await asyncio.sleep(0)
|
||||
assert not second.done()
|
||||
assert state.get_stats().queued == 1
|
||||
finish_stream.set()
|
||||
assert (await first)[0]["status"] == 200
|
||||
assert (await second)[0]["status"] == 200
|
||||
assert state.get_stats().admitted == 0
|
||||
assert state.get_stats().queued == 0
|
||||
|
||||
|
||||
class _FakeGauge:
|
||||
def __init__(self) -> None:
|
||||
self.value = 0.0
|
||||
|
||||
def inc(self, amount: float = 1) -> None:
|
||||
self.value += amount
|
||||
|
||||
def dec(self, amount: float = 1) -> None:
|
||||
self.value -= amount
|
||||
|
||||
|
||||
class _FakeCounter:
|
||||
def __init__(self) -> None:
|
||||
self.by_reason: Final[dict[str, _FakeGauge]] = {}
|
||||
|
||||
def labels(self, reason: str) -> _FakeGauge:
|
||||
return self.by_reason.setdefault(reason, _FakeGauge())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metrics_track_admitted_queued_and_rejected() -> None:
|
||||
admitted: Final = _FakeGauge()
|
||||
queued: Final = _FakeGauge()
|
||||
rejected: Final = _FakeCounter()
|
||||
state: Final = AdmissionControlState(
|
||||
lambda: AdmissionControlMetrics(admitted_gauge=admitted, queued_gauge=queued, rejected_counter=rejected)
|
||||
)
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
middleware: Final = AdmissionControlMiddleware(
|
||||
_handler_with_release(started, release),
|
||||
lambda: AdmissionControlSettings(1, 1, 0.05),
|
||||
state,
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(_call(middleware))
|
||||
await started.wait()
|
||||
second: Final = asyncio.create_task(_call(middleware))
|
||||
await asyncio.sleep(0)
|
||||
assert (admitted.value, queued.value) == (1.0, 1.0)
|
||||
await _call(middleware)
|
||||
assert rejected.by_reason["queue_full"].value == 1.0
|
||||
await second
|
||||
assert rejected.by_reason["queue_timeout"].value == 1.0
|
||||
release.set()
|
||||
await first
|
||||
assert (admitted.value, queued.value) == (0.0, 0.0)
|
||||
|
||||
|
||||
def test_create_prometheus_admission_metrics_registers_named_metrics() -> None:
|
||||
from prometheus_client import REGISTRY
|
||||
|
||||
metrics: Final = create_prometheus_admission_metrics()
|
||||
if metrics is not None:
|
||||
metrics.admitted_gauge.inc()
|
||||
metrics.queued_gauge.inc()
|
||||
metrics.rejected_counter.labels(reason="queue_full").inc()
|
||||
assert REGISTRY.get_sample_value("litellm_admission_admitted_requests") == 1.0
|
||||
assert REGISTRY.get_sample_value("litellm_admission_queued_requests") == 1.0
|
||||
assert REGISTRY.get_sample_value("litellm_admission_rejected_requests_total", {"reason": "queue_full"}) is not None
|
||||
assert create_prometheus_admission_metrics() is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("settings", "expected"),
|
||||
(
|
||||
({}, None),
|
||||
({"max_in_flight_requests_per_worker": None}, None),
|
||||
({"max_in_flight_requests_per_worker": 0}, None),
|
||||
({"max_in_flight_requests_per_worker": "many"}, None),
|
||||
({"max_in_flight_requests_per_worker": 3, "max_queued_requests_per_worker": -1}, None),
|
||||
({"max_in_flight_requests_per_worker": 3, "admission_queue_timeout_seconds": 0}, None),
|
||||
({"max_in_flight_requests_per_worker": 3, "admission_queue_timeout_seconds": -0.5}, None),
|
||||
(
|
||||
{"max_in_flight_requests_per_worker": 3, "max_queued_requests_per_worker": 0},
|
||||
AdmissionControlSettings(3, 0, 1.0),
|
||||
),
|
||||
(
|
||||
{"max_in_flight_requests_per_worker": 3},
|
||||
AdmissionControlSettings(3, 3, 1.0),
|
||||
),
|
||||
(
|
||||
{
|
||||
"max_in_flight_requests_per_worker": 3,
|
||||
"max_queued_requests_per_worker": 5,
|
||||
"admission_queue_timeout_seconds": 0.25,
|
||||
},
|
||||
AdmissionControlSettings(3, 5, 0.25),
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_get_admission_control_settings(
|
||||
settings: dict[str, object],
|
||||
expected: AdmissionControlSettings | None,
|
||||
) -> None:
|
||||
assert get_admission_control_settings(settings) == expected
|
||||
|
||||
|
||||
def test_invalid_admission_control_settings_logs_once(caplog: pytest.LogCaptureFixture) -> None:
|
||||
_parse_admission_control_settings.cache_clear()
|
||||
caplog.set_level("ERROR")
|
||||
settings: Final = {"max_in_flight_requests_per_worker": [1]}
|
||||
|
||||
assert get_admission_control_settings(settings) is None
|
||||
assert get_admission_control_settings(settings) is None
|
||||
|
||||
messages: Final = tuple(
|
||||
record.message
|
||||
for record in caplog.records
|
||||
if record.message.startswith("Ignoring invalid admission control settings")
|
||||
)
|
||||
assert len(messages) == 1
|
||||
|
|
@ -319,9 +319,6 @@ class TestVertexAIPassThroughHandler:
|
|||
mock_handler.get_default_base_target_url.return_value = (
|
||||
f"https://{test_location}-aiplatform.googleapis.com/"
|
||||
)
|
||||
mock_handler.update_base_target_url_with_credential_location = Mock(
|
||||
return_value=f"https://{test_location}-aiplatform.googleapis.com/"
|
||||
)
|
||||
mock_get_handler.return_value = mock_handler
|
||||
|
||||
# Mock create_pass_through_route to return a function that returns a mock response
|
||||
|
|
@ -427,9 +424,6 @@ class TestVertexAIPassThroughHandler:
|
|||
mock_handler.get_default_base_target_url.return_value = (
|
||||
"https://aiplatform.googleapis.com/"
|
||||
)
|
||||
mock_handler.update_base_target_url_with_credential_location = Mock(
|
||||
return_value="https://aiplatform.googleapis.com/"
|
||||
)
|
||||
mock_get_handler.return_value = mock_handler
|
||||
|
||||
# Mock create_pass_through_route to return a function that returns a mock response
|
||||
|
|
@ -530,9 +524,6 @@ class TestVertexAIPassThroughHandler:
|
|||
mock_handler.get_default_base_target_url.return_value = (
|
||||
f"https://{default_location}-aiplatform.googleapis.com/"
|
||||
)
|
||||
mock_handler.update_base_target_url_with_credential_location = Mock(
|
||||
return_value=f"https://{default_location}-aiplatform.googleapis.com/"
|
||||
)
|
||||
mock_get_handler.return_value = mock_handler
|
||||
|
||||
# Mock create_pass_through_route to return a function that returns a mock response
|
||||
|
|
@ -1308,9 +1299,6 @@ class TestVertexAIDiscoveryPassThroughHandler:
|
|||
mock_handler.get_default_base_target_url.return_value = (
|
||||
"https://discoveryengine.googleapis.com"
|
||||
)
|
||||
mock_handler.update_base_target_url_with_credential_location = Mock(
|
||||
return_value="https://discoveryengine.googleapis.com"
|
||||
)
|
||||
mock_get_handler.return_value = mock_handler
|
||||
|
||||
# Mock create_pass_through_route to return a function that returns a mock response
|
||||
|
|
@ -3650,7 +3638,6 @@ class TestVertexRawPredictStreamingClassification:
|
|||
base_url = "https://us-east5-aiplatform.googleapis.com/"
|
||||
mock_handler = Mock()
|
||||
mock_handler.get_default_base_target_url.return_value = base_url
|
||||
mock_handler.update_base_target_url_with_credential_location = Mock(return_value=base_url)
|
||||
|
||||
module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
|
||||
with (
|
||||
|
|
@ -4234,6 +4221,126 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak:
|
|||
assert "sk-master-1234" not in " ".join(f"{name}:{value}" for name, value in forwarded.items())
|
||||
|
||||
|
||||
class TestVertexPassthroughDefaultLocationOnShortRoutes:
|
||||
PROJECT = "test-project"
|
||||
SHORT_ROUTE = "publishers/google/models/gemini-2.5-flash:generateContent"
|
||||
|
||||
@staticmethod
|
||||
def _forwarder() -> Mock:
|
||||
return Mock(return_value=AsyncMock(return_value={"status": "success"}))
|
||||
|
||||
async def _forward(
|
||||
self,
|
||||
monkeypatch,
|
||||
endpoint: str,
|
||||
default_config: dict | None,
|
||||
headers: list[tuple[bytes, bytes]],
|
||||
forwarder: Mock,
|
||||
) -> None:
|
||||
from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import (
|
||||
PassthroughEndpointRouter,
|
||||
)
|
||||
|
||||
async def receive():
|
||||
return {"type": "http.request", "body": b"{}", "more_body": False}
|
||||
|
||||
request: Final = Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": f"/vertex_ai/{endpoint}",
|
||||
"headers": headers,
|
||||
"query_string": b"",
|
||||
},
|
||||
receive=receive,
|
||||
)
|
||||
router: Final = PassthroughEndpointRouter()
|
||||
if default_config is not None:
|
||||
router.set_default_vertex_config(dict(default_config))
|
||||
module: Final = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
|
||||
monkeypatch.setattr(f"{module}.passthrough_endpoint_router", router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
mock_credentials: Final = Mock()
|
||||
mock_credentials.token = "test-token"
|
||||
caller: Final = UserAPIKeyAuth(api_key="test-key")
|
||||
with (
|
||||
mock.patch( # test-quality-ok: the route mints its Google token through its own VertexBase, nothing injects the credential loader
|
||||
"litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth",
|
||||
return_value=(mock_credentials, self.PROJECT),
|
||||
),
|
||||
mock.patch(f"{module}.create_pass_through_route", new=forwarder),
|
||||
mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=caller)),
|
||||
):
|
||||
await vertex_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
fastapi_response=Response(),
|
||||
user_api_key_dict=caller,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("endpoint", "location", "expected_target"),
|
||||
[
|
||||
(
|
||||
SHORT_ROUTE,
|
||||
"global",
|
||||
"https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/" + SHORT_ROUTE,
|
||||
),
|
||||
(
|
||||
f"v1/{SHORT_ROUTE}",
|
||||
"global",
|
||||
"https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/" + SHORT_ROUTE,
|
||||
),
|
||||
(
|
||||
f"v1beta1/{SHORT_ROUTE}",
|
||||
"global",
|
||||
"https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/" + SHORT_ROUTE,
|
||||
),
|
||||
(
|
||||
SHORT_ROUTE,
|
||||
"us-central1",
|
||||
"https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/"
|
||||
+ SHORT_ROUTE,
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_default_vertex_config_location_fills_routes_without_project_and_location(
|
||||
self, monkeypatch, endpoint, location, expected_target
|
||||
):
|
||||
forwarder: Final = self._forwarder()
|
||||
await self._forward(
|
||||
monkeypatch,
|
||||
endpoint,
|
||||
{"vertex_project": self.PROJECT, "vertex_location": location, "vertex_credentials": "test-creds"},
|
||||
[(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")],
|
||||
forwarder,
|
||||
)
|
||||
forwarded: Final = forwarder.call_args.kwargs
|
||||
assert str(forwarded["target"]) == expected_target
|
||||
assert forwarded["custom_headers"]["Authorization"] == "Bearer test-token"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("default_config", "headers"),
|
||||
[
|
||||
(None, [(b"content-type", b"application/json"), (b"authorization", b"Bearer ya29.byo-google-oauth")]),
|
||||
(
|
||||
{"vertex_project": PROJECT, "vertex_credentials": "test-creds"},
|
||||
[(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")],
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_no_location_anywhere_is_a_400_not_a_500(self, monkeypatch, default_config, headers):
|
||||
forwarder: Final = self._forwarder()
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await self._forward(monkeypatch, self.SHORT_ROUTE, default_config, headers, forwarder)
|
||||
forwarder.assert_not_called()
|
||||
assert raised.value.status_code == 400
|
||||
assert "/projects/<project>/locations/<location>/" in str(raised.value.detail)
|
||||
assert "default_vertex_config" in str(raised.value.detail)
|
||||
|
||||
|
||||
class TestGetAzureAISearchIndexFromEndpoint:
|
||||
"""The operable index is only the segment right after ``indexes``.
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import pytest
|
|||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
VertexAIPassThroughHandler,
|
||||
_base_vertex_proxy_route,
|
||||
_upstream_headers_for_vertex_route,
|
||||
)
|
||||
|
|
@ -20,6 +21,7 @@ async def test_vertex_passthrough_load_balancing():
|
|||
mock_request = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_handler = MagicMock()
|
||||
mock_handler.get_default_base_target_url.return_value = "https://test.url"
|
||||
|
||||
# Mock the router
|
||||
mock_router = MagicMock()
|
||||
|
|
@ -68,7 +70,6 @@ async def test_vertex_passthrough_load_balancing():
|
|||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
mock_prep_headers.return_value = (
|
||||
{},
|
||||
"https://test.url",
|
||||
False,
|
||||
"test-project-lb",
|
||||
"us-central1-lb",
|
||||
|
|
@ -290,12 +291,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
|
|||
mock_vertex_credentials.vertex_location = "us-central1"
|
||||
mock_vertex_credentials.vertex_credentials = "test-credentials"
|
||||
|
||||
# Create mock handler
|
||||
mock_handler = MagicMock()
|
||||
mock_handler.update_base_target_url_with_credential_location.return_value = (
|
||||
"https://us-central1-aiplatform.googleapis.com"
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
VertexBase,
|
||||
|
|
@ -313,7 +308,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
|
|||
# Call the function
|
||||
(
|
||||
headers,
|
||||
base_target_url,
|
||||
headers_passed_through,
|
||||
vertex_project,
|
||||
vertex_location,
|
||||
|
|
@ -323,8 +317,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
|
|||
router_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
base_target_url="https://us-central1-aiplatform.googleapis.com",
|
||||
get_vertex_pass_through_handler=mock_handler,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"),
|
||||
)
|
||||
|
||||
|
|
@ -394,7 +386,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
|
|||
"content-type": "application/json",
|
||||
"Authorization": "Bearer vertex-access-token",
|
||||
},
|
||||
"https://aiplatform.googleapis.com",
|
||||
False,
|
||||
"test-project",
|
||||
"global",
|
||||
|
|
@ -406,7 +397,7 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
|
|||
endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}",
|
||||
request=MagicMock(),
|
||||
fastapi_response=MagicMock(),
|
||||
get_vertex_pass_through_handler=MagicMock(),
|
||||
get_vertex_pass_through_handler=VertexAIPassThroughHandler(),
|
||||
)
|
||||
|
||||
upstream_headers = mock_create_route.call_args.kwargs["custom_headers"]
|
||||
|
|
@ -473,12 +464,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
|
|||
mock_vertex_credentials.vertex_location = "us-central1"
|
||||
mock_vertex_credentials.vertex_credentials = "test-credentials"
|
||||
|
||||
# Create mock handler
|
||||
mock_handler = MagicMock()
|
||||
mock_handler.update_base_target_url_with_credential_location.return_value = (
|
||||
"https://us-central1-aiplatform.googleapis.com"
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
VertexBase,
|
||||
|
|
@ -495,7 +480,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
|
|||
|
||||
(
|
||||
headers,
|
||||
_base_target_url,
|
||||
_headers_passed_through,
|
||||
_vertex_project,
|
||||
_vertex_location,
|
||||
|
|
@ -505,8 +489,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
|
|||
router_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
base_target_url="https://us-central1-aiplatform.googleapis.com",
|
||||
get_vertex_pass_through_handler=mock_handler,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"),
|
||||
)
|
||||
|
||||
|
|
@ -742,7 +724,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
|
|||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
mock_prep_headers.return_value = (
|
||||
{},
|
||||
"https://global-aiplatform.googleapis.com",
|
||||
False,
|
||||
"nv-gcpllmgwit-20250411173346",
|
||||
"global",
|
||||
|
|
|
|||
|
|
@ -5,14 +5,12 @@ import hashlib
|
|||
import json
|
||||
import re
|
||||
from datetime import timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import litellm
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
|
|
@ -3325,7 +3323,7 @@ def _compare_nested_dicts(
|
|||
return differences
|
||||
|
||||
# Check for keys in actual but not in expected
|
||||
for key in actual.keys():
|
||||
for key in actual:
|
||||
current_path = f"{path}.{key}" if path else key
|
||||
if current_path not in ignore_keys and key not in expected:
|
||||
differences.append(f"Extra key in actual: {current_path}")
|
||||
|
|
@ -3495,24 +3493,22 @@ async def test_view_spend_logs_summarize_parameter(client, monkeypatch):
|
|||
# Return individual log entries when summarize=false
|
||||
return mock_spend_logs
|
||||
|
||||
async def group_by(self, *args, **kwargs):
|
||||
# Return grouped data when summarize=true
|
||||
# Simplified mock response for grouped data
|
||||
async def query_raw(self, sql_query, *params):
|
||||
yesterday = datetime.datetime.now(timezone.utc) - timedelta(days=1)
|
||||
return [
|
||||
{
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"model": "gpt-3.5-turbo",
|
||||
"startTime": yesterday.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
|
||||
"_sum": {"spend": 0.05},
|
||||
"day": yesterday.date().isoformat(),
|
||||
"spend": 0.05,
|
||||
},
|
||||
{
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"model": "gpt-4",
|
||||
"startTime": yesterday.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
|
||||
"_sum": {"spend": 0.10},
|
||||
"day": yesterday.date().isoformat(),
|
||||
"spend": 0.10,
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -3850,47 +3846,30 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
|
|||
"""
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
# This simulates the summarized data that Prisma's `group_by` would return.
|
||||
mock_summarized_response = [
|
||||
{
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"model": "gpt-4",
|
||||
"startTime": (datetime.now(timezone.utc) - timedelta(days=1)).strftime(
|
||||
"%Y-%m-%dT%H:%M:%S.%fZ"
|
||||
),
|
||||
"_sum": {"spend": 0.15},
|
||||
"day": (datetime.now(timezone.utc) - timedelta(days=1)).date().isoformat(),
|
||||
"spend": 0.15,
|
||||
}
|
||||
]
|
||||
|
||||
# This mock class will replace the real Prisma client.
|
||||
class MockDB:
|
||||
def __init__(self):
|
||||
self.litellm_spendlogs = self
|
||||
|
||||
async def group_by(self, *args, **kwargs):
|
||||
# We assert that the `gte` and `lte` values are strings in ISO format.
|
||||
# If they were datetime objects, this test would fail.
|
||||
where_clause = kwargs.get("where", {})
|
||||
start_time_filter = where_clause.get("startTime", {})
|
||||
|
||||
assert "gte" in start_time_filter
|
||||
assert "lte" in start_time_filter
|
||||
assert isinstance(start_time_filter["gte"], str)
|
||||
assert isinstance(start_time_filter["lte"], str)
|
||||
assert "T" in start_time_filter["gte"] # Check for ISO format 'T' separator
|
||||
|
||||
# If the assertions pass, return the mock response.
|
||||
async def query_raw(self, sql_query, *params):
|
||||
assert isinstance(params[0], str)
|
||||
assert isinstance(params[1], str)
|
||||
assert "T" in params[0]
|
||||
assert "T" in params[1]
|
||||
return mock_summarized_response
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
# Apply the monkeypatch to replace the real prisma_client with our mock.
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient())
|
||||
|
||||
# Define a date range for the test.
|
||||
start_date = (datetime.now(timezone.utc) - timedelta(days=2)).strftime("%Y-%m-%d")
|
||||
end_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||||
|
||||
|
|
@ -3898,8 +3877,6 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
|
|||
user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
try:
|
||||
# Call the endpoint with both start and end dates.
|
||||
# We don't need `summarize=true` as it's the default.
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={
|
||||
|
|
@ -3909,11 +3886,9 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
|
|||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
# ASSERTIONS
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Check that the response is not empty and has the summarized structure.
|
||||
assert isinstance(data, list)
|
||||
assert len(data) > 0
|
||||
assert "startTime" in data[0]
|
||||
|
|
@ -3924,6 +3899,183 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
|
|||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_view_spend_logs_summarize_groups_by_day_in_sql(client, monkeypatch):
|
||||
mock_rows = [
|
||||
{
|
||||
"day": "2024-01-01",
|
||||
"api_key": "hashed::sk-abc",
|
||||
"user": "u1",
|
||||
"model": "gpt-4",
|
||||
"spend": 0.1,
|
||||
},
|
||||
{
|
||||
"day": "2024-01-01",
|
||||
"api_key": "hashed::sk-abc",
|
||||
"user": "u1",
|
||||
"model": "gpt-4o",
|
||||
"spend": 0.2,
|
||||
},
|
||||
]
|
||||
|
||||
class MockDB:
|
||||
def __init__(self):
|
||||
self.captured_sql = None
|
||||
self.captured_params = None
|
||||
|
||||
async def query_raw(self, sql_query, *params):
|
||||
self.captured_sql = sql_query
|
||||
self.captured_params = params
|
||||
return mock_rows
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
def hash_token(self, token):
|
||||
return "hashed::" + token
|
||||
|
||||
mock_prisma_client = MockPrismaClient()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={
|
||||
"start_date": "2024-01-01",
|
||||
"end_date": "2024-01-03",
|
||||
"api_key": "sk-abc",
|
||||
"request_id": "req-123",
|
||||
"user_id": "u1",
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
sql = mock_prisma_client.db.captured_sql
|
||||
assert "date_trunc('day'" in sql
|
||||
assert "GROUP BY" in sql
|
||||
assert "find_many" not in sql
|
||||
assert not hasattr(mock_prisma_client.db, "group_by")
|
||||
assert mock_prisma_client.db.captured_params == (
|
||||
"2024-01-01T00:00:00+00:00",
|
||||
"2024-01-03T00:00:00+00:00",
|
||||
"hashed::sk-abc",
|
||||
"req-123",
|
||||
"u1",
|
||||
)
|
||||
assert len(data) == 3
|
||||
assert data[0]["startTime"] == "2024-01-01"
|
||||
assert data[0]["spend"] == pytest.approx(0.3)
|
||||
assert data[0]["models"] == {"gpt-4": 0.1, "gpt-4o": 0.2}
|
||||
assert data[0]["users"] == {"u1": pytest.approx(0.3)}
|
||||
assert data[0]["hashed::sk-abc"] == pytest.approx(0.3)
|
||||
assert data[1] == {
|
||||
"startTime": "2024-01-02",
|
||||
"spend": 0,
|
||||
"users": {},
|
||||
"models": {},
|
||||
}
|
||||
assert data[2] == {
|
||||
"startTime": "2024-01-03",
|
||||
"spend": 0,
|
||||
"users": {},
|
||||
"models": {},
|
||||
}
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_view_spend_logs_summarize_empty_rows(client, monkeypatch):
|
||||
class MockDB:
|
||||
async def query_raw(self, sql_query, *params):
|
||||
return []
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient())
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={"start_date": "2024-01-01", "end_date": "2024-01-01"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == []
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_view_spend_logs_summarize_unhashed_api_key_without_padding(client, monkeypatch):
|
||||
mock_rows = [
|
||||
{
|
||||
"day": "2024-01-01",
|
||||
"api_key": "plain-key",
|
||||
"user": "u1",
|
||||
"model": "gpt-4",
|
||||
"spend": 0.4,
|
||||
}
|
||||
]
|
||||
|
||||
class MockDB:
|
||||
def __init__(self):
|
||||
self.captured_params = None
|
||||
|
||||
async def query_raw(self, sql_query, *params):
|
||||
self.captured_params = params
|
||||
return mock_rows
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
mock_prisma_client = MockPrismaClient()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={
|
||||
"start_date": "2024-01-01",
|
||||
"end_date": "2024-01-01",
|
||||
"api_key": "plain-key",
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert mock_prisma_client.db.captured_params == (
|
||||
"2024-01-01T00:00:00+00:00",
|
||||
"2024-01-01T00:00:00+00:00",
|
||||
"plain-key",
|
||||
)
|
||||
assert data == [
|
||||
{
|
||||
"startTime": "2024-01-01",
|
||||
"spend": pytest.approx(0.4),
|
||||
"plain-key": pytest.approx(0.4),
|
||||
"users": {"u1": pytest.approx(0.4)},
|
||||
"models": {"gpt-4": pytest.approx(0.4)},
|
||||
}
|
||||
]
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_with_error_code(client):
|
||||
"""Test filtering spend logs by error code"""
|
||||
|
|
@ -4832,13 +4984,14 @@ class _CaptureFilterDB:
|
|||
def __init__(self):
|
||||
self.litellm_spendlogs = self
|
||||
self.captured_where = None
|
||||
self.captured_params = None
|
||||
|
||||
async def find_many(self, *args, **kwargs):
|
||||
self.captured_where = kwargs.get("where")
|
||||
return []
|
||||
|
||||
async def group_by(self, *args, **kwargs):
|
||||
self.captured_where = kwargs.get("where")
|
||||
async def query_raw(self, sql_query, *params):
|
||||
self.captured_params = params
|
||||
return []
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -5101,6 +5101,40 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment():
|
|||
)
|
||||
|
||||
|
||||
def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete():
|
||||
"""
|
||||
Regression for #29064: upsert_deployment removed the old deployment from
|
||||
model_list but left it in the global pattern_router, so wildcard requests
|
||||
round-robined between the stale and the corrected deployment.
|
||||
"""
|
||||
from litellm.types.router import Deployment, LiteLLM_Params
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/openai/*", "api_key": "sk-old"},
|
||||
"model_info": {"id": "global-wildcard"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
router.upsert_deployment(
|
||||
Deployment(
|
||||
model_name="openai/*",
|
||||
litellm_params=LiteLLM_Params(model="openai/*", api_key="sk-new"),
|
||||
model_info={"id": "global-wildcard"},
|
||||
)
|
||||
)
|
||||
|
||||
matches = router.pattern_router.route("openai/gpt-5.2")
|
||||
assert matches is not None
|
||||
assert [m["litellm_params"]["api_key"] for m in matches] == ["sk-new"]
|
||||
|
||||
router.delete_deployment(id="global-wildcard")
|
||||
assert router.pattern_router.patterns == {}
|
||||
|
||||
|
||||
def test_pattern_match_router_remove_deployment():
|
||||
"""
|
||||
remove_deployment must drop only the deployment with the given model id and
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
"limit": 22328
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 26758
|
||||
"limit": 26750
|
||||
},
|
||||
"LIT003": {
|
||||
"limit": 261
|
||||
|
|
@ -27,10 +27,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT010": {
|
||||
"limit": 16470
|
||||
"limit": 16468
|
||||
},
|
||||
"LIT011": {
|
||||
"limit": 5516
|
||||
"limit": 5514
|
||||
},
|
||||
"LIT012": {
|
||||
"limit": 4489
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ const TagTable: React.FC<TagTableProps> = ({ data, onEdit, onDelete, onSelectTag
|
|||
data={data}
|
||||
columns={columns}
|
||||
getRowId={(tag, index) => tag.name || String(index)}
|
||||
fillHeight
|
||||
sortingMode="client"
|
||||
sorting={sorting}
|
||||
onSortingChange={setSorting}
|
||||
|
|
|
|||
|
|
@ -126,7 +126,7 @@ const TagManagement: React.FC<TagProps> = ({ accessToken, userID, userRole }) =>
|
|||
}, [accessToken]);
|
||||
|
||||
return (
|
||||
<div className="mx-4 h-[75vh]">
|
||||
<div className="mx-4 h-full">
|
||||
{selectedTagId ? (
|
||||
<TagInfoView
|
||||
tagId={selectedTagId}
|
||||
|
|
@ -139,7 +139,7 @@ const TagManagement: React.FC<TagProps> = ({ accessToken, userID, userRole }) =>
|
|||
editTag={editTag}
|
||||
/>
|
||||
) : (
|
||||
<div className="mt-2 h-[75vh] w-full gap-2 p-8">
|
||||
<div className="flex h-full w-full flex-col p-8 pt-10">
|
||||
<div className="mt-2 mb-4 flex w-full items-center justify-between">
|
||||
<h1>Tag Management</h1>
|
||||
<div className="flex items-center space-x-2">
|
||||
|
|
@ -162,23 +162,21 @@ const TagManagement: React.FC<TagProps> = ({ accessToken, userID, userRole }) =>
|
|||
</p>
|
||||
</div>
|
||||
|
||||
<Button className="mb-4" onClick={() => setIsCreateModalVisible(true)}>
|
||||
<Button className="mb-4 self-start" onClick={() => setIsCreateModalVisible(true)}>
|
||||
+ Create New Tag
|
||||
</Button>
|
||||
|
||||
<div className="mt-2 grid h-[75vh] w-full grid-cols-1 gap-2 pt-2 pb-2">
|
||||
<div>
|
||||
<TagTable
|
||||
data={tags}
|
||||
isLoading={isLoadingTags}
|
||||
onEdit={(tag) => {
|
||||
setSelectedTagId(tag.name);
|
||||
setEditTag(true);
|
||||
}}
|
||||
onDelete={handleDelete}
|
||||
onSelectTag={setSelectedTagId}
|
||||
/>
|
||||
</div>
|
||||
<div className="mt-2 flex min-h-0 flex-1 flex-col">
|
||||
<TagTable
|
||||
data={tags}
|
||||
isLoading={isLoadingTags}
|
||||
onEdit={(tag) => {
|
||||
setSelectedTagId(tag.name);
|
||||
setEditTag(true);
|
||||
}}
|
||||
onDelete={handleDelete}
|
||||
onSelectTag={setSelectedTagId}
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* Create Tag Modal */}
|
||||
|
|
|
|||
|
|
@ -137,8 +137,8 @@ const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID
|
|||
/>
|
||||
</div>
|
||||
) : (
|
||||
<div className="mx-4 h-[75vh]">
|
||||
<div className="gap-2 p-8 h-[75vh] w-full mt-2">
|
||||
<div className="mx-4">
|
||||
<div className="gap-2 p-8 w-full mt-2">
|
||||
<div className="flex justify-between mt-2 w-full items-center mb-4">
|
||||
<h1 className="text-xl font-semibold tracking-tight text-foreground">Vector Store Management</h1>
|
||||
<div className="flex items-center space-x-2">
|
||||
|
|
|
|||
|
|
@ -400,7 +400,7 @@ const ModelHubTable: React.FC<ModelHubTableProps> = ({ accessToken, publicPage,
|
|||
}
|
||||
|
||||
return (
|
||||
<div className="mx-4 h-[75vh]">
|
||||
<div className="mx-4">
|
||||
{publicPage == false ? (
|
||||
<div className="w-full m-2 mt-2 p-8">
|
||||
{/* Header with Title, Description and URL */}
|
||||
|
|
|
|||
|
|
@ -559,6 +559,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
|
|||
{
|
||||
key: "your-teams",
|
||||
label: "Your Teams",
|
||||
className: "flex min-h-0 flex-1 flex-col",
|
||||
children: (
|
||||
<>
|
||||
<TeamsTable
|
||||
|
|
@ -608,6 +609,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
|
|||
{
|
||||
key: "available-teams",
|
||||
label: "Available Teams",
|
||||
className: "min-h-0 flex-1 overflow-y-auto",
|
||||
children: <AvailableTeamsPanel accessToken={accessToken} userID={userID} />,
|
||||
},
|
||||
...(isProxyAdminRole(userRole || "")
|
||||
|
|
@ -615,6 +617,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
|
|||
{
|
||||
key: "default-settings",
|
||||
label: "Default Team Settings",
|
||||
className: "min-h-0 flex-1 overflow-y-auto",
|
||||
children: <TeamSSOSettings accessToken={accessToken} userID={userID || ""} userRole={userRole || ""} />,
|
||||
},
|
||||
]
|
||||
|
|
@ -622,7 +625,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
|
|||
];
|
||||
|
||||
return (
|
||||
<main className={selectedTeamId ? "px-12 py-6" : "p-8"}>
|
||||
<main className={selectedTeamId ? "px-12 py-6" : "flex h-full flex-col p-8"}>
|
||||
{selectedTeamId ? (
|
||||
<TeamInfoView
|
||||
teamId={selectedTeamId}
|
||||
|
|
@ -642,7 +645,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
|
|||
premiumUser={premiumUser}
|
||||
/>
|
||||
) : (
|
||||
<Tabs defaultValue={tabItems[0].key} className="gap-6">
|
||||
<Tabs defaultValue={tabItems[0].key} className="min-h-0 flex-1 gap-6">
|
||||
<PageHeader
|
||||
icon={<Users />}
|
||||
title="Teams"
|
||||
|
|
@ -674,7 +677,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
|
|||
)}
|
||||
/>
|
||||
{tabItems.map((item) => (
|
||||
<TabsContent key={item.key} value={item.key}>
|
||||
<TabsContent key={item.key} value={item.key} className={item.className}>
|
||||
{item.children}
|
||||
</TabsContent>
|
||||
))}
|
||||
|
|
|
|||
|
|
@ -164,7 +164,7 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet
|
|||
isLoading={isLoading}
|
||||
loadingMessage="Loading teams..."
|
||||
noDataMessage="No teams found"
|
||||
maxBodyHeight="calc(75vh - 210px)"
|
||||
fillHeight
|
||||
size="compact"
|
||||
toolbar={(table) => (
|
||||
<>
|
||||
|
|
|
|||
|
|
@ -256,7 +256,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
|
|||
}
|
||||
|
||||
return (
|
||||
<div className="flex h-full flex-col gap-6 overflow-hidden">
|
||||
<div className="flex min-h-0 flex-1 flex-col gap-6">
|
||||
<PageHeader
|
||||
icon={<KeyRound />}
|
||||
title="Virtual Keys"
|
||||
|
|
@ -283,7 +283,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
|
|||
isLoading={isLoading}
|
||||
loadingMessage="Loading keys..."
|
||||
noDataMessage="No keys found"
|
||||
maxBodyHeight="calc(75vh - 210px)"
|
||||
fillHeight
|
||||
size="compact"
|
||||
toolbar={(table) => (
|
||||
<>
|
||||
|
|
|
|||
|
|
@ -1,26 +1,58 @@
|
|||
import React from "react";
|
||||
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
|
||||
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
import { DEFAULT_DEPLOYMENT_AFFINITY } from "./ComplexityRouterConfig";
|
||||
import { DEFAULT_DEPLOYMENT_AFFINITY, DEFAULT_SESSION_AFFINITY_TTL_SECONDS } from "./ComplexityRouterConfig";
|
||||
|
||||
export const AffinityControls: React.FC<{
|
||||
value: ComplexityRouterConfigValue;
|
||||
onChange: (value: ComplexityRouterConfigValue) => void;
|
||||
}> = ({ value, onChange }) => (
|
||||
<>
|
||||
<div className="flex items-center gap-2 mb-2">
|
||||
<Switch
|
||||
checked={value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY}
|
||||
onCheckedChange={(deploymentAffinity) => onChange({ ...value, deployment_affinity: deploymentAffinity })}
|
||||
aria-label="Pin a session to one deployment per model group"
|
||||
/>
|
||||
<strong className="font-semibold">Pin a session to one deployment per model group</strong>
|
||||
</div>
|
||||
<span className="block text-xs text-muted-foreground">
|
||||
Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to
|
||||
load-balance every turn.
|
||||
</span>
|
||||
</>
|
||||
);
|
||||
}> = ({ value, onChange }) => {
|
||||
const [ttlDraft, setTtlDraft] = React.useState<string | null>(null);
|
||||
const commitTtl = (raw: string) => {
|
||||
setTtlDraft(null);
|
||||
if (raw.trim() === "") {
|
||||
onChange({ ...value, session_affinity_ttl_seconds: undefined });
|
||||
return;
|
||||
}
|
||||
const parsed = Number(raw);
|
||||
if (!Number.isFinite(parsed)) return;
|
||||
onChange({ ...value, session_affinity_ttl_seconds: Math.max(1, Math.round(parsed)) });
|
||||
};
|
||||
|
||||
return (
|
||||
<>
|
||||
<div className="flex items-center gap-2 mb-2">
|
||||
<Switch
|
||||
checked={value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY}
|
||||
onCheckedChange={(deploymentAffinity) => onChange({ ...value, deployment_affinity: deploymentAffinity })}
|
||||
aria-label="Pin a session to one deployment per model group"
|
||||
/>
|
||||
<strong className="font-semibold">Pin a session to one deployment per model group</strong>
|
||||
</div>
|
||||
<span className="block text-xs mb-3 text-muted-foreground">
|
||||
Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to
|
||||
load-balance every turn.
|
||||
</span>
|
||||
<div style={{ maxWidth: 320 }}>
|
||||
<label className="block text-sm font-medium mb-1" htmlFor="session-affinity-ttl">
|
||||
How long a pin survives idle (seconds)
|
||||
</label>
|
||||
<Input
|
||||
id="session-affinity-ttl"
|
||||
inputMode="numeric"
|
||||
value={ttlDraft ?? value.session_affinity_ttl_seconds ?? ""}
|
||||
placeholder={String(DEFAULT_SESSION_AFFINITY_TTL_SECONDS)}
|
||||
onChange={(event) => setTtlDraft(event.target.value)}
|
||||
onBlur={(event) => commitTtl(event.target.value)}
|
||||
/>
|
||||
<span className="block text-xs mt-1 text-muted-foreground">
|
||||
Refreshes after every request that reuses a pin. Empty tracks the backend default of{" "}
|
||||
{DEFAULT_SESSION_AFFINITY_TTL_SECONDS} seconds.
|
||||
</span>
|
||||
</div>
|
||||
</>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -947,6 +947,46 @@ describe("ComplexityRouterConfig affinity panel", () => {
|
|||
|
||||
expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).not.toBeChecked();
|
||||
});
|
||||
|
||||
it("writes an idle TTL on blur and keeps the partial input as a draft while typing", () => {
|
||||
const onChange = vi.fn();
|
||||
renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} />);
|
||||
fireEvent.click(screen.getByText("Advanced: Affinity"));
|
||||
|
||||
const ttl = screen.getByLabelText("How long a pin survives idle (seconds)");
|
||||
expect(ttl).toHaveAttribute("placeholder", "3600");
|
||||
fireEvent.change(ttl, { target: { value: "300" } });
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
fireEvent.blur(ttl);
|
||||
|
||||
expect(onChange).toHaveBeenCalledWith({ ...defaultValue, session_affinity_ttl_seconds: 300 });
|
||||
});
|
||||
|
||||
it("clearing the idle TTL returns the router to its backend default", () => {
|
||||
const onChange = vi.fn();
|
||||
const value = { ...defaultValue, session_affinity_ttl_seconds: 300 };
|
||||
renderWithProviders(<ComplexityRouterConfig {...baseProps} value={value} onChange={onChange} />);
|
||||
fireEvent.click(screen.getByText("Advanced: Affinity"));
|
||||
|
||||
const ttl = screen.getByLabelText("How long a pin survives idle (seconds)");
|
||||
expect(ttl).toHaveValue("300");
|
||||
fireEvent.change(ttl, { target: { value: "" } });
|
||||
fireEvent.blur(ttl);
|
||||
|
||||
expect(onChange).toHaveBeenCalledWith({ ...value, session_affinity_ttl_seconds: undefined });
|
||||
});
|
||||
|
||||
it("clamps a non-positive idle TTL to the backend's minimum", () => {
|
||||
const onChange = vi.fn();
|
||||
renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} />);
|
||||
fireEvent.click(screen.getByText("Advanced: Affinity"));
|
||||
|
||||
const ttl = screen.getByLabelText("How long a pin survives idle (seconds)");
|
||||
fireEvent.change(ttl, { target: { value: "0" } });
|
||||
fireEvent.blur(ttl);
|
||||
|
||||
expect(onChange).toHaveBeenCalledWith({ ...defaultValue, session_affinity_ttl_seconds: 1 });
|
||||
});
|
||||
});
|
||||
|
||||
describe("ComplexityRouterConfig default model", () => {
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3;
|
|||
export const DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS = 8000;
|
||||
export const MIN_QUOTED_CONTEXT_TURN_CHARS = 120;
|
||||
export const DEFAULT_SESSION_AFFINITY = false;
|
||||
export const DEFAULT_SESSION_AFFINITY_TTL_SECONDS = 3600;
|
||||
export const DEFAULT_DEPLOYMENT_AFFINITY = true;
|
||||
|
||||
export type ClassificationMode = "every_request" | "user_turn";
|
||||
|
|
@ -411,6 +412,7 @@ export interface ComplexityRouterConfigValue {
|
|||
hybrid_boundary_margin?: number;
|
||||
classification_mode?: ClassificationMode;
|
||||
session_affinity?: boolean;
|
||||
session_affinity_ttl_seconds?: number;
|
||||
modality_routing?: boolean;
|
||||
modality_pin_override?: boolean;
|
||||
deployment_affinity?: boolean;
|
||||
|
|
|
|||
|
|
@ -493,7 +493,7 @@ describe("AddAutoRouterTab", () => {
|
|||
});
|
||||
});
|
||||
|
||||
it("carries session affinity turned on through to the create payload", async () => {
|
||||
it("carries session affinity turned on and its idle window through to the create payload", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getMissingTiersError).mockReturnValue(null);
|
||||
|
||||
|
|
@ -503,12 +503,17 @@ describe("AddAutoRouterTab", () => {
|
|||
expandDetailedConfiguration();
|
||||
await user.click(screen.getByText("Advanced: Classification Method"));
|
||||
await user.click(await screen.findByRole("radio", { name: /Once per session/ }));
|
||||
await user.click(screen.getByText("Advanced: Affinity"));
|
||||
const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)");
|
||||
fireEvent.change(ttl, { target: { value: "300" } });
|
||||
fireEvent.blur(ttl);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /add auto router/i }));
|
||||
|
||||
await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled());
|
||||
expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({
|
||||
session_affinity: true,
|
||||
session_affinity_ttl_seconds: 300,
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -388,6 +388,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
reasoningOverrideMinScore: complexityRouterConfig.reasoning_override_min_score,
|
||||
enableContextWindowEscalation: complexityRouterConfig.enable_context_window_escalation,
|
||||
contextWindowEscalationBuffer: complexityRouterConfig.context_window_escalation_buffer,
|
||||
sessionAffinityTtlSeconds: complexityRouterConfig.session_affinity_ttl_seconds,
|
||||
};
|
||||
|
||||
const submitRecommendedRouter = async (name: string) => {
|
||||
|
|
|
|||
|
|
@ -73,6 +73,15 @@ describe("buildComplexityRouterConfig", () => {
|
|||
expect(config.context_window_escalation_buffer).toBe(0.9);
|
||||
});
|
||||
|
||||
it("omits session_affinity_ttl_seconds when untouched, so the router tracks the backend default", () => {
|
||||
expect(buildComplexityRouterConfig(baseParams)).not.toHaveProperty("session_affinity_ttl_seconds");
|
||||
});
|
||||
|
||||
it("emits an explicit session affinity idle window", () => {
|
||||
const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinityTtlSeconds: 300 });
|
||||
expect(config.session_affinity_ttl_seconds).toBe(300);
|
||||
});
|
||||
|
||||
it("trims escalation keywords and drops blank entries", () => {
|
||||
const config = buildComplexityRouterConfig({
|
||||
...baseParams,
|
||||
|
|
|
|||
|
|
@ -138,6 +138,7 @@ export interface BuildComplexityRouterConfigParams {
|
|||
tierModelParams?: TierModelParamsByTier;
|
||||
enableContextWindowEscalation?: boolean;
|
||||
contextWindowEscalationBuffer?: number;
|
||||
sessionAffinityTtlSeconds?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -175,6 +176,7 @@ export interface ComplexityRouterConfigPayload {
|
|||
hybrid_boundary_margin?: number;
|
||||
classification_mode: ClassificationMode;
|
||||
session_affinity: boolean;
|
||||
session_affinity_ttl_seconds?: number;
|
||||
deployment_affinity: boolean;
|
||||
modality_routing: boolean;
|
||||
modality_pin_override: boolean;
|
||||
|
|
@ -456,6 +458,7 @@ export const buildComplexityRouterConfig = ({
|
|||
tierModelParams,
|
||||
enableContextWindowEscalation,
|
||||
contextWindowEscalationBuffer,
|
||||
sessionAffinityTtlSeconds,
|
||||
}: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => {
|
||||
const serializedTierModelConfigs = customTierSet
|
||||
? serializeTierModelConfigs(
|
||||
|
|
@ -522,6 +525,9 @@ export const buildComplexityRouterConfig = ({
|
|||
...(contextWindowEscalationBuffer !== undefined && {
|
||||
context_window_escalation_buffer: contextWindowEscalationBuffer,
|
||||
}),
|
||||
...(sessionAffinityTtlSeconds !== undefined && {
|
||||
session_affinity_ttl_seconds: sessionAffinityTtlSeconds,
|
||||
}),
|
||||
...scorerKnobs,
|
||||
};
|
||||
if (!customTierSet) return payload;
|
||||
|
|
|
|||
|
|
@ -257,6 +257,43 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("buildUpdatedComplexityRouterConfig session affinity ttl", () => {
|
||||
it("writes an edited idle window", () => {
|
||||
const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, session_affinity_ttl_seconds: 300 });
|
||||
expect(result.session_affinity_ttl_seconds).toBe(300);
|
||||
});
|
||||
|
||||
it("carries a stored idle window through an untouched open-and-save", () => {
|
||||
const stored = { ...STORED, session_affinity_ttl_seconds: 900 };
|
||||
const result = buildUpdatedComplexityRouterConfig(stored, hydrateComplexityRouterConfig(stored, undefined));
|
||||
expect(result.session_affinity_ttl_seconds).toBe(900);
|
||||
});
|
||||
|
||||
it("drops the key when the field is cleared, so the router goes back to tracking the backend default", () => {
|
||||
const result = buildUpdatedComplexityRouterConfig(
|
||||
{ ...STORED, session_affinity_ttl_seconds: 900 },
|
||||
{ ...FORM_VALUE, session_affinity_ttl_seconds: undefined },
|
||||
);
|
||||
expect(result).not.toHaveProperty("session_affinity_ttl_seconds");
|
||||
});
|
||||
|
||||
it("keeps the idle window on a custom tier set, whose deployment pin still uses it", () => {
|
||||
const result = buildUpdatedComplexityRouterConfig(STORED, {
|
||||
...FORM_VALUE,
|
||||
session_affinity_ttl_seconds: 300,
|
||||
custom_tier_set: {
|
||||
tiers: [
|
||||
{ id: "a", name: "CASUAL", definition: "small talk", models: ["gpt-4o-mini"] },
|
||||
{ id: "b", name: "AUDIT", definition: "security review", models: ["o1"] },
|
||||
],
|
||||
fallback_tier_id: "a",
|
||||
},
|
||||
});
|
||||
expect(result.session_affinity).toBe(false);
|
||||
expect(result.session_affinity_ttl_seconds).toBe(300);
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildUpdatedComplexityRouterConfig modality pin override", () => {
|
||||
it("writes modality_pin_override explicitly both ways", () => {
|
||||
expect(
|
||||
|
|
@ -529,6 +566,7 @@ describe("managed keys survive an untouched open-and-save", () => {
|
|||
classifier_fallback: "default_model",
|
||||
classification_mode: "user_turn",
|
||||
session_affinity: true,
|
||||
session_affinity_ttl_seconds: 300,
|
||||
modality_routing: true,
|
||||
modality_pin_override: true,
|
||||
deployment_affinity: false,
|
||||
|
|
|
|||
|
|
@ -550,6 +550,49 @@ describe("EditAutoRouterModal deployment affinity", () => {
|
|||
expect(savedConfig().deployment_affinity).toBe(false);
|
||||
});
|
||||
|
||||
it("preserves an idle TTL through an untouched save", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithStoredConfig({ ...STORED_CONFIG, session_affinity_ttl_seconds: 300 });
|
||||
|
||||
await user.click(await screen.findByText("Advanced: Affinity"));
|
||||
expect(await screen.findByLabelText("How long a pin survives idle (seconds)")).toHaveValue("300");
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
|
||||
expect(savedConfig().session_affinity_ttl_seconds).toBe(300);
|
||||
});
|
||||
|
||||
it("persists an edited idle TTL", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithStoredConfig(STORED_CONFIG);
|
||||
|
||||
await user.click(await screen.findByText("Advanced: Affinity"));
|
||||
const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)");
|
||||
fireEvent.change(ttl, { target: { value: "300" } });
|
||||
fireEvent.blur(ttl);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
|
||||
expect(savedConfig().session_affinity_ttl_seconds).toBe(300);
|
||||
});
|
||||
|
||||
it("removes the idle TTL when cleared", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithStoredConfig({ ...STORED_CONFIG, session_affinity_ttl_seconds: 300 });
|
||||
|
||||
await user.click(await screen.findByText("Advanced: Affinity"));
|
||||
const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)");
|
||||
fireEvent.change(ttl, { target: { value: "" } });
|
||||
fireEvent.blur(ttl);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
|
||||
expect(savedConfig()).not.toHaveProperty("session_affinity_ttl_seconds");
|
||||
});
|
||||
|
||||
// modality_pin_override is a managed key, so the modal rewrites it from form state on save. A
|
||||
// hydration gap would silently turn a stored override off on the next untouched save.
|
||||
it("shows a stored modality_pin_override=true as on and preserves it through an untouched save", async () => {
|
||||
|
|
|
|||
|
|
@ -104,6 +104,7 @@ export interface StoredComplexityRouterConfig {
|
|||
dimension_weights?: unknown;
|
||||
reasoning_override_min_score?: unknown;
|
||||
session_affinity?: unknown;
|
||||
session_affinity_ttl_seconds?: unknown;
|
||||
modality_routing?: unknown;
|
||||
modality_pin_override?: unknown;
|
||||
deployment_affinity?: unknown;
|
||||
|
|
@ -182,6 +183,11 @@ export const hydrateComplexityRouterConfig = (
|
|||
reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score),
|
||||
session_affinity:
|
||||
typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY,
|
||||
session_affinity_ttl_seconds:
|
||||
typeof parsedConfig.session_affinity_ttl_seconds === "number" &&
|
||||
Number.isFinite(parsedConfig.session_affinity_ttl_seconds)
|
||||
? parsedConfig.session_affinity_ttl_seconds
|
||||
: undefined,
|
||||
modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false,
|
||||
modality_pin_override:
|
||||
typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false,
|
||||
|
|
@ -224,6 +230,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
|
|||
"hybrid_boundary_margin",
|
||||
"classification_mode",
|
||||
"session_affinity",
|
||||
"session_affinity_ttl_seconds",
|
||||
"modality_routing",
|
||||
"modality_pin_override",
|
||||
"deployment_affinity",
|
||||
|
|
@ -322,6 +329,7 @@ export const buildUpdatedComplexityRouterConfig = (
|
|||
classifierContextIncludeAssistantTurns: value.classifier_context_include_assistant_turns,
|
||||
classifierFallback: value.classifier_fallback,
|
||||
sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY,
|
||||
sessionAffinityTtlSeconds: value.session_affinity_ttl_seconds,
|
||||
modalityRouting: value.modality_routing ?? false,
|
||||
modalityPinOverride: value.modality_pin_override ?? false,
|
||||
deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
|
||||
|
|
|
|||
|
|
@ -613,8 +613,11 @@ describe("DataTable layout", () => {
|
|||
|
||||
it("makes the header sticky and constrains body height when maxBodyHeight is set", () => {
|
||||
render(<DataTable data={CHARLIE_ALICE_BOB} columns={nameEmailColumns} maxBodyHeight={240} />);
|
||||
expect(screen.getByTestId("data-table-head")).toHaveClass("sticky");
|
||||
expect(screen.getByTestId("data-table-scroller")).toHaveStyle({ maxHeight: "240px" });
|
||||
const scroller = screen.getByTestId("data-table-scroller");
|
||||
expect(scroller).toHaveStyle({ maxHeight: "240px" });
|
||||
expect(scroller).toHaveClass("overflow-auto");
|
||||
expect(scroller).toHaveClass("[&_[data-slot=table-container]]:overflow-visible");
|
||||
expect(screen.getByTestId("data-table-head")).toHaveClass("sticky", "bg-background");
|
||||
});
|
||||
|
||||
it("caps fillHeight at the parent's height instead of stretching to it, so a short table stays short", () => {
|
||||
|
|
|
|||
|
|
@ -59,18 +59,22 @@ const noop = () => {};
|
|||
/**
|
||||
* Height-filling mode. The table still sizes to its rows; the parent's height is only a ceiling, so
|
||||
* a short table keeps its footer under the last row and a long one scrolls its rows instead of the
|
||||
* page. `table-container` is the Table primitive's own overflow-x wrapper; left as a scroll box it
|
||||
* captures the sticky header and the header scrolls away with the rows. And rows pass under that
|
||||
* header, which the semi-transparent header row tint alone would not hide.
|
||||
* page.
|
||||
*/
|
||||
const FILL_CLASSES = {
|
||||
outer: "flex max-h-full min-h-0 flex-col",
|
||||
frame: "flex min-h-0 flex-col",
|
||||
body: "min-h-0 [&_[data-slot=table-container]]:overflow-visible",
|
||||
body: "min-h-0",
|
||||
} as const;
|
||||
|
||||
const NO_FILL_CLASSES = { outer: "", frame: "", body: "" } as const;
|
||||
|
||||
const STICKY_CLASSES = {
|
||||
body: "[&_[data-slot=table-container]]:overflow-visible",
|
||||
header: "bg-background",
|
||||
} as const;
|
||||
|
||||
const NO_FILL_CLASSES = { outer: "", frame: "", body: "", header: "" } as const;
|
||||
const NO_STICKY_CLASSES = { body: "", header: "" } as const;
|
||||
|
||||
function columnDefId<TData, TValue>(column: ColumnDef<TData, TValue>): string | undefined {
|
||||
if ("id" in column && typeof column.id === "string") {
|
||||
|
|
@ -533,6 +537,7 @@ export function DataTable<TData extends RowData, TValue>(props: DataTableProps<T
|
|||
const visibleColumnCount = table.getVisibleLeafColumns().length;
|
||||
const stickyHeader = maxBodyHeight !== undefined || fillHeight;
|
||||
const fill = fillHeight ? FILL_CLASSES : NO_FILL_CLASSES;
|
||||
const sticky = stickyHeader ? STICKY_CLASSES : NO_STICKY_CLASSES;
|
||||
const tableStyle = enableColumnResizing ? { width: table.getTotalSize(), minWidth: "100%" } : undefined;
|
||||
|
||||
const renderPagination = (): React.ReactNode => {
|
||||
|
|
@ -593,13 +598,13 @@ export function DataTable<TData extends RowData, TValue>(props: DataTableProps<T
|
|||
{toolbar !== undefined && <div className="shrink-0 border-b border-border px-4 py-3">{toolbar(table)}</div>}
|
||||
<div
|
||||
data-testid="data-table-scroller"
|
||||
className={cn(stickyHeader ? "overflow-auto" : "overflow-x-auto", fill.body)}
|
||||
className={cn(stickyHeader ? "overflow-auto" : "overflow-x-auto", sticky.body, fill.body)}
|
||||
style={maxBodyHeight !== undefined ? { maxHeight: maxBodyHeight } : undefined}
|
||||
>
|
||||
<TableRoot className={enableColumnResizing ? "table-fixed" : ""} style={tableStyle}>
|
||||
<TableHeader
|
||||
data-testid="data-table-head"
|
||||
className={cn(stickyHeader ? "sticky top-0 z-sticky" : "", fill.header)}
|
||||
className={cn(stickyHeader ? "sticky top-0 z-sticky" : "", sticky.header)}
|
||||
>
|
||||
{table.getHeaderGroups().map((headerGroup) => (
|
||||
<TableRow key={headerGroup.id} className="bg-muted/50">
|
||||
|
|
|
|||
|
|
@ -443,7 +443,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
|
|||
}, []);
|
||||
|
||||
return (
|
||||
<div className="w-full h-full overflow-hidden">
|
||||
<div className="w-full">
|
||||
{selectedKey ? (
|
||||
<KeyInfoView
|
||||
keyId={selectedKey.token}
|
||||
|
|
@ -453,7 +453,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
|
|||
onDelete={refetch}
|
||||
/>
|
||||
) : (
|
||||
<div className="py-4 flex-1 overflow-hidden">
|
||||
<div className="py-4">
|
||||
<DataTable
|
||||
data={displayKeys}
|
||||
columns={columns}
|
||||
|
|
@ -471,7 +471,6 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
|
|||
columnResizeMode="onChange"
|
||||
isLoading={isLoading || isFetching}
|
||||
loadingMessage="Loading keys..."
|
||||
maxBodyHeight="75vh"
|
||||
size="compact"
|
||||
toolbar={(table) => (
|
||||
<>
|
||||
|
|
|
|||
|
|
@ -216,24 +216,22 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
|
|||
const canCreateKey = userRole !== "Admin Viewer" && userRole !== "proxy_admin_viewer";
|
||||
|
||||
return (
|
||||
<main className="h-[75vh] p-8">
|
||||
<div className="flex h-full flex-col">
|
||||
<VirtualKeysTable
|
||||
headerActions={
|
||||
canCreateKey ? (
|
||||
<CreateKey
|
||||
key={selectedTeam ? selectedTeam.team_id : null}
|
||||
team={selectedTeam as Team | null}
|
||||
teams={teams as Team[]}
|
||||
data={keys}
|
||||
addKey={addKey}
|
||||
autoOpenCreate={autoOpenCreate}
|
||||
prefillData={prefillData}
|
||||
/>
|
||||
) : undefined
|
||||
}
|
||||
/>
|
||||
</div>
|
||||
<main className="flex h-full flex-col p-8">
|
||||
<VirtualKeysTable
|
||||
headerActions={
|
||||
canCreateKey ? (
|
||||
<CreateKey
|
||||
key={selectedTeam ? selectedTeam.team_id : null}
|
||||
team={selectedTeam as Team | null}
|
||||
teams={teams as Team[]}
|
||||
data={keys}
|
||||
addKey={addKey}
|
||||
autoOpenCreate={autoOpenCreate}
|
||||
prefillData={prefillData}
|
||||
/>
|
||||
) : undefined
|
||||
}
|
||||
/>
|
||||
</main>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -86,6 +86,7 @@ export function RequestLogsTable({
|
|||
data={data}
|
||||
columns={columns}
|
||||
getRowId={(row) => row.request_id}
|
||||
fillHeight
|
||||
sortingMode="server"
|
||||
sorting={sorting}
|
||||
onSortingChange={onSortingChange}
|
||||
|
|
|
|||
|
|
@ -27,6 +27,9 @@ const AUDIT_LOGS_TAB: LogsTab = { id: "audit logs", label: "Audit Logs" };
|
|||
const DELETED_KEYS_TAB: LogsTab = { id: "deleted keys", label: "Deleted Keys" };
|
||||
const DELETED_TEAMS_TAB: LogsTab = { id: "deleted teams", label: "Deleted Teams" };
|
||||
|
||||
const tabContentClassName = (tabId: LogsTabId): string =>
|
||||
tabId === REQUEST_LOGS_TAB.id ? "flex min-h-0 flex-1 flex-col" : "min-h-0 flex-1 overflow-y-auto";
|
||||
|
||||
export default function SpendLogsTable({ accessToken, token, userRole, userID, premiumUser }: SpendLogsTableProps) {
|
||||
const [activeTab, setActiveTab] = useState<LogsTabId>(REQUEST_LOGS_TAB.id);
|
||||
const canViewAuditLogs = useCan("viewAuditLogs");
|
||||
|
|
@ -78,8 +81,8 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p
|
|||
};
|
||||
|
||||
return (
|
||||
<div className="box-border w-full overflow-x-hidden p-6">
|
||||
<Tabs value={activeTab} onValueChange={(value) => setActiveTab(value as LogsTabId)}>
|
||||
<div className="flex h-full w-full flex-col p-6">
|
||||
<Tabs value={activeTab} onValueChange={(value) => setActiveTab(value as LogsTabId)} className="min-h-0 flex-1">
|
||||
<TabsList variant="line">
|
||||
{tabs.map((tab) => (
|
||||
<TabsTrigger key={tab.id} value={tab.id} className="flex-none">
|
||||
|
|
@ -88,7 +91,7 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p
|
|||
))}
|
||||
</TabsList>
|
||||
{tabs.map((tab) => (
|
||||
<TabsContent key={tab.id} value={tab.id} keepMounted>
|
||||
<TabsContent key={tab.id} value={tab.id} keepMounted className={tabContentClassName(tab.id)}>
|
||||
{renderPanel(tab.id)}
|
||||
</TabsContent>
|
||||
))}
|
||||
|
|
|
|||
|
|
@ -83,9 +83,19 @@ describe("autorouter_presets", () => {
|
|||
expect(config.tier_boundaries).toBeUndefined();
|
||||
expect(config.token_thresholds).toBeUndefined();
|
||||
expect(config.dimension_weights).toBeUndefined();
|
||||
expect(config.session_affinity_ttl_seconds).toBeUndefined();
|
||||
}
|
||||
});
|
||||
|
||||
it("carries a preset's session affinity idle window into the prefilled form state", () => {
|
||||
const config = getPresetByKey("anthropic_family")!.complexity_router_config;
|
||||
const prefill = buildPresetPrefill({ ...config, session_affinity_ttl_seconds: 300 }, groupsOnly([]));
|
||||
expect(prefill.complexityRouterConfig.session_affinity_ttl_seconds).toBe(300);
|
||||
expect(
|
||||
buildPresetPrefill(config, groupsOnly([])).complexityRouterConfig.session_affinity_ttl_seconds,
|
||||
).toBeUndefined();
|
||||
});
|
||||
|
||||
it("keeps the model-family presets on the heuristic classifier", () => {
|
||||
for (const key of ["anthropic_family", "gemini_family", "openai_family"]) {
|
||||
expect(getPresetByKey(key)!.complexity_router_config.classifier_type).toBe("heuristic");
|
||||
|
|
|
|||
|
|
@ -284,6 +284,7 @@ export const buildPresetPrefill = (
|
|||
classifier_context_include_assistant_turns: config.classifier_context_include_assistant_turns,
|
||||
classification_mode: config.classification_mode ?? DEFAULT_CLASSIFICATION_MODE,
|
||||
session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY,
|
||||
session_affinity_ttl_seconds: config.session_affinity_ttl_seconds,
|
||||
deployment_affinity: config.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
|
||||
modality_routing: config.modality_routing ?? false,
|
||||
modality_pin_override: config.modality_pin_override ?? false,
|
||||
|
|
|
|||
16
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
16
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -25501,6 +25501,12 @@ export interface components {
|
|||
* @description Documents all the fields supported by `general_settings` in config.yaml
|
||||
*/
|
||||
ConfigGeneralSettings: {
|
||||
/**
|
||||
* Admission Queue Timeout Seconds
|
||||
* @description maximum time a request waits for a worker slot
|
||||
* @default 1
|
||||
*/
|
||||
admission_queue_timeout_seconds: number;
|
||||
/**
|
||||
* Alert To Webhook Url
|
||||
* @description Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}`
|
||||
|
|
@ -25709,11 +25715,21 @@ export interface components {
|
|||
* @description max file size in MB for /v1/files uploads, for any purpose, if a file is larger than this size it will be rejected before being forwarded to the provider
|
||||
*/
|
||||
max_file_size_mb?: number | null;
|
||||
/**
|
||||
* Max In Flight Requests Per Worker
|
||||
* @description maximum concurrent requests handled by each worker
|
||||
*/
|
||||
max_in_flight_requests_per_worker?: number | null;
|
||||
/**
|
||||
* Max Parallel Requests
|
||||
* @description maximum parallel requests for each api key
|
||||
*/
|
||||
max_parallel_requests?: number | null;
|
||||
/**
|
||||
* Max Queued Requests Per Worker
|
||||
* @description maximum requests waiting for a worker slot
|
||||
*/
|
||||
max_queued_requests_per_worker?: number | null;
|
||||
/**
|
||||
* Max Request Size Mb
|
||||
* @description max request size in MB, if a request is larger than this size it will be rejected
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue