mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
Merge remote-tracking branch 'origin/main' into litellm_decrease_anys_opus5_r5
This commit is contained in:
commit
596783c257
68 changed files with 3561 additions and 301 deletions
|
|
@ -51,6 +51,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/cache_settings",
|
||||
"/coordination_redis/",
|
||||
"/cost_tracking",
|
||||
"/cost_optimization/",
|
||||
"/cost/",
|
||||
"/credentials",
|
||||
"/credential",
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "is_default" BOOLEAN NOT NULL DEFAULT false;
|
||||
|
|
@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable {
|
|||
models String[] @default([]) // Model names or patterns
|
||||
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
|
||||
priority Int? // Explicit execution order
|
||||
is_default Boolean @default(false) // Applied only when no non-default attachment matches
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
|
|
|
|||
|
|
@ -1684,6 +1684,9 @@ if TYPE_CHECKING:
|
|||
from .llms.bedrock.messages.mantle_transformation import (
|
||||
AmazonMantleMessagesConfig as AmazonMantleMessagesConfig,
|
||||
)
|
||||
from .llms.bedrock_mantle.messages.transformation import (
|
||||
BedrockMantleAnthropicMessagesConfig as BedrockMantleAnthropicMessagesConfig,
|
||||
)
|
||||
from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig
|
||||
from .llms.together_ai.chat.transformation import (
|
||||
TogetherAIChatConfig as TogetherAIChatConfig,
|
||||
|
|
|
|||
|
|
@ -176,6 +176,7 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"BedrockClaudePlatformMessagesConfig",
|
||||
"AmazonAnthropicClaudeMessagesConfig",
|
||||
"AmazonMantleMessagesConfig",
|
||||
"BedrockMantleAnthropicMessagesConfig",
|
||||
"TogetherAIConfig",
|
||||
"TogetherAIChatConfig",
|
||||
"NLPCloudConfig",
|
||||
|
|
@ -746,6 +747,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
".llms.bedrock.messages.mantle_transformation",
|
||||
"AmazonMantleMessagesConfig",
|
||||
),
|
||||
"BedrockMantleAnthropicMessagesConfig": (
|
||||
".llms.bedrock_mantle.messages.transformation",
|
||||
"BedrockMantleAnthropicMessagesConfig",
|
||||
),
|
||||
"TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"),
|
||||
"TogetherAIChatConfig": (
|
||||
".llms.together_ai.chat.transformation",
|
||||
|
|
|
|||
|
|
@ -131,6 +131,41 @@
|
|||
"web-fetch-2025-09-10": null,
|
||||
"web-search-2025-03-05": null
|
||||
},
|
||||
"bedrock_mantle": {
|
||||
"advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19",
|
||||
"advisor-tool-2026-03-01": null,
|
||||
"bash_20241022": null,
|
||||
"bash_20250124": null,
|
||||
"claude-code-20250219": "claude-code-20250219",
|
||||
"code-execution-2025-08-25": null,
|
||||
"compact-2026-01-12": "compact-2026-01-12",
|
||||
"computer-use-2025-01-24": "computer-use-2025-01-24",
|
||||
"computer-use-2025-11-24": "computer-use-2025-11-24",
|
||||
"context-1m-2025-08-07": "context-1m-2025-08-07",
|
||||
"context-management-2025-06-27": "context-management-2025-06-27",
|
||||
"effort-2025-11-24": "effort-2025-11-24",
|
||||
"fast-mode-2026-02-01": null,
|
||||
"files-api-2025-04-14": null,
|
||||
"fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
|
||||
"interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
|
||||
"mcp-client-2025-04-04": null,
|
||||
"mcp-client-2025-11-20": null,
|
||||
"mcp-servers-2025-12-04": null,
|
||||
"output-128k-2025-02-19": "output-128k-2025-02-19",
|
||||
"per-turn-control-2026-07-01": "per-turn-control-2026-07-01",
|
||||
"prompt-caching-scope-2026-01-05": null,
|
||||
"skills-2025-10-02": null,
|
||||
"structured-output-2024-03-01": null,
|
||||
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
|
||||
"text_editor_20241022": null,
|
||||
"text_editor_20250124": null,
|
||||
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
|
||||
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
|
||||
"tool-examples-2025-10-29": "tool-examples-2025-10-29",
|
||||
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
|
||||
"web-fetch-2025-09-10": null,
|
||||
"web-search-2025-03-05": "web-search-2025-03-05"
|
||||
},
|
||||
"vertex_ai": {
|
||||
"advisor-tool-2026-03-01": null,
|
||||
"advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19",
|
||||
|
|
|
|||
|
|
@ -334,7 +334,7 @@ def update_headers_with_filtered_beta(
|
|||
Updated headers dict
|
||||
"""
|
||||
existing_beta: Final = headers.get("anthropic-beta")
|
||||
if not existing_beta:
|
||||
if existing_beta is None:
|
||||
return headers
|
||||
|
||||
# Parse existing beta headers
|
||||
|
|
|
|||
|
|
@ -402,6 +402,7 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = (
|
|||
if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None
|
||||
else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT
|
||||
)
|
||||
PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20
|
||||
DEFAULT_TRIM_RATIO: Final = float(
|
||||
os.getenv("DEFAULT_TRIM_RATIO", 0.75)
|
||||
) # default ratio of tokens to trim from the end of a prompt
|
||||
|
|
|
|||
|
|
@ -7,12 +7,13 @@ import base64
|
|||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Awaitable, Callable, Generator
|
||||
from collections.abc import Awaitable, Callable, Generator, Sequence
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, TypeAlias, TypeVar
|
||||
|
||||
import anyio
|
||||
import httpx2
|
||||
from httpx2._client import UseClientDefault
|
||||
from httpx2._types import AuthTypes
|
||||
|
|
@ -38,6 +39,8 @@ from mcp.types import (
|
|||
ListPromptsResult,
|
||||
ListResourcesResult,
|
||||
ListResourceTemplatesResult,
|
||||
PaginatedRequestParams,
|
||||
PaginatedResult,
|
||||
Prompt,
|
||||
ResourceTemplate,
|
||||
ServerNotification,
|
||||
|
|
@ -49,7 +52,12 @@ from mcp.types import Tool as MCPTool
|
|||
from pydantic import AnyUrl
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT
|
||||
from litellm.constants import (
|
||||
MCP_CLIENT_TIMEOUT,
|
||||
MCP_NPM_CACHE_DIR,
|
||||
MCP_TOOL_LISTING_MAX_PAGES,
|
||||
MCP_TOOL_LISTING_TIMEOUT,
|
||||
)
|
||||
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
|
||||
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response
|
||||
|
|
@ -147,6 +155,8 @@ def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None:
|
|||
|
||||
|
||||
TSessionResult = TypeVar("TSessionResult")
|
||||
_ListPage = TypeVar("_ListPage", bound=PaginatedResult)
|
||||
_ListItem = TypeVar("_ListItem")
|
||||
|
||||
|
||||
class _MCPHTTPClient(httpx2.AsyncClient):
|
||||
|
|
@ -793,6 +803,33 @@ class MCPClient:
|
|||
# Return a default error result instead of raising
|
||||
return self.error_tool_result(e)
|
||||
|
||||
async def _list_optional_pages(
|
||||
self,
|
||||
fetch_page: Callable[[PaginatedRequestParams | None], Awaitable[_ListPage]],
|
||||
items_of: Callable[[_ListPage], Sequence[_ListItem]],
|
||||
) -> list[_ListItem]: # mutable-ok: existing list discovery API
|
||||
items: Final[list[_ListItem]] = [] # mutable-ok: bounded iterative page accumulation
|
||||
cursors: Final[set[str]] = set() # mutable-ok: constant-time detection of cursor cycles
|
||||
cursor: str | None = None # rebind-ok: iterative traversal avoids recursion at the existing page cap
|
||||
with anyio.fail_after(max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)):
|
||||
for page_index in range(MCP_TOOL_LISTING_MAX_PAGES):
|
||||
try:
|
||||
page = await fetch_page( # rebind-ok: each SDK page replaces the previous one
|
||||
None if cursor is None else PaginatedRequestParams(cursor=cursor)
|
||||
)
|
||||
except MCPError as error:
|
||||
if page_index > 0 and error.error.code == METHOD_NOT_FOUND:
|
||||
raise RuntimeError("MCP list operation became unavailable during pagination") from error
|
||||
raise
|
||||
items.extend(items_of(page))
|
||||
if not page.next_cursor:
|
||||
return items
|
||||
if page.next_cursor in cursors:
|
||||
raise RuntimeError("MCP list pagination repeated a cursor")
|
||||
cursors.add(page.next_cursor)
|
||||
cursor = page.next_cursor
|
||||
raise RuntimeError(f"MCP list pagination exceeded {MCP_TOOL_LISTING_MAX_PAGES} pages")
|
||||
|
||||
async def list_prompts(self, *, raise_on_error: bool = False) -> list[Prompt]:
|
||||
"""List available prompts from the server."""
|
||||
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
|
||||
|
|
@ -802,7 +839,11 @@ class MCPClient:
|
|||
if capabilities is not None and capabilities.prompts is None:
|
||||
return ListPromptsResult(prompts=[])
|
||||
try:
|
||||
return await session.list_prompts()
|
||||
return ListPromptsResult(
|
||||
prompts=await self._list_optional_pages(
|
||||
lambda params: session.list_prompts(params=params), lambda page: page.prompts
|
||||
)
|
||||
)
|
||||
except MCPError as error:
|
||||
if error.error.code != METHOD_NOT_FOUND:
|
||||
raise
|
||||
|
|
@ -892,7 +933,11 @@ class MCPClient:
|
|||
if capabilities is not None and capabilities.resources is None:
|
||||
return ListResourcesResult(resources=[])
|
||||
try:
|
||||
return await session.list_resources()
|
||||
return ListResourcesResult(
|
||||
resources=await self._list_optional_pages(
|
||||
lambda params: session.list_resources(params=params), lambda page: page.resources
|
||||
)
|
||||
)
|
||||
except MCPError as error:
|
||||
if error.error.code != METHOD_NOT_FOUND:
|
||||
raise
|
||||
|
|
@ -941,7 +986,12 @@ class MCPClient:
|
|||
if capabilities is not None and capabilities.resources is None:
|
||||
return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload
|
||||
try:
|
||||
return await session.list_resource_templates()
|
||||
return ListResourceTemplatesResult(
|
||||
resource_templates=await self._list_optional_pages(
|
||||
lambda params: session.list_resource_templates(params=params),
|
||||
lambda page: page.resource_templates,
|
||||
)
|
||||
)
|
||||
except MCPError as error:
|
||||
if error.error.code != METHOD_NOT_FOUND:
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -46,6 +46,8 @@ from litellm.types.llms.openai import (
|
|||
AllMessageValues,
|
||||
ChatCompletionDocumentObject,
|
||||
ChatCompletionNamedToolChoiceParam,
|
||||
ChatCompletionRedactedThinkingBlock,
|
||||
ChatCompletionThinkingBlock,
|
||||
ChatCompletionToolParam,
|
||||
OpenAIMessageContentListBlock,
|
||||
)
|
||||
|
|
@ -854,6 +856,8 @@ def _count_content_list(
|
|||
content_list: str
|
||||
| Iterable[
|
||||
OpenAIMessageContentListBlock
|
||||
| ChatCompletionThinkingBlock
|
||||
| ChatCompletionRedactedThinkingBlock
|
||||
| AnthropicMessagesTextParam
|
||||
| AnthropicMessagesImageParam
|
||||
| AnthropicMessagesDocumentParam
|
||||
|
|
@ -898,9 +902,9 @@ def _count_content_list(
|
|||
use_default_image_token_count,
|
||||
default_token_count,
|
||||
)
|
||||
elif c["type"] == "thinking":
|
||||
elif c["type"] in ("thinking", "redacted_thinking"):
|
||||
# Claude extended thinking content block
|
||||
# Count the thinking text and skip signature (opaque signature blob)
|
||||
# Count the thinking text and skip the opaque blobs (signature, redacted data)
|
||||
thinking_text = str(c.get("thinking", ""))
|
||||
if thinking_text:
|
||||
num_tokens += count_function(thinking_text)
|
||||
|
|
@ -920,7 +924,8 @@ def _count_content_list(
|
|||
raise ValueError(
|
||||
f"Invalid content item type: {content_type}. "
|
||||
f"Expected str or dict with 'type' field "
|
||||
f"(text, image_url, image, document, file, tool_use, tool_result, thinking, tool_reference)."
|
||||
f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, "
|
||||
f"tool_reference)."
|
||||
)
|
||||
return num_tokens
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -651,6 +651,11 @@ def anthropic_messages_handler(
|
|||
"display": "summarized",
|
||||
}
|
||||
|
||||
resolved_api_base: Final = (
|
||||
dynamic_api_base
|
||||
if dynamic_api_base is not None and anthropic_messages_provider_config.uses_get_llm_provider_api_base()
|
||||
else api_base
|
||||
)
|
||||
return base_llm_http_handler.anthropic_messages_handler(
|
||||
model=model,
|
||||
messages=strip_provider_specific_fields_from_anthropic_messages(messages),
|
||||
|
|
@ -662,7 +667,7 @@ def anthropic_messages_handler(
|
|||
litellm_params=litellm_params,
|
||||
logging_obj=litellm_logging_obj,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
api_base=resolved_api_base,
|
||||
stream=stream,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -128,6 +128,9 @@ class BaseAnthropicMessagesConfig(ABC):
|
|||
"""
|
||||
return True
|
||||
|
||||
def uses_get_llm_provider_api_base(self) -> bool:
|
||||
return False
|
||||
|
||||
def get_async_streaming_response_iterator(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from collections.abc import AsyncIterator
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
|
|
@ -445,13 +445,16 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
# Bedrock InvokeModel DOES support ``clear_tool_uses_20250919`` under the
|
||||
# ``context-management-2025-06-27`` beta. AWS docs:
|
||||
# https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-tool-use.md
|
||||
_BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: dict[str, str] = {
|
||||
"compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
|
||||
"clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
|
||||
}
|
||||
_BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType(
|
||||
{
|
||||
"compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
|
||||
"clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@classmethod
|
||||
def _filter_context_management_for_bedrock_invoke(
|
||||
cls,
|
||||
anthropic_messages_request: dict,
|
||||
beta_set: set,
|
||||
) -> None:
|
||||
|
|
@ -481,7 +484,7 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
anthropic_messages_request.pop("context_management", None)
|
||||
return
|
||||
|
||||
supported: Final = AmazonAnthropicClaudeMessagesConfig._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS
|
||||
supported: Final = cls._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS
|
||||
retained_edits: Final = [e for e in edits if isinstance(e, dict) and e.get("type") in supported]
|
||||
if not retained_edits:
|
||||
anthropic_messages_request.pop("context_management", None)
|
||||
|
|
@ -546,15 +549,16 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
if "tool-search-tool-2025-10-19" in beta_set:
|
||||
beta_set.add("tool-examples-2025-10-29")
|
||||
|
||||
beta_provider: Final = self.custom_llm_provider or "bedrock"
|
||||
filtered_betas: Final = sorted(
|
||||
filter_and_transform_beta_headers(
|
||||
beta_headers=list(beta_set),
|
||||
provider="bedrock",
|
||||
provider=beta_provider,
|
||||
)
|
||||
)
|
||||
|
||||
dropped_user_betas: Final = sorted(
|
||||
b for b in user_beta_set if not filter_and_transform_beta_headers([b], provider="bedrock")
|
||||
b for b in user_beta_set if not filter_and_transform_beta_headers([b], provider=beta_provider)
|
||||
)
|
||||
if dropped_user_betas:
|
||||
verbose_logger.warning(
|
||||
|
|
|
|||
0
litellm/llms/bedrock_mantle/messages/__init__.py
Normal file
0
litellm/llms/bedrock_mantle/messages/__init__.py
Normal file
127
litellm/llms/bedrock_mantle/messages/transformation.py
Normal file
127
litellm/llms/bedrock_mantle/messages/transformation.py
Normal file
|
|
@ -0,0 +1,127 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
|
||||
DEFAULT_ANTHROPIC_API_VERSION,
|
||||
)
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.common_utils import MANTLE_MESSAGES_PATH
|
||||
from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig
|
||||
from litellm.llms.bedrock_mantle.common_utils import (
|
||||
MANTLE_HOST_RE,
|
||||
BedrockMantleAuthMixin,
|
||||
resolve_mantle_region,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
_BASE_SUFFIXES_TO_STRIP: Final = (
|
||||
MANTLE_MESSAGES_PATH,
|
||||
"/v1/messages",
|
||||
"/messages",
|
||||
"/anthropic/v1",
|
||||
"/openai/v1",
|
||||
"/v1",
|
||||
)
|
||||
_BODY_FIELDS_MANTLE_READS_FROM_HEADERS: Final = frozenset({"anthropic_version", "anthropic_beta"})
|
||||
_ANTHROPIC_BETAS: Final = TypeAdapter(tuple[str, ...])
|
||||
_MANTLE_REQUEST: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mapping[str, object]) -> str:
|
||||
region: Final = resolve_mantle_region(MappingProxyType({**litellm_params, "api_base": api_base}))
|
||||
configured: Final = (
|
||||
api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") or f"https://bedrock-mantle.{region}.api.aws"
|
||||
).rstrip("/")
|
||||
stripped: Final = next(
|
||||
(configured[: -len(suffix)] for suffix in _BASE_SUFFIXES_TO_STRIP if configured.endswith(suffix)),
|
||||
configured,
|
||||
)
|
||||
host: Final = f"https://bedrock-mantle.{region}.api.aws" if MANTLE_HOST_RE.match(stripped) else stripped
|
||||
return f"{host}{MANTLE_MESSAGES_PATH}"
|
||||
|
||||
|
||||
class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleMessagesConfig):
|
||||
_BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType(
|
||||
{
|
||||
**AmazonMantleMessagesConfig._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS,
|
||||
"clear_thinking_20251015": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
|
||||
}
|
||||
)
|
||||
|
||||
def __init__(self, aws_signer: BaseAWSLLM | None = None) -> None:
|
||||
AmazonMantleMessagesConfig.__init__(self)
|
||||
self._aws_signer = aws_signer or self
|
||||
|
||||
@property
|
||||
def custom_llm_provider(self) -> str | None:
|
||||
return "bedrock_mantle"
|
||||
|
||||
def uses_get_llm_provider_api_base(self) -> bool:
|
||||
return True
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
return build_mantle_native_messages_url(api_base=api_base, litellm_params=litellm_params)
|
||||
|
||||
def validate_anthropic_messages_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> tuple[dict, str | None]:
|
||||
merged_headers, resolved_api_base = super().validate_anthropic_messages_environment(
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
)
|
||||
if any(name.lower() == "anthropic-version" for name in merged_headers):
|
||||
return merged_headers, resolved_api_base
|
||||
return { # mutable-ok: the base class contract returns a dict the handler signs into in place
|
||||
**merged_headers,
|
||||
"anthropic-version": DEFAULT_ANTHROPIC_API_VERSION,
|
||||
}, resolved_api_base
|
||||
|
||||
def transform_anthropic_messages_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
request: Final = _MANTLE_REQUEST.validate_python(
|
||||
super().transform_anthropic_messages_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
),
|
||||
)
|
||||
betas: Final = request.get("anthropic_beta")
|
||||
if betas is not None:
|
||||
header_betas: Final = ",".join(_ANTHROPIC_BETAS.validate_python(betas))
|
||||
headers["anthropic-beta"] = header_betas # rebind-ok: the handler signs and sends this same dict
|
||||
return { # mutable-ok: the base class contract returns the dict the handler serializes as the body
|
||||
key: value for key, value in request.items() if key not in _BODY_FIELDS_MANTLE_READS_FROM_HEADERS
|
||||
}
|
||||
|
|
@ -42971,21 +42971,21 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/deepseek/deepseek-v4-pro": {
|
||||
"input_cost_per_token": 9.22722e-07,
|
||||
"input_cost_per_token": 9.19242e-07,
|
||||
"input_cost_per_token_cache_hit": 4.4e-08,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 384000,
|
||||
"max_tokens": 384000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.845444e-06,
|
||||
"output_cost_per_token": 1.838484e-06,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 7.68935e-08,
|
||||
"cache_read_input_token_cost": 7.66035e-08,
|
||||
"supports_audio_input": false,
|
||||
"supports_pdf_input": false,
|
||||
"supports_vision": false,
|
||||
|
|
@ -54441,16 +54441,19 @@
|
|||
"zai.glm-4.7": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"max_input_tokens": 203000,
|
||||
"max_output_tokens": 4000,
|
||||
"max_tokens": 4000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.2e-06,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/"
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/",
|
||||
"supports_audio_input": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"zai.glm-5": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
|
|
@ -54470,16 +54473,19 @@
|
|||
"zai.glm-4.7-flash": {
|
||||
"input_cost_per_token": 7e-08,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"max_input_tokens": 203000,
|
||||
"max_output_tokens": 4000,
|
||||
"max_tokens": 4000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4e-07,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/"
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/",
|
||||
"supports_audio_input": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"zai/glm-5": {
|
||||
"cache_creation_input_token_cost": 0,
|
||||
|
|
@ -60558,6 +60564,34 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"bedrock_mantle/anthropic.claude-haiku-4-5": {
|
||||
"cache_creation_input_token_cost": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-06,
|
||||
"cache_read_input_token_cost": 1e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 64000,
|
||||
"max_tokens": 64000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-06,
|
||||
"source": "https://aws.amazon.com/about-aws/whats-new/2025/10/claude-4-5-haiku-anthropic-amazon-bedrock",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_parallel_tool_use_config": true,
|
||||
"prompt_cache_min_tokens": 4096,
|
||||
"input_cost_per_token_batches": 5e-07,
|
||||
"output_cost_per_token_batches": 2.5e-06
|
||||
},
|
||||
"us.xai.grok-4.6": {
|
||||
"input_cost_per_token": 2.2e-06,
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
|
|
@ -76803,7 +76837,7 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3.html",
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
|
|
|
|||
|
|
@ -3865,7 +3865,7 @@ if MCP_AVAILABLE:
|
|||
try:
|
||||
data: Final = json.loads(body)
|
||||
return isinstance(data, dict) and data.get("method") == "initialize"
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
|
||||
return False
|
||||
|
||||
def _extract_initialize_client_info(body: bytes) -> Implementation | None:
|
||||
|
|
@ -4791,7 +4791,7 @@ if MCP_AVAILABLE:
|
|||
"MCP: detected JSON-RPC response POST (id=%s), skipping session lock to avoid deadlock",
|
||||
_peeked.get("id"),
|
||||
)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
|
||||
# Peek cap truncated the body, so it can't be fully parsed.
|
||||
# Scan the top-level keys (depth-aware) instead of a flat
|
||||
# substring search: a response's result payload may nest a
|
||||
|
|
|
|||
|
|
@ -34982,6 +34982,12 @@
|
|||
"PolicyAttachmentCreateRequest": {
|
||||
"description": "Request body for creating a policy attachment.",
|
||||
"properties": {
|
||||
"default": {
|
||||
"default": false,
|
||||
"description": "Apply this attachment only when no non-default attachment matches the request.",
|
||||
"title": "Default",
|
||||
"type": "boolean"
|
||||
},
|
||||
"keys": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -35113,6 +35119,12 @@
|
|||
"description": "Who created the attachment.",
|
||||
"title": "Created By"
|
||||
},
|
||||
"default": {
|
||||
"default": false,
|
||||
"description": "Apply this attachment only when no non-default attachment matches the request.",
|
||||
"title": "Default",
|
||||
"type": "boolean"
|
||||
},
|
||||
"definition_location": {
|
||||
"default": "db",
|
||||
"description": "Where this attachment is defined: 'db' (database) or 'config' (config.yaml).",
|
||||
|
|
@ -37141,6 +37153,12 @@
|
|||
"PolicyAttachmentCreateRequest": {
|
||||
"description": "Request body for creating a policy attachment.",
|
||||
"properties": {
|
||||
"default": {
|
||||
"default": false,
|
||||
"description": "Apply this attachment only when no non-default attachment matches the request.",
|
||||
"title": "Default",
|
||||
"type": "boolean"
|
||||
},
|
||||
"keys": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -3216,7 +3216,9 @@ def _match_and_track_policies(
|
|||
attachment_registry: Final = (
|
||||
attachment_registry_override if attachment_registry_override is not None else get_attachment_registry()
|
||||
)
|
||||
matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons(context)
|
||||
matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons(
|
||||
context, PolicyMatcher.policy_applies(context, policies_override)
|
||||
)
|
||||
matching_policy_names: Final = [m["policy_name"] for m in matches_with_reasons]
|
||||
policy_reasons: Final = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons}
|
||||
|
||||
|
|
@ -3418,7 +3420,12 @@ async def add_guardrails_from_policy_engine(
|
|||
|
||||
|
||||
_ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join(
|
||||
(LlmProviders.ANTHROPIC.value, LlmProviders.BEDROCK.value, LlmProviders.VERTEX_AI.value)
|
||||
(
|
||||
LlmProviders.ANTHROPIC.value,
|
||||
LlmProviders.BEDROCK.value,
|
||||
LlmProviders.BEDROCK_MANTLE.value,
|
||||
LlmProviders.VERTEX_AI.value,
|
||||
)
|
||||
)
|
||||
_ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS: Final = LlmProviders.ANTHROPIC.value
|
||||
|
||||
|
|
|
|||
184
litellm/proxy/management_endpoints/prompt_caching_requests.py
Normal file
184
litellm/proxy/management_endpoints/prompt_caching_requests.py
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
from collections.abc import Callable, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from pydantic import BaseModel, Json, TypeAdapter
|
||||
|
||||
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth, user_api_key_has_admin_view
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.spend_tracking.savings import (
|
||||
extract_cache_creation_tokens,
|
||||
extract_cache_read_tokens,
|
||||
marks_gateway_injection,
|
||||
prompt_caching_savings_for_request,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
||||
_query_raw_rows, # pyright: ignore[reportPrivateUsage] # existing typed spend-query adapter; rows validated below
|
||||
)
|
||||
from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY
|
||||
from litellm.types.management_endpoints.prompt_caching_requests import (
|
||||
PromptCachingRequest,
|
||||
PromptCachingRequestCursor,
|
||||
PromptCachingRequestFilter,
|
||||
PromptCachingRequestsResponse,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
||||
def _numeric_token_sql(path: str) -> str:
|
||||
value: Final = f"metadata #> '{{usage_object,{path}}}'"
|
||||
return (
|
||||
f"CASE WHEN jsonb_typeof({value}) = 'number' THEN ({value} #>> '{{}}')::numeric "
|
||||
f"WHEN {value} = 'true'::jsonb THEN 1 WHEN {value} = 'false'::jsonb THEN 0 END"
|
||||
)
|
||||
|
||||
|
||||
def _cache_tokens_sql(*paths: str) -> str:
|
||||
candidates: Final = ", ".join(f"NULLIF(({_numeric_token_sql(path)}), 0)" for path in paths)
|
||||
return f"TRUNC(COALESCE({candidates}, 0))"
|
||||
|
||||
|
||||
_CACHE_READ_SQL: Final = _cache_tokens_sql("cache_read_input_tokens", "prompt_tokens_details,cached_tokens")
|
||||
_CACHE_CREATION_SQL: Final = _cache_tokens_sql(
|
||||
"cache_creation_input_tokens",
|
||||
"prompt_tokens_details,cache_write_tokens",
|
||||
"prompt_tokens_details,cache_creation_tokens",
|
||||
)
|
||||
_GATEWAY_INJECTED_SQL: Final = (
|
||||
f"(jsonb_typeof(metadata->'{GATEWAY_INJECTED_CACHE_METADATA_KEY}') = 'string' "
|
||||
f"AND (metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' = '' "
|
||||
f"OR metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' = model_id))"
|
||||
)
|
||||
_FILTER_SQL: Final = MappingProxyType(
|
||||
{
|
||||
"all": f"({_GATEWAY_INJECTED_SQL} OR {_CACHE_READ_SQL} > 0 OR {_CACHE_CREATION_SQL} > 0)",
|
||||
"injected": _GATEWAY_INJECTED_SQL,
|
||||
"hits": f"{_CACHE_READ_SQL} > 0",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def prompt_caching_requests_sql(filter: PromptCachingRequestFilter) -> str:
|
||||
return f"""
|
||||
SELECT request_id, "startTime" AS start_time, "endTime" AS end_time,
|
||||
model, model_id, custom_llm_provider, spend,
|
||||
CASE WHEN jsonb_typeof(metadata->'usage_object') = 'object'
|
||||
THEN metadata->'usage_object' END AS usage_object,
|
||||
CASE WHEN jsonb_typeof(metadata->'cost_breakdown') = 'object'
|
||||
THEN metadata->'cost_breakdown' END AS cost_breakdown,
|
||||
CASE WHEN jsonb_typeof(metadata->'{GATEWAY_INJECTED_CACHE_METADATA_KEY}') = 'string'
|
||||
THEN metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' END AS gateway_marker
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= ($1::text::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" <= ($2::text::timestamptz AT TIME ZONE 'UTC')
|
||||
AND COALESCE(LOWER(cache_hit), 'false') != 'true'
|
||||
AND {_FILTER_SQL[filter]}
|
||||
AND ($4::text::timestamptz IS NULL OR
|
||||
("startTime", request_id) < (($4::text::timestamptz AT TIME ZONE 'UTC'), $5::text))
|
||||
ORDER BY "startTime" DESC, request_id DESC
|
||||
LIMIT $3::integer
|
||||
"""
|
||||
|
||||
|
||||
class _PromptCachingRow(BaseModel):
|
||||
request_id: str
|
||||
start_time: datetime
|
||||
end_time: datetime
|
||||
model: str
|
||||
model_id: str | None
|
||||
custom_llm_provider: str | None
|
||||
spend: float
|
||||
usage_object: Json[Mapping[str, object]] | Mapping[str, object] | None
|
||||
cost_breakdown: Json[Mapping[str, object]] | Mapping[str, object] | None
|
||||
gateway_marker: str | None
|
||||
|
||||
|
||||
_REQUEST_ROWS: Final = TypeAdapter(tuple[_PromptCachingRow, ...])
|
||||
|
||||
|
||||
def _request_result(row: _PromptCachingRow, llm_router: "Callable[[], Router | None]") -> PromptCachingRequest:
|
||||
return PromptCachingRequest(
|
||||
request_id=row.request_id,
|
||||
start_time=row.start_time.replace(tzinfo=timezone.utc) if row.start_time.tzinfo is None else row.start_time,
|
||||
model=row.model,
|
||||
gateway_injected=marks_gateway_injection(
|
||||
MappingProxyType({GATEWAY_INJECTED_CACHE_METADATA_KEY: row.gateway_marker}), row.model_id
|
||||
),
|
||||
cache_read_tokens=extract_cache_read_tokens(row.usage_object),
|
||||
cache_creation_tokens=extract_cache_creation_tokens(row.usage_object),
|
||||
spend=row.spend,
|
||||
net_savings=prompt_caching_savings_for_request(
|
||||
model=row.model,
|
||||
custom_llm_provider=row.custom_llm_provider,
|
||||
usage_object=row.usage_object,
|
||||
model_id=row.model_id,
|
||||
llm_router=llm_router,
|
||||
cost_breakdown=row.cost_breakdown,
|
||||
billed_at=row.end_time,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/cost_optimization/prompt_caching/requests",
|
||||
tags=["Cost Optimization"], # mutable-ok: FastAPI's route API requires a list
|
||||
response_model=PromptCachingRequestsResponse,
|
||||
)
|
||||
async def get_prompt_caching_requests(
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
start_date: datetime,
|
||||
end_date: datetime,
|
||||
page_size: Annotated[int, Query(ge=1, le=100)] = 50,
|
||||
filter: PromptCachingRequestFilter = "all",
|
||||
cursor_start_time: datetime | None = None,
|
||||
cursor_request_id: Annotated[str | None, Query(min_length=1)] = None,
|
||||
) -> PromptCachingRequestsResponse:
|
||||
from litellm.proxy.proxy_server import llm_router, prisma_client
|
||||
|
||||
if not user_api_key_has_admin_view(user_api_key_dict):
|
||||
raise HTTPException(status_code=403, detail="Only proxy admin roles can view prompt caching requests")
|
||||
if (cursor_start_time is None) != (cursor_request_id is None):
|
||||
raise HTTPException(status_code=400, detail="cursor_start_time and cursor_request_id must be provided together")
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
start: Final = start_date.replace(tzinfo=timezone.utc) if start_date.tzinfo is None else start_date
|
||||
end: Final = end_date.replace(tzinfo=timezone.utc) if end_date.tzinfo is None else end_date
|
||||
if end < start:
|
||||
raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date")
|
||||
cursor_time: Final = (
|
||||
cursor_start_time.replace(tzinfo=timezone.utc)
|
||||
if cursor_start_time is not None and cursor_start_time.tzinfo is None
|
||||
else cursor_start_time
|
||||
)
|
||||
rows: Final = _REQUEST_ROWS.validate_python(
|
||||
await _query_raw_rows(
|
||||
prisma_client,
|
||||
prompt_caching_requests_sql(filter),
|
||||
start.isoformat(),
|
||||
end.isoformat(),
|
||||
page_size + 1,
|
||||
cursor_time.isoformat() if cursor_time is not None else None,
|
||||
cursor_request_id,
|
||||
)
|
||||
or ()
|
||||
)
|
||||
|
||||
def current_router() -> "Router | None":
|
||||
return llm_router
|
||||
|
||||
requests: Final = tuple(_request_result(row, current_router) for row in rows[:page_size])
|
||||
has_more: Final = len(rows) > page_size
|
||||
return PromptCachingRequestsResponse(
|
||||
requests=requests,
|
||||
page_size=page_size,
|
||||
has_more=has_more,
|
||||
next_cursor=PromptCachingRequestCursor(start_time=requests[-1].start_time, request_id=requests[-1].request_id)
|
||||
if has_more
|
||||
else None,
|
||||
)
|
||||
|
|
@ -5,6 +5,7 @@ Attachments define WHERE policies apply, separate from the policy definitions.
|
|||
This allows the same policy to be attached to multiple scopes.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, TypedDict
|
||||
|
|
@ -119,35 +120,49 @@ class AttachmentRegistry:
|
|||
models=attachment_data.get("models"),
|
||||
tags=attachment_data.get("tags"),
|
||||
priority=attachment_data.get("priority"),
|
||||
default=attachment_data.get("default", False),
|
||||
)
|
||||
|
||||
def get_attached_policies(self, context: PolicyMatchContext) -> list[str]:
|
||||
def get_attached_policies(
|
||||
self,
|
||||
context: PolicyMatchContext,
|
||||
policy_applies: Callable[[str], bool] | None = None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Get list of policy names attached to the given context.
|
||||
|
||||
Args:
|
||||
context: The request context to match against
|
||||
policy_applies: Optional predicate; attachments whose policy does not apply are ignored
|
||||
|
||||
Returns:
|
||||
List of policy names that are attached to matching scopes
|
||||
"""
|
||||
return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context)]
|
||||
return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context, policy_applies)]
|
||||
|
||||
def get_attached_policies_with_reasons(self, context: PolicyMatchContext) -> list[PolicyAttachmentMatch]:
|
||||
def get_attached_policies_with_reasons(
|
||||
self,
|
||||
context: PolicyMatchContext,
|
||||
policy_applies: Callable[[str], bool] | None = None,
|
||||
) -> list[PolicyAttachmentMatch]:
|
||||
"""
|
||||
Get list of policy names and match reasons for the given context.
|
||||
|
||||
Returns a list of dicts with 'policy_name' and 'matched_via' keys.
|
||||
The 'matched_via' describes which dimension caused the match.
|
||||
Attachments whose policy fails `policy_applies` are dropped before defaults are considered.
|
||||
"""
|
||||
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
|
||||
|
||||
in_scope: Final = tuple(
|
||||
attachment
|
||||
for attachment in self._attachments
|
||||
if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context)
|
||||
and (policy_applies is None or policy_applies(attachment.policy))
|
||||
)
|
||||
non_default: Final = tuple(attachment for attachment in in_scope if not attachment.default)
|
||||
matching_attachments: Final = sorted(
|
||||
(
|
||||
attachment
|
||||
for attachment in self._attachments
|
||||
if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context)
|
||||
),
|
||||
non_default or tuple(attachment for attachment in in_scope if attachment.default),
|
||||
key=_attachment_sort_key,
|
||||
)
|
||||
broadest_attachment_by_policy: Final = MappingProxyType(
|
||||
|
|
@ -169,6 +184,11 @@ class AttachmentRegistry:
|
|||
@staticmethod
|
||||
def _describe_match_reason(attachment: PolicyAttachment, context: PolicyMatchContext) -> str:
|
||||
"""Describe why an attachment matched the context."""
|
||||
reason: Final = AttachmentRegistry._describe_scope_match(attachment, context)
|
||||
return f"default:{reason}" if attachment.default else reason
|
||||
|
||||
@staticmethod
|
||||
def _describe_scope_match(attachment: PolicyAttachment, context: PolicyMatchContext) -> str:
|
||||
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
|
||||
|
||||
if attachment.is_global():
|
||||
|
|
@ -324,6 +344,7 @@ class AttachmentRegistry:
|
|||
"models": attachment_request.models or [],
|
||||
"tags": attachment_request.tags or [],
|
||||
"priority": attachment_request.priority,
|
||||
"is_default": attachment_request.default,
|
||||
"created_at": datetime.now(timezone.utc),
|
||||
"updated_at": datetime.now(timezone.utc),
|
||||
"created_by": created_by,
|
||||
|
|
@ -340,6 +361,7 @@ class AttachmentRegistry:
|
|||
models=attachment_request.models,
|
||||
tags=attachment_request.tags,
|
||||
priority=attachment_request.priority,
|
||||
default=attachment_request.default,
|
||||
)
|
||||
self.add_attachment(attachment)
|
||||
|
||||
|
|
@ -352,6 +374,7 @@ class AttachmentRegistry:
|
|||
models=created_attachment.models or [],
|
||||
tags=created_attachment.tags or [],
|
||||
priority=created_attachment.priority,
|
||||
default=created_attachment.is_default,
|
||||
created_at=created_attachment.created_at,
|
||||
updated_at=created_attachment.updated_at,
|
||||
created_by=created_attachment.created_by,
|
||||
|
|
@ -429,6 +452,7 @@ class AttachmentRegistry:
|
|||
models=attachment.models or [],
|
||||
tags=attachment.tags or [],
|
||||
priority=attachment.priority,
|
||||
default=attachment.is_default,
|
||||
created_at=attachment.created_at,
|
||||
updated_at=attachment.updated_at,
|
||||
created_by=attachment.created_by,
|
||||
|
|
@ -468,6 +492,7 @@ class AttachmentRegistry:
|
|||
models=a.models or [],
|
||||
tags=a.tags or [],
|
||||
priority=a.priority,
|
||||
default=a.is_default,
|
||||
created_at=a.created_at,
|
||||
updated_at=a.updated_at,
|
||||
created_by=a.created_by,
|
||||
|
|
@ -502,6 +527,7 @@ class AttachmentRegistry:
|
|||
models=(attachment_response.models if attachment_response.models else None),
|
||||
tags=attachment_response.tags if attachment_response.tags else None,
|
||||
priority=attachment_response.priority,
|
||||
default=attachment_response.default,
|
||||
)
|
||||
for attachment_response in attachments
|
||||
]
|
||||
|
|
|
|||
|
|
@ -61,6 +61,7 @@ def _config_attachment_to_db_response(index: int, attachment: PolicyAttachment)
|
|||
models=attachment.models or [],
|
||||
tags=attachment.tags or [],
|
||||
priority=attachment.priority,
|
||||
default=attachment.default,
|
||||
definition_location="config",
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ apply to a given request based on team alias, key alias, and model.
|
|||
Policies are matched via policy_attachments which define WHERE each policy applies.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -113,7 +114,7 @@ class PolicyMatcher:
|
|||
verbose_proxy_logger.debug("AttachmentRegistry not initialized, returning empty list")
|
||||
return []
|
||||
|
||||
return registry.get_attached_policies(context)
|
||||
return registry.get_attached_policies(context, PolicyMatcher.policy_applies(context))
|
||||
|
||||
@staticmethod
|
||||
def get_matching_policies_from_registry(
|
||||
|
|
@ -130,9 +131,31 @@ class PolicyMatcher:
|
|||
"""
|
||||
return PolicyMatcher.get_matching_policies(context=context)
|
||||
|
||||
@staticmethod
|
||||
def policy_applies(
|
||||
context: PolicyMatchContext,
|
||||
policies: dict[str, Policy] | None = None,
|
||||
) -> Callable[[str], bool]:
|
||||
"""Predicate telling whether a policy exists and its condition matches the context."""
|
||||
resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies()
|
||||
return lambda policy_name: bool(
|
||||
PolicyMatcher.get_policies_with_matching_conditions(
|
||||
policy_names=(policy_name,),
|
||||
context=context,
|
||||
policies=resolved,
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _registry_policies() -> dict[str, Policy]:
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
|
||||
registry: Final = get_policy_registry()
|
||||
return registry.get_all_policies() if registry.is_initialized() else {}
|
||||
|
||||
@staticmethod
|
||||
def get_policies_with_matching_conditions(
|
||||
policy_names: list[str],
|
||||
policy_names: Sequence[str],
|
||||
context: PolicyMatchContext,
|
||||
policies: dict[str, Policy] | None = None,
|
||||
) -> list[str]:
|
||||
|
|
@ -152,17 +175,12 @@ class PolicyMatcher:
|
|||
List of policy names whose conditions match the context
|
||||
"""
|
||||
from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
|
||||
if policies is None:
|
||||
registry: Final = get_policy_registry()
|
||||
if not registry.is_initialized():
|
||||
return []
|
||||
policies = registry.get_all_policies()
|
||||
resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies()
|
||||
|
||||
matching_policies: Final = []
|
||||
for policy_name in policy_names:
|
||||
policy = policies.get(policy_name)
|
||||
policy = resolved.get(policy_name)
|
||||
if policy is None:
|
||||
continue
|
||||
# Policy matches if it has no condition OR condition evaluates to True
|
||||
|
|
|
|||
|
|
@ -265,7 +265,9 @@ async def resolve_policies_for_context(
|
|||
)
|
||||
|
||||
# Get matching policies with reasons
|
||||
match_results: Final = get_attachment_registry().get_attached_policies_with_reasons(context=context)
|
||||
match_results: Final = get_attachment_registry().get_attached_policies_with_reasons(
|
||||
context=context, policy_applies=PolicyMatcher.policy_applies(context)
|
||||
)
|
||||
|
||||
if not match_results:
|
||||
return PolicyResolveResponse(
|
||||
|
|
|
|||
|
|
@ -84,7 +84,9 @@ def _retrieval_context(
|
|||
|
||||
|
||||
def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]:
|
||||
matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context)
|
||||
matches: Final = get_attachment_registry().get_attached_policies_with_reasons(
|
||||
context, PolicyMatcher.policy_applies(context)
|
||||
)
|
||||
if not matches:
|
||||
return (), MappingProxyType({})
|
||||
applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions(
|
||||
|
|
|
|||
|
|
@ -601,6 +601,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
|
|||
from litellm.proxy.management_endpoints.organization_endpoints import (
|
||||
router as organization_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.prompt_caching_requests import (
|
||||
router as prompt_caching_requests_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.router_settings_endpoints import (
|
||||
router as router_settings_router,
|
||||
)
|
||||
|
|
@ -19274,6 +19277,7 @@ app.include_router(workflow_management_router)
|
|||
app.include_router(memory_router)
|
||||
app.include_router(plugin_router)
|
||||
app.include_router(cost_tracking_settings_router)
|
||||
app.include_router(prompt_caching_requests_router)
|
||||
app.include_router(router_settings_router)
|
||||
app.include_router(fallback_management_router)
|
||||
app.include_router(cache_settings_router)
|
||||
|
|
|
|||
|
|
@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable {
|
|||
models String[] @default([]) // Model names or patterns
|
||||
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
|
||||
priority Int? // Explicit execution order
|
||||
is_default Boolean @default(false) // Applied only when no non-default attachment matches
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
|
|
|
|||
|
|
@ -578,6 +578,56 @@ def autorouter_savings_for_logging_payload(
|
|||
)
|
||||
|
||||
|
||||
def _request_savings_pricing(
|
||||
model: str | None,
|
||||
custom_llm_provider: str | None,
|
||||
model_id: str | None,
|
||||
llm_router: "Callable[[], Router | None] | None",
|
||||
) -> tuple[str | None, ModelInfo | None]:
|
||||
router_instance: Final = llm_router() if llm_router else None
|
||||
identity: Final = _resolve_model(model, custom_llm_provider)
|
||||
pricing: Final = _effective_model_info(router_instance, model_id, model or "") or (
|
||||
_model_info(identity) if identity else None
|
||||
)
|
||||
return identity.provider if identity else custom_llm_provider, pricing
|
||||
|
||||
|
||||
def _prompt_caching_savings(
|
||||
pricing: ModelInfo | None,
|
||||
provider: str | None,
|
||||
usage_object: Mapping[str, object] | None,
|
||||
cost_breakdown: Mapping[str, object] | None,
|
||||
billed_at: datetime | str | None,
|
||||
) -> float | None:
|
||||
usage: Final = _usage_from_spend_log(usage_object)
|
||||
if pricing is None or usage is None:
|
||||
return None
|
||||
basis: Final = _pricing_basis(cost_breakdown)
|
||||
result: Final = calculate_prompt_caching_savings(
|
||||
model_info=pricing,
|
||||
usage=usage,
|
||||
custom_llm_provider=provider,
|
||||
service_tier=basis.service_tier,
|
||||
data_residency=basis.data_residency,
|
||||
vertex_location=basis.vertex_location,
|
||||
billed_at=_coerce_billed_at(billed_at),
|
||||
)
|
||||
return result if isfinite(result) else None
|
||||
|
||||
|
||||
def prompt_caching_savings_for_request(
|
||||
model: str | None,
|
||||
custom_llm_provider: str | None,
|
||||
usage_object: Mapping[str, object] | None,
|
||||
model_id: str | None = None,
|
||||
llm_router: "Callable[[], Router | None] | None" = None,
|
||||
cost_breakdown: Mapping[str, object] | None = None,
|
||||
billed_at: datetime | str | None = None,
|
||||
) -> float | None:
|
||||
request_pricing: Final = _request_savings_pricing(model, custom_llm_provider, model_id, llm_router)
|
||||
return _prompt_caching_savings(request_pricing[1], request_pricing[0], usage_object, cost_breakdown, billed_at)
|
||||
|
||||
|
||||
def compute_savings_spend(
|
||||
model: str | None,
|
||||
custom_llm_provider: str | None,
|
||||
|
|
@ -639,29 +689,12 @@ def compute_savings_spend(
|
|||
# Deployment rates when the request came through one, public rates otherwise --
|
||||
# `_effective_model_info` merges a deployment's configured prices over the built-in
|
||||
# map, so a negotiated price is not silently replaced by the list rate.
|
||||
router_instance: Router | None = llm_router() if llm_router else None
|
||||
identity: Final = _resolve_model(model, custom_llm_provider)
|
||||
pricing: Final = _effective_model_info(router_instance, model_id, model or "") or (
|
||||
_model_info(identity) if identity else None
|
||||
)
|
||||
request_pricing: Final = _request_savings_pricing(model, custom_llm_provider, model_id, llm_router)
|
||||
provider: Final = request_pricing[0]
|
||||
pricing: Final = request_pricing[1]
|
||||
input_cost: Final = (_get_cost_per_unit(pricing, "input_cost_per_token") or 0.0) if pricing else 0.0
|
||||
compression: Final = max(compression_saved_tokens, 0) * input_cost
|
||||
usage: Final = _usage_from_spend_log(usage_object)
|
||||
basis: Final = _pricing_basis(cost_breakdown)
|
||||
billed_at_datetime: Final = _coerce_billed_at(billed_at)
|
||||
prompt_caching: Final = (
|
||||
calculate_prompt_caching_savings(
|
||||
model_info=pricing,
|
||||
usage=usage,
|
||||
custom_llm_provider=identity.provider if identity else custom_llm_provider,
|
||||
service_tier=basis.service_tier,
|
||||
data_residency=basis.data_residency,
|
||||
vertex_location=basis.vertex_location,
|
||||
billed_at=billed_at_datetime,
|
||||
)
|
||||
if pricing is not None and usage is not None
|
||||
else 0.0
|
||||
)
|
||||
prompt_caching: Final = _prompt_caching_savings(pricing, provider, usage_object, cost_breakdown, billed_at) or 0.0
|
||||
gateway_injected_caching: Final = prompt_caching if gateway_injected_cache else 0.0
|
||||
|
||||
# The figure the logging path recorded wins, before the usage gate on purpose: a row
|
||||
|
|
|
|||
|
|
@ -4,12 +4,19 @@ Wrapper around router cache. Meant to store model id when prompt caching support
|
|||
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from pydantic_core import to_jsonable_python
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import PROMPT_CACHE_LOOKBACK_POSITIONS
|
||||
from litellm.litellm_core_utils.logging_utils import truncate_base64_in_messages
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -28,27 +35,102 @@ class PromptCachingCacheValue(TypedDict):
|
|||
model_id: str
|
||||
|
||||
|
||||
PROMPT_CACHE_PIN_TTL_SECONDS: Final = 300
|
||||
_TOOL_RUN_BLOCK_TYPES: Final = frozenset({"tool_use", "tool_result"})
|
||||
_PREFIX_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...])
|
||||
_TOOLS_ADAPTER: Final = TypeAdapter(tuple[JsonValue, ...])
|
||||
_PINS_ADAPTER: Final[TypeAdapter[tuple[JsonValue, ...] | None]] = TypeAdapter(tuple[JsonValue, ...] | None)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PrefixPosition:
|
||||
cache_key: str
|
||||
position: int
|
||||
|
||||
|
||||
def _sorted_pairs(pairs: Iterable[tuple[str, JsonValue]]) -> tuple[tuple[str, JsonValue], ...]:
|
||||
return tuple(sorted(pairs, key=lambda pair: pair[0]))
|
||||
|
||||
|
||||
def _canonical_bytes(value: object) -> bytes:
|
||||
return json.dumps(value, sort_keys=True, separators=(",", ":")).encode()
|
||||
|
||||
|
||||
def _block_unit(
|
||||
envelope: tuple[tuple[str, JsonValue], ...], message_run_type: str | None, block: JsonValue
|
||||
) -> tuple[bytes, str | None]:
|
||||
if not isinstance(block, dict):
|
||||
return _canonical_bytes((envelope, block)), message_run_type
|
||||
block_type: Final = block.get("type")
|
||||
block_run_type: Final = block_type if isinstance(block_type, str) and block_type in _TOOL_RUN_BLOCK_TYPES else None
|
||||
stripped: Final = _sorted_pairs(item for item in block.items() if item[0] != "cache_control")
|
||||
return _canonical_bytes((envelope, stripped)), message_run_type or block_run_type
|
||||
|
||||
|
||||
def _message_units(message: Mapping[str, JsonValue]) -> tuple[tuple[bytes, str | None], ...]:
|
||||
envelope: Final = _sorted_pairs(item for item in message.items() if item[0] not in ("content", "cache_control"))
|
||||
message_run_type: Final = "tool_result" if message.get("role") == "tool" else None
|
||||
content: Final = message.get("content")
|
||||
if isinstance(content, list) and content:
|
||||
return tuple(_block_unit(envelope, message_run_type, block) for block in content)
|
||||
if isinstance(content, str) and content:
|
||||
return ((_canonical_bytes((envelope, (("text", content), ("type", "text")))), message_run_type),)
|
||||
return ((_canonical_bytes((envelope, None)), message_run_type),)
|
||||
|
||||
|
||||
def _chain_digest(digest: bytes, unit: bytes) -> bytes:
|
||||
return hashlib.sha256(digest + unit).digest()
|
||||
|
||||
|
||||
def _seed(tools: Sequence[ChatCompletionToolParam] | None) -> bytes:
|
||||
if tools is None:
|
||||
return hashlib.sha256(b"").digest()
|
||||
return hashlib.sha256(
|
||||
_canonical_bytes(
|
||||
_TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True, bytes_mode="base64"))
|
||||
)
|
||||
).digest()
|
||||
|
||||
|
||||
def _positions_of(
|
||||
prefix: tuple[Mapping[str, JsonValue], ...], tools: Sequence[ChatCompletionToolParam] | None
|
||||
) -> tuple[PrefixPosition, ...]:
|
||||
units: Final = tuple(unit for message in prefix for unit in _message_units(message))
|
||||
digests: Final = tuple(accumulate((unit_bytes for unit_bytes, _ in units), _chain_digest, initial=_seed(tools)))[1:]
|
||||
run_types: Final = tuple(run_type for _, run_type in units)
|
||||
positions: Final = accumulate(
|
||||
0 if run_type is not None and run_type == previous else 1
|
||||
for run_type, previous in zip(run_types, (None, *run_types[:-1]))
|
||||
)
|
||||
return tuple(
|
||||
PrefixPosition(cache_key=f"deployment:{digest.hex()}:prompt_caching", position=position)
|
||||
for digest, position in zip(digests, positions)
|
||||
)
|
||||
|
||||
|
||||
def _lookback_keys(positions: tuple[PrefixPosition, ...]) -> tuple[str, ...]:
|
||||
if not positions:
|
||||
return ()
|
||||
oldest_probed_position: Final = positions[-1].position - PROMPT_CACHE_LOOKBACK_POSITIONS
|
||||
return tuple(entry.cache_key for entry in reversed(positions) if entry.position > oldest_probed_position)
|
||||
|
||||
|
||||
def _pinned_value(value: JsonValue) -> PromptCachingCacheValue | None:
|
||||
if not isinstance(value, dict):
|
||||
return None
|
||||
model_id: Final = value.get("model_id")
|
||||
return PromptCachingCacheValue(model_id=model_id) if isinstance(model_id, str) else None
|
||||
|
||||
|
||||
def _first_pin(values: tuple[JsonValue, ...] | None) -> PromptCachingCacheValue | None:
|
||||
if values is None:
|
||||
return None
|
||||
return next((pin for pin in map(_pinned_value, values) if pin is not None), None)
|
||||
|
||||
|
||||
class PromptCachingCache:
|
||||
def __init__(self, cache: DualCache):
|
||||
self.cache = cache
|
||||
self.in_memory_cache = InMemoryCache()
|
||||
|
||||
@staticmethod
|
||||
def serialize_object(obj: Any) -> object:
|
||||
"""Helper function to serialize Pydantic objects, dictionaries, or fallback to string."""
|
||||
if hasattr(obj, "dict"):
|
||||
# If the object is a Pydantic model, use its `dict()` method
|
||||
return obj.dict()
|
||||
elif isinstance(obj, dict):
|
||||
# If the object is a dictionary, serialize it with sorted keys
|
||||
return json.dumps(obj, sort_keys=True, separators=(",", ":")) # Standardize serialization
|
||||
|
||||
elif isinstance(obj, list):
|
||||
# Serialize lists by ensuring each element is handled properly
|
||||
return [PromptCachingCache.serialize_object(item) for item in obj]
|
||||
elif isinstance(obj, (int, float, bool)):
|
||||
return obj # Keep primitive types as-is
|
||||
return str(obj)
|
||||
|
||||
@staticmethod
|
||||
def extract_cacheable_prefix(
|
||||
|
|
@ -140,114 +222,116 @@ class PromptCachingCache:
|
|||
return cacheable_prefix
|
||||
|
||||
@staticmethod
|
||||
def get_prompt_caching_cache_key(
|
||||
def prefix_positions(
|
||||
messages: list[AllMessageValues] | None,
|
||||
tools: list[ChatCompletionToolParam] | None,
|
||||
) -> str | None:
|
||||
if messages is None and tools is None:
|
||||
return None
|
||||
tools: Sequence[ChatCompletionToolParam] | None,
|
||||
) -> tuple[PrefixPosition, ...]:
|
||||
"""
|
||||
One cache key per content block of the cacheable prefix, oldest block first.
|
||||
|
||||
# Extract cacheable prefix from messages (only include up to last cache_control block)
|
||||
cacheable_messages = None
|
||||
if messages is not None:
|
||||
cacheable_messages = PromptCachingCache.extract_cacheable_prefix(messages)
|
||||
# If no cacheable prefix found, return None (can't cache)
|
||||
if not cacheable_messages:
|
||||
return None
|
||||
Each key hashes the prefix content up to and including that block, with cache_control markers
|
||||
left out, so the key of a block is the same whichever turn's breakpoint the prefix ends at.
|
||||
String content hashes like a single text block, which is how the provider treats it and how
|
||||
Claude Code re-sends a previously marked message. `position` counts a run of consecutive
|
||||
tool_use (or tool_result) blocks as one, matching the provider's lookback window.
|
||||
|
||||
# Use serialize_object for consistent and stable serialization
|
||||
data_to_hash: Final = {}
|
||||
if cacheable_messages is not None:
|
||||
serialized_messages: Final = PromptCachingCache.serialize_object(cacheable_messages)
|
||||
data_to_hash["messages"] = serialized_messages
|
||||
if tools is not None:
|
||||
serialized_tools: Final = PromptCachingCache.serialize_object(tools)
|
||||
data_to_hash["tools"] = serialized_tools
|
||||
|
||||
# Combine serialized data into a single string
|
||||
data_to_hash_str: Final = json.dumps(
|
||||
data_to_hash,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
The prefix is hashed in the shape the success event sees it, with long base64 data URIs
|
||||
already replaced by their size placeholder, so a request carrying the raw image bytes
|
||||
derives the same keys the write side stored.
|
||||
"""
|
||||
if not messages:
|
||||
return ()
|
||||
return _positions_of(
|
||||
_PREFIX_ADAPTER.validate_python(
|
||||
to_jsonable_python(
|
||||
truncate_base64_in_messages(PromptCachingCache.extract_cacheable_prefix(messages)),
|
||||
serialize_unknown=True,
|
||||
bytes_mode="base64",
|
||||
)
|
||||
),
|
||||
tools,
|
||||
)
|
||||
|
||||
# Create a hash of the serialized data for a stable cache key
|
||||
hashed_data: Final = hashlib.sha256(data_to_hash_str.encode()).hexdigest()
|
||||
return f"deployment:{hashed_data}:prompt_caching"
|
||||
@staticmethod
|
||||
async def async_prefix_positions(
|
||||
messages: list[AllMessageValues] | None,
|
||||
tools: Sequence[ChatCompletionToolParam] | None,
|
||||
) -> tuple[PrefixPosition, ...]:
|
||||
if not messages:
|
||||
return ()
|
||||
return await offload_token_count(PromptCachingCache.prefix_positions)(messages, tools)
|
||||
|
||||
@staticmethod
|
||||
def get_prompt_caching_cache_key(
|
||||
messages: list[AllMessageValues] | None,
|
||||
tools: Sequence[ChatCompletionToolParam] | None,
|
||||
) -> str | None:
|
||||
positions: Final = PromptCachingCache.prefix_positions(messages, tools)
|
||||
return positions[-1].cache_key if positions else None
|
||||
|
||||
def add_model_id(
|
||||
self,
|
||||
model_id: str,
|
||||
messages: list[AllMessageValues] | None,
|
||||
tools: list[ChatCompletionToolParam] | None,
|
||||
tools: Sequence[ChatCompletionToolParam] | None,
|
||||
) -> None:
|
||||
if messages is None and tools is None:
|
||||
return
|
||||
|
||||
cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
|
||||
# If no cacheable prefix found, don't cache (can't generate cache key)
|
||||
if cache_key is None:
|
||||
return
|
||||
|
||||
self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=300)
|
||||
return
|
||||
self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=PROMPT_CACHE_PIN_TTL_SECONDS)
|
||||
|
||||
async def async_add_model_id(
|
||||
self,
|
||||
model_id: str,
|
||||
messages: list[AllMessageValues] | None,
|
||||
tools: list[ChatCompletionToolParam] | None,
|
||||
tools: Sequence[ChatCompletionToolParam] | None,
|
||||
) -> None:
|
||||
if messages is None and tools is None:
|
||||
return
|
||||
|
||||
cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
|
||||
# If no cacheable prefix found, don't cache (can't generate cache key)
|
||||
if cache_key is None:
|
||||
positions: Final = await PromptCachingCache.async_prefix_positions(messages, tools)
|
||||
if not positions:
|
||||
return
|
||||
|
||||
await self.cache.async_set_cache(
|
||||
cache_key,
|
||||
positions[-1].cache_key,
|
||||
PromptCachingCacheValue(model_id=model_id),
|
||||
ttl=300, # store for 5 minutes
|
||||
ttl=PROMPT_CACHE_PIN_TTL_SECONDS,
|
||||
)
|
||||
return
|
||||
|
||||
async def async_get_model_id(
|
||||
self,
|
||||
messages: list[AllMessageValues] | None,
|
||||
tools: list[ChatCompletionToolParam] | None,
|
||||
tools: Sequence[ChatCompletionToolParam] | None,
|
||||
) -> PromptCachingCacheValue | None:
|
||||
"""
|
||||
Get model ID from cache using the cacheable prefix.
|
||||
|
||||
The cache key is based on the cacheable prefix (everything up to and including
|
||||
the last cache_control block), so requests with the same cacheable prefix but
|
||||
different user messages will have the same cache key.
|
||||
Find the deployment that last served this prefix, walking back from the breakpoint the
|
||||
same way the provider cache does, so a breakpoint that moved forward since the last
|
||||
turn still lands on the deployment whose cache holds the earlier prefix.
|
||||
"""
|
||||
if messages is None and tools is None:
|
||||
cache_keys: Final = _lookback_keys(await PromptCachingCache.async_prefix_positions(messages, tools))
|
||||
if not cache_keys:
|
||||
return None
|
||||
|
||||
# Generate cache key using cacheable prefix
|
||||
cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
|
||||
if cache_key is None:
|
||||
return None
|
||||
|
||||
# Perform cache lookup
|
||||
cache_result: Final = await self.cache.async_get_cache(key=cache_key)
|
||||
return cache_result
|
||||
return _first_pin(
|
||||
_PINS_ADAPTER.validate_python(
|
||||
await self.cache.async_batch_get_cache(
|
||||
keys=list(cache_keys), # mutable-ok: DualCache.async_batch_get_cache only takes a list
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
def get_model_id(
|
||||
self,
|
||||
messages: list[AllMessageValues] | None,
|
||||
tools: list[ChatCompletionToolParam] | None,
|
||||
tools: Sequence[ChatCompletionToolParam] | None,
|
||||
) -> PromptCachingCacheValue | None:
|
||||
if messages is None and tools is None:
|
||||
cache_keys: Final = _lookback_keys(PromptCachingCache.prefix_positions(messages, tools))
|
||||
if not cache_keys:
|
||||
return None
|
||||
|
||||
cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
|
||||
# If no cacheable prefix found, return None (can't cache)
|
||||
if cache_key is None:
|
||||
return None
|
||||
|
||||
return self.cache.get_cache(cache_key)
|
||||
return _first_pin(
|
||||
_PINS_ADAPTER.validate_python(
|
||||
self.cache.batch_get_cache(
|
||||
keys=list(cache_keys), # mutable-ok: DualCache.batch_get_cache only takes a list
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,35 @@
|
|||
from datetime import datetime
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
PromptCachingRequestFilter: TypeAlias = Literal["all", "injected", "hits"]
|
||||
|
||||
|
||||
class PromptCachingRequest(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
request_id: str
|
||||
start_time: datetime
|
||||
model: str
|
||||
gateway_injected: bool
|
||||
cache_read_tokens: int
|
||||
cache_creation_tokens: int
|
||||
spend: float
|
||||
net_savings: float | None
|
||||
|
||||
|
||||
class PromptCachingRequestCursor(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
start_time: datetime
|
||||
request_id: str
|
||||
|
||||
|
||||
class PromptCachingRequestsResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
requests: tuple[PromptCachingRequest, ...]
|
||||
page_size: int
|
||||
has_more: bool
|
||||
next_cursor: PromptCachingRequestCursor | None
|
||||
|
|
@ -294,6 +294,10 @@ class PolicyAttachment(BaseModel):
|
|||
le=2147483647,
|
||||
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
|
||||
)
|
||||
default: bool = Field(
|
||||
default=False,
|
||||
description="Apply this attachment only when no non-default attachment matches the request.",
|
||||
)
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
|
|
|||
|
|
@ -311,6 +311,10 @@ class PolicyAttachmentCreateRequest(BaseModel):
|
|||
le=2147483647,
|
||||
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
|
||||
)
|
||||
default: bool = Field(
|
||||
default=False,
|
||||
description="Apply this attachment only when no non-default attachment matches the request.",
|
||||
)
|
||||
|
||||
|
||||
class PolicyAttachmentDBResponse(BaseModel):
|
||||
|
|
@ -327,6 +331,10 @@ class PolicyAttachmentDBResponse(BaseModel):
|
|||
default=None,
|
||||
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
|
||||
)
|
||||
default: bool = Field(
|
||||
default=False,
|
||||
description="Apply this attachment only when no non-default attachment matches the request.",
|
||||
)
|
||||
created_at: datetime | None = Field(default=None, description="When the attachment was created.")
|
||||
updated_at: datetime | None = Field(default=None, description="When the attachment was last updated.")
|
||||
created_by: str | None = Field(default=None, description="Who created the attachment.")
|
||||
|
|
|
|||
|
|
@ -5624,6 +5624,12 @@ def _get_model_info_from_generalization(
|
|||
return None
|
||||
|
||||
|
||||
def _strip_mantle_region_prefix(model: str) -> str:
|
||||
from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix
|
||||
|
||||
return split_mantle_region_prefix(model)[1]
|
||||
|
||||
|
||||
def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> PotentialModelNamesAndCustomLLMProvider:
|
||||
if custom_llm_provider is None:
|
||||
# Get custom_llm_provider
|
||||
|
|
@ -5656,20 +5662,30 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P
|
|||
|
||||
split_model = strip_bedrock_routing_prefix(split_model)
|
||||
|
||||
region_free_split_model: Final = (
|
||||
_strip_mantle_region_prefix(split_model) if custom_llm_provider == "bedrock_mantle" else split_model
|
||||
)
|
||||
region_free_combined_stripped_model_name: Final = (
|
||||
f"bedrock_mantle/{_strip_model_name(model=region_free_split_model, custom_llm_provider=custom_llm_provider)}"
|
||||
if custom_llm_provider == "bedrock_mantle"
|
||||
else combined_stripped_model_name
|
||||
)
|
||||
provider_model_info: Final = (
|
||||
ProviderConfigManager.get_provider_model_info(model=split_model, provider=LlmProviders(custom_llm_provider))
|
||||
ProviderConfigManager.get_provider_model_info(
|
||||
model=region_free_split_model, provider=LlmProviders(custom_llm_provider)
|
||||
)
|
||||
if custom_llm_provider in LlmProvidersSet
|
||||
else None
|
||||
)
|
||||
provider_cost_key: Final = (
|
||||
provider_model_info.get_model_cost_key(split_model) if provider_model_info is not None else None
|
||||
provider_model_info.get_model_cost_key(region_free_split_model) if provider_model_info is not None else None
|
||||
)
|
||||
|
||||
return PotentialModelNamesAndCustomLLMProvider(
|
||||
split_model=split_model,
|
||||
split_model=region_free_split_model,
|
||||
combined_model_name=combined_model_name,
|
||||
stripped_model_name=stripped_model_name,
|
||||
combined_stripped_model_name=combined_stripped_model_name,
|
||||
combined_stripped_model_name=region_free_combined_stripped_model_name,
|
||||
provider_prefixed_model_name=provider_cost_key or provider_prefixed_model_name,
|
||||
custom_llm_provider=cast(str, custom_llm_provider),
|
||||
)
|
||||
|
|
@ -8681,6 +8697,13 @@ class ProviderConfigManager:
|
|||
from litellm.llms.bedrock.common_utils import BedrockModelInfo
|
||||
|
||||
return BedrockModelInfo.get_bedrock_provider_config_for_messages_api(model)
|
||||
elif litellm.LlmProviders.BEDROCK_MANTLE == provider:
|
||||
if "claude" in model_lower:
|
||||
from litellm.llms.bedrock_mantle.messages.transformation import (
|
||||
BedrockMantleAnthropicMessagesConfig,
|
||||
)
|
||||
|
||||
return BedrockMantleAnthropicMessagesConfig()
|
||||
elif litellm.LlmProviders.VERTEX_AI == provider:
|
||||
if "claude" in model_lower:
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import (
|
||||
|
|
|
|||
|
|
@ -42971,21 +42971,21 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/deepseek/deepseek-v4-pro": {
|
||||
"input_cost_per_token": 9.22722e-07,
|
||||
"input_cost_per_token": 9.19242e-07,
|
||||
"input_cost_per_token_cache_hit": 4.4e-08,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 384000,
|
||||
"max_tokens": 384000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.845444e-06,
|
||||
"output_cost_per_token": 1.838484e-06,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 7.68935e-08,
|
||||
"cache_read_input_token_cost": 7.66035e-08,
|
||||
"supports_audio_input": false,
|
||||
"supports_pdf_input": false,
|
||||
"supports_vision": false,
|
||||
|
|
@ -54441,16 +54441,19 @@
|
|||
"zai.glm-4.7": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"max_input_tokens": 203000,
|
||||
"max_output_tokens": 4000,
|
||||
"max_tokens": 4000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.2e-06,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/"
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/",
|
||||
"supports_audio_input": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"zai.glm-5": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
|
|
@ -54470,16 +54473,19 @@
|
|||
"zai.glm-4.7-flash": {
|
||||
"input_cost_per_token": 7e-08,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"max_input_tokens": 203000,
|
||||
"max_output_tokens": 4000,
|
||||
"max_tokens": 4000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4e-07,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/"
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/",
|
||||
"supports_audio_input": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"zai/glm-5": {
|
||||
"cache_creation_input_token_cost": 0,
|
||||
|
|
@ -60558,6 +60564,34 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"bedrock_mantle/anthropic.claude-haiku-4-5": {
|
||||
"cache_creation_input_token_cost": 1.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-06,
|
||||
"cache_read_input_token_cost": 1e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"supports_tool_search": true,
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 64000,
|
||||
"max_tokens": 64000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-06,
|
||||
"source": "https://aws.amazon.com/about-aws/whats-new/2025/10/claude-4-5-haiku-anthropic-amazon-bedrock",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_parallel_tool_use_config": true,
|
||||
"prompt_cache_min_tokens": 4096,
|
||||
"input_cost_per_token_batches": 5e-07,
|
||||
"output_cost_per_token_batches": 2.5e-06
|
||||
},
|
||||
"us.xai.grok-4.6": {
|
||||
"input_cost_per_token": 2.2e-06,
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
|
|
@ -76803,7 +76837,7 @@
|
|||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3.html",
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
|
|
|
|||
|
|
@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable {
|
|||
models String[] @default([]) // Model names or patterns
|
||||
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
|
||||
priority Int? // Explicit execution order
|
||||
is_default Boolean @default(false) // Applied only when no non-default attachment matches
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
|
|
|
|||
|
|
@ -2003,7 +2003,7 @@ def test_provider_specific_header():
|
|||
)
|
||||
# Verify multi-provider support: anthropic headers work across multiple providers
|
||||
assert data["provider_specific_header"] == {
|
||||
"custom_llm_provider": "anthropic,bedrock,vertex_ai",
|
||||
"custom_llm_provider": "anthropic,bedrock,bedrock_mantle,vertex_ai",
|
||||
"extra_headers": {
|
||||
"anthropic-beta": "prompt-caching-2024-07-31",
|
||||
},
|
||||
|
|
@ -2075,7 +2075,7 @@ def test_provider_specific_header_multi_provider():
|
|||
assert "provider_specific_header" in data
|
||||
assert (
|
||||
data["provider_specific_header"]["custom_llm_provider"]
|
||||
== "anthropic,bedrock,vertex_ai"
|
||||
== "anthropic,bedrock,bedrock_mantle,vertex_ai"
|
||||
)
|
||||
assert data["provider_specific_header"]["extra_headers"] == {
|
||||
"anthropic-beta": "context-1m-2025-08-07",
|
||||
|
|
|
|||
|
|
@ -11,57 +11,9 @@ from unittest.mock import patch, MagicMock, AsyncMock
|
|||
from create_mock_standard_logging_payload import create_standard_logging_payload
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
import unittest
|
||||
from pydantic import BaseModel
|
||||
from litellm.router_utils.prompt_caching_cache import PromptCachingCache
|
||||
|
||||
|
||||
class ExampleModel(BaseModel):
|
||||
field1: str
|
||||
field2: int
|
||||
|
||||
|
||||
def test_serialize_pydantic_object():
|
||||
model = ExampleModel(field1="value", field2=42)
|
||||
serialized = PromptCachingCache.serialize_object(model)
|
||||
assert serialized == {"field1": "value", "field2": 42}
|
||||
|
||||
|
||||
def test_serialize_dict():
|
||||
obj = {"b": 2, "a": 1}
|
||||
serialized = PromptCachingCache.serialize_object(obj)
|
||||
assert serialized == '{"a":1,"b":2}' # JSON string with sorted keys
|
||||
|
||||
|
||||
def test_serialize_nested_dict():
|
||||
obj = {"z": {"b": 2, "a": 1}, "x": [1, 2, {"c": 3}]}
|
||||
serialized = PromptCachingCache.serialize_object(obj)
|
||||
expected = '{"x":[1,2,{"c":3}],"z":{"a":1,"b":2}}' # JSON string with sorted keys
|
||||
assert serialized == expected
|
||||
|
||||
|
||||
def test_serialize_list():
|
||||
obj = ["item1", {"a": 1, "b": 2}, 42]
|
||||
serialized = PromptCachingCache.serialize_object(obj)
|
||||
expected = ["item1", '{"a":1,"b":2}', 42]
|
||||
assert serialized == expected
|
||||
|
||||
|
||||
def test_serialize_fallback():
|
||||
obj = 12345 # Simple non-serializable object
|
||||
serialized = PromptCachingCache.serialize_object(obj)
|
||||
assert serialized == 12345
|
||||
|
||||
|
||||
def test_serialize_non_serializable():
|
||||
class CustomClass:
|
||||
def __str__(self):
|
||||
return "custom_object"
|
||||
|
||||
obj = CustomClass()
|
||||
serialized = PromptCachingCache.serialize_object(obj)
|
||||
assert serialized == "custom_object" # Fallback to string conversion
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_prompt_caching_same_cacheable_prefix_routes_to_same_deployment():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -2036,6 +2036,15 @@ async def test_optional_discovery_preserves_cancellation(method: str) -> None:
|
|||
},
|
||||
},
|
||||
)
|
||||
if not (payload.params or {}).get("cursor"):
|
||||
field: Final = {
|
||||
"prompts/list": "prompts",
|
||||
"resources/list": "resources",
|
||||
"resources/templates/list": "resourceTemplates",
|
||||
}[method]
|
||||
return httpx2.Response(
|
||||
200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [], "nextCursor": "pending-page"}}
|
||||
)
|
||||
ready.set()
|
||||
await pending.wait()
|
||||
return httpx2.Response(202)
|
||||
|
|
@ -2055,6 +2064,255 @@ async def test_optional_discovery_preserves_cancellation(method: str) -> None:
|
|||
await asyncio.wait_for(task, timeout=3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list"))
|
||||
@pytest.mark.parametrize("session_id", (None, "pagination-session"))
|
||||
@pytest.mark.parametrize("empty_middle", (False, True))
|
||||
async def test_optional_discovery_collects_all_pages(method: str, session_id: str | None, empty_middle: bool) -> None:
|
||||
from mcp.types import Prompt, PromptArgument, Resource, ResourceTemplate
|
||||
|
||||
field: Final = {
|
||||
"prompts/list": "prompts",
|
||||
"resources/list": "resources",
|
||||
"resources/templates/list": "resourceTemplates",
|
||||
}[method]
|
||||
entries: Final = tuple(
|
||||
{
|
||||
"prompts/list": Prompt(
|
||||
name=f"item-{index}",
|
||||
description="prompt description",
|
||||
arguments=[PromptArgument(name="query", required=True)],
|
||||
),
|
||||
"resources/list": Resource(
|
||||
name=f"item-{index}",
|
||||
uri=f"test://item/{index}",
|
||||
mime_type="text/plain",
|
||||
description="resource description",
|
||||
),
|
||||
"resources/templates/list": ResourceTemplate(
|
||||
name=f"item-{index}", uri_template=f"test://item/{index}/{{query}}", mime_type="text/plain"
|
||||
),
|
||||
}[method]
|
||||
for index in range(5)
|
||||
)
|
||||
|
||||
def respond(request: httpx2.Request) -> httpx2.Response:
|
||||
if request.method == "GET":
|
||||
return httpx2.Response(405)
|
||||
if request.method == "DELETE":
|
||||
return httpx2.Response(200)
|
||||
payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content)
|
||||
if not isinstance(payload, JSONRPCRequest):
|
||||
return httpx2.Response(202)
|
||||
if payload.method == "initialize":
|
||||
return httpx2.Response(
|
||||
200,
|
||||
headers={"mcp-session-id": session_id} if session_id else {},
|
||||
json={
|
||||
"jsonrpc": "2.0",
|
||||
"id": payload.id,
|
||||
"result": {
|
||||
"protocolVersion": payload.params["protocolVersion"],
|
||||
"capabilities": {"prompts": {}, "resources": {}},
|
||||
"serverInfo": {"name": "paged", "version": "1"},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert payload.method == method
|
||||
assert request.headers.get("mcp-session-id") == session_id
|
||||
cursor: Final = (payload.params or {}).get("cursor")
|
||||
assert cursor in (None, "opaque:/second+page", "opaque:/last+page")
|
||||
page: Final = (
|
||||
entries[:3] if cursor is None else (() if empty_middle and cursor == "opaque:/second+page" else entries[3:])
|
||||
)
|
||||
next_cursor: Final = (
|
||||
"opaque:/second+page"
|
||||
if cursor is None
|
||||
else "opaque:/last+page"
|
||||
if empty_middle and cursor == "opaque:/second+page"
|
||||
else ""
|
||||
)
|
||||
return httpx2.Response(
|
||||
200,
|
||||
json={
|
||||
"jsonrpc": "2.0",
|
||||
"id": payload.id,
|
||||
"result": {
|
||||
field: [item.model_dump(mode="json", by_alias=True) for item in page],
|
||||
"nextCursor": next_cursor,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
responder: Final = Mock(side_effect=respond)
|
||||
client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp")
|
||||
operation: Final = {
|
||||
"prompts/list": client.list_prompts,
|
||||
"resources/list": client.list_resources,
|
||||
"resources/templates/list": client.list_resource_templates,
|
||||
}[method]
|
||||
assert await operation(raise_on_error=True) == list(entries)
|
||||
requests: Final = tuple(
|
||||
_JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content)
|
||||
for call in responder.call_args_list
|
||||
if call.args[0].method == "POST"
|
||||
)
|
||||
assert sum(isinstance(request, JSONRPCRequest) and request.method == "initialize" for request in requests) == 1
|
||||
assert tuple(
|
||||
(request.params or {}).get("cursor")
|
||||
for request in requests
|
||||
if isinstance(request, JSONRPCRequest) and request.method == method
|
||||
) == ((None, "opaque:/second+page", "opaque:/last+page") if empty_middle else (None, "opaque:/second+page"))
|
||||
assert sum(call.args[0].method == "DELETE" for call in responder.call_args_list) == (1 if session_id else 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list"))
|
||||
@pytest.mark.parametrize(
|
||||
"failure", ("repeat", "cycle", "cap", "method_not_found", "internal_error", "unauthorized", "deadline")
|
||||
)
|
||||
@pytest.mark.parametrize("strict", (False, True))
|
||||
async def test_optional_discovery_rejects_incomplete_walks(
|
||||
method: str, failure: str, strict: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_MAX_PAGES", 3 if failure == "cycle" else 2, raising=False)
|
||||
monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_TIMEOUT", 0.05)
|
||||
field: Final = {
|
||||
"prompts/list": "prompts",
|
||||
"resources/list": "resources",
|
||||
"resources/templates/list": "resourceTemplates",
|
||||
}[method]
|
||||
entry: Final = {
|
||||
"prompts/list": {"name": "first"},
|
||||
"resources/list": {"name": "first", "uri": "test://first"},
|
||||
"resources/templates/list": {"name": "first", "uriTemplate": "test://{name}"},
|
||||
}[method]
|
||||
cancelled: Final = asyncio.Event()
|
||||
|
||||
async def respond(request: httpx2.Request) -> httpx2.Response:
|
||||
payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content)
|
||||
if not isinstance(payload, JSONRPCRequest):
|
||||
return httpx2.Response(202)
|
||||
if payload.method == "initialize":
|
||||
return httpx2.Response(
|
||||
200,
|
||||
json={
|
||||
"jsonrpc": "2.0",
|
||||
"id": payload.id,
|
||||
"result": {
|
||||
"protocolVersion": payload.params["protocolVersion"],
|
||||
"capabilities": {"prompts": {}, "resources": {}},
|
||||
"serverInfo": {"name": "interrupted", "version": "1"},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert payload.method == method
|
||||
cursor: Final = (payload.params or {}).get("cursor")
|
||||
if cursor is not None:
|
||||
if failure == "deadline":
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
cancelled.set()
|
||||
if failure == "unauthorized":
|
||||
return httpx2.Response(401)
|
||||
if failure in ("method_not_found", "internal_error"):
|
||||
return httpx2.Response(
|
||||
200,
|
||||
json={
|
||||
"jsonrpc": "2.0",
|
||||
"id": payload.id,
|
||||
"error": {
|
||||
"code": -32601 if failure == "method_not_found" else -32603,
|
||||
"message": "Later page unavailable",
|
||||
},
|
||||
},
|
||||
)
|
||||
next_cursor: Final = (
|
||||
"private-cursor-2" if cursor == "private-cursor-1" and failure != "repeat" else "private-cursor-1"
|
||||
)
|
||||
return httpx2.Response(
|
||||
200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry], "nextCursor": next_cursor}}
|
||||
)
|
||||
|
||||
responder: Final = AsyncMock(side_effect=respond)
|
||||
client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp", timeout=0.2)
|
||||
operation: Final = {
|
||||
"prompts/list": client.list_prompts,
|
||||
"resources/list": client.list_resources,
|
||||
"resources/templates/list": client.list_resource_templates,
|
||||
}[method]
|
||||
if strict:
|
||||
error_type: Final = {
|
||||
"internal_error": MCPError,
|
||||
"unauthorized": httpx2.HTTPStatusError,
|
||||
"deadline": TimeoutError,
|
||||
}.get(failure, RuntimeError)
|
||||
with pytest.raises(error_type):
|
||||
await operation(raise_on_error=True)
|
||||
else:
|
||||
assert await operation() == []
|
||||
assert len(
|
||||
tuple(
|
||||
payload
|
||||
for call in responder.call_args_list
|
||||
if isinstance(payload := _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content), JSONRPCRequest)
|
||||
and payload.method == method
|
||||
)
|
||||
) == (3 if failure == "cycle" else 2)
|
||||
assert "private-cursor" not in caplog.text
|
||||
if failure == "deadline":
|
||||
assert cancelled.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list"))
|
||||
async def test_optional_discovery_allows_exhaustion_at_page_cap(method: str, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_MAX_PAGES", 2, raising=False)
|
||||
field: Final = {
|
||||
"prompts/list": "prompts",
|
||||
"resources/list": "resources",
|
||||
"resources/templates/list": "resourceTemplates",
|
||||
}[method]
|
||||
|
||||
def respond(request: httpx2.Request) -> httpx2.Response:
|
||||
payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content)
|
||||
if not isinstance(payload, JSONRPCRequest):
|
||||
return httpx2.Response(202)
|
||||
if payload.method == "initialize":
|
||||
result: Final = {
|
||||
"protocolVersion": payload.params["protocolVersion"],
|
||||
"capabilities": {"prompts": {}, "resources": {}},
|
||||
"serverInfo": {"name": "empty-pages", "version": "1"},
|
||||
}
|
||||
return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result})
|
||||
assert payload.method == method
|
||||
return httpx2.Response(
|
||||
200,
|
||||
json={
|
||||
"jsonrpc": "2.0",
|
||||
"id": payload.id,
|
||||
"result": {field: [], "nextCursor": None if (payload.params or {}).get("cursor") else "last-page"},
|
||||
},
|
||||
)
|
||||
|
||||
responder: Final = Mock(side_effect=respond)
|
||||
client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp")
|
||||
operation: Final = {
|
||||
"prompts/list": client.list_prompts,
|
||||
"resources/list": client.list_resources,
|
||||
"resources/templates/list": client.list_resource_templates,
|
||||
}[method]
|
||||
assert await operation(raise_on_error=True) == []
|
||||
assert (
|
||||
sum(
|
||||
isinstance(payload := _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content), JSONRPCRequest)
|
||||
and payload.method == method
|
||||
for call in responder.call_args_list
|
||||
)
|
||||
== 2
|
||||
)
|
||||
|
||||
|
||||
def test_client_import_before_proxy_credentials_succeeds_in_fresh_process():
|
||||
import subprocess
|
||||
|
|
|
|||
|
|
@ -1257,6 +1257,25 @@ def test_token_counter_with_thinking_content():
|
|||
), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}"
|
||||
|
||||
|
||||
|
||||
def test_token_counter_with_redacted_thinking_content():
|
||||
"""
|
||||
A replayed redacted_thinking block (Anthropic redacted reasoning, or the /v1/messages bridge's stand-in
|
||||
for a reasoning item with no summary) counts zero tokens for its encrypted payload, like a thinking
|
||||
block with no text. It used to raise, which made is_prompt_caching_valid_prompt return False and the
|
||||
prompt_caching pre-call check stop pinning the deployment that held the cached prefix.
|
||||
"""
|
||||
model = "anthropic/claude-sonnet-4-5-20250929"
|
||||
reply = {"type": "text", "text": "Draw from the box labeled Mixed, because that label must be wrong."}
|
||||
redacted_block = {"type": "redacted_thinking", "data": "EqQBCkYIBRgCKkBjZ2xhc3M" * 30}
|
||||
user_turn = {"role": "user", "content": [{"type": "text", "text": "Which box do you draw from?"}]}
|
||||
follow_up = {"role": "user", "content": [{"type": "text", "text": "Restate that in one sentence."}]}
|
||||
|
||||
without_block = [user_turn, {"role": "assistant", "content": [reply]}, follow_up]
|
||||
with_block = [user_turn, {"role": "assistant", "content": [redacted_block, reply]}, follow_up]
|
||||
|
||||
assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block)
|
||||
|
||||
def test_token_counter_with_tool_reference_block():
|
||||
"""
|
||||
Regression test: a message containing an Anthropic tool-search
|
||||
|
|
|
|||
|
|
@ -1440,6 +1440,46 @@ async def test_anthropic_messages_leaves_non_provider_failures_unmapped():
|
|||
assert "Traceback" not in str(excinfo.value)
|
||||
|
||||
|
||||
def _recording_client(seen_urls: list[str]) -> AsyncHTTPHandler:
|
||||
def record_and_answer(request: httpx.Request) -> httpx.Response:
|
||||
seen_urls.append(str(request.url))
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_test",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "deepseek-chat",
|
||||
"content": [{"type": "text", "text": "pong"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 3, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
upstream = AsyncHTTPHandler()
|
||||
upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(record_and_answer))
|
||||
return upstream
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_messages_api_base_env_is_not_shadowed_by_the_chat_default(monkeypatch):
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages import handler
|
||||
|
||||
monkeypatch.delenv("DEEPSEEK_API_BASE", raising=False)
|
||||
monkeypatch.setenv("DEEPSEEK_ANTHROPIC_API_BASE", "https://deepseek.internal.example/anthropic")
|
||||
seen_urls: list[str] = []
|
||||
|
||||
await handler.anthropic_messages(
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
model="deepseek/deepseek-chat",
|
||||
api_key="sk-test",
|
||||
client=_recording_client(seen_urls),
|
||||
)
|
||||
|
||||
assert seen_urls == ["https://deepseek.internal.example/anthropic/v1/messages"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_forwards_safeguards_and_unknown_beta_to_anthropic():
|
||||
"""Shapes are what Claude Code 2.1.278 sends and api.anthropic.com returns, captured 2026-09-21."""
|
||||
|
|
|
|||
|
|
@ -0,0 +1,484 @@
|
|||
"""
|
||||
Unit tests for the bedrock_mantle native Anthropic Messages route.
|
||||
|
||||
Mantle serves its Claude models only on `/anthropic/v1/messages` (the OpenAI
|
||||
paths reject them), so `bedrock_mantle/anthropic.claude-*` requests on
|
||||
/v1/messages must hit that endpoint directly instead of the chat-completions
|
||||
bridge. These tests lock the dispatcher gate, the URL derivation from the
|
||||
OpenAI-surface base that get_llm_provider pre-fills, the version header, the
|
||||
Bearer/SigV4 auth chain, and the wire request through the public entrypoint.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.caching.llm_caching_handler import LLMClientCache
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock_mantle.messages.transformation import (
|
||||
BedrockMantleAnthropicMessagesConfig,
|
||||
build_mantle_native_messages_url,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
MESSAGES_PATH = "/anthropic/v1/messages"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _httpx_transport_with_fresh_clients(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_ambient_mantle_env(monkeypatch):
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
|
||||
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
|
||||
monkeypatch.delenv("AWS_REGION", raising=False)
|
||||
|
||||
|
||||
def _anthropic_response() -> httpx.Response:
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
json={
|
||||
"id": "msg_test",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "anthropic.claude-sonnet-5",
|
||||
"content": [{"type": "text", "text": "pong"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 3, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
_SSE_EVENTS = (
|
||||
(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_stream",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "anthropic.claude-sonnet-5",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 3, "output_tokens": 1},
|
||||
},
|
||||
},
|
||||
),
|
||||
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
|
||||
(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}),
|
||||
("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
|
||||
|
||||
def _sse_response() -> httpx.Response:
|
||||
body = "".join(f"event: {event}\ndata: {json.dumps(payload)}\n\n" for event, payload in _SSE_EVENTS).encode()
|
||||
return httpx.Response(status_code=200, content=body, headers={"content-type": "text/event-stream"})
|
||||
|
||||
|
||||
def _mantle_messages_route(region: str) -> respx.Route:
|
||||
return respx.post(f"https://bedrock-mantle.{region}.api.aws{MESSAGES_PATH}")
|
||||
|
||||
|
||||
def _sent_body(route: respx.Route) -> dict:
|
||||
return json.loads(route.calls.last.request.content)
|
||||
|
||||
|
||||
class TestDispatch:
|
||||
def test_claude_models_get_the_native_messages_config(self):
|
||||
config = ProviderConfigManager.get_provider_anthropic_messages_config(
|
||||
model="anthropic.claude-sonnet-5", provider=litellm.LlmProviders.BEDROCK_MANTLE
|
||||
)
|
||||
assert isinstance(config, BedrockMantleAnthropicMessagesConfig)
|
||||
assert config.custom_llm_provider == "bedrock_mantle"
|
||||
|
||||
@pytest.mark.parametrize("model", ["openai.gpt-5.6-sol", "openai.gpt-oss-120b-1:0", "google.gemma-4-31b"])
|
||||
def test_non_claude_models_keep_the_bridge(self, model):
|
||||
assert (
|
||||
ProviderConfigManager.get_provider_anthropic_messages_config(
|
||||
model=model, provider=litellm.LlmProviders.BEDROCK_MANTLE
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
class TestURL:
|
||||
@pytest.mark.parametrize(
|
||||
"api_base",
|
||||
[
|
||||
"https://bedrock-mantle.us-east-1.api.aws/v1",
|
||||
"https://bedrock-mantle.us-east-1.api.aws/openai/v1",
|
||||
"https://bedrock-mantle.us-east-1.api.aws/openai/v1/",
|
||||
"https://bedrock-mantle.us-east-1.api.aws",
|
||||
"https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages",
|
||||
],
|
||||
)
|
||||
def test_prefilled_openai_base_becomes_the_messages_endpoint(self, api_base):
|
||||
url = build_mantle_native_messages_url(api_base, {"aws_region_name": "us-east-1"})
|
||||
assert url == f"https://bedrock-mantle.us-east-1.api.aws{MESSAGES_PATH}"
|
||||
|
||||
def test_aws_region_name_wins_over_the_prefilled_host_region(self):
|
||||
url = build_mantle_native_messages_url(
|
||||
"https://bedrock-mantle.us-east-1.api.aws/v1", {"aws_region_name": "us-east-2"}
|
||||
)
|
||||
assert url == f"https://bedrock-mantle.us-east-2.api.aws{MESSAGES_PATH}"
|
||||
|
||||
def test_host_region_is_used_when_no_region_param(self):
|
||||
url = build_mantle_native_messages_url("https://bedrock-mantle.eu-west-1.api.aws/v1", {})
|
||||
assert url == f"https://bedrock-mantle.eu-west-1.api.aws{MESSAGES_PATH}"
|
||||
|
||||
def test_custom_host_is_preserved(self):
|
||||
url = build_mantle_native_messages_url("https://vpce-abc.bedrock-mantle.example.com/v1", {})
|
||||
assert url == f"https://vpce-abc.bedrock-mantle.example.com{MESSAGES_PATH}"
|
||||
|
||||
def test_env_base_is_used_without_api_base(self, monkeypatch):
|
||||
monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://mantle-proxy.internal/openai/v1")
|
||||
assert build_mantle_native_messages_url(None, {}) == f"https://mantle-proxy.internal{MESSAGES_PATH}"
|
||||
|
||||
def test_default_host_comes_from_mantle_region_env(self, monkeypatch):
|
||||
monkeypatch.setenv("BEDROCK_MANTLE_REGION", "ap-northeast-1")
|
||||
assert (
|
||||
build_mantle_native_messages_url(None, {})
|
||||
== f"https://bedrock-mantle.ap-northeast-1.api.aws{MESSAGES_PATH}"
|
||||
)
|
||||
|
||||
def test_config_get_complete_url_reads_litellm_params(self):
|
||||
config = BedrockMantleAnthropicMessagesConfig()
|
||||
url = config.get_complete_url(
|
||||
api_base="https://bedrock-mantle.us-east-1.api.aws/v1",
|
||||
api_key=None,
|
||||
model="anthropic.claude-sonnet-5",
|
||||
optional_params={},
|
||||
litellm_params={"aws_region_name": "us-west-2"},
|
||||
)
|
||||
assert url == f"https://bedrock-mantle.us-west-2.api.aws{MESSAGES_PATH}"
|
||||
|
||||
|
||||
class TestEnvironment:
|
||||
def _validate(self, headers: dict, litellm_params: dict) -> dict:
|
||||
config = BedrockMantleAnthropicMessagesConfig()
|
||||
merged, _ = config.validate_anthropic_messages_environment(
|
||||
headers=headers,
|
||||
model="anthropic.claude-sonnet-5",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
return merged
|
||||
|
||||
def test_adds_the_anthropic_version_header(self):
|
||||
assert self._validate({}, {})["anthropic-version"] == "2023-06-01"
|
||||
|
||||
def test_keeps_a_caller_supplied_version_header(self):
|
||||
merged = self._validate({"Anthropic-Version": "2024-01-01"}, {})
|
||||
assert merged["Anthropic-Version"] == "2024-01-01"
|
||||
assert "anthropic-version" not in merged
|
||||
|
||||
def test_project_id_becomes_the_workspace_header(self):
|
||||
assert self._validate({}, {"aws_bedrock_project_id": "proj_123"})["anthropic-workspace"] == "proj_123"
|
||||
|
||||
|
||||
class TestRequestBody:
|
||||
def test_body_carries_model_and_stream_but_not_the_invoke_version(self):
|
||||
config = BedrockMantleAnthropicMessagesConfig()
|
||||
body = config.transform_anthropic_messages_request(
|
||||
model="anthropic.claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
anthropic_messages_optional_request_params={"max_tokens": 8, "stream": True},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
assert body["model"] == "anthropic.claude-sonnet-5"
|
||||
assert body["stream"] is True
|
||||
assert body["max_tokens"] == 8
|
||||
assert "anthropic_version" not in body
|
||||
|
||||
def test_body_omits_stream_when_not_streaming(self):
|
||||
config = BedrockMantleAnthropicMessagesConfig()
|
||||
body = config.transform_anthropic_messages_request(
|
||||
model="anthropic.claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
anthropic_messages_optional_request_params={"max_tokens": 8},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
assert "stream" not in body
|
||||
|
||||
|
||||
class TestAuth:
|
||||
def test_bearer_from_api_key_skips_aws_credentials(self):
|
||||
signer = BaseAWSLLM()
|
||||
signer.get_credentials = MagicMock(side_effect=AssertionError("must not resolve AWS credentials"))
|
||||
config = BedrockMantleAnthropicMessagesConfig(aws_signer=signer)
|
||||
headers, signed = config.sign_request(
|
||||
headers={"anthropic-version": "2023-06-01"},
|
||||
optional_params={},
|
||||
request_data={"model": "anthropic.claude-sonnet-5"},
|
||||
api_base=f"https://bedrock-mantle.us-east-1.api.aws{MESSAGES_PATH}",
|
||||
api_key="arg-bearer",
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer arg-bearer"
|
||||
assert headers["anthropic-version"] == "2023-06-01"
|
||||
assert signed == b'{"model": "anthropic.claude-sonnet-5"}'
|
||||
|
||||
def test_bearer_from_mantle_env_key(self, monkeypatch):
|
||||
monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer")
|
||||
config = BedrockMantleAnthropicMessagesConfig()
|
||||
headers, _ = config.sign_request(
|
||||
headers={},
|
||||
optional_params={},
|
||||
request_data={},
|
||||
api_base=f"https://bedrock-mantle.us-east-1.api.aws{MESSAGES_PATH}",
|
||||
api_key=None,
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer env-bearer"
|
||||
|
||||
def test_sigv4_scope_is_pinned_to_the_url_host_region(self):
|
||||
config = BedrockMantleAnthropicMessagesConfig()
|
||||
headers, signed = config.sign_request(
|
||||
headers={"anthropic-version": "2023-06-01"},
|
||||
optional_params={
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
request_data={"model": "anthropic.claude-sonnet-5"},
|
||||
api_base=f"https://bedrock-mantle.us-west-2.api.aws{MESSAGES_PATH}",
|
||||
api_key=None,
|
||||
)
|
||||
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert "/us-west-2/bedrock/aws4_request" in headers["Authorization"]
|
||||
assert signed == b'{"model": "anthropic.claude-sonnet-5"}'
|
||||
|
||||
|
||||
class TestWireRequest:
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_claude_request_hits_the_native_messages_endpoint(self):
|
||||
route = _mantle_messages_route("us-east-1").mock(return_value=_anthropic_response())
|
||||
|
||||
response = await litellm.anthropic_messages(
|
||||
model="bedrock_mantle/anthropic.claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
max_tokens=8,
|
||||
api_key="test-bearer",
|
||||
aws_region_name="us-east-1",
|
||||
)
|
||||
|
||||
assert response["content"][0]["text"] == "pong"
|
||||
assert route.call_count == 1
|
||||
sent = route.calls.last.request
|
||||
assert sent.headers["authorization"] == "Bearer test-bearer"
|
||||
assert sent.headers["anthropic-version"] == "2023-06-01"
|
||||
assert "x-api-key" not in sent.headers
|
||||
body = _sent_body(route)
|
||||
assert body["model"] == "anthropic.claude-sonnet-5"
|
||||
assert body["messages"] == [{"role": "user", "content": "ping"}]
|
||||
assert "anthropic_version" not in body
|
||||
assert "stream" not in body
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_region_prefix_selects_the_host_and_is_not_sent_as_model(self):
|
||||
route = _mantle_messages_route("us-east-2").mock(return_value=_anthropic_response())
|
||||
|
||||
await litellm.anthropic_messages(
|
||||
model="bedrock_mantle/us-east-2/anthropic.claude-haiku-4-5",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
max_tokens=8,
|
||||
api_key="test-bearer",
|
||||
)
|
||||
|
||||
assert route.call_count == 1
|
||||
assert _sent_body(route)["model"] == "anthropic.claude-haiku-4-5"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_streaming_sends_stream_and_passes_the_sse_through(self):
|
||||
route = _mantle_messages_route("us-east-1").mock(return_value=_sse_response())
|
||||
|
||||
response = await litellm.anthropic_messages(
|
||||
model="bedrock_mantle/anthropic.claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
max_tokens=8,
|
||||
stream=True,
|
||||
api_key="test-bearer",
|
||||
aws_region_name="us-east-1",
|
||||
)
|
||||
raw = b"".join([chunk async for chunk in response])
|
||||
|
||||
assert route.call_count == 1
|
||||
assert _sent_body(route)["stream"] is True
|
||||
text = raw.decode()
|
||||
assert "event: message_start" in text
|
||||
assert '"text": "pong"' in text
|
||||
assert "event: message_stop" in text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_sigv4_request_signs_against_the_messages_url(self):
|
||||
route = _mantle_messages_route("us-east-1").mock(return_value=_anthropic_response())
|
||||
|
||||
await litellm.anthropic_messages(
|
||||
model="bedrock_mantle/anthropic.claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
max_tokens=8,
|
||||
aws_access_key_id="AKIAEXAMPLE",
|
||||
aws_secret_access_key="c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
|
||||
aws_region_name="us-east-1",
|
||||
)
|
||||
|
||||
assert route.call_count == 1
|
||||
authorization = route.calls.last.request.headers["authorization"]
|
||||
assert authorization.startswith("AWS4-HMAC-SHA256")
|
||||
assert "/us-east-1/bedrock/aws4_request" in authorization
|
||||
|
||||
|
||||
def _sent_betas(route: respx.Route) -> list[str]:
|
||||
return route.calls.last.request.headers["anthropic-beta"].split(",")
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_beta_headers_config")
|
||||
class TestBetaHeadersOnTheWire:
|
||||
async def _send(self, **request_params) -> respx.Route:
|
||||
route = _mantle_messages_route("us-east-1").mock(return_value=_anthropic_response())
|
||||
await litellm.anthropic_messages(
|
||||
model="bedrock_mantle/anthropic.claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
max_tokens=8,
|
||||
api_key="test-bearer",
|
||||
aws_region_name="us-east-1",
|
||||
**request_params,
|
||||
)
|
||||
return route
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_betas_mantle_accepts_reach_it_in_the_header(self):
|
||||
route = await self._send(
|
||||
extra_headers={
|
||||
"anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14,context-management-2025-06-27"
|
||||
}
|
||||
)
|
||||
|
||||
assert _sent_betas(route) == [
|
||||
"claude-code-20250219",
|
||||
"context-management-2025-06-27",
|
||||
"interleaved-thinking-2025-05-14",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_betas_a_proxy_client_sends_reach_mantle_filtered(self):
|
||||
from litellm.proxy.litellm_pre_call_utils import add_provider_specific_headers_to_request
|
||||
|
||||
proxy_request_data: dict = {}
|
||||
add_provider_specific_headers_to_request(
|
||||
data=proxy_request_data,
|
||||
headers={
|
||||
"anthropic-beta": "claude-code-20250219,fast-mode-2026-02-01,interleaved-thinking-2025-05-14",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"user-agent": "claude-cli/2.1.239",
|
||||
},
|
||||
)
|
||||
|
||||
route = await self._send(**proxy_request_data)
|
||||
|
||||
assert _sent_betas(route) == ["claude-code-20250219", "interleaved-thinking-2025-05-14"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_betas_mantle_rejects_are_dropped_before_the_request(self):
|
||||
route = await self._send(
|
||||
extra_headers={"anthropic-beta": "code-execution-2025-08-25,context-1m-2025-08-07,files-api-2025-04-14"}
|
||||
)
|
||||
|
||||
assert _sent_betas(route) == ["context-1m-2025-08-07"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_no_beta_header_is_sent_when_every_value_is_rejected(self):
|
||||
route = await self._send(extra_headers={"anthropic-beta": "code-execution-2025-08-25"})
|
||||
|
||||
assert "anthropic-beta" not in route.calls.last.request.headers
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_advanced_tool_use_is_renamed_to_the_beta_mantle_knows(self):
|
||||
route = await self._send(extra_headers={"anthropic-beta": "advanced-tool-use-2025-11-20"})
|
||||
|
||||
assert "tool-search-tool-2025-10-19" in _sent_betas(route)
|
||||
assert "advanced-tool-use-2025-11-20" not in _sent_betas(route)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_a_feature_beta_joins_the_callers_betas_in_the_header(self):
|
||||
route = await self._send(
|
||||
extra_headers={"anthropic-beta": "context-1m-2025-08-07"},
|
||||
context_management={"edits": [{"type": "clear_tool_uses_20250919"}]},
|
||||
)
|
||||
|
||||
assert _sent_betas(route) == ["context-1m-2025-08-07", "context-management-2025-06-27"]
|
||||
assert _sent_body(route)["context_management"] == {"edits": [{"type": "clear_tool_uses_20250919"}]}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_betas_and_version_never_travel_in_the_body(self):
|
||||
route = await self._send(
|
||||
extra_headers={"anthropic-beta": "context-1m-2025-08-07"},
|
||||
context_management={"edits": [{"type": "clear_tool_uses_20250919"}]},
|
||||
anthropic_version="bedrock-2023-05-31",
|
||||
)
|
||||
|
||||
body = _sent_body(route)
|
||||
assert "anthropic_beta" not in body
|
||||
assert "anthropic_version" not in body
|
||||
assert route.calls.last.request.headers["anthropic-version"] == "2023-06-01"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_clear_thinking_edit_is_forwarded_with_thinking_on(self):
|
||||
edits = [{"type": "clear_thinking_20251015", "keep": "all"}, {"type": "clear_tool_uses_20250919"}]
|
||||
route = await self._send(
|
||||
context_management={"edits": edits},
|
||||
thinking={"type": "adaptive"},
|
||||
)
|
||||
|
||||
body = _sent_body(route)
|
||||
assert body["context_management"] == {"edits": edits}
|
||||
assert body["thinking"] == {"type": "adaptive"}
|
||||
assert "context-management-2025-06-27" in _sent_betas(route)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_tools_reach_mantle_unchanged(self):
|
||||
tools = [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Look up the weather",
|
||||
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
|
||||
}
|
||||
]
|
||||
route = await self._send(tools=tools, tool_choice={"type": "auto"})
|
||||
|
||||
body = _sent_body(route)
|
||||
assert body["tools"] == tools
|
||||
assert body["tool_choice"] == {"type": "auto"}
|
||||
|
|
@ -1992,6 +1992,7 @@ async def test_streamable_http_session_manager_is_stateless():
|
|||
(
|
||||
("POST", b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', True),
|
||||
("POST", b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}', False),
|
||||
("POST", b"", False),
|
||||
("GET", b"", False),
|
||||
("DELETE", b"", False),
|
||||
),
|
||||
|
|
@ -2465,6 +2466,68 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body():
|
|||
assert total_streamed == len(first_chunk) + sum(len(b) for b in oversized_tail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ("initialize", "tools/call"))
|
||||
@pytest.mark.parametrize("chunked", (False, True))
|
||||
@pytest.mark.parametrize(
|
||||
("character", "bytes_before_cap"),
|
||||
(("é", 0), ("é", 1), ("中", 1), ("中", 2), ("😀", 1), ("😀", 2), ("😀", 3)),
|
||||
)
|
||||
async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap(
|
||||
method: str, chunked: bool, character: str, bytes_before_cap: int
|
||||
) -> None:
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
params: Final = (
|
||||
{
|
||||
"protocolVersion": LATEST_HANDSHAKE_VERSION,
|
||||
"capabilities": {},
|
||||
"clientInfo": {"name": "<<text>>", "version": "1"},
|
||||
}
|
||||
if method == "initialize"
|
||||
else {"name": "update_full_document", "arguments": {"markdown": "<<text>>"}}
|
||||
)
|
||||
template: Final = json.dumps({"jsonrpc": "2.0", "id": 1, "method": method, "params": params}).encode()
|
||||
prefix, suffix = template.split(b"<<text>>")
|
||||
cap: Final = mcp_module._MCP_ROUTING_PEEK_MAX_BYTES
|
||||
body: Final = prefix + b"x" * (cap - bytes_before_cap - len(prefix)) + character.encode() + b"tail" + suffix
|
||||
chunks: Final = (body[: cap - 1], body[cap - 1 : cap], body[cap:]) if chunked else (body,)
|
||||
messages: Final[tuple[Message, ...]] = tuple(
|
||||
{"type": "http.request", "body": chunk, "more_body": index < len(chunks) - 1}
|
||||
for index, chunk in enumerate(chunks)
|
||||
)
|
||||
receive: Final = AsyncMock(side_effect=messages)
|
||||
send: Final = AsyncMock()
|
||||
received: Final[asyncio.Future[bytes]] = asyncio.get_running_loop().create_future()
|
||||
|
||||
async def handle_request(_: Scope, downstream_receive: Receive, outgoing: Send) -> None:
|
||||
assert receive.await_count == (2 if chunked else 1)
|
||||
received.set_result(await _drain_body(downstream_receive))
|
||||
await outgoing({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await outgoing({"type": "http.response.body", "body": b"{}"})
|
||||
|
||||
stateless_handle: Final = AsyncMock(side_effect=handle_request)
|
||||
stateful_handle: Final = AsyncMock()
|
||||
scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
with (
|
||||
_client_allowlist_patches({}, None),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.session_manager_stateless",
|
||||
SimpleNamespace(handle_request=stateless_handle),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.session_manager_stateful",
|
||||
SimpleNamespace(handle_request=stateful_handle),
|
||||
),
|
||||
):
|
||||
await mcp_module.handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert send.call_args_list[0].args[0]["status"] == 200
|
||||
assert received.result() == body
|
||||
stateless_handle.assert_awaited_once()
|
||||
stateful_handle.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforce_stateful_session_cap_evicts_oldest_idle_then_rejects():
|
||||
"""
|
||||
|
|
@ -4016,7 +4079,12 @@ def test_jsonrpc_text_has_top_level_method_ignores_nested_method():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_truncated_jsonrpc_response_with_nested_method_skips_lock():
|
||||
@pytest.mark.parametrize("response_field", ("result", "error"))
|
||||
@pytest.mark.parametrize(("character", "bytes_before_cap"), (("", 0), ("x", 0), ("é", 1), ("中", 2), ("😀", 3)))
|
||||
@pytest.mark.parametrize("cancel_request", (False, True))
|
||||
async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(
|
||||
response_field: str, character: str, bytes_before_cap: int, cancel_request: bool
|
||||
) -> None:
|
||||
"""Regression: a large JSON-RPC *response* POST whose ``result`` payload
|
||||
nests a ``method`` key must skip the per-session lock so it does not
|
||||
deadlock behind the in-flight request POST that is holding the lock while
|
||||
|
|
@ -4044,7 +4112,7 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock():
|
|||
async def handle(s, r, se):
|
||||
msg = await r()
|
||||
body = msg.get("body", b"") or b""
|
||||
if b'"result"' in body:
|
||||
if body == response_body:
|
||||
response_handled.set()
|
||||
else:
|
||||
request_in_handle.set()
|
||||
|
|
@ -4071,9 +4139,16 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock():
|
|||
# A JSON-RPC response larger than the routing peek cap so it can't be fully
|
||||
# parsed, with a nested "method" key in the first bytes to trip a flat
|
||||
# substring heuristic.
|
||||
response_body = (
|
||||
'{"jsonrpc":"2.0","id":99,"result":{"toolResult":{"method":"GET","payload":"' + ("x" * 5000) + '"}}}'
|
||||
response_prefix: Final = (
|
||||
'{"jsonrpc":"2.0","id":99,"' + response_field
|
||||
+ '":{"code":-32000,"message":"test","data":{"method":"GET","payload":"'
|
||||
).encode()
|
||||
response_body: Final = (
|
||||
response_prefix
|
||||
+ b"x" * (mcp_server._MCP_ROUTING_PEEK_MAX_BYTES - bytes_before_cap - len(response_prefix) if character else 0)
|
||||
+ character.encode()
|
||||
+ b'tail"}}}'
|
||||
)
|
||||
|
||||
try:
|
||||
with (
|
||||
|
|
@ -4101,8 +4176,17 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock():
|
|||
# lock held by req_task and this wait would time out (deadlock).
|
||||
await asyncio.wait_for(response_handled.wait(), timeout=1.0)
|
||||
|
||||
gate.set()
|
||||
await asyncio.gather(req_task, resp_task)
|
||||
await resp_task
|
||||
assert not req_task.done()
|
||||
if cancel_request:
|
||||
req_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await req_task
|
||||
else:
|
||||
gate.set()
|
||||
await req_task
|
||||
assert not mcp_server._stateful_session_locks[session_id].locked()
|
||||
assert session_id not in mcp_server._stateful_session_active_request_counts
|
||||
finally:
|
||||
gate.set()
|
||||
mcp_server._stateful_session_auth_contexts.pop(session_id, None)
|
||||
|
|
|
|||
|
|
@ -13182,6 +13182,8 @@ class _DiscoveryUpstream:
|
|||
await self.release.wait()
|
||||
if self.outcome == "failure":
|
||||
return httpx2.Response(503)
|
||||
if self.outcome == "paged_failure" and (payload.params or {}).get("cursor"):
|
||||
return httpx2.Response(503)
|
||||
if self.outcome == "cancelled":
|
||||
raise asyncio.CancelledError()
|
||||
if self.outcome == "rejected":
|
||||
|
|
@ -13196,7 +13198,12 @@ class _DiscoveryUpstream:
|
|||
},
|
||||
"tools/list": {"tools": []},
|
||||
}[payload.method]
|
||||
return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result})
|
||||
continuation: Final = (
|
||||
{"nextCursor": "last-page"}
|
||||
if self.outcome in ("paged", "paged_failure") and not (payload.params or {}).get("cursor")
|
||||
else {}
|
||||
)
|
||||
return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {**result, **continuation}})
|
||||
|
||||
@property
|
||||
def initializes(self) -> int:
|
||||
|
|
@ -13262,6 +13269,29 @@ async def test_discovery_cache_empty_results_and_failures(kind: str, outcome: st
|
|||
assert upstream.initializes == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ("prompts", "resources", "templates"))
|
||||
async def test_discovery_cache_retries_failed_pagination_before_caching_complete_list(kind: str) -> None:
|
||||
manager: Final = MCPServerManager()
|
||||
upstream: Final = _DiscoveryUpstream()
|
||||
upstream.outcome = "paged_failure"
|
||||
operation: Final = {
|
||||
"prompts": manager.get_prompts_from_server,
|
||||
"resources": manager.get_resources_from_server,
|
||||
"templates": manager.get_resource_templates_from_server,
|
||||
}[kind]
|
||||
with _mcp_upstream(upstream.respond):
|
||||
assert await operation(_discovery_server(), None) == []
|
||||
assert upstream.initializes == 1
|
||||
upstream.outcome = "paged"
|
||||
recovered: Final = await operation(_discovery_server(), None)
|
||||
assert [item.name for item in recovered] == ["discovery-example", "discovery-example"]
|
||||
assert upstream.initializes == 2
|
||||
requests_after_recovery: Final = upstream.requests
|
||||
assert await operation(_discovery_server(), None) == recovered
|
||||
assert upstream.requests == requests_after_recovery
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discovery_cache_isolates_forwarded_credentials_and_shares_static_auth() -> None:
|
||||
import respx
|
||||
|
|
|
|||
|
|
@ -0,0 +1,321 @@
|
|||
import json
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import psycopg
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from fastapi import FastAPI
|
||||
from prisma import Prisma
|
||||
from pydantic import TypeAdapter
|
||||
from pytest_postgresql import factories
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.prompt_caching_requests import router
|
||||
from litellm.proxy.spend_tracking.savings import (
|
||||
extract_cache_creation_tokens,
|
||||
extract_cache_read_tokens,
|
||||
marks_gateway_injection,
|
||||
)
|
||||
from litellm.types.management_endpoints.prompt_caching_requests import (
|
||||
PromptCachingRequestFilter,
|
||||
PromptCachingRequestsResponse,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
_cache_postgresql_proc: Final = factories.postgresql_proc() # pyright: ignore[reportUnknownMemberType] # third-party fixture factory has incomplete callable types
|
||||
_cache_postgresql: Final = factories.postgresql("_cache_postgresql_proc")
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
|
||||
_JSON_ROWS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
_START: Final = "2026-09-01T00:00:00Z"
|
||||
_END: Final = "2026-09-02T00:00:00Z"
|
||||
_URL: Final = "/cost_optimization/prompt_caching/requests"
|
||||
_MODEL: Final = "claude-sonnet-5"
|
||||
_MARKER: Final = "litellm_gateway_injected_cache"
|
||||
_DDL: Final = """
|
||||
CREATE TABLE "LiteLLM_SpendLogs" (
|
||||
request_id TEXT PRIMARY KEY, "startTime" TIMESTAMP, "endTime" TIMESTAMP,
|
||||
model TEXT, model_id TEXT, custom_llm_provider TEXT, spend DOUBLE PRECISION,
|
||||
metadata JSONB, cache_hit TEXT
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Case:
|
||||
request_id: str
|
||||
metadata: Mapping[str, object]
|
||||
cache_hit: str | None = None
|
||||
start_time: datetime = datetime(2026, 9, 1, 12, 0, 0, 123456)
|
||||
|
||||
def matches(self, filter: PromptCachingRequestFilter) -> bool:
|
||||
if self.cache_hit is not None and self.cache_hit.lower() == "true":
|
||||
return False
|
||||
if not datetime(2026, 9, 1) <= self.start_time <= datetime(2026, 9, 2):
|
||||
return False
|
||||
usage: Final = self.metadata.get("usage_object")
|
||||
normalized: Final = _JSON_OBJECT.validate_python(usage) if isinstance(usage, Mapping) else None
|
||||
injected: Final = marks_gateway_injection(self.metadata, "dep-a")
|
||||
reads: Final = extract_cache_read_tokens(normalized)
|
||||
writes: Final = extract_cache_creation_tokens(normalized)
|
||||
match filter:
|
||||
case "injected":
|
||||
return injected
|
||||
case "hits":
|
||||
return reads > 0
|
||||
case "all":
|
||||
return injected or reads > 0 or writes > 0
|
||||
|
||||
|
||||
_CASES: Final = (
|
||||
_Case("injected-empty", {_MARKER: ""}),
|
||||
_Case("injected-deployment", {_MARKER: "dep-a"}),
|
||||
_Case("wrong-deployment", {_MARKER: "dep-b"}),
|
||||
_Case("legacy-read", {"usage_object": {"cache_read_input_tokens": 100}}),
|
||||
_Case("nested-read", {"usage_object": {"prompt_tokens_details": {"cached_tokens": 100}}}),
|
||||
_Case("write", {"usage_object": {"cache_creation_input_tokens": 100}}),
|
||||
_Case("nested-write", {"usage_object": {"prompt_tokens_details": {"cache_write_tokens": 100}}}),
|
||||
_Case("nested-creation", {"usage_object": {"prompt_tokens_details": {"cache_creation_tokens": 100}}}),
|
||||
_Case(
|
||||
"top-precedence",
|
||||
{"usage_object": {"cache_read_input_tokens": -2, "prompt_tokens_details": {"cached_tokens": 100}}},
|
||||
),
|
||||
_Case(
|
||||
"zero-fallback",
|
||||
{"usage_object": {"cache_read_input_tokens": 0, "prompt_tokens_details": {"cached_tokens": 100}}},
|
||||
),
|
||||
_Case(
|
||||
"fractional-precedence",
|
||||
{"usage_object": {"cache_read_input_tokens": 0.5, "prompt_tokens_details": {"cached_tokens": 100}}},
|
||||
),
|
||||
_Case("malformed-number", {"usage_object": {"cache_read_input_tokens": "100"}}),
|
||||
_Case("malformed-container", {"usage_object": [100]}),
|
||||
_Case("boolean-number", {"usage_object": {"cache_read_input_tokens": True}}),
|
||||
_Case("boolean-marker", {_MARKER: True}),
|
||||
_Case("response-cache", {_MARKER: "", "usage_object": {"cache_read_input_tokens": 100}}, "True"),
|
||||
_Case("outside-before", {_MARKER: ""}, start_time=datetime(2026, 8, 31, 23, 59, 59)),
|
||||
_Case(
|
||||
"outside-after", {"usage_object": {"cache_read_input_tokens": 100}}, start_time=datetime(2026, 9, 2, 0, 0, 1)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(loop_scope="function")
|
||||
async def _cache_prisma(
|
||||
_cache_postgresql: psycopg.Connection[tuple[object, ...]],
|
||||
) -> AsyncIterator[Prisma]:
|
||||
info: Final = _cache_postgresql.info
|
||||
database: Final = Prisma(datasource={
|
||||
"url": f"postgresql://{info.user}@{info.host}:{info.port}/{info.dbname}?connection_limit=1",
|
||||
})
|
||||
await database.connect()
|
||||
try:
|
||||
yield database
|
||||
finally:
|
||||
await database.disconnect()
|
||||
|
||||
|
||||
def _seed(connection: psycopg.Connection[tuple[object, ...]], cases: tuple[_Case, ...] = _CASES) -> None:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(_DDL)
|
||||
cursor.executemany(
|
||||
"""INSERT INTO "LiteLLM_SpendLogs"
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s)""",
|
||||
tuple(
|
||||
(
|
||||
case.request_id,
|
||||
case.start_time,
|
||||
datetime(2026, 9, 1, 12, 0, 1),
|
||||
_MODEL,
|
||||
"dep-a",
|
||||
"anthropic",
|
||||
0.01,
|
||||
json.dumps(dict(case.metadata)),
|
||||
case.cache_hit,
|
||||
)
|
||||
for case in cases
|
||||
),
|
||||
)
|
||||
connection.commit()
|
||||
|
||||
|
||||
def _app(role: LitellmUserRoles | None) -> FastAPI:
|
||||
application: Final = FastAPI()
|
||||
application.include_router(router)
|
||||
|
||||
def caller() -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(user_role=role)
|
||||
|
||||
application.dependency_overrides[user_api_key_auth] = caller
|
||||
return application
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("filter", ["all", "injected", "hits"])
|
||||
@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
|
||||
async def test_request_filters_match_accounting_and_paginate_before_projection(
|
||||
_cache_postgresql: psycopg.Connection[tuple[object, ...]],
|
||||
_cache_prisma: Prisma,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
filter: PromptCachingRequestFilter,
|
||||
role: LitellmUserRoles,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
_seed(_cache_postgresql)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=_cache_prisma))
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
expected: Final = tuple(sorted((case.request_id for case in _CASES if case.matches(filter)), reverse=True))
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=_app(role)), base_url="http://test") as client:
|
||||
first: Final = await client.get(
|
||||
_URL, params={"start_date": _START, "end_date": _END, "filter": filter, "page_size": 2}
|
||||
)
|
||||
assert first.status_code == 200
|
||||
first_page: Final = PromptCachingRequestsResponse.model_validate_json(first.content)
|
||||
assert tuple(row.request_id for row in first_page.requests) == expected[:2]
|
||||
assert first_page.has_more is (len(expected) > 2)
|
||||
assert (first_page.next_cursor is not None) is first_page.has_more
|
||||
if first_page.next_cursor is not None:
|
||||
assert first_page.next_cursor.request_id == expected[1]
|
||||
assert first_page.next_cursor.start_time == first_page.requests[-1].start_time
|
||||
next_response: Final = await client.get(
|
||||
_URL, params={
|
||||
"start_date": _START, "end_date": _END, "filter": filter, "page_size": 2,
|
||||
"cursor_start_time": first_page.next_cursor.start_time.astimezone(
|
||||
timezone(timedelta(hours=-7))
|
||||
).isoformat(),
|
||||
"cursor_request_id": first_page.next_cursor.request_id,
|
||||
}
|
||||
)
|
||||
assert next_response.status_code == 200
|
||||
next_page: Final = PromptCachingRequestsResponse.model_validate_json(next_response.content)
|
||||
assert tuple(row.request_id for row in next_page.requests) == expected[2:4]
|
||||
assert next_page.has_more is (len(expected) > 4)
|
||||
assert (next_page.next_cursor is not None) is next_page.has_more
|
||||
second: Final = await client.get(
|
||||
_URL, params={"start_date": _START, "end_date": _END, "filter": filter, "page_size": 100}
|
||||
)
|
||||
assert second.status_code == 200
|
||||
complete: Final = PromptCachingRequestsResponse.model_validate_json(second.content)
|
||||
assert tuple(row.request_id for row in complete.requests) == expected
|
||||
assert complete.has_more is False
|
||||
assert complete.next_cursor is None
|
||||
assert all(row.start_time.tzinfo == timezone.utc for row in complete.requests)
|
||||
payload: Final = _JSON_OBJECT.validate_json(second.content)
|
||||
assert set(payload) == {"requests", "page_size", "has_more", "next_cursor"}
|
||||
serialized_rows: Final = _JSON_ROWS.validate_python(payload["requests"])
|
||||
assert set(serialized_rows[0]) == {
|
||||
"request_id",
|
||||
"start_time",
|
||||
"model",
|
||||
"gateway_injected",
|
||||
"cache_read_tokens",
|
||||
"cache_creation_tokens",
|
||||
"spend",
|
||||
"net_savings",
|
||||
}
|
||||
by_id: Final = {row.request_id: row for row in complete.requests}
|
||||
if filter == "all":
|
||||
assert by_id["injected-empty"].gateway_injected is True
|
||||
assert by_id["injected-empty"].net_savings is None
|
||||
assert by_id["legacy-read"].gateway_injected is False
|
||||
assert by_id["legacy-read"].net_savings is not None and by_id["legacy-read"].net_savings > 0
|
||||
assert by_id["write"].net_savings is not None and by_id["write"].net_savings < 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("role", [None, LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY])
|
||||
async def test_non_admin_is_denied_before_database_access(
|
||||
role: LitellmUserRoles | None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=_app(role)), base_url="http://test") as client:
|
||||
response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END})
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("params", [
|
||||
{"filter": "savings"}, {"page_size": 0}, {"page_size": 101}, {"start_date": "invalid"},
|
||||
{"cursor_start_time": "invalid", "cursor_request_id": "request"},
|
||||
{"cursor_start_time": _START, "cursor_request_id": ""},
|
||||
])
|
||||
async def test_invalid_request_is_rejected(params: Mapping[str, str | int]) -> None:
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test"
|
||||
) as client:
|
||||
response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, **params})
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("params", [{"cursor_start_time": _START}, {"cursor_request_id": "request"}])
|
||||
async def test_incomplete_cursor_is_rejected(
|
||||
params: Mapping[str, str], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test"
|
||||
) as client:
|
||||
response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, **params})
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("delete_before_cursor", [False, True])
|
||||
async def test_cursor_keeps_remaining_requests_once_during_insertions_and_deletions(
|
||||
_cache_postgresql: psycopg.Connection[tuple[object, ...]],
|
||||
_cache_prisma: Prisma,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
delete_before_cursor: bool,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
cases: Final = (*_CASES, _Case(
|
||||
"older-cache-read", {"usage_object": {"cache_read_input_tokens": 100}}, start_time=datetime(2026, 9, 1, 11),
|
||||
))
|
||||
_seed(_cache_postgresql, cases)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=_cache_prisma))
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
expected: Final = (*sorted((case.request_id for case in _CASES if case.matches("all")), reverse=True), "older-cache-read")
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test"
|
||||
) as client:
|
||||
first: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, "page_size": 2})
|
||||
assert first.status_code == 200
|
||||
first_page: Final = PromptCachingRequestsResponse.model_validate_json(first.content)
|
||||
assert tuple(row.request_id for row in first_page.requests) == expected[:2]
|
||||
assert first_page.next_cursor is not None
|
||||
with _cache_postgresql.cursor() as cursor:
|
||||
cursor.executemany(
|
||||
"""INSERT INTO "LiteLLM_SpendLogs"
|
||||
SELECT %s, %s, "endTime", model, model_id, custom_llm_provider, spend, metadata, cache_hit
|
||||
FROM "LiteLLM_SpendLogs" WHERE request_id = %s""",
|
||||
(
|
||||
("newer-request", datetime(2026, 9, 1, 13), expected[0]),
|
||||
("zz-higher-id", cases[0].start_time, expected[0]),
|
||||
),
|
||||
)
|
||||
if delete_before_cursor:
|
||||
cursor.execute('DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (expected[0],))
|
||||
_cache_postgresql.commit()
|
||||
following: Final = await client.get(_URL, params={
|
||||
"start_date": _START, "end_date": _END, "page_size": 100,
|
||||
"cursor_start_time": first_page.next_cursor.start_time.isoformat(),
|
||||
"cursor_request_id": first_page.next_cursor.request_id,
|
||||
})
|
||||
assert following.status_code == 200
|
||||
following_page: Final = PromptCachingRequestsResponse.model_validate_json(following.content)
|
||||
assert tuple(row.request_id for row in following_page.requests) == expected[2:]
|
||||
assert following_page.has_more is False
|
||||
assert following_page.next_cursor is None
|
||||
|
|
@ -14,7 +14,8 @@ from litellm.proxy.policy_engine.attachment_registry import (
|
|||
AttachmentRegistry,
|
||||
get_attachment_registry,
|
||||
)
|
||||
from litellm.types.proxy.policy_engine import PolicyMatchContext
|
||||
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
|
||||
from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyGuardrails, PolicyMatchContext
|
||||
|
||||
|
||||
class TestGetAttachedPolicies:
|
||||
|
|
@ -30,9 +31,7 @@ class TestGetAttachedPolicies:
|
|||
)
|
||||
|
||||
# Should match any context
|
||||
context = PolicyMatchContext(
|
||||
team_alias="any-team", key_alias="any-key", model="any-model"
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="any-team", key_alias="any-key", model="any-model")
|
||||
attached = registry.get_attached_policies(context)
|
||||
assert "global-baseline" in attached
|
||||
|
||||
|
|
@ -46,15 +45,11 @@ class TestGetAttachedPolicies:
|
|||
)
|
||||
|
||||
# Match
|
||||
context = PolicyMatchContext(
|
||||
team_alias="healthcare-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4")
|
||||
assert "healthcare-policy" in registry.get_attached_policies(context)
|
||||
|
||||
# No match - different team
|
||||
context_other = PolicyMatchContext(
|
||||
team_alias="finance-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
context_other = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4")
|
||||
assert "healthcare-policy" not in registry.get_attached_policies(context_other)
|
||||
|
||||
def test_key_wildcard_pattern_attachment(self):
|
||||
|
|
@ -67,15 +62,11 @@ class TestGetAttachedPolicies:
|
|||
)
|
||||
|
||||
# Match - key starts with dev-key-
|
||||
context = PolicyMatchContext(
|
||||
team_alias="team", key_alias="dev-key-123", model="gpt-4"
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="dev-key-123", model="gpt-4")
|
||||
assert "dev-policy" in registry.get_attached_policies(context)
|
||||
|
||||
# No match - different prefix
|
||||
context_prod = PolicyMatchContext(
|
||||
team_alias="team", key_alias="prod-key-123", model="gpt-4"
|
||||
)
|
||||
context_prod = PolicyMatchContext(team_alias="team", key_alias="prod-key-123", model="gpt-4")
|
||||
assert "dev-policy" not in registry.get_attached_policies(context_prod)
|
||||
|
||||
def test_model_specific_attachment(self):
|
||||
|
|
@ -92,9 +83,7 @@ class TestGetAttachedPolicies:
|
|||
assert "gpt4-policy" in registry.get_attached_policies(context)
|
||||
|
||||
# No match
|
||||
context_other = PolicyMatchContext(
|
||||
team_alias="team", key_alias="key", model="gpt-3.5"
|
||||
)
|
||||
context_other = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-3.5")
|
||||
assert "gpt4-policy" not in registry.get_attached_policies(context_other)
|
||||
|
||||
def test_model_wildcard_pattern(self):
|
||||
|
|
@ -107,15 +96,11 @@ class TestGetAttachedPolicies:
|
|||
)
|
||||
|
||||
# Match
|
||||
context = PolicyMatchContext(
|
||||
team_alias="team", key_alias="key", model="bedrock/claude-3"
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="key", model="bedrock/claude-3")
|
||||
assert "bedrock-policy" in registry.get_attached_policies(context)
|
||||
|
||||
# No match
|
||||
context_other = PolicyMatchContext(
|
||||
team_alias="team", key_alias="key", model="openai/gpt-4"
|
||||
)
|
||||
context_other = PolicyMatchContext(team_alias="team", key_alias="key", model="openai/gpt-4")
|
||||
assert "bedrock-policy" not in registry.get_attached_policies(context_other)
|
||||
|
||||
def test_multiple_attachments_match_same_context(self):
|
||||
|
|
@ -129,9 +114,7 @@ class TestGetAttachedPolicies:
|
|||
]
|
||||
)
|
||||
|
||||
context = PolicyMatchContext(
|
||||
team_alias="healthcare-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4")
|
||||
attached = registry.get_attached_policies(context)
|
||||
|
||||
# All three should match
|
||||
|
|
@ -277,9 +260,7 @@ class TestGetAttachedPolicies:
|
|||
]
|
||||
)
|
||||
|
||||
context = PolicyMatchContext(
|
||||
team_alias="healthcare-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4")
|
||||
attached = registry.get_attached_policies(context)
|
||||
|
||||
# Should only appear once
|
||||
|
|
@ -288,9 +269,7 @@ class TestGetAttachedPolicies:
|
|||
def test_many_distinct_policies_resolve_in_linear_time(self):
|
||||
policy_count = 20_000
|
||||
registry = AttachmentRegistry()
|
||||
registry.load_attachments(
|
||||
[{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)]
|
||||
)
|
||||
registry.load_attachments([{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)])
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4")
|
||||
|
||||
started = time.perf_counter()
|
||||
|
|
@ -318,9 +297,7 @@ class TestGetAttachedPolicies:
|
|||
]
|
||||
)
|
||||
|
||||
context = PolicyMatchContext(
|
||||
team_alias="finance-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4")
|
||||
attached = registry.get_attached_policies(context)
|
||||
assert attached == []
|
||||
|
||||
|
|
@ -338,23 +315,15 @@ class TestGetAttachedPolicies:
|
|||
)
|
||||
|
||||
# Match - both team and model match
|
||||
context = PolicyMatchContext(
|
||||
team_alias="healthcare-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4")
|
||||
assert "strict-policy" in registry.get_attached_policies(context)
|
||||
|
||||
# No match - team matches but model doesn't
|
||||
context_wrong_model = PolicyMatchContext(
|
||||
team_alias="healthcare-team", key_alias="key", model="gpt-3.5"
|
||||
)
|
||||
assert "strict-policy" not in registry.get_attached_policies(
|
||||
context_wrong_model
|
||||
)
|
||||
context_wrong_model = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-3.5")
|
||||
assert "strict-policy" not in registry.get_attached_policies(context_wrong_model)
|
||||
|
||||
# No match - model matches but team doesn't
|
||||
context_wrong_team = PolicyMatchContext(
|
||||
team_alias="finance-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
context_wrong_team = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4")
|
||||
assert "strict-policy" not in registry.get_attached_policies(context_wrong_team)
|
||||
|
||||
|
||||
|
|
@ -527,6 +496,111 @@ class TestMatchAttribution:
|
|||
assert "catch-all" in attached
|
||||
|
||||
|
||||
class TestDefaultAttachments:
|
||||
"""`default: true` attachments apply only when no non-default attachment matches."""
|
||||
|
||||
@staticmethod
|
||||
def _registry() -> AttachmentRegistry:
|
||||
registry = AttachmentRegistry()
|
||||
registry.load_attachments(
|
||||
[
|
||||
{"policy": "guardrail-y", "scope": "*", "default": True},
|
||||
{"policy": "guardrail-x", "tags": ["opt-in"]},
|
||||
]
|
||||
)
|
||||
return registry
|
||||
|
||||
def test_opted_in_request_gets_only_the_opt_in_policy(self):
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"])
|
||||
|
||||
assert self._registry().get_attached_policies(context) == ["guardrail-x"]
|
||||
|
||||
def test_request_without_opt_in_falls_back_to_default_policy(self):
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2")
|
||||
|
||||
assert self._registry().get_attached_policies(context) == ["guardrail-y"]
|
||||
|
||||
def test_default_attachment_still_honors_its_own_scope(self):
|
||||
registry = AttachmentRegistry()
|
||||
registry.load_attachments([{"policy": "team-default", "teams": ["team-a"], "default": True}])
|
||||
|
||||
assert registry.get_attached_policies(PolicyMatchContext(team_alias="team-a", key_alias="k", model="m")) == [
|
||||
"team-default"
|
||||
]
|
||||
assert registry.get_attached_policies(PolicyMatchContext(team_alias="team-b", key_alias="k", model="m")) == []
|
||||
|
||||
def test_all_matching_defaults_apply_when_nothing_else_matches(self):
|
||||
registry = AttachmentRegistry()
|
||||
registry.load_attachments(
|
||||
[
|
||||
{"policy": "default-a", "scope": "*", "default": True},
|
||||
{"policy": "default-b", "teams": ["team-a"], "default": True},
|
||||
{"policy": "opt-in", "tags": ["opt-in"]},
|
||||
]
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="team-a", key_alias="k", model="m")
|
||||
|
||||
assert registry.get_attached_policies(context) == ["default-a", "default-b"]
|
||||
|
||||
def test_non_default_attachments_remain_additive(self):
|
||||
registry = AttachmentRegistry()
|
||||
registry.load_attachments(
|
||||
[
|
||||
{"policy": "baseline", "scope": "*"},
|
||||
{"policy": "opt-in", "tags": ["opt-in"]},
|
||||
{"policy": "fallback", "scope": "*", "default": True},
|
||||
]
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="m", tags=["opt-in"])
|
||||
|
||||
assert registry.get_attached_policies(context) == ["baseline", "opt-in"]
|
||||
|
||||
def test_default_match_reason_is_labelled(self):
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="m")
|
||||
|
||||
results = self._registry().get_attached_policies_with_reasons(context)
|
||||
|
||||
assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}]
|
||||
|
||||
def test_inapplicable_opt_in_policy_does_not_suppress_default(self):
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"])
|
||||
policies = {
|
||||
"guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"])),
|
||||
"guardrail-x": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="claude.*")),
|
||||
}
|
||||
|
||||
results = self._registry().get_attached_policies_with_reasons(
|
||||
context, PolicyMatcher.policy_applies(context, policies)
|
||||
)
|
||||
|
||||
assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}]
|
||||
|
||||
def test_attachment_to_missing_policy_does_not_suppress_default(self):
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"])
|
||||
policies = {"guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"]))}
|
||||
|
||||
assert self._registry().get_attached_policies(context, PolicyMatcher.policy_applies(context, policies)) == [
|
||||
"guardrail-y"
|
||||
]
|
||||
|
||||
def test_applicable_opt_in_policy_still_wins_with_predicate(self):
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"])
|
||||
policies = {
|
||||
"guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"])),
|
||||
"guardrail-x": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="gpt.*")),
|
||||
}
|
||||
|
||||
assert self._registry().get_attached_policies(context, PolicyMatcher.policy_applies(context, policies)) == [
|
||||
"guardrail-x"
|
||||
]
|
||||
|
||||
def test_default_defaults_to_false_when_omitted(self):
|
||||
registry = AttachmentRegistry()
|
||||
registry.load_attachments([{"policy": "p"}])
|
||||
|
||||
assert registry.get_all_attachments()[0].default is False
|
||||
|
||||
|
||||
class TestAttachmentRegistrySingleton:
|
||||
"""Test global singleton behavior."""
|
||||
|
||||
|
|
@ -557,6 +631,7 @@ def _make_db_attachment_row(
|
|||
scope: str | None = None,
|
||||
teams: list[str] | None = None,
|
||||
priority: int | None = None,
|
||||
is_default: bool = False,
|
||||
) -> MagicMock:
|
||||
row = MagicMock()
|
||||
row.attachment_id = attachment_id
|
||||
|
|
@ -567,6 +642,7 @@ def _make_db_attachment_row(
|
|||
row.models = []
|
||||
row.tags = []
|
||||
row.priority = priority
|
||||
row.is_default = is_default
|
||||
row.created_at = datetime.now(timezone.utc)
|
||||
row.updated_at = datetime.now(timezone.utc)
|
||||
row.created_by = None
|
||||
|
|
@ -576,9 +652,7 @@ def _make_db_attachment_row(
|
|||
|
||||
def _prisma_with_attachment_rows(rows: list[MagicMock]) -> MagicMock:
|
||||
prisma = MagicMock()
|
||||
prisma.configure_mock(
|
||||
**{"db.litellm_policyattachmenttable.find_many": AsyncMock(return_value=rows)}
|
||||
)
|
||||
prisma.configure_mock(**{"db.litellm_policyattachmenttable.find_many": AsyncMock(return_value=rows)})
|
||||
return prisma
|
||||
|
||||
|
||||
|
|
@ -629,6 +703,15 @@ class TestConfigAttachmentsPreservedAcrossDbSync:
|
|||
|
||||
assert registry.get_all_attachments()[0].priority == 7
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_round_trips_db_attachment_default_flag(self):
|
||||
registry = AttachmentRegistry()
|
||||
db_row = _make_db_attachment_row(is_default=True)
|
||||
|
||||
await registry.sync_attachments_from_db(_prisma_with_attachment_rows([db_row]))
|
||||
|
||||
assert registry.get_all_attachments()[0].default is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self):
|
||||
registry = AttachmentRegistry()
|
||||
|
|
|
|||
|
|
@ -8,8 +8,11 @@ Tests:
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm.proxy.policy_engine.attachment_registry as attachment_registry_module
|
||||
import litellm.proxy.policy_engine.policy_registry as policy_registry_module
|
||||
from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry
|
||||
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
|
||||
from litellm.proxy.policy_engine.policy_registry import PolicyRegistry
|
||||
from litellm.types.proxy.policy_engine import (
|
||||
PolicyMatchContext,
|
||||
PolicyScope,
|
||||
|
|
@ -196,3 +199,48 @@ class TestPolicyMatcherWithAttachments:
|
|||
attached = registry.get_attached_policies(context)
|
||||
|
||||
assert "healthcare-policy" not in attached
|
||||
|
||||
|
||||
def _global_registries(monkeypatch):
|
||||
policies = PolicyRegistry()
|
||||
policies.load_policies(
|
||||
{
|
||||
"guardrail-y": {"guardrails": {"add": ["y"]}},
|
||||
"guardrail-x": {"guardrails": {"add": ["x"]}, "condition": {"model": "claude.*"}},
|
||||
}
|
||||
)
|
||||
attachments = AttachmentRegistry()
|
||||
attachments.load_attachments(
|
||||
[
|
||||
{"policy": "guardrail-x", "tags": ["opt-in"]},
|
||||
{"policy": "guardrail-y", "scope": "*", "default": True},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(policy_registry_module, "get_policy_registry", lambda: policies)
|
||||
monkeypatch.setattr(attachment_registry_module, "get_attachment_registry", lambda: attachments)
|
||||
return policies
|
||||
|
||||
|
||||
class TestGetMatchingPoliciesFallback:
|
||||
def test_condition_failing_opt_in_falls_back_to_default(self, monkeypatch):
|
||||
_global_registries(monkeypatch)
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5", tags=["opt-in"])
|
||||
|
||||
assert PolicyMatcher.get_matching_policies(context=context) == ["guardrail-y"]
|
||||
|
||||
def test_condition_passing_opt_in_suppresses_default(self, monkeypatch):
|
||||
_global_registries(monkeypatch)
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="claude-haiku", tags=["opt-in"])
|
||||
|
||||
assert PolicyMatcher.get_matching_policies(context=context) == ["guardrail-x"]
|
||||
|
||||
def test_policy_applies_reads_registry_once(self, monkeypatch):
|
||||
policies = _global_registries(monkeypatch)
|
||||
calls = []
|
||||
original = policies.get_all_policies
|
||||
monkeypatch.setattr(policies, "get_all_policies", lambda: calls.append(1) or original())
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5", tags=["opt-in"])
|
||||
|
||||
PolicyMatcher.get_matching_policies(context=context)
|
||||
|
||||
assert len(calls) == 1
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from litellm.proxy.spend_tracking.savings import (
|
|||
compute_autorouter_savings,
|
||||
compute_savings_spend,
|
||||
marks_gateway_injection,
|
||||
prompt_caching_savings_for_request,
|
||||
)
|
||||
from litellm.router import Router
|
||||
from litellm.types.utils import Usage
|
||||
|
|
@ -18,6 +19,42 @@ from litellm.types.utils import Usage
|
|||
pytestmark = pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model,usage", [
|
||||
(None, {"cache_read_input_tokens": 100}),
|
||||
("claude-sonnet-5", None),
|
||||
("claude-sonnet-5", {"prompt_tokens": "invalid"}),
|
||||
])
|
||||
def test_prompt_cache_estimate_distinguishes_unknown_from_zero(model: str | None, usage: dict[str, object] | None) -> None:
|
||||
assert prompt_caching_savings_for_request(model, "anthropic", usage) is None
|
||||
assert compute_savings_spend(model, "anthropic", 0, False, usage_object=usage).prompt_caching == 0
|
||||
assert prompt_caching_savings_for_request("claude-sonnet-5", "anthropic", {"prompt_tokens": 100}) == 0
|
||||
|
||||
|
||||
def test_prompt_cache_estimate_uses_the_rollup_pricing_and_retains_write_premiums() -> None:
|
||||
router: Final = Router(model_list=[{
|
||||
"model_name": "negotiated",
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-5", "input_cost_per_token": 1e-6,
|
||||
"cache_creation_input_token_cost": 1.25e-6, "cache_read_input_token_cost": 1e-7,
|
||||
},
|
||||
"model_info": {"id": "negotiated-cache-prices"},
|
||||
}])
|
||||
|
||||
def current_router() -> Router:
|
||||
return router
|
||||
|
||||
usage: Final = {"cache_read_input_tokens": 1000, "cache_creation_input_tokens": 20000}
|
||||
estimate: Final = prompt_caching_savings_for_request(
|
||||
"claude-sonnet-5", "anthropic", usage, model_id="negotiated-cache-prices", llm_router=current_router,
|
||||
)
|
||||
rollup: Final = compute_savings_spend(
|
||||
"claude-sonnet-5", "anthropic", 0, True, usage_object=usage,
|
||||
model_id="negotiated-cache-prices", llm_router=current_router,
|
||||
)
|
||||
assert estimate == pytest.approx(1000 * (1e-6 - 1e-7) - 20000 * (1.25e-6 - 1e-6))
|
||||
assert estimate == rollup.prompt_caching == rollup.gateway_injected_caching
|
||||
|
||||
|
||||
@pytest.mark.parametrize("modifier", [{"speed": "fast"}, {"inference_geo": "us"}])
|
||||
@pytest.mark.parametrize("continuing", [False, True])
|
||||
def test_baseline_preserves_anthropic_pricing_fields(modifier: dict[str, str], continuing: bool) -> None:
|
||||
|
|
|
|||
|
|
@ -7249,7 +7249,7 @@ CROSS_ACCOUNT_AUTHORIZATION = "Bearer deliberately-configured-pass-through-token
|
|||
|
||||
SIGV4_PREFIX = "AWS4-HMAC-SHA256"
|
||||
AUTHORIZATION_HEADER_CASINGS = ["authorization", "Authorization", "AUTHORIZATION"]
|
||||
LEAK_TARGET_PROVIDERS = ["bedrock", "bedrock_converse", "vertex_ai"]
|
||||
LEAK_TARGET_PROVIDERS = ["bedrock", "bedrock_converse", "bedrock_mantle", "vertex_ai"]
|
||||
|
||||
BEDROCK_ENDPOINT = (
|
||||
"https://bedrock-runtime.us-west-2.amazonaws.com/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke"
|
||||
|
|
@ -7342,6 +7342,28 @@ def test_oauth_credential_entry_is_scoped_to_anthropic_alone():
|
|||
assert [entry["custom_llm_provider"] for entry in credential_entries] == ["anthropic"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"])
|
||||
def test_client_anthropic_api_headers_reach_every_anthropic_messages_provider(custom_llm_provider):
|
||||
client_headers = {
|
||||
"anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"user-agent": "claude-cli/2.1.239",
|
||||
}
|
||||
|
||||
forwarded = _headers_forwarded_to(client_headers, custom_llm_provider)
|
||||
|
||||
assert forwarded == {
|
||||
"anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14",
|
||||
"anthropic-version": "2023-06-01",
|
||||
}
|
||||
|
||||
|
||||
def test_client_anthropic_api_headers_stay_off_openai_compatible_providers():
|
||||
forwarded = _headers_forwarded_to({"anthropic-beta": "claude-code-20250219"}, "openai")
|
||||
|
||||
assert forwarded == {}
|
||||
|
||||
|
||||
def test_no_provider_specific_header_when_client_sends_nothing_anthropic():
|
||||
data: dict = {}
|
||||
add_provider_specific_headers_to_request(
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import pytest
|
|||
import litellm
|
||||
from litellm.anthropic_beta_headers_manager import (
|
||||
filter_and_transform_beta_headers,
|
||||
update_headers_with_filtered_beta,
|
||||
update_request_with_filtered_beta,
|
||||
)
|
||||
|
||||
|
|
@ -511,3 +512,20 @@ class TestAnthropicBetaHeadersFiltering:
|
|||
assert (
|
||||
"unknown-header-123" not in filtered
|
||||
), f"Unknown header should not be in result for {provider}"
|
||||
|
||||
@pytest.mark.parametrize("provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"])
|
||||
def test_blank_anthropic_beta_header_is_removed(self, provider):
|
||||
headers = {"anthropic-beta": "", "anthropic-version": "2023-06-01"}
|
||||
|
||||
assert update_headers_with_filtered_beta(headers, provider) == {"anthropic-version": "2023-06-01"}
|
||||
|
||||
@pytest.mark.parametrize("provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"])
|
||||
def test_whitespace_only_anthropic_beta_header_is_removed(self, provider):
|
||||
headers = {"anthropic-beta": " , ", "anthropic-version": "2023-06-01"}
|
||||
|
||||
assert update_headers_with_filtered_beta(headers, provider) == {"anthropic-version": "2023-06-01"}
|
||||
|
||||
def test_absent_anthropic_beta_header_is_left_alone(self):
|
||||
headers = {"anthropic-version": "2023-06-01"}
|
||||
|
||||
assert update_headers_with_filtered_beta(headers, "bedrock_mantle") == {"anthropic-version": "2023-06-01"}
|
||||
|
|
|
|||
|
|
@ -3522,6 +3522,58 @@ def test_cost_per_token_region_name_applies_to_provider_prefixed_model(_local_mo
|
|||
)
|
||||
|
||||
|
||||
def test_completion_cost_mantle_native_messages_prices_claude_from_the_bedrock_row(_local_model_cost_map):
|
||||
"""Mantle's native Messages API answers with Anthropic's canonical model name and the proxy
|
||||
resolves a Mantle region for every call, so the first cost candidate is
|
||||
bedrock_mantle/<region>/claude-sonnet-5. That name has no row of its own and must fall through to
|
||||
the deployment's bare Bedrock row instead of stopping on an unpriced capability rule at $0."""
|
||||
|
||||
response = litellm.ModelResponse(
|
||||
id="msg_x",
|
||||
choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
model="claude-sonnet-5",
|
||||
usage={"prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110},
|
||||
)
|
||||
row = litellm.model_cost["anthropic.claude-sonnet-5"]
|
||||
expected = 100 * row["input_cost_per_token"] + 10 * row["output_cost_per_token"]
|
||||
assert expected > 0
|
||||
|
||||
for region_name in ("us-east-1", None):
|
||||
assert litellm.completion_cost(
|
||||
completion_response=response,
|
||||
model="bedrock_mantle/anthropic.claude-sonnet-5",
|
||||
custom_llm_provider="bedrock_mantle",
|
||||
region_name=region_name,
|
||||
) == pytest.approx(expected)
|
||||
|
||||
|
||||
def test_completion_cost_mantle_native_messages_prices_haiku_from_the_mantle_row(_local_model_cost_map):
|
||||
"""Mantle serves Anthropic's un-versioned haiku id, which has no bare Bedrock row (Bedrock's carries
|
||||
the -20251001-v1:0 suffix), and Claude Code sends every small-fast-model call to it. Both the plain
|
||||
and the region-prefixed deployment names must price from bedrock_mantle/anthropic.claude-haiku-4-5
|
||||
instead of billing $0."""
|
||||
|
||||
response = litellm.ModelResponse(
|
||||
id="msg_x",
|
||||
choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
model="claude-haiku-4-5",
|
||||
usage={"prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110},
|
||||
)
|
||||
row = litellm.model_cost["bedrock_mantle/anthropic.claude-haiku-4-5"]
|
||||
expected = 100 * row["input_cost_per_token"] + 10 * row["output_cost_per_token"]
|
||||
assert expected > 0
|
||||
|
||||
for model in (
|
||||
"bedrock_mantle/anthropic.claude-haiku-4-5",
|
||||
"bedrock_mantle/us-east-2/anthropic.claude-haiku-4-5",
|
||||
):
|
||||
assert litellm.completion_cost(
|
||||
completion_response=response,
|
||||
model=model,
|
||||
custom_llm_provider="bedrock_mantle",
|
||||
) == pytest.approx(expected), model
|
||||
|
||||
|
||||
def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map):
|
||||
"""An explicit base_model keeps pricing on that model's own key even when the request carries a
|
||||
region with different regional rates, so the private provider model never widens region pricing."""
|
||||
|
|
|
|||
|
|
@ -1163,6 +1163,21 @@ def test_get_model_info_bedrock_regional_inference_profile_pricing(local_model_c
|
|||
assert control["key"] == "au.anthropic.claude-opus-4-8"
|
||||
|
||||
|
||||
def test_get_model_info_bedrock_mantle_region_prefix_falls_back_to_the_mantle_row(local_model_cost_map):
|
||||
"""A Mantle deployment name may carry the region as a prefix (bedrock_mantle/us-east-2/<model>).
|
||||
That name has no cost row of its own, so pricing must fall through to the region-free
|
||||
bedrock_mantle/<model> row instead of raising, while a region that has its own row keeps it."""
|
||||
for model, expected_key in (
|
||||
("bedrock_mantle/us-east-2/anthropic.claude-haiku-4-5", "bedrock_mantle/anthropic.claude-haiku-4-5"),
|
||||
("bedrock_mantle/us-east-2/openai.gpt-5.6-sol", "bedrock_mantle/openai.gpt-5.6-sol"),
|
||||
("bedrock_mantle/us-gov-west-1/openai.gpt-5.4", "bedrock_mantle/us-gov-west-1/openai.gpt-5.4"),
|
||||
):
|
||||
info = litellm.get_model_info(model=model, custom_llm_provider="bedrock_mantle")
|
||||
assert info["key"] == expected_key, model
|
||||
assert info["input_cost_per_token"] == litellm.model_cost[expected_key]["input_cost_per_token"], model
|
||||
assert info["input_cost_per_token"] > 0, model
|
||||
|
||||
|
||||
def test_openai_models_in_model_info(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
|
|
@ -3646,6 +3661,28 @@ class TestGetOptionalParamsTencent:
|
|||
assert isinstance(config, TencentAnthropicMessagesConfig)
|
||||
assert config.custom_llm_provider == "tencent"
|
||||
|
||||
def test_bedrock_mantle_claude_messages_config_routing(self):
|
||||
import litellm
|
||||
from litellm.llms.bedrock_mantle.messages.transformation import (
|
||||
BedrockMantleAnthropicMessagesConfig,
|
||||
)
|
||||
|
||||
config = ProviderConfigManager.get_provider_anthropic_messages_config(
|
||||
model="anthropic.claude-sonnet-5",
|
||||
provider=litellm.LlmProviders.BEDROCK_MANTLE,
|
||||
)
|
||||
assert isinstance(config, BedrockMantleAnthropicMessagesConfig)
|
||||
assert config.custom_llm_provider == "bedrock_mantle"
|
||||
|
||||
def test_bedrock_mantle_openai_models_keep_the_messages_bridge(self):
|
||||
import litellm
|
||||
|
||||
config = ProviderConfigManager.get_provider_anthropic_messages_config(
|
||||
model="openai.gpt-5.6-sol",
|
||||
provider=litellm.LlmProviders.BEDROCK_MANTLE,
|
||||
)
|
||||
assert config is None
|
||||
|
||||
|
||||
class TestValidateEnvironmentTencent:
|
||||
"""Tests that validate_environment resolves TENCENT_API_KEY for the tencent provider."""
|
||||
|
|
|
|||
|
|
@ -1,12 +1,13 @@
|
|||
import asyncio
|
||||
import copy
|
||||
from typing import cast
|
||||
import functools
|
||||
from typing import Final, cast
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT
|
||||
from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, PROMPT_CACHE_LOOKBACK_POSITIONS
|
||||
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import (
|
||||
|
|
@ -19,6 +20,23 @@ from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_p
|
|||
|
||||
MODEL_GROUP_ALIAS = "my-claude-group"
|
||||
OPUS_4_6_MIN_TOKENS = 4096
|
||||
CALLBACK_REGISTRIES: Final = (
|
||||
"input_callback",
|
||||
"success_callback",
|
||||
"failure_callback",
|
||||
"_async_success_callback",
|
||||
"_async_failure_callback",
|
||||
"callbacks",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fresh_callback_registries(monkeypatch):
|
||||
"""`litellm.logging_callback_manager` keeps one callback per class, so a
|
||||
`PromptCachingDeploymentCheck` or `_SentMessagesCapture` left behind by an
|
||||
earlier test would swallow the next test's success events."""
|
||||
for registry in CALLBACK_REGISTRIES:
|
||||
monkeypatch.setattr(litellm, registry, [])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -210,6 +228,58 @@ async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is
|
|||
AUTO_CACHING_MODEL = "anthropic/claude-sonnet-4-5"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replayed_redacted_thinking_block_still_records_and_pins():
|
||||
"""
|
||||
A model that returns no reasoning summary (gpt-5.x through the /v1/messages bridge, Anthropic with
|
||||
redacted reasoning) hands the client a `redacted_thinking` block, and the client replays it on every
|
||||
later turn. The token count behind `is_prompt_caching_valid_prompt` raised on that block, the helper
|
||||
swallowed it to False, and the check neither recorded the serving deployment nor pinned it, so the
|
||||
conversation bounced across the group and paid a cache write on each deployment.
|
||||
"""
|
||||
cache = DualCache()
|
||||
check = PromptCachingDeploymentCheck(cache=cache)
|
||||
model = "openai/gpt-5.6-sol"
|
||||
deployments = _deployments(model, model, model)
|
||||
messages = cast(
|
||||
list[AllMessageValues],
|
||||
[
|
||||
*_messages(word_count=3000),
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:" + "Z" * 400},
|
||||
{"type": "text", "text": "Draw from the box labeled Mixed."},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Restate that in one sentence."},
|
||||
],
|
||||
)
|
||||
|
||||
assert is_prompt_caching_valid_prompt(model=model, messages=messages) is True
|
||||
|
||||
await check.async_log_success_event(
|
||||
kwargs={
|
||||
"standard_logging_object": {
|
||||
"call_type": "anthropic_messages",
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"model_id": "dep-2",
|
||||
}
|
||||
},
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
filtered = await check.async_filter_deployments(
|
||||
model=MODEL_GROUP_ALIAS,
|
||||
healthy_deployments=deployments,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
assert filtered == [deployments[1]]
|
||||
|
||||
|
||||
def _auto_caching_messages() -> list[AllMessageValues]:
|
||||
"""A prompt over the model minimum that carries no client cache_control."""
|
||||
return cast(
|
||||
|
|
@ -552,3 +622,292 @@ async def test_async_log_success_event_counts_the_prompt_off_the_event_loop():
|
|||
"model_id": "dep-1"
|
||||
}
|
||||
assert_loop_stayed_free(took, lags)
|
||||
|
||||
|
||||
LONG_PROMPT = "word " * 3000
|
||||
ONE_PIXEL_PNG = (
|
||||
"data:image/png;base64,"
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
|
||||
)
|
||||
|
||||
|
||||
def _turn(*messages: dict) -> list[AllMessageValues]:
|
||||
return cast(list[AllMessageValues], list(messages))
|
||||
|
||||
|
||||
def _text(text: str) -> dict:
|
||||
return {"type": "text", "text": text}
|
||||
|
||||
|
||||
def _marked(text: str) -> dict:
|
||||
return {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pin_survives_the_breakpoint_moving_to_the_next_turn():
|
||||
"""
|
||||
The regression. Claude Code marks only the newest user message each turn, so the last breakpoint
|
||||
moves forward every turn. The key hashed the prefix up to that moving breakpoint, markers
|
||||
included, so no turn after the first ever found the pin the previous turn wrote, and a
|
||||
multi-deployment group re-rolled the deployment mid-session, paying a cache write on a
|
||||
deployment whose provider cache held nothing of the conversation.
|
||||
"""
|
||||
cache = DualCache()
|
||||
check = PromptCachingDeploymentCheck(cache=cache)
|
||||
deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)
|
||||
turn_one = _turn({"role": "user", "content": [_marked(LONG_PROMPT)]})
|
||||
turn_two = _turn(
|
||||
{"role": "user", "content": [_text(LONG_PROMPT)]},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
{"role": "user", "content": [_marked("next")]},
|
||||
)
|
||||
|
||||
await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=turn_one, tools=None)
|
||||
|
||||
filtered = await check.async_filter_deployments(
|
||||
model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two
|
||||
)
|
||||
|
||||
assert filtered == [deployments[1]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pin_survives_the_marked_message_coming_back_as_string_content():
|
||||
"""
|
||||
Claude Code sends the message that carries a breakpoint as a one-block content list and re-sends
|
||||
it next turn as plain string content once the marker has moved on. The provider caches both
|
||||
shapes identically, so the key has to as well, or the walk-back never lands on the turn-one write.
|
||||
"""
|
||||
cache = DualCache()
|
||||
check = PromptCachingDeploymentCheck(cache=cache)
|
||||
deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)
|
||||
turn_one = _turn(
|
||||
{"role": "system", "content": [_marked(LONG_PROMPT)]},
|
||||
{"role": "user", "content": [_marked("hello")]},
|
||||
)
|
||||
turn_two = _turn(
|
||||
{"role": "system", "content": LONG_PROMPT},
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "content": "hi"},
|
||||
{"role": "user", "content": [_marked("again")]},
|
||||
)
|
||||
|
||||
await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-1", messages=turn_one, tools=None)
|
||||
|
||||
filtered = await check.async_filter_deployments(
|
||||
model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two
|
||||
)
|
||||
|
||||
assert filtered == [deployments[0]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lookback_stops_where_the_provider_cache_stops():
|
||||
"""
|
||||
Anthropic finds a cached prefix at most PROMPT_CACHE_LOOKBACK_POSITIONS block positions behind a
|
||||
breakpoint, the breakpoint block included. Probing further would pin to a deployment whose cache
|
||||
the provider will not consult, and probing less would drop pins the provider still honors.
|
||||
"""
|
||||
prompt_cache = PromptCachingCache(cache=DualCache())
|
||||
await prompt_cache.async_add_model_id(
|
||||
model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("block 0")]}), tools=None
|
||||
)
|
||||
|
||||
def turn_with_blocks_after(count: int) -> list[AllMessageValues]:
|
||||
later = [_text(f"block {index}") for index in range(1, count)] + [_marked(f"block {count}")]
|
||||
return _turn({"role": "user", "content": [_text("block 0"), *later]})
|
||||
|
||||
inside_window = turn_with_blocks_after(PROMPT_CACHE_LOOKBACK_POSITIONS - 1)
|
||||
past_window = turn_with_blocks_after(PROMPT_CACHE_LOOKBACK_POSITIONS)
|
||||
|
||||
assert await prompt_cache.async_get_model_id(messages=inside_window, tools=None) == {"model_id": "dep-1"}
|
||||
assert prompt_cache.get_model_id(messages=inside_window, tools=None) == {"model_id": "dep-1"}
|
||||
assert await prompt_cache.async_get_model_id(messages=past_window, tools=None) is None
|
||||
assert prompt_cache.get_model_id(messages=past_window, tools=None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_run_of_tool_blocks_counts_as_one_lookback_position():
|
||||
"""
|
||||
The provider counts consecutive tool_use blocks as one lookback position, and consecutive
|
||||
tool_result blocks as one, in both the Anthropic and the OpenAI message shapes. An agent turn that
|
||||
fans out into many tool calls would otherwise push the previous breakpoint out of the window
|
||||
after a single turn, which is exactly when the conversation is longest and the cache matters most.
|
||||
"""
|
||||
prompt_cache = PromptCachingCache(cache=DualCache())
|
||||
await prompt_cache.async_add_model_id(
|
||||
model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("task")]}), tools=None
|
||||
)
|
||||
fan_out = PROMPT_CACHE_LOOKBACK_POSITIONS + 5
|
||||
|
||||
def anthropic_shaped(tool_use_type: str, tool_result_type: str) -> list[AllMessageValues]:
|
||||
return _turn(
|
||||
{"role": "user", "content": [_text("task")]},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": tool_use_type, "id": f"call-{index}", "name": "read", "input": {"index": index}}
|
||||
for index in range(fan_out)
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
*(
|
||||
{"type": tool_result_type, "tool_use_id": f"call-{index}", "content": "ok"}
|
||||
for index in range(fan_out)
|
||||
),
|
||||
_marked("continue"),
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
openai_shaped = _turn(
|
||||
{"role": "user", "content": [_text("task")]},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{"id": f"call-{index}", "type": "function", "function": {"name": "read", "arguments": "{}"}}
|
||||
for index in range(fan_out)
|
||||
],
|
||||
},
|
||||
*({"role": "tool", "tool_call_id": f"call-{index}", "content": "ok"} for index in range(fan_out)),
|
||||
{"role": "user", "content": [_marked("continue")]},
|
||||
)
|
||||
|
||||
assert await prompt_cache.async_get_model_id(messages=anthropic_shaped("tool_use", "tool_result"), tools=None) == {
|
||||
"model_id": "dep-1"
|
||||
}
|
||||
assert await prompt_cache.async_get_model_id(messages=openai_shaped, tools=None) == {"model_id": "dep-1"}
|
||||
assert await prompt_cache.async_get_model_id(messages=anthropic_shaped("text", "text"), tools=None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_edited_earlier_block_does_not_inherit_the_pin():
|
||||
"""
|
||||
Every key must bind the whole prefix before its block, not the block alone, or a conversation
|
||||
that repeats a pinned block after an edit walks back onto a cache the provider no longer holds.
|
||||
"""
|
||||
prompt_cache = PromptCachingCache(cache=DualCache())
|
||||
await prompt_cache.async_add_model_id(
|
||||
model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("original")]}), tools=None
|
||||
)
|
||||
edited = _turn(
|
||||
{"role": "user", "content": [_text("edited")]},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
{"role": "user", "content": [_marked("original")]},
|
||||
)
|
||||
|
||||
assert await prompt_cache.async_get_model_id(messages=edited, tools=None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_swapped_roles_do_not_inherit_the_pin():
|
||||
"""The message envelope is part of what the provider caches, so the same blocks under other roles key apart."""
|
||||
prompt_cache = PromptCachingCache(cache=DualCache())
|
||||
pinned = _turn(
|
||||
{"role": "user", "content": [_text("question")]},
|
||||
{"role": "assistant", "content": [_marked("answer")]},
|
||||
)
|
||||
swapped = _turn(
|
||||
{"role": "assistant", "content": [_text("question")]},
|
||||
{"role": "user", "content": [_marked("answer")]},
|
||||
)
|
||||
await prompt_cache.async_add_model_id(model_id="dep-1", messages=pinned, tools=None)
|
||||
|
||||
assert await prompt_cache.async_get_model_id(messages=pinned, tools=None) == {"model_id": "dep-1"}
|
||||
assert await prompt_cache.async_get_model_id(messages=swapped, tools=None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_bytes_in_a_block_hash_instead_of_failing_the_request():
|
||||
"""A block carrying raw bytes must key like any other block rather than raising out of the router filter."""
|
||||
prompt_cache = PromptCachingCache(cache=DualCache())
|
||||
binary_block = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": b"\xff\xfe"}}
|
||||
turn = _turn({"role": "user", "content": [binary_block, _marked("describe")]})
|
||||
await prompt_cache.async_add_model_id(model_id="dep-1", messages=turn, tools=None)
|
||||
|
||||
assert await prompt_cache.async_get_model_id(messages=turn, tools=None) == {"model_id": "dep-1"}
|
||||
|
||||
|
||||
class _BrokenBatchReadCache(DualCache):
|
||||
async def async_batch_get_cache(self, keys, parent_otel_span=None, local_only=False, **kwargs):
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_batch_read_pins_nothing():
|
||||
"""DualCache answers None rather than a list when the batch read raises, and routing must fall through."""
|
||||
prompt_cache = PromptCachingCache(cache=_BrokenBatchReadCache())
|
||||
|
||||
assert (
|
||||
await prompt_cache.async_get_model_id(messages=_turn({"role": "user", "content": [_marked("x")]}), tools=None)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pin_matches_when_the_success_event_truncated_an_image_payload(monkeypatch, local_model_cost_map):
|
||||
"""
|
||||
The success event only ever sees the standard logging payload, whose long base64 data URIs are
|
||||
replaced by size placeholders, while routing sees the raw request. Hashing the raw bytes on the
|
||||
read side would key every image-carrying session past its own pin.
|
||||
"""
|
||||
capture = _SentMessagesCapture()
|
||||
monkeypatch.setattr(litellm, "callbacks", [capture])
|
||||
image = {"type": "image_url", "image_url": {"url": ONE_PIXEL_PNG}}
|
||||
turn_one = _turn({"role": "user", "content": [image, _marked(LONG_PROMPT)]})
|
||||
|
||||
await litellm.acompletion(
|
||||
model=AUTO_CACHING_MODEL, messages=copy.deepcopy(turn_one), mock_response="ok", api_key="sk-fake"
|
||||
)
|
||||
logged = await _eventually(lambda: capture.messages)
|
||||
assert logged is not None
|
||||
assert logged != turn_one
|
||||
|
||||
cache = DualCache()
|
||||
await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=logged, tools=None)
|
||||
turn_two = _turn(
|
||||
{"role": "user", "content": [image, _text(LONG_PROMPT)]},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
{"role": "user", "content": [_marked("next")]},
|
||||
)
|
||||
deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)
|
||||
|
||||
filtered = await PromptCachingDeploymentCheck(cache=cache).async_filter_deployments(
|
||||
model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two
|
||||
)
|
||||
|
||||
assert filtered == [deployments[1]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claude_code_style_session_stays_on_one_deployment_across_turns(local_model_cost_map):
|
||||
"""
|
||||
End to end over the router with a client that marks only the newest user message each turn, the
|
||||
way Claude Code does. Every turn has to land on the deployment that served the first one.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": MODEL_GROUP_ALIAS,
|
||||
"litellm_params": {"model": AUTO_CACHING_MODEL, "api_key": "sk-fake"},
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
for model_id in (f"dep-{number}" for number in range(1, 7))
|
||||
],
|
||||
optional_pre_call_checks=["prompt_caching"],
|
||||
)
|
||||
user_turns = [LONG_PROMPT, *(f"follow-up {number}" for number in range(1, 9))]
|
||||
history: list[AllMessageValues] = []
|
||||
served: list[str] = []
|
||||
for text in user_turns:
|
||||
request = cast(list[AllMessageValues], [*history, {"role": "user", "content": [_marked(text)]}])
|
||||
response = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=request, mock_response="ok")
|
||||
served.append(response._hidden_params["model_id"])
|
||||
pin_key = PromptCachingCache.get_prompt_caching_cache_key(request, None)
|
||||
assert await _eventually(functools.partial(router.cache.get_cache, key=pin_key)) is not None
|
||||
history = [*history, {"role": "user", "content": [_text(text)]}, {"role": "assistant", "content": "ok"}]
|
||||
|
||||
assert served == [served[0]] * len(user_turns)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@
|
|||
import React, { useMemo, useState } from "react";
|
||||
import { ArrowDown, ArrowUp, ArrowUpDown, Info } from "lucide-react";
|
||||
|
||||
import AdvancedDatePicker from "@/components/shared/advanced_date_picker";
|
||||
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
|
||||
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
|
||||
import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
|
|
@ -81,7 +80,7 @@ const SortableHead = ({
|
|||
};
|
||||
|
||||
const CacheLeakageCard: React.FC<CacheLeakageCardProps> = ({ activity }) => {
|
||||
const { dateValue, onDateChange, results, loading, isFetchingMore, apiKeyTruncation } = activity;
|
||||
const { results, loading, isFetchingMore, apiKeyTruncation } = activity;
|
||||
const [dimension, setDimension] = useState<CacheLeakageDimension>("key");
|
||||
const [sort, setSort] = useState<SortState>({ column: "potentialSavings", dir: "desc" });
|
||||
const leakage = useMemo(() => computeCacheLeakage(results, dimension), [results, dimension]);
|
||||
|
|
@ -111,9 +110,6 @@ const CacheLeakageCard: React.FC<CacheLeakageCardProps> = ({ activity }) => {
|
|||
cached token, after cache-write premiums.
|
||||
</p>
|
||||
</div>
|
||||
<div className="shrink-0">
|
||||
<AdvancedDatePicker value={dateValue} onValueChange={onDateChange} />
|
||||
</div>
|
||||
</div>
|
||||
<Tabs value={dimension} onValueChange={(value) => setDimension(value === "model" ? "model" : "key")}>
|
||||
<TabsList>
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () =>
|
|||
}));
|
||||
|
||||
vi.mock("./PromptCompressionTab", () => ({ __esModule: true, default: () => <div /> }));
|
||||
vi.mock("./PromptCachingRequestsTable", () => ({ default: () => <div /> }));
|
||||
|
||||
import CostOptimizationView from "./CostOptimizationView";
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,248 @@
|
|||
import { Profiler } from "react";
|
||||
import { act, fireEvent, renderWithProviders, screen, testQueryClient, waitFor, within } from "@/../tests/test-utils";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
import type { components } from "@/lib/http/schema";
|
||||
import PromptCachingRequestsTable from "./PromptCachingRequestsTable";
|
||||
import type { DateRange } from "./useDailyActivityRange";
|
||||
|
||||
type CacheRequest = components["schemas"]["PromptCachingRequest"];
|
||||
type RequestsResponse = components["schemas"]["PromptCachingRequestsResponse"];
|
||||
const firstCursor = { start_time: "2026-09-01T11:59:59.123456Z", request_id: "first-boundary?&" };
|
||||
const secondCursor = { start_time: firstCursor.start_time, request_id: "second-boundary" };
|
||||
const fetchMock = vi.fn<typeof fetch>();
|
||||
const dates = { from: new Date(2026, 8, 1, 12), to: new Date(2026, 8, 2, 12) };
|
||||
const request = (overrides: Partial<CacheRequest> = {}): CacheRequest => ({
|
||||
request_id: "request-default",
|
||||
start_time: "2026-09-01T12:00:00Z",
|
||||
model: "cache-test-model",
|
||||
gateway_injected: true,
|
||||
cache_read_tokens: 0,
|
||||
cache_creation_tokens: 1000,
|
||||
spend: 0.0375,
|
||||
net_savings: -0.0075,
|
||||
...overrides,
|
||||
});
|
||||
const response = (requests: CacheRequest[], nextCursor: RequestsResponse["next_cursor"] = null) => {
|
||||
const body: RequestsResponse = { requests, has_more: nextCursor !== null, next_cursor: nextCursor, page_size: 50 };
|
||||
return Response.json(body);
|
||||
};
|
||||
const lastQuery = () => new URL(String(fetchMock.mock.calls.at(-1)?.[0]), "http://localhost").searchParams;
|
||||
|
||||
describe("PromptCachingRequestsTable", () => {
|
||||
beforeEach(() => {
|
||||
fetchMock.mockReset();
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
testQueryClient.clear();
|
||||
vi.unstubAllGlobals();
|
||||
vi.unstubAllEnvs();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it("separates recorded injection from cache hits, retains write premiums and unknown savings, and links each request", async () => {
|
||||
const clientHit = {
|
||||
request_id: "client-hit",
|
||||
gateway_injected: false,
|
||||
cache_read_tokens: 10000,
|
||||
cache_creation_tokens: 0,
|
||||
net_savings: 0.27,
|
||||
};
|
||||
fetchMock.mockResolvedValue(
|
||||
response([
|
||||
request({ request_id: "injected/write?&", net_savings: -0.0075 }),
|
||||
request(clientHit),
|
||||
request({ request_id: "unknown-price", net_savings: null }),
|
||||
request({ request_id: "no-benefit", net_savings: 0 }),
|
||||
]),
|
||||
);
|
||||
renderWithProviders(<PromptCachingRequestsTable accessToken="token-a" dateValue={dates} />);
|
||||
|
||||
const table = await screen.findByRole("table", { name: "Prompt caching requests" });
|
||||
const write = within(table).getByRole("row", { name: /injected\/write/ });
|
||||
expect(within(write).getByText("Recorded")).toBeInTheDocument();
|
||||
expect(within(write).getByText("1,000")).toBeInTheDocument();
|
||||
expect(within(write).getByText("$0.0375")).toBeInTheDocument();
|
||||
expect(within(write).getByText("-$0.0075")).toBeInTheDocument();
|
||||
expect(within(write).getByText(new Date("2026-09-01T12:00:00Z").toLocaleString())).toBeInTheDocument();
|
||||
expect(within(write).getByText("cache-test-model")).toHaveAttribute("title", "cache-test-model");
|
||||
expect(within(write).getByRole("link")).toHaveAttribute("href", "/ui/logs?log_id=injected%2Fwrite%3F%26");
|
||||
|
||||
const hit = within(table).getByRole("row", { name: /client-hit/ });
|
||||
expect(within(hit).getByText("Not recorded")).toBeInTheDocument();
|
||||
expect(within(hit).getByText("10,000")).toBeInTheDocument();
|
||||
expect(within(hit).getByText("$0.2700")).toBeInTheDocument();
|
||||
expect(within(table).getByRole("row", { name: /unknown-price/ })).toHaveTextContent("Unavailable");
|
||||
expect(within(table).getByRole("row", { name: /no-benefit/ })).toHaveTextContent("$0.00");
|
||||
expect(screen.getByText(/after cache-write premiums/)).toBeInTheDocument();
|
||||
expect(lastQuery().get("start_date")).toBe("2026-09-01T00:00:00.000Z");
|
||||
expect(lastQuery().get("end_date")).toBe("2026-09-02T23:59:59.999Z");
|
||||
expect(fetchMock.mock.calls[0][1]?.headers).toEqual(expect.objectContaining({ Authorization: "Bearer token-a" }));
|
||||
});
|
||||
|
||||
it("forwards complete server cursors, goes back to prior cursors, and clears them for each caching filter", async () => {
|
||||
fetchMock.mockImplementation(async (input) => {
|
||||
const query = new URL(String(input), "http://localhost").searchParams;
|
||||
const pages = new Map([
|
||||
[null, 1],
|
||||
[firstCursor.request_id, 2],
|
||||
[secondCursor.request_id, 3],
|
||||
]);
|
||||
const page = pages.get(query.get("cursor_request_id"));
|
||||
const nextCursor =
|
||||
new Map([
|
||||
[1, firstCursor],
|
||||
[2, secondCursor],
|
||||
]).get(page ?? 0) ?? null;
|
||||
return response([request({ request_id: `${query.get("filter")}-${page}` })], nextCursor);
|
||||
});
|
||||
renderWithProviders(<PromptCachingRequestsTable accessToken="token-a" dateValue={dates} />);
|
||||
await screen.findByRole("link", { name: "all-1" });
|
||||
expect(screen.getByRole("button", { name: "Previous" })).toBeDisabled();
|
||||
expect(lastQuery().has("page")).toBe(false);
|
||||
expect(lastQuery().has("cursor_request_id")).toBe(false);
|
||||
|
||||
fireEvent.click(screen.getByRole("button", { name: "Next" }));
|
||||
await screen.findByRole("link", { name: "all-2" });
|
||||
expect(screen.getByText("Page 2")).toBeInTheDocument();
|
||||
expect(lastQuery().get("cursor_start_time")).toBe(firstCursor.start_time);
|
||||
expect(lastQuery().get("cursor_request_id")).toBe(firstCursor.request_id);
|
||||
fireEvent.click(screen.getByRole("button", { name: "Next" }));
|
||||
await screen.findByRole("link", { name: "all-3" });
|
||||
expect(screen.getByText("Page 3")).toBeInTheDocument();
|
||||
expect(lastQuery().get("cursor_start_time")).toBe(secondCursor.start_time);
|
||||
expect(lastQuery().get("cursor_request_id")).toBe(secondCursor.request_id);
|
||||
expect(screen.getByRole("button", { name: "Next" })).toBeDisabled();
|
||||
|
||||
await testQueryClient.invalidateQueries({ refetchType: "none" });
|
||||
fireEvent.click(screen.getByRole("button", { name: "Previous" }));
|
||||
await screen.findByRole("link", { name: "all-2" });
|
||||
await waitFor(() => expect(lastQuery().get("cursor_request_id")).toBe(firstCursor.request_id));
|
||||
expect(lastQuery().get("cursor_start_time")).toBe(firstCursor.start_time);
|
||||
expect(screen.getByText("Page 2")).toBeInTheDocument();
|
||||
fireEvent.click(screen.getByRole("button", { name: "Previous" }));
|
||||
await screen.findByRole("link", { name: "all-1" });
|
||||
await waitFor(() => expect(lastQuery().has("cursor_request_id")).toBe(false));
|
||||
expect(lastQuery().has("cursor_start_time")).toBe(false);
|
||||
fireEvent.click(screen.getByRole("button", { name: "Next" }));
|
||||
await screen.findByRole("link", { name: "all-2" });
|
||||
|
||||
fireEvent.click(screen.getByRole("tab", { name: "LiteLLM injected" }));
|
||||
await screen.findByRole("link", { name: "injected-1" });
|
||||
expect(screen.queryByRole("link", { name: "all-2" })).not.toBeInTheDocument();
|
||||
expect(lastQuery().get("filter")).toBe("injected");
|
||||
expect(lastQuery().has("cursor_request_id")).toBe(false);
|
||||
expect(lastQuery().has("cursor_start_time")).toBe(false);
|
||||
|
||||
fireEvent.click(screen.getByRole("button", { name: "Next" }));
|
||||
await screen.findByRole("link", { name: "injected-2" });
|
||||
fireEvent.click(screen.getByRole("tab", { name: "Cache hits" }));
|
||||
await screen.findByRole("link", { name: "hits-1" });
|
||||
expect(lastQuery().get("filter")).toBe("hits");
|
||||
expect(lastQuery().get("page_size")).toBe("50");
|
||||
expect(screen.getByText("Page 1")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("includes the current UTC day for a range ending today, matching the activity totals", async () => {
|
||||
vi.stubEnv("TZ", "America/Los_Angeles");
|
||||
vi.setSystemTime(new Date("2026-09-20T03:00:00Z"));
|
||||
fetchMock.mockResolvedValue(response([]));
|
||||
const today = { from: new Date(2026, 8, 19), to: new Date() };
|
||||
renderWithProviders(<PromptCachingRequestsTable accessToken="token-a" dateValue={today} />);
|
||||
|
||||
await screen.findByText("No matching prompt caching requests in this range");
|
||||
expect(lastQuery().get("start_date")).toBe("2026-09-19T00:00:00.000Z");
|
||||
expect(lastQuery().get("end_date")).toBe("2026-09-20T23:59:59.999Z");
|
||||
});
|
||||
|
||||
it.each(["date", "authentication"])(
|
||||
"hides every old-scope frame and resets pagination when %s changes",
|
||||
async (change) => {
|
||||
fetchMock.mockResolvedValueOnce(response([request({ request_id: "old-first" })], firstCursor));
|
||||
fetchMock.mockResolvedValueOnce(response([request({ request_id: "old-second" })]));
|
||||
const committedOldRows: boolean[] = [];
|
||||
const snapshot = () => {
|
||||
committedOldRows.push(screen.queryByRole("link", { name: "old-second" }) !== null);
|
||||
};
|
||||
const tree = (accessToken: string, dateValue: DateRange) => (
|
||||
<Profiler id="request-scope" onRender={snapshot}>
|
||||
<PromptCachingRequestsTable accessToken={accessToken} dateValue={dateValue} />
|
||||
</Profiler>
|
||||
);
|
||||
const { rerender } = renderWithProviders(tree("token-a", dates));
|
||||
await screen.findByRole("link", { name: "old-first" });
|
||||
fireEvent.click(screen.getByRole("button", { name: "Next" }));
|
||||
await screen.findByRole("link", { name: "old-second" });
|
||||
|
||||
const pending = Promise.withResolvers<Response>();
|
||||
fetchMock.mockReturnValueOnce(pending.promise);
|
||||
committedOldRows.length = 0;
|
||||
rerender(
|
||||
tree(
|
||||
change === "authentication" ? "token-b" : "token-a",
|
||||
change === "date" ? { ...dates, to: new Date(2026, 8, 3) } : dates,
|
||||
),
|
||||
);
|
||||
|
||||
expect(screen.getByRole("status")).toHaveTextContent("Loading requests");
|
||||
expect(committedOldRows.length).toBeGreaterThan(0);
|
||||
expect(committedOldRows.every((visible) => !visible)).toBe(true);
|
||||
expect(lastQuery().has("cursor_request_id")).toBe(false);
|
||||
expect(lastQuery().has("cursor_start_time")).toBe(false);
|
||||
if (change === "date") {
|
||||
expect(lastQuery().get("end_date")).toBe("2026-09-03T23:59:59.999Z");
|
||||
} else {
|
||||
expect(fetchMock.mock.calls.at(-1)?.[1]?.headers).toEqual(
|
||||
expect.objectContaining({ Authorization: "Bearer token-b" }),
|
||||
);
|
||||
}
|
||||
|
||||
pending.resolve(response([request({ request_id: "new-first" })]));
|
||||
await screen.findByRole("link", { name: "new-first" });
|
||||
expect(screen.getByText("Page 1")).toBeInTheDocument();
|
||||
expect(committedOldRows.every((visible) => !visible)).toBe(true);
|
||||
},
|
||||
);
|
||||
|
||||
it("ignores a delayed response from the previous caching filter", async () => {
|
||||
const stale = Promise.withResolvers<Response>();
|
||||
const current = Promise.withResolvers<Response>();
|
||||
fetchMock.mockReturnValueOnce(stale.promise).mockReturnValueOnce(current.promise);
|
||||
renderWithProviders(<PromptCachingRequestsTable accessToken="token-a" dateValue={dates} />);
|
||||
fireEvent.click(screen.getByRole("tab", { name: "Cache hits" }));
|
||||
expect(lastQuery().get("filter")).toBe("hits");
|
||||
|
||||
current.resolve(response([request({ request_id: "current-hit" })]));
|
||||
await screen.findByRole("link", { name: "current-hit" });
|
||||
await act(async () => {
|
||||
stale.resolve(response([request({ request_id: "stale-all" })], firstCursor));
|
||||
await stale.promise;
|
||||
});
|
||||
|
||||
expect(screen.getByRole("link", { name: "current-hit" })).toBeInTheDocument();
|
||||
expect(screen.queryByRole("link", { name: "stale-all" })).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Next" })).toBeDisabled();
|
||||
});
|
||||
|
||||
it("offers retry after a failed read and shows the empty state after it succeeds", async () => {
|
||||
fetchMock.mockRejectedValueOnce(new Error("offline"));
|
||||
fetchMock.mockResolvedValueOnce(response([]));
|
||||
renderWithProviders(<PromptCachingRequestsTable accessToken="token-a" dateValue={dates} />);
|
||||
|
||||
expect(await screen.findByRole("alert")).toHaveTextContent("Could not load prompt caching requests");
|
||||
fireEvent.click(screen.getByRole("button", { name: "Retry" }));
|
||||
expect(await screen.findByText("No matching prompt caching requests in this range")).toBeInTheDocument();
|
||||
expect(screen.queryByRole("alert")).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Next" })).toBeDisabled();
|
||||
expect(fetchMock).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("does not request data for an incomplete date range", async () => {
|
||||
renderWithProviders(<PromptCachingRequestsTable accessToken="token-a" dateValue={{ from: dates.from }} />);
|
||||
expect(screen.getByText("Select a date range to view requests")).toBeInTheDocument();
|
||||
expect(screen.queryByRole("status")).not.toBeInTheDocument();
|
||||
await waitFor(() => expect(fetchMock).not.toHaveBeenCalled());
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,186 @@
|
|||
"use client";
|
||||
|
||||
import { useQuery, type UseQueryOptions } from "@tanstack/react-query";
|
||||
import Link from "next/link";
|
||||
import { useState } from "react";
|
||||
|
||||
import { apiClient } from "@/components/networking";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
|
||||
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
|
||||
import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import { LOG_ID_QUERY_PARAM } from "@/components/view_logs/logDetailRouting";
|
||||
import type { paths } from "@/lib/http/schema";
|
||||
import { formatNumberWithCommas } from "@/utils/dataUtils";
|
||||
import { uiHref } from "@/utils/uiHref";
|
||||
import { usd } from "./costOptimizationUtils";
|
||||
import { benchmarksWindow as activityWindow } from "./useAutoRouterBenchmarks";
|
||||
import type { DateRange } from "./useDailyActivityRange";
|
||||
|
||||
const REQUESTS_PATH = "/cost_optimization/prompt_caching/requests";
|
||||
type RequestsEndpoint = paths[typeof REQUESTS_PATH]["get"];
|
||||
type RequestsResponse = RequestsEndpoint["responses"][200]["content"]["application/json"];
|
||||
type RequestsQuery = NonNullable<RequestsEndpoint["parameters"]["query"]>;
|
||||
type RequestFilter = NonNullable<RequestsQuery["filter"]>;
|
||||
type RequestCursor = RequestsResponse["next_cursor"];
|
||||
|
||||
interface PromptCachingRequestsTableProps {
|
||||
accessToken: string;
|
||||
dateValue: DateRange;
|
||||
}
|
||||
|
||||
export default function PromptCachingRequestsTable({ accessToken, dateValue }: PromptCachingRequestsTableProps) {
|
||||
const [filter, setFilter] = useState<RequestFilter>("all");
|
||||
const window = activityWindow(dateValue, new Date());
|
||||
const startDate = window.start_date ? `${window.start_date}T00:00:00.000Z` : "";
|
||||
const endDate = window.end_date ? `${window.end_date}T23:59:59.999Z` : "";
|
||||
const scope = JSON.stringify([accessToken, startDate, endDate, filter]);
|
||||
const [pagination, setPagination] = useState<{ scope: string; cursors: readonly RequestCursor[] }>({
|
||||
scope,
|
||||
cursors: [null],
|
||||
});
|
||||
const cursors = pagination.scope === scope ? pagination.cursors : [null];
|
||||
const cursor = cursors.at(-1);
|
||||
const page = cursors.length;
|
||||
|
||||
if (pagination.scope !== scope) {
|
||||
setPagination({ scope, cursors: [null] });
|
||||
}
|
||||
|
||||
const enabled = Boolean(accessToken && startDate && endDate);
|
||||
const query: RequestsQuery = {
|
||||
start_date: startDate,
|
||||
end_date: endDate,
|
||||
filter,
|
||||
page_size: 50,
|
||||
cursor_start_time: cursor?.start_time,
|
||||
cursor_request_id: cursor?.request_id,
|
||||
};
|
||||
const queryOptions: UseQueryOptions<RequestsResponse> = {
|
||||
queryKey: [REQUESTS_PATH, accessToken, query],
|
||||
queryFn: ({ signal }) => apiClient.get<RequestsResponse>(REQUESTS_PATH, { accessToken, query, signal }),
|
||||
enabled,
|
||||
retry: false,
|
||||
};
|
||||
const requests = useQuery(queryOptions);
|
||||
const nextCursor = requests.data?.next_cursor;
|
||||
|
||||
const changeFilter = (value: unknown) => {
|
||||
if (value === "all" || value === "injected" || value === "hits") {
|
||||
setFilter(value);
|
||||
}
|
||||
};
|
||||
|
||||
return (
|
||||
<Card>
|
||||
<CardHeader className="gap-3">
|
||||
<div>
|
||||
<CardTitle>Prompt caching requests</CardTitle>
|
||||
<p className="mt-1 text-sm text-muted-foreground">
|
||||
Requests with recorded LiteLLM injection or provider cache reads or writes. A cache hit alone does not
|
||||
establish LiteLLM injection; older logs may not record it.
|
||||
</p>
|
||||
<p className="mt-1 text-sm text-muted-foreground">
|
||||
Net savings are estimated from logged usage and current configured pricing, after cache-write premiums.
|
||||
Negative values mean caching cost more; unavailable means the request could not be priced.
|
||||
</p>
|
||||
</div>
|
||||
<Tabs value={filter} onValueChange={changeFilter}>
|
||||
<TabsList aria-label="Prompt caching request filters">
|
||||
<TabsTrigger value="all">All caching</TabsTrigger>
|
||||
<TabsTrigger value="injected">LiteLLM injected</TabsTrigger>
|
||||
<TabsTrigger value="hits">Cache hits</TabsTrigger>
|
||||
</TabsList>
|
||||
</Tabs>
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
{!enabled && <p className="py-8 text-center text-muted-foreground">Select a date range to view requests</p>}
|
||||
{enabled && requests.isPending && (
|
||||
<p role="status" className="py-8 text-center text-muted-foreground">
|
||||
Loading requests...
|
||||
</p>
|
||||
)}
|
||||
{enabled && requests.isError && (
|
||||
<div role="alert" className="flex items-center justify-center gap-3 py-8">
|
||||
<p>Could not load prompt caching requests</p>
|
||||
<Button variant="outline" onClick={() => void requests.refetch()} disabled={requests.isFetching}>
|
||||
Retry
|
||||
</Button>
|
||||
</div>
|
||||
)}
|
||||
{enabled && requests.isSuccess && (
|
||||
<>
|
||||
{requests.data.requests.length === 0 ? (
|
||||
<p className="py-8 text-center text-muted-foreground">
|
||||
No matching prompt caching requests in this range
|
||||
</p>
|
||||
) : (
|
||||
<Table aria-label="Prompt caching requests">
|
||||
<TableHeader>
|
||||
<TableRow>
|
||||
<TableHead>Request</TableHead>
|
||||
<TableHead>Model</TableHead>
|
||||
<TableHead>LiteLLM injection</TableHead>
|
||||
<TableHead className="text-right">Cache reads</TableHead>
|
||||
<TableHead className="text-right">Cache writes</TableHead>
|
||||
<TableHead className="text-right">Actual cost</TableHead>
|
||||
<TableHead className="text-right">Net savings</TableHead>
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
{requests.data.requests.map((request) => (
|
||||
<TableRow key={request.request_id}>
|
||||
<TableCell>
|
||||
<Link
|
||||
href={uiHref(`logs?${new URLSearchParams({ [LOG_ID_QUERY_PARAM]: request.request_id })}`)}
|
||||
className="block max-w-40 truncate text-primary underline underline-offset-2"
|
||||
title={request.request_id}
|
||||
>
|
||||
{request.request_id}
|
||||
</Link>
|
||||
<time dateTime={request.start_time} className="mt-1 block text-xs text-muted-foreground">
|
||||
{new Date(request.start_time).toLocaleString()}
|
||||
</time>
|
||||
</TableCell>
|
||||
<TableCell>
|
||||
<span className="block max-w-36 truncate" title={request.model}>
|
||||
{request.model}
|
||||
</span>
|
||||
</TableCell>
|
||||
<TableCell>{request.gateway_injected ? "Recorded" : "Not recorded"}</TableCell>
|
||||
<TableCell className="text-right">{formatNumberWithCommas(request.cache_read_tokens)}</TableCell>
|
||||
<TableCell className="text-right">
|
||||
{formatNumberWithCommas(request.cache_creation_tokens)}
|
||||
</TableCell>
|
||||
<TableCell className="text-right">{usd(request.spend)}</TableCell>
|
||||
<TableCell className="text-right">
|
||||
{request.net_savings === null ? "Unavailable" : usd(request.net_savings)}
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
))}
|
||||
</TableBody>
|
||||
</Table>
|
||||
)}
|
||||
<div className="mt-4 flex items-center justify-end gap-3">
|
||||
<Button
|
||||
variant="outline"
|
||||
disabled={page === 1}
|
||||
onClick={() => setPagination({ scope, cursors: cursors.slice(0, -1) })}
|
||||
>
|
||||
Previous
|
||||
</Button>
|
||||
<span className="text-sm text-muted-foreground">Page {page}</span>
|
||||
<Button
|
||||
variant="outline"
|
||||
disabled={!requests.data.has_more || !nextCursor}
|
||||
onClick={() => nextCursor && setPagination({ scope, cursors: [...cursors, nextCursor] })}
|
||||
>
|
||||
Next
|
||||
</Button>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
);
|
||||
}
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
import { render, waitFor, screen } from "@testing-library/react";
|
||||
import { fireEvent, render, waitFor, screen } from "@testing-library/react";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
|
||||
const mockGetGeneralSettingsCall = vi.fn();
|
||||
|
|
@ -12,6 +12,21 @@ vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () =>
|
|||
}));
|
||||
|
||||
const mockCacheLeakageCard = vi.fn();
|
||||
const mockRequestsTable = vi.fn();
|
||||
const nextDateRange = { from: new Date(2026, 8, 1), to: new Date(2026, 8, 2) };
|
||||
|
||||
vi.mock("./PromptCachingRequestsTable", () => ({
|
||||
default: (props: unknown) => {
|
||||
mockRequestsTable(props);
|
||||
return <div data-testid="caching-requests" />;
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock("@/components/shared/advanced_date_picker", () => ({
|
||||
default: ({ onValueChange }: { onValueChange: (range: typeof nextDateRange) => void }) => (
|
||||
<button onClick={() => onValueChange(nextDateRange)}>Change caching dates</button>
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("./CacheLeakageCard", () => ({
|
||||
__esModule: true,
|
||||
|
|
@ -24,7 +39,7 @@ vi.mock("./CacheLeakageCard", () => ({
|
|||
import PromptCachingTab from "./PromptCachingTab";
|
||||
|
||||
describe("PromptCachingTab", () => {
|
||||
it("renders the cache leakage table alongside the caching settings", async () => {
|
||||
it("shares the selected dates between requests and cache leakage alongside caching settings", async () => {
|
||||
mockGetGeneralSettingsCall.mockResolvedValue([]);
|
||||
|
||||
const activity = {
|
||||
|
|
@ -42,6 +57,10 @@ describe("PromptCachingTab", () => {
|
|||
|
||||
expect(screen.getByTestId("caching-settings")).toBeInTheDocument();
|
||||
expect(screen.getByTestId("cache-leakage-card")).toBeInTheDocument();
|
||||
expect(screen.getByTestId("caching-requests")).toBeInTheDocument();
|
||||
expect(mockRequestsTable).toHaveBeenCalledWith({ accessToken: "test-token", dateValue: activity.dateValue });
|
||||
fireEvent.click(screen.getByRole("button", { name: "Change caching dates" }));
|
||||
expect(activity.onDateChange).toHaveBeenCalledWith(nextDateRange);
|
||||
await waitFor(() => expect(mockCacheLeakageCard).toHaveBeenCalledWith(expect.objectContaining({ activity })));
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -3,12 +3,14 @@
|
|||
import React, { useCallback, useEffect, useState } from "react";
|
||||
|
||||
import { getGeneralSettingsCall } from "@/components/networking";
|
||||
import AdvancedDatePicker from "@/components/shared/advanced_date_picker";
|
||||
import { toast } from "@/lib/toast";
|
||||
import {
|
||||
PromptCachingPanel,
|
||||
generalSettingsItem,
|
||||
} from "@/app/(dashboard)/router-settings/_components/general_settings";
|
||||
import CacheLeakageCard from "./CacheLeakageCard";
|
||||
import PromptCachingRequestsTable from "./PromptCachingRequestsTable";
|
||||
import { DailyActivityRange } from "./useDailyActivityRange";
|
||||
|
||||
interface PromptCachingTabProps {
|
||||
|
|
@ -48,6 +50,11 @@ const PromptCachingTab: React.FC<PromptCachingTabProps> = ({ accessToken, activi
|
|||
return (
|
||||
<div className="w-full space-y-6">
|
||||
<PromptCachingPanel accessToken={accessToken} settings={settings} onChange={handleChange} />
|
||||
<div className="flex flex-wrap items-center justify-between gap-3">
|
||||
<p className="text-sm text-muted-foreground">Date range for requests and cache leakage</p>
|
||||
<AdvancedDatePicker value={activity.dateValue} onValueChange={activity.onDateChange} />
|
||||
</div>
|
||||
<PromptCachingRequestsTable accessToken={accessToken} dateValue={activity.dateValue} />
|
||||
<CacheLeakageCard activity={activity} />
|
||||
</div>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -65,6 +65,19 @@ describe("AttachmentTable", () => {
|
|||
);
|
||||
});
|
||||
|
||||
it("should show a Default badge only for default attachments", () => {
|
||||
const attachments = [
|
||||
makeAttachment({ attachment_id: "att-def00001", policy_name: "fallback", default: true }),
|
||||
makeAttachment({ attachment_id: "att-def00002", policy_name: "regular" }),
|
||||
];
|
||||
renderWithProviders(<AttachmentTable {...defaultProps} attachments={attachments} />);
|
||||
const rows = screen.getAllByRole("row").slice(1);
|
||||
const fallbackRow = rows.find((row) => within(row).queryByText("fallback"));
|
||||
const regularRow = rows.find((row) => within(row).queryByText("regular"));
|
||||
expect(within(fallbackRow!).getByText("Default")).toBeInTheDocument();
|
||||
expect(within(regularRow!).queryByText("Default")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should show skeleton rows when isLoading is true", () => {
|
||||
renderWithProviders(<AttachmentTable {...defaultProps} isLoading />);
|
||||
expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0);
|
||||
|
|
|
|||
|
|
@ -181,6 +181,20 @@ export const getAttachmentTableColumns = ({
|
|||
<span className="font-mono text-xs">{row.original.priority}</span>
|
||||
),
|
||||
},
|
||||
{
|
||||
id: "default",
|
||||
accessorFn: (row) => (row.default ? 1 : 0),
|
||||
meta: { title: "Default" },
|
||||
header: ({ column }) => <DataTableSortHeader column={column} title="Default" />,
|
||||
size: 100,
|
||||
enableSorting: true,
|
||||
cell: ({ row }) =>
|
||||
row.original.default ? (
|
||||
<StatusBadge tone="info" label="Default" tooltip="Applied only when no non-default attachment matches" />
|
||||
) : (
|
||||
<span className="text-muted-foreground">-</span>
|
||||
),
|
||||
},
|
||||
{
|
||||
id: "created_at",
|
||||
accessorFn: (row) => row.created_at ?? "",
|
||||
|
|
|
|||
|
|
@ -237,6 +237,21 @@ describe("AddAttachmentForm", () => {
|
|||
expect(createAttachment).toHaveBeenCalledWith("test-token", { policy_name: "policy-alpha", scope: "*" });
|
||||
});
|
||||
|
||||
it("sends default: true when the Default switch is turned on", async () => {
|
||||
const user = userEvent.setup();
|
||||
const createAttachment = vi.fn().mockResolvedValue({});
|
||||
renderWithProviders(<AddAttachmentForm {...defaultProps} createAttachment={createAttachment} />);
|
||||
await selectPolicy(user, "policy-alpha");
|
||||
await user.click(screen.getByRole("switch", { name: /default/i }));
|
||||
await submit(user);
|
||||
await waitFor(() => expect(createAttachment).toHaveBeenCalledTimes(1));
|
||||
expect(createAttachment).toHaveBeenCalledWith("test-token", {
|
||||
policy_name: "policy-alpha",
|
||||
scope: "*",
|
||||
default: true,
|
||||
});
|
||||
});
|
||||
|
||||
it.each([
|
||||
["2147483648", /at most 2147483647/i],
|
||||
["-2147483649", /at least -2147483648/i],
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import { Button } from "@/components/ui/button";
|
|||
import { Input } from "@/components/ui/input";
|
||||
import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group";
|
||||
import { Separator } from "@/components/ui/separator";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
|
||||
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
|
||||
import { useZodForm } from "@/lib/forms/useZodForm";
|
||||
|
|
@ -38,6 +39,7 @@ interface AttachmentFormValues {
|
|||
models: string[];
|
||||
tags: string[];
|
||||
priority: number | null;
|
||||
default: boolean;
|
||||
}
|
||||
|
||||
const EMPTY_VALUES: AttachmentFormValues = {
|
||||
|
|
@ -47,6 +49,7 @@ const EMPTY_VALUES: AttachmentFormValues = {
|
|||
models: [],
|
||||
tags: [],
|
||||
priority: null,
|
||||
default: false,
|
||||
};
|
||||
|
||||
const INT32_MIN = -2147483648;
|
||||
|
|
@ -64,6 +67,7 @@ const attachmentShape = {
|
|||
.min(INT32_MIN, `Priority must be at least ${INT32_MIN}`)
|
||||
.max(INT32_MAX, `Priority must be at most ${INT32_MAX}`)
|
||||
.nullable(),
|
||||
default: z.boolean(),
|
||||
};
|
||||
|
||||
const buildAttachmentSchema = (scopeType: ScopeType, teamsLoaded: boolean, availableTeams: string[]) =>
|
||||
|
|
@ -453,9 +457,23 @@ const AddAttachmentForm: React.FC<AddAttachmentFormProps> = ({
|
|||
/>
|
||||
)}
|
||||
</FormField>
|
||||
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="default"
|
||||
label={labelWithHint(
|
||||
"Default (fallback)",
|
||||
"A default attachment is applied only when no non-default attachment matches the request.",
|
||||
)}
|
||||
description="Use this for the guardrail everyone gets unless they opt in to another attachment."
|
||||
>
|
||||
{({ value, onChange, ref, ...field }) => (
|
||||
<Switch {...field} inputRef={ref} checked={value === true} onCheckedChange={onChange} />
|
||||
)}
|
||||
</FormField>
|
||||
</FieldGroup>
|
||||
|
||||
{impactResult && <ImpactPreviewAlert impactResult={impactResult} />}
|
||||
{impactResult && <ImpactPreviewAlert impactResult={impactResult} isDefault={form.watch("default")} />}
|
||||
|
||||
<div className="flex justify-end space-x-2 mt-4">
|
||||
<Button type="button" variant="secondary" onClick={handleClose}>
|
||||
|
|
|
|||
|
|
@ -80,6 +80,16 @@ describe("buildAttachmentData", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("default", () => {
|
||||
it.each(["global", "specific"] as const)("should send default: true for a %s scope", (scopeType) => {
|
||||
expect(buildAttachmentData({ policy_name: "p", default: true }, scopeType).default).toBe(true);
|
||||
});
|
||||
|
||||
it.each([undefined, false])("should omit default when it is %s", (value) => {
|
||||
expect(buildAttachmentData({ policy_name: "p", default: value }, "specific")).not.toHaveProperty("default");
|
||||
});
|
||||
});
|
||||
|
||||
describe("priority", () => {
|
||||
it.each(["global", "specific"] as const)("should include priority for a %s scope", (scopeType) => {
|
||||
expect(buildAttachmentData({ policy_name: "p", priority: 0 }, scopeType).priority).toBe(0);
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ export interface AttachmentFormInput {
|
|||
models?: string[];
|
||||
tags?: string[];
|
||||
priority?: number | null;
|
||||
default?: boolean;
|
||||
}
|
||||
|
||||
export function buildAttachmentData(
|
||||
|
|
@ -25,5 +26,6 @@ export function buildAttachmentData(
|
|||
if (formValues.tags && formValues.tags.length > 0) data.tags = formValues.tags;
|
||||
}
|
||||
if (typeof formValues.priority === "number") data.priority = formValues.priority;
|
||||
if (formValues.default === true) data.default = true;
|
||||
return data;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -69,6 +69,17 @@ describe("ImpactPreviewAlert", () => {
|
|||
expect(screen.getByText(/1 key\b/i)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should present the counts as an upper bound for a default attachment", () => {
|
||||
renderWithProviders(<ImpactPreviewAlert impactResult={specificImpact} isDefault />);
|
||||
expect(screen.getByText(/would affect up to/i)).toBeInTheDocument();
|
||||
expect(screen.getByText(/no non-default attachment matches/i)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should not qualify the counts for a non-default attachment", () => {
|
||||
renderWithProviders(<ImpactPreviewAlert impactResult={specificImpact} />);
|
||||
expect(screen.queryByText(/up to/i)).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should not show a key section when there are no sample keys", () => {
|
||||
const noKeys = { affected_keys_count: 0, affected_teams_count: 2, sample_keys: [], sample_teams: ["t1", "t2"] };
|
||||
renderWithProviders(<ImpactPreviewAlert impactResult={noKeys} />);
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ interface ImpactResult {
|
|||
|
||||
interface ImpactPreviewAlertProps {
|
||||
impactResult: ImpactResult;
|
||||
isDefault?: boolean;
|
||||
}
|
||||
|
||||
interface SampleListProps {
|
||||
|
|
@ -32,8 +33,9 @@ const SampleList: React.FC<SampleListProps> = ({ label, samples, totalCount }) =
|
|||
</div>
|
||||
);
|
||||
|
||||
const ImpactPreviewAlert: React.FC<ImpactPreviewAlertProps> = ({ impactResult }) => {
|
||||
const ImpactPreviewAlert: React.FC<ImpactPreviewAlertProps> = ({ impactResult, isDefault = false }) => {
|
||||
const isGlobal = impactResult.affected_keys_count === -1;
|
||||
const qualifier = isDefault ? "up to " : "";
|
||||
|
||||
return (
|
||||
<Alert className="mb-4">
|
||||
|
|
@ -47,7 +49,7 @@ const ImpactPreviewAlert: React.FC<ImpactPreviewAlertProps> = ({ impactResult })
|
|||
) : (
|
||||
<div>
|
||||
<span>
|
||||
This attachment would affect{" "}
|
||||
This attachment would affect {qualifier}
|
||||
<strong>
|
||||
{impactResult.affected_keys_count} key{impactResult.affected_keys_count !== 1 ? "s" : ""}
|
||||
</strong>{" "}
|
||||
|
|
@ -57,6 +59,11 @@ const ImpactPreviewAlert: React.FC<ImpactPreviewAlertProps> = ({ impactResult })
|
|||
</strong>
|
||||
.
|
||||
</span>
|
||||
{isDefault && (
|
||||
<div className="text-xs text-muted-foreground">
|
||||
Default attachments only apply to requests no non-default attachment matches, so fewer may be affected.
|
||||
</div>
|
||||
)}
|
||||
{impactResult.sample_keys.length > 0 && (
|
||||
<SampleList
|
||||
label="Keys"
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@ export interface PolicyAttachment {
|
|||
models: string[];
|
||||
tags: string[];
|
||||
priority?: number | null;
|
||||
default?: boolean;
|
||||
created_at?: string;
|
||||
updated_at?: string;
|
||||
created_by?: string;
|
||||
|
|
@ -80,6 +81,7 @@ export interface PolicyAttachmentCreateRequest {
|
|||
models?: string[];
|
||||
tags?: string[];
|
||||
priority?: number;
|
||||
default?: boolean;
|
||||
}
|
||||
|
||||
export interface PolicyListResponse {
|
||||
|
|
|
|||
107
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
107
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -3534,6 +3534,23 @@ export interface paths {
|
|||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/cost_optimization/prompt_caching/requests": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
/** Get Prompt Caching Requests */
|
||||
get: operations["get_prompt_caching_requests_cost_optimization_prompt_caching_requests_get"];
|
||||
put?: never;
|
||||
post?: never;
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/credentials": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -35222,6 +35239,12 @@ export interface components {
|
|||
* @description Request body for creating a policy attachment.
|
||||
*/
|
||||
PolicyAttachmentCreateRequest: {
|
||||
/**
|
||||
* Default
|
||||
* @description Apply this attachment only when no non-default attachment matches the request.
|
||||
* @default false
|
||||
*/
|
||||
default: boolean;
|
||||
/**
|
||||
* Keys
|
||||
* @description Key aliases or patterns this attachment applies to.
|
||||
|
|
@ -35278,6 +35301,12 @@ export interface components {
|
|||
* @description Who created the attachment.
|
||||
*/
|
||||
created_by?: string | null;
|
||||
/**
|
||||
* Default
|
||||
* @description Apply this attachment only when no non-default attachment matches the request.
|
||||
* @default false
|
||||
*/
|
||||
default: boolean;
|
||||
/**
|
||||
* Definition Location
|
||||
* @description Where this attachment is defined: 'db' (database) or 'config' (config.yaml).
|
||||
|
|
@ -35878,6 +35907,48 @@ export interface components {
|
|||
prompt_id: string;
|
||||
prompt_info?: components["schemas"]["PromptInfo"] | null;
|
||||
};
|
||||
/** PromptCachingRequest */
|
||||
PromptCachingRequest: {
|
||||
/** Cache Creation Tokens */
|
||||
cache_creation_tokens: number;
|
||||
/** Cache Read Tokens */
|
||||
cache_read_tokens: number;
|
||||
/** Gateway Injected */
|
||||
gateway_injected: boolean;
|
||||
/** Model */
|
||||
model: string;
|
||||
/** Net Savings */
|
||||
net_savings: number | null;
|
||||
/** Request Id */
|
||||
request_id: string;
|
||||
/** Spend */
|
||||
spend: number;
|
||||
/**
|
||||
* Start Time
|
||||
* Format: date-time
|
||||
*/
|
||||
start_time: string;
|
||||
};
|
||||
/** PromptCachingRequestCursor */
|
||||
PromptCachingRequestCursor: {
|
||||
/** Request Id */
|
||||
request_id: string;
|
||||
/**
|
||||
* Start Time
|
||||
* Format: date-time
|
||||
*/
|
||||
start_time: string;
|
||||
};
|
||||
/** PromptCachingRequestsResponse */
|
||||
PromptCachingRequestsResponse: {
|
||||
/** Has More */
|
||||
has_more: boolean;
|
||||
next_cursor: components["schemas"]["PromptCachingRequestCursor"] | null;
|
||||
/** Page Size */
|
||||
page_size: number;
|
||||
/** Requests */
|
||||
requests: components["schemas"]["PromptCachingRequest"][];
|
||||
};
|
||||
/** PromptInfo */
|
||||
PromptInfo: {
|
||||
/**
|
||||
|
|
@ -47333,6 +47404,42 @@ export interface operations {
|
|||
};
|
||||
};
|
||||
};
|
||||
get_prompt_caching_requests_cost_optimization_prompt_caching_requests_get: {
|
||||
parameters: {
|
||||
query: {
|
||||
start_date: string;
|
||||
end_date: string;
|
||||
page_size?: number;
|
||||
filter?: "all" | "injected" | "hits";
|
||||
cursor_start_time?: string | null;
|
||||
cursor_request_id?: string | null;
|
||||
};
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody?: never;
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["PromptCachingRequestsResponse"];
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
422: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["HTTPValidationError"];
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
get_credentials_credentials_get: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue