mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge origin/litellm_internal_staging into litellm_mantle_daybreak_blue_cost_map
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
commit
71191de8a5
102 changed files with 2375 additions and 3192 deletions
3
.github/workflows/test-rust.yml
vendored
3
.github/workflows/test-rust.yml
vendored
|
|
@ -117,6 +117,9 @@ jobs:
|
|||
|
||||
- run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
|
||||
|
||||
- name: Run pytest tests/test_litellm_rust with the compiled extension
|
||||
run: make test-rust-extension
|
||||
|
||||
- run: >-
|
||||
uv build --wheel --out-dir panic-dist
|
||||
--config-setting "maturin.build-args=--features panic-test,extension-module"
|
||||
|
|
|
|||
13
Makefile
13
Makefile
|
|
@ -4,6 +4,7 @@
|
|||
.PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \
|
||||
test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
|
||||
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
|
||||
test-rust-extension \
|
||||
info lint lint-inner lint-dev lint-checks format \
|
||||
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
|
||||
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
|
||||
|
|
@ -54,6 +55,7 @@ help:
|
|||
@echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)"
|
||||
@echo " make test-integration - Run integration tests"
|
||||
@echo " make test-unit-helm - Run helm unit tests"
|
||||
@echo " make test-rust-extension - Build the Rust extension and run its public Python tests"
|
||||
@echo ""
|
||||
@echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide"
|
||||
@echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine."
|
||||
|
|
@ -289,6 +291,17 @@ pre-commit:
|
|||
@$(MAKE) check
|
||||
|
||||
# Testing targets
|
||||
test-rust-extension:
|
||||
@temporary=$$(mktemp -d) && \
|
||||
trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \
|
||||
$(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \
|
||||
set -- "$$temporary"/wheels/*.whl && \
|
||||
[ "$$#" -eq 1 ] && \
|
||||
UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
|
||||
$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
|
||||
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
|
||||
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
|
||||
|
||||
test: install-test-deps
|
||||
$(UV_RUN) pytest tests/
|
||||
|
||||
|
|
|
|||
|
|
@ -55,6 +55,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16
|
|||
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
|
||||
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
|
||||
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
|
||||
budget_reservation_disabled_info_emitted = False
|
||||
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
|
||||
DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512))
|
||||
SQS_SEND_MESSAGE_ACTION: Final = "SendMessage"
|
||||
|
|
|
|||
|
|
@ -16,6 +16,8 @@ from collections.abc import Iterable, Mapping, Sequence
|
|||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
|
|
@ -23,6 +25,7 @@ from litellm.integrations.prompt_management_base import PromptManagementClient
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
with_prompt_cache_breakpoint,
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request
|
||||
from litellm.types.integrations.anthropic_cache_control_hook import (
|
||||
GATEWAY_INJECTED_CACHE_METADATA_KEY,
|
||||
GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
|
||||
|
|
@ -62,10 +65,26 @@ OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset(
|
|||
)
|
||||
OPENAI_API_HOST: Final = "api.openai.com"
|
||||
OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
|
||||
_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
|
||||
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
|
||||
AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
|
||||
|
||||
|
||||
def _validated_object_mapping(value: object) -> dict[object, object] | None:
|
||||
try:
|
||||
return _OBJECT_MAPPING_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _validated_object_list(value: object) -> list[object] | None:
|
||||
try:
|
||||
return _OBJECT_LIST_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def supports_openai_prompt_cache_breakpoint(model: str) -> bool:
|
||||
model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model)
|
||||
if model_map_flag is not None:
|
||||
|
|
@ -114,6 +133,36 @@ CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_
|
|||
|
||||
|
||||
class AnthropicCacheControlHook(CustomPromptManagement):
|
||||
@staticmethod
|
||||
def _request_value(request_kwargs: object, key: str) -> object:
|
||||
request_mapping: Final = _validated_object_mapping(request_kwargs)
|
||||
if request_mapping is None:
|
||||
return None
|
||||
return request_mapping.get(key)
|
||||
|
||||
@staticmethod
|
||||
def _request_user_agent(request_kwargs: object) -> str | None:
|
||||
proxy_server_request: Final = AnthropicCacheControlHook._request_value(request_kwargs, "proxy_server_request")
|
||||
proxy_server_request_mapping: Final = _validated_object_mapping(proxy_server_request)
|
||||
if proxy_server_request_mapping is None:
|
||||
return None
|
||||
headers: Final = proxy_server_request_mapping.get("headers")
|
||||
headers_mapping: Final = _validated_object_mapping(headers)
|
||||
if headers_mapping is None:
|
||||
return None
|
||||
user_agent: Final = next(
|
||||
(value for key, value in headers_mapping.items() if isinstance(key, str) and key.lower() == "user-agent"),
|
||||
None,
|
||||
)
|
||||
return user_agent if isinstance(user_agent, str) else None
|
||||
|
||||
@staticmethod
|
||||
def _request_system(request_kwargs: object) -> str | list[object] | None:
|
||||
system: Final = AnthropicCacheControlHook._request_value(request_kwargs, "system")
|
||||
if isinstance(system, str):
|
||||
return system
|
||||
return _validated_object_list(system)
|
||||
|
||||
def get_chat_completion_prompt(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -520,12 +569,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
points: Sequence[CacheControlInjectionPoint],
|
||||
messages: list[AllMessageValues],
|
||||
tools: list[object] | None,
|
||||
cache_control: object,
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
api_base: object,
|
||||
prompt_cache_options: object,
|
||||
) -> Sequence[Mapping[str, object]] | None:
|
||||
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools):
|
||||
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control):
|
||||
return None
|
||||
return AnthropicCacheControlHook._stamped_with_dialect(
|
||||
points, model, custom_llm_provider, api_base, prompt_cache_options
|
||||
|
|
@ -561,6 +611,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
messages: list[AllMessageValues],
|
||||
system: str | list | None,
|
||||
tools: list | None,
|
||||
cache_control: object = None,
|
||||
) -> bool:
|
||||
"""Whether configured injection points must yield to client-set cache_control.
|
||||
|
||||
|
|
@ -573,13 +624,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
"""
|
||||
if all(point.get("_litellm_judged") for point in points):
|
||||
return False
|
||||
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools)
|
||||
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control)
|
||||
|
||||
@staticmethod
|
||||
def _request_has_cache_control(
|
||||
messages: list[AllMessageValues],
|
||||
system: str | list | None,
|
||||
tools: list | None = None,
|
||||
cache_control: object = None,
|
||||
) -> bool:
|
||||
"""Return True if the request already carries any client-supplied cache_control.
|
||||
|
||||
|
|
@ -591,6 +643,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
carry the mark either at the top level (Anthropic shape) or nested under
|
||||
``function`` (OpenAI shape); the Anthropic chat transform accepts both.
|
||||
"""
|
||||
if cache_control is not None:
|
||||
return True
|
||||
if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0:
|
||||
return True
|
||||
if tools is not None:
|
||||
|
|
@ -612,6 +666,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
custom_llm_provider: str | None,
|
||||
tools: list | None = None,
|
||||
enable_prompt_caching: bool | None = None,
|
||||
cache_control: object = None,
|
||||
request_kwargs: object = None,
|
||||
) -> list[CacheControlInjectionPoint]:
|
||||
"""Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on.
|
||||
|
||||
|
|
@ -649,7 +705,12 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
if not supports_prompt_caching(model=model, custom_llm_provider=provider):
|
||||
return []
|
||||
|
||||
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools):
|
||||
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control):
|
||||
return []
|
||||
|
||||
if is_claude_code_one_shot_subagent_request(
|
||||
messages, system, tools, AnthropicCacheControlHook._request_user_agent(request_kwargs)
|
||||
):
|
||||
return []
|
||||
|
||||
control: Final = AnthropicCacheControlHook._default_control()
|
||||
|
|
@ -665,6 +726,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
models: Iterable[str],
|
||||
tools: list[AllToolParamValues] | None = None,
|
||||
enable_prompt_caching: bool | None = None,
|
||||
request_kwargs: object = None,
|
||||
) -> list[AllMessageValues]:
|
||||
"""Return the messages auto prompt caching will send, default breakpoints included.
|
||||
|
||||
|
|
@ -681,11 +743,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
for candidate in (
|
||||
AnthropicCacheControlHook.get_default_injection_points(
|
||||
messages=messages,
|
||||
system=None,
|
||||
model=model,
|
||||
custom_llm_provider=None,
|
||||
tools=tools,
|
||||
enable_prompt_caching=enable_prompt_caching,
|
||||
system=AnthropicCacheControlHook._request_system(request_kwargs),
|
||||
cache_control=AnthropicCacheControlHook._request_value(request_kwargs, "cache_control"),
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
for model in models
|
||||
)
|
||||
|
|
@ -730,6 +794,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
non_default_params["cache_control_injection_points"],
|
||||
messages,
|
||||
tools,
|
||||
non_default_params.get("cache_control"),
|
||||
model,
|
||||
custom_llm_provider,
|
||||
api_base,
|
||||
|
|
@ -747,6 +812,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
tools=tools,
|
||||
enable_prompt_caching=enable_prompt_caching,
|
||||
cache_control=non_default_params.get("cache_control"),
|
||||
request_kwargs=non_default_params,
|
||||
)
|
||||
if points:
|
||||
non_default_params["cache_control_injection_points"] = points
|
||||
|
|
@ -853,10 +920,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
enable_prompt_caching: Final = cast( # cast-ok: kwargs is untyped; key stamped as bool by the proxy
|
||||
bool | None, kwargs.pop("enable_prompt_caching", None)
|
||||
)
|
||||
cache_control: Final = kwargs.get("cache_control")
|
||||
configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list
|
||||
list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
|
||||
)
|
||||
if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools):
|
||||
if configured and AnthropicCacheControlHook._should_stand_down(
|
||||
configured, typed_messages, system, tools, cache_control
|
||||
):
|
||||
return messages, system
|
||||
injection_points: list[CacheControlInjectionPoint] = configured or []
|
||||
if not injection_points and model is not None:
|
||||
|
|
@ -867,6 +937,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
enable_prompt_caching=enable_prompt_caching,
|
||||
cache_control=cache_control,
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
if not injection_points:
|
||||
return messages, system
|
||||
|
|
|
|||
|
|
@ -67,6 +67,96 @@ _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$")
|
|||
_INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$")
|
||||
_DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$")
|
||||
_DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)")
|
||||
_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:"
|
||||
_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
|
||||
_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
|
||||
|
||||
def is_claude_code_user_agent(user_agent: str) -> bool:
|
||||
return user_agent.startswith("claude-cli/")
|
||||
|
||||
|
||||
def _validated_claude_code_mapping(value: object) -> dict[object, object] | None:
|
||||
try:
|
||||
return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _validated_claude_code_list(value: object) -> list[object] | None:
|
||||
try:
|
||||
return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None:
|
||||
stripped: Final = text.strip()
|
||||
if "\n" in stripped or "\r" in stripped or not stripped.startswith(_CLAUDE_CODE_BILLING_HEADER_PREFIX):
|
||||
return None
|
||||
fields: Final = tuple(
|
||||
field
|
||||
for raw_field in stripped.removeprefix(_CLAUDE_CODE_BILLING_HEADER_PREFIX).split(";")
|
||||
if (field := raw_field.strip())
|
||||
)
|
||||
if not fields or any("=" not in field for field in fields):
|
||||
return None
|
||||
parsed_fields: Final = tuple(
|
||||
(parts[0].strip(), parts[1].strip()) for field in fields for parts in (field.split("=", 1),)
|
||||
)
|
||||
if any(not key or not value for key, value in parsed_fields):
|
||||
return None
|
||||
return parsed_fields
|
||||
|
||||
|
||||
def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None:
|
||||
if isinstance(system, str):
|
||||
return (system,)
|
||||
blocks: Final = _validated_claude_code_list(system)
|
||||
if blocks is None:
|
||||
return None
|
||||
block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks)
|
||||
if any(block is None for block in block_mappings):
|
||||
return None
|
||||
text_values: Final = tuple(
|
||||
block.get("text") for block in block_mappings if block is not None and block.get("type") == "text"
|
||||
)
|
||||
if len(text_values) != len(blocks) or any(not isinstance(text, str) for text in text_values):
|
||||
return None
|
||||
meaningful_text: Final = tuple(text for text in text_values if isinstance(text, str) and text.strip())
|
||||
return meaningful_text or None
|
||||
|
||||
|
||||
def _is_claude_code_subagent_billing_system(system: object) -> bool:
|
||||
billing_texts: Final = _claude_code_billing_texts(system)
|
||||
if billing_texts is None:
|
||||
return False
|
||||
billing_fields: Final = tuple(
|
||||
fields for text in billing_texts if (fields := _claude_code_billing_fields(text)) is not None
|
||||
)
|
||||
if len(billing_fields) != len(billing_texts):
|
||||
return False
|
||||
subagent_values: Final = tuple(
|
||||
value for fields in billing_fields for key, value in fields if key == "cc_is_subagent"
|
||||
)
|
||||
return subagent_values == ("true",)
|
||||
|
||||
|
||||
def is_claude_code_one_shot_subagent_request(
|
||||
messages: list[AllMessageValues],
|
||||
system: object,
|
||||
tools: object,
|
||||
user_agent: str | None,
|
||||
) -> bool:
|
||||
only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None
|
||||
return (
|
||||
user_agent is not None
|
||||
and is_claude_code_user_agent(user_agent)
|
||||
and not tools
|
||||
and only_message is not None
|
||||
and only_message.get("role") == "user"
|
||||
and _is_claude_code_subagent_billing_system(system)
|
||||
)
|
||||
|
||||
|
||||
def _strip_bedrock_id_suffixes(model: str) -> str:
|
||||
|
|
|
|||
|
|
@ -82,7 +82,7 @@ async def anthropic_messages_with_mcp(
|
|||
LiteLLM_Proxy_MCP_Handler,
|
||||
)
|
||||
|
||||
mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
|
||||
mcp_references, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
|
||||
|
||||
if not mcp_references:
|
||||
return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
|
||||
|
|
|
|||
|
|
@ -3108,15 +3108,17 @@ class MCPRequestHandler:
|
|||
@staticmethod
|
||||
async def _get_allowed_mcp_servers_for_agent(
|
||||
user_api_key_auth: UserAPIKeyAuth | None = None,
|
||||
agent_object_permission=None,
|
||||
agent_object_permission: LiteLLM_ObjectPermissionTable | None = None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Get allowed MCP servers for an agent (from the agent's object_permission).
|
||||
|
||||
Returns the MCP servers from the agent's object_permission.
|
||||
If agent has no object_permission, returns [] (no extra restriction). An entitlement the
|
||||
agent LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here so the
|
||||
resolver denies.
|
||||
Returns the agent's direct servers, the servers in its access groups, and the servers reached
|
||||
through its toolsets, exactly as the key, team, and org levels count theirs. If agent has no
|
||||
object_permission, returns [] (no extra restriction). An entitlement the agent LINKS but that
|
||||
cannot be read, or a declared toolset that resolves to no grants, raises
|
||||
``UnloadableEntitlementError`` out of here so the resolver denies instead of reading the
|
||||
agent as unrestricted.
|
||||
|
||||
Args:
|
||||
user_api_key_auth: User auth with agent_id
|
||||
|
|
@ -3126,31 +3128,30 @@ class MCPRequestHandler:
|
|||
if not user_api_key_auth or not user_api_key_auth.agent_id:
|
||||
return []
|
||||
|
||||
obj_perm = agent_object_permission
|
||||
if obj_perm is None:
|
||||
obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
|
||||
obj_perm: Final = (
|
||||
agent_object_permission
|
||||
if agent_object_permission is not None
|
||||
else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
|
||||
)
|
||||
if obj_perm is None:
|
||||
return []
|
||||
|
||||
try:
|
||||
direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or []
|
||||
if isinstance(direct_mcp_servers, str):
|
||||
direct_mcp_servers = []
|
||||
mcp_access_groups = getattr(obj_perm, "mcp_access_groups", None) or []
|
||||
if isinstance(mcp_access_groups, str):
|
||||
mcp_access_groups = []
|
||||
|
||||
# Permission entries may be server_ids OR names/aliases — expand to ids.
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(list(direct_mcp_servers))
|
||||
|
||||
access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(mcp_access_groups)
|
||||
all_servers: Final = expanded_direct_servers + access_group_servers
|
||||
return list(set(all_servers))
|
||||
expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(
|
||||
obj_perm.mcp_servers or []
|
||||
)
|
||||
access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(
|
||||
obj_perm.mcp_access_groups or []
|
||||
)
|
||||
toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(obj_perm)
|
||||
return list({*expanded_direct_servers, *access_group_servers, *toolset_grants})
|
||||
except Exception as e:
|
||||
if isinstance(e, UnloadableEntitlementError):
|
||||
raise
|
||||
verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e)
|
||||
return []
|
||||
|
||||
|
|
@ -3158,13 +3159,15 @@ class MCPRequestHandler:
|
|||
async def _get_agent_tool_permissions_for_server(
|
||||
server_id: str,
|
||||
user_api_key_auth: UserAPIKeyAuth | None = None,
|
||||
agent_object_permission=None,
|
||||
agent_object_permission: LiteLLM_ObjectPermissionTable | None = None,
|
||||
) -> list[str] | None:
|
||||
"""
|
||||
Get allowed tool names for a server from the agent's object_permission.
|
||||
Returns None if agent has no tool restrictions for this server. An entitlement the agent
|
||||
LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here, which the
|
||||
tool resolver turns into deny-all for the server rather than an unrestricted tool list.
|
||||
Get allowed tool names for a server from the agent's object_permission: the union of its
|
||||
direct tool permissions and the tools its toolsets grant on that server, mirroring the key and
|
||||
team levels. Returns None if agent has no tool restrictions for this server. An entitlement the
|
||||
agent LINKS but that cannot be read, or a declared toolset that resolves to no grants, raises
|
||||
``UnloadableEntitlementError`` out of here, which the tool resolver turns into deny-all for the
|
||||
server rather than an unrestricted tool list.
|
||||
|
||||
Args:
|
||||
server_id: Server ID to check permissions for
|
||||
|
|
@ -3175,24 +3178,30 @@ class MCPRequestHandler:
|
|||
if not user_api_key_auth or not user_api_key_auth.agent_id:
|
||||
return None
|
||||
|
||||
obj_perm = agent_object_permission
|
||||
if obj_perm is None:
|
||||
obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
|
||||
obj_perm: Final = (
|
||||
agent_object_permission
|
||||
if agent_object_permission is not None
|
||||
else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
|
||||
)
|
||||
if obj_perm is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
mcp_tool_permissions: Final = getattr(obj_perm, "mcp_tool_permissions", None)
|
||||
if not mcp_tool_permissions or not isinstance(mcp_tool_permissions, dict):
|
||||
return None
|
||||
# Dict keys may be server_ids OR names/aliases; normalize before lookup.
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
tools: Final = global_mcp_server_manager.expand_tool_permissions(mcp_tool_permissions).get(server_id)
|
||||
return list(tools) if tools else None
|
||||
direct_tools: Final = (
|
||||
global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id)
|
||||
if obj_perm.mcp_tool_permissions
|
||||
else None
|
||||
)
|
||||
toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(obj_perm, server_id)
|
||||
agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools)
|
||||
return list(agent_tools) if agent_tools else None
|
||||
except Exception as e:
|
||||
if isinstance(e, UnloadableEntitlementError):
|
||||
raise
|
||||
verbose_logger.warning("Failed to get agent tool permissions for server: %s", e)
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from pydantic import AnyUrl, ConfigDict
|
|||
from starlette.requests import Request as StarletteRequest
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.types import Message, Receive, Scope, Send
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
|
||||
|
|
@ -816,6 +817,11 @@ if MCP_AVAILABLE:
|
|||
}
|
||||
}
|
||||
return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta})
|
||||
except HTTPException as e:
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.types import INVALID_REQUEST, ErrorData
|
||||
|
||||
raise McpError(ErrorData(code=INVALID_REQUEST, message=_http_detail_message(e.detail))) from e
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error in list_tools endpoint: %s", e)
|
||||
# Return empty list instead of failing completely
|
||||
|
|
@ -1095,6 +1101,7 @@ if MCP_AVAILABLE:
|
|||
mcp_server_auth_headers=mcp_server_auth_headers,
|
||||
oauth2_headers=oauth2_headers,
|
||||
raw_headers=raw_headers,
|
||||
client_ip=_client_ip,
|
||||
host_progress_callback=host_progress_callback,
|
||||
**data, # for logging
|
||||
)
|
||||
|
|
@ -1128,7 +1135,7 @@ if MCP_AVAILABLE:
|
|||
except HTTPException as e:
|
||||
verbose_logger.error("HTTPException in MCP tool call: %s", e)
|
||||
return CallToolResult(
|
||||
content=[TextContent(text=f"Error: {e.detail}", type="text")],
|
||||
content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")],
|
||||
isError=True,
|
||||
)
|
||||
except MCPUpstreamAuthError as e:
|
||||
|
|
@ -1392,7 +1399,7 @@ if MCP_AVAILABLE:
|
|||
########################################################
|
||||
|
||||
async def _get_allowed_mcp_servers_from_mcp_server_names(
|
||||
mcp_servers: list[str] | None,
|
||||
mcp_servers: Sequence[str] | None,
|
||||
allowed_mcp_servers: list[MCPServer],
|
||||
) -> list[MCPServer]:
|
||||
"""
|
||||
|
|
@ -1413,13 +1420,10 @@ if MCP_AVAILABLE:
|
|||
server_name_matched = False
|
||||
|
||||
for server in allowed_mcp_servers:
|
||||
if server:
|
||||
match_list = [s.lower() for s in iter_known_server_prefixes(server) if s]
|
||||
|
||||
if server_or_group.lower() in match_list:
|
||||
filtered_server[server.server_id] = server
|
||||
server_name_matched = True
|
||||
break
|
||||
if server and _server_answers_to(server, server_or_group):
|
||||
filtered_server[server.server_id] = server
|
||||
server_name_matched = True
|
||||
break
|
||||
|
||||
if not server_name_matched:
|
||||
try:
|
||||
|
|
@ -1449,6 +1453,72 @@ if MCP_AVAILABLE:
|
|||
|
||||
return allowed_mcp_servers
|
||||
|
||||
def _http_detail_message(detail: object) -> str:
|
||||
return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail)
|
||||
|
||||
def _server_answers_to(server: MCPServer, name: str) -> bool:
|
||||
requested: Final = name.lower()
|
||||
return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known)
|
||||
|
||||
class _McpDeniedDetail(TypedDict):
|
||||
error: ReadOnly[str]
|
||||
|
||||
async def raise_denied_scoped_mcp_access(
|
||||
requested_names: Sequence[str],
|
||||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
client_ip: str | None = None,
|
||||
) -> None:
|
||||
"""A scoped request (``/mcp/<name>`` path or ``x-mcp-servers`` header) resolved to zero
|
||||
allowed servers, so the denial must be loud: a silent 200 with no tools reads as a healthy
|
||||
server with no tools. Unknown, unauthorized, and access-group names all share one generic
|
||||
error so scoping cannot probe which servers exist; the agent variant fires only when the
|
||||
same request resolves once the agent binding is stripped, proving the binding caused the veto."""
|
||||
agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None
|
||||
if user_api_key_auth is not None and agent_id:
|
||||
resolved_without_agent: Final = await _get_allowed_mcp_servers(
|
||||
user_api_key_auth=user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None})),
|
||||
mcp_servers=requested_names,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
|
||||
def _resolved_to_server(name: str) -> bool:
|
||||
return any(_server_answers_to(server, name) for server in resolved_without_agent)
|
||||
|
||||
vetoed_server: Final = next((name for name in requested_names if _resolved_to_server(name)), None)
|
||||
if vetoed_server is not None:
|
||||
agent_denial: Final[_McpDeniedDetail] = {
|
||||
"error": (
|
||||
f"MCP server '{vetoed_server}' is not available to this key: the key is bound to "
|
||||
f"agent '{agent_id}', whose MCP grants do not include this server. Add the server "
|
||||
f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or "
|
||||
f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
|
||||
)
|
||||
}
|
||||
raise HTTPException(status_code=403, detail=agent_denial)
|
||||
vetoed_group: Final = next(
|
||||
(
|
||||
name
|
||||
for name in requested_names
|
||||
if not _resolved_to_server(name)
|
||||
and any(name in (server.access_groups or ()) for server in resolved_without_agent)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if vetoed_group is not None:
|
||||
group_denial: Final[_McpDeniedDetail] = {
|
||||
"error": (
|
||||
f"MCP access group '{vetoed_group}' is not available to this key: the key is bound to "
|
||||
f"agent '{agent_id}', whose MCP grants do not include it. Add the group to the "
|
||||
f"agent's object_permission.mcp_access_groups (edit the agent in the Admin UI or "
|
||||
f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
|
||||
)
|
||||
}
|
||||
raise HTTPException(status_code=403, detail=group_denial)
|
||||
generic_denial: Final[_McpDeniedDetail] = {
|
||||
"error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}"
|
||||
}
|
||||
raise HTTPException(status_code=403, detail=generic_denial)
|
||||
|
||||
def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool:
|
||||
"""
|
||||
Check if a tool name matches any name in the filter list.
|
||||
|
|
@ -1541,7 +1611,7 @@ if MCP_AVAILABLE:
|
|||
|
||||
async def _get_allowed_mcp_servers(
|
||||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
mcp_servers: list[str] | None,
|
||||
mcp_servers: Sequence[str] | None,
|
||||
client_ip: str | None = None,
|
||||
) -> list[MCPServer]:
|
||||
"""Return allowed MCP servers for a request after applying filters.
|
||||
|
|
@ -1977,6 +2047,12 @@ if MCP_AVAILABLE:
|
|||
mcp_servers=mcp_servers,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
if mcp_servers and not allowed_mcp_servers:
|
||||
await raise_denied_scoped_mcp_access(
|
||||
requested_names=mcp_servers,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
|
||||
# Pre-fetch OAuth credentials only when at least one server uses OAuth2,
|
||||
# to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
|
||||
|
|
@ -2404,6 +2480,8 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools))
|
||||
return listing
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error getting tools from managed MCP servers: %s", e)
|
||||
# Continue with an empty listing instead of failing completely
|
||||
|
|
@ -3086,6 +3164,7 @@ if MCP_AVAILABLE:
|
|||
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
|
||||
oauth2_headers: dict[str, str] | None = None,
|
||||
raw_headers: dict[str, str] | None = None,
|
||||
client_ip: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> CallToolResult:
|
||||
"""
|
||||
|
|
@ -3116,6 +3195,12 @@ if MCP_AVAILABLE:
|
|||
mcp_servers=mcp_servers,
|
||||
allowed_mcp_servers=allowed_mcp_servers,
|
||||
)
|
||||
if mcp_servers and not allowed_mcp_servers:
|
||||
await raise_denied_scoped_mcp_access(
|
||||
requested_names=mcp_servers,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
if not allowed_mcp_servers:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
|
|
|
|||
|
|
@ -366,6 +366,7 @@ async def handle_mcp_tool_call(
|
|||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_get_allowed_mcp_servers,
|
||||
execute_mcp_tool,
|
||||
raise_denied_scoped_mcp_access,
|
||||
)
|
||||
|
||||
allowed_mcp_servers: Final = await _get_allowed_mcp_servers(
|
||||
|
|
@ -373,6 +374,12 @@ async def handle_mcp_tool_call(
|
|||
mcp_servers=mcp_servers,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
if mcp_servers and not allowed_mcp_servers:
|
||||
await raise_denied_scoped_mcp_access(
|
||||
requested_names=mcp_servers,
|
||||
user_api_key_auth=user_api_key_dict,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
|
||||
# Reject before dispatch when the key has no accessible servers; otherwise an
|
||||
# unprefixed local tool name would fall through to the local registry in
|
||||
|
|
|
|||
|
|
@ -2634,6 +2634,20 @@
|
|||
],
|
||||
"title": "Mcp Tool Permissions"
|
||||
},
|
||||
"mcp_toolsets": {
|
||||
"anyOf": [
|
||||
{
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Mcp Toolsets"
|
||||
},
|
||||
"models": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -2833,7 +2833,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
"Enable only if your deployment is experiencing phantom "
|
||||
"BudgetExceededError responses caused by leaked reservations "
|
||||
"(see GitHub issue #27639). "
|
||||
"A proxy-level WARNING is logged on every request while this flag "
|
||||
"An INFO notice is logged once per worker at config load while this flag "
|
||||
"is active as a reminder that hard enforcement is relaxed."
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ from fastapi import HTTPException, Request, status
|
|||
from pydantic import PositiveInt, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm import Router, provider_list
|
||||
from litellm import Router, constants, provider_list
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY,
|
||||
|
|
@ -1390,6 +1390,24 @@ def warn_once_if_custom_auth_skips_common_checks(
|
|||
_custom_auth_common_checks_warning_emitted = True
|
||||
|
||||
|
||||
def log_once_if_budget_reservation_disabled(
|
||||
*,
|
||||
disabled: bool,
|
||||
logger: Logger = verbose_proxy_logger,
|
||||
) -> None:
|
||||
if constants.budget_reservation_disabled_info_emitted or not disabled:
|
||||
return
|
||||
logger.info(
|
||||
"disable_budget_reservation is enabled: skipping optimistic budget "
|
||||
"reservation. Budget enforcement is read-time only. Concurrent "
|
||||
"requests can each pass the spend check before their cost is recorded, "
|
||||
"so a configured budget may be briefly exceeded under high concurrency. "
|
||||
"Set disable_budget_reservation to False or remove it to restore "
|
||||
"hard per-request budget enforcement."
|
||||
)
|
||||
constants.budget_reservation_disabled_info_emitted = True # rebind-ok: process-wide one-shot sentinel
|
||||
|
||||
|
||||
def is_pass_through_provider_route(route: str) -> bool:
|
||||
PROVIDER_SPECIFIC_PASS_THROUGH_ROUTES: Final = [
|
||||
"vertex-ai",
|
||||
|
|
|
|||
|
|
@ -2706,14 +2706,6 @@ async def _reserve_budget_after_common_checks(
|
|||
if skip_budget_checks:
|
||||
return
|
||||
if general_settings.get("disable_budget_reservation") is True:
|
||||
verbose_proxy_logger.warning(
|
||||
"disable_budget_reservation is enabled: skipping optimistic budget "
|
||||
"reservation. Budget enforcement is read-time only — concurrent "
|
||||
"requests can each pass the spend check before their cost is recorded, "
|
||||
"so a configured budget may be briefly exceeded under high concurrency. "
|
||||
"Set disable_budget_reservation to False or remove it to restore "
|
||||
"hard per-request budget enforcement."
|
||||
)
|
||||
return
|
||||
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ Pre-call hook that filters MCP tools semantically before LLM inference.
|
|||
Reduces context window size and improves tool selection accuracy.
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final, Optional
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -104,7 +104,7 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
)
|
||||
|
||||
# Parse to separate MCP tools from other tools
|
||||
mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
|
||||
mcp_tools, _ = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
|
||||
|
||||
if not mcp_tools:
|
||||
return []
|
||||
|
|
@ -173,7 +173,11 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
return [name for name in names if name]
|
||||
|
||||
@staticmethod
|
||||
def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]:
|
||||
async def _narrow_mcp_references(
|
||||
tools: Sequence[Mapping[str, object]],
|
||||
selected_tool_names: list[str],
|
||||
served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] | None = None,
|
||||
) -> list[object]:
|
||||
"""
|
||||
Restrict each litellm_proxy MCP reference to the semantically selected tools.
|
||||
|
||||
|
|
@ -192,13 +196,14 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
LiteLLM_Proxy_MCP_Handler,
|
||||
)
|
||||
|
||||
via_gateway: Final = await (
|
||||
LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools, served_names)
|
||||
if served_names is not None
|
||||
else LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools)
|
||||
)
|
||||
return [
|
||||
(
|
||||
{**tool, "allowed_tools": selected_tool_names}
|
||||
if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool])
|
||||
else tool
|
||||
)
|
||||
for tool in tools
|
||||
{**tool, "allowed_tools": selected_tool_names} if isinstance(tool, dict) and routed else tool
|
||||
for tool, routed in zip(tools, via_gateway, strict=True)
|
||||
]
|
||||
|
||||
def _is_mcp_tool(self, tool: object) -> bool:
|
||||
|
|
@ -325,7 +330,7 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools)
|
||||
|
||||
selected_tool_names: Final = self._selected_tool_names(filtered_expanded_tools)
|
||||
narrowed_tools: Final = self._narrow_mcp_references(tools, selected_tool_names)
|
||||
narrowed_tools: Final = await self._narrow_mcp_references(tools, selected_tool_names)
|
||||
data["tools"] = narrowed_tools
|
||||
self._emit_filter_metadata_safe(
|
||||
data=data,
|
||||
|
|
|
|||
|
|
@ -789,12 +789,6 @@ def apply_missing_session_id_policy(
|
|||
)
|
||||
|
||||
|
||||
def is_claude_code_user_agent(user_agent: str) -> bool:
|
||||
"""Claude Code identifies itself as ``claude-cli/<version> ...``; the IDE
|
||||
extensions and the Agent SDK run through the same CLI and share that prefix."""
|
||||
return user_agent.startswith("claude-cli/")
|
||||
|
||||
|
||||
def is_codex_user_agent(user_agent: str) -> bool:
|
||||
"""Codex builds its user agent as ``<originator>/<version> ...`` and ships
|
||||
several first-party originators: ``codex-tui``, ``codex_cli_rs``,
|
||||
|
|
@ -811,6 +805,8 @@ def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_c
|
|||
requests routed to providers that reject them. An explicit drop_params
|
||||
from the caller or in the operator's ``litellm_settings`` always wins
|
||||
over this default."""
|
||||
from litellm.llms.anthropic.common_utils import is_claude_code_user_agent
|
||||
|
||||
if not (is_claude_code_user_agent(user_agent) or is_codex_user_agent(user_agent)):
|
||||
return False
|
||||
if "drop_params" in data:
|
||||
|
|
|
|||
|
|
@ -306,6 +306,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
from litellm.proxy.auth.auth_utils import (
|
||||
check_response_size_is_safe,
|
||||
is_request_body_safe,
|
||||
log_once_if_budget_reservation_disabled,
|
||||
warn_once_if_custom_auth_skips_common_checks,
|
||||
)
|
||||
from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
|
||||
|
|
@ -5653,6 +5654,10 @@ class ProxyConfig:
|
|||
run_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)),
|
||||
)
|
||||
|
||||
log_once_if_budget_reservation_disabled(
|
||||
disabled=general_settings.get("disable_budget_reservation") is True,
|
||||
)
|
||||
|
||||
custom_key_generate: Final = general_settings.get("custom_key_generate", None)
|
||||
if custom_key_generate is not None:
|
||||
user_custom_key_generate = get_instance_fn(value=custom_key_generate, config_file_path=config_file_path)
|
||||
|
|
|
|||
|
|
@ -177,7 +177,7 @@ async def aresponses_api_with_mcp(
|
|||
(
|
||||
mcp_tools_with_litellm_proxy,
|
||||
other_tools,
|
||||
) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
|
||||
) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
|
||||
|
||||
# Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform)
|
||||
# Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata)
|
||||
|
|
@ -236,6 +236,7 @@ async def aresponses_api_with_mcp(
|
|||
"timeout": timeout,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
**kwargs,
|
||||
"_skip_mcp_handler": True,
|
||||
}
|
||||
|
||||
# Handle MCP streaming if requested
|
||||
|
|
@ -898,13 +899,14 @@ def _responses_try_dispatch_mcp_gateway(
|
|||
custom_llm_provider: str | None,
|
||||
kwargs: dict[str, object],
|
||||
_is_async: bool,
|
||||
skip_mcp_handler: bool,
|
||||
) -> Any | None:
|
||||
"""Return a response when MCP gateway handles the call; otherwise None."""
|
||||
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
)
|
||||
|
||||
if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools):
|
||||
if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools):
|
||||
return None
|
||||
mcp_call_kwargs: Final = {
|
||||
"input": input,
|
||||
|
|
@ -1074,6 +1076,7 @@ def responses(
|
|||
litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj")
|
||||
litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None)
|
||||
_is_async: Final = kwargs.pop("aresponses", False) is True
|
||||
skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False)
|
||||
use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs)
|
||||
|
||||
client_headers: Final = kwargs.get("headers")
|
||||
|
|
@ -1168,6 +1171,7 @@ def responses(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
_is_async=_is_async,
|
||||
skip_mcp_handler=skip_mcp_handler,
|
||||
)
|
||||
if _mcp_dispatch is not None:
|
||||
return _mcp_dispatch
|
||||
|
|
|
|||
|
|
@ -106,7 +106,7 @@ async def acompletion_with_mcp(
|
|||
(
|
||||
mcp_tools_with_litellm_proxy,
|
||||
other_tools,
|
||||
) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
|
||||
) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
|
||||
|
||||
if not mcp_tools_with_litellm_proxy:
|
||||
# No MCP tools, proceed with regular completion
|
||||
|
|
@ -114,6 +114,7 @@ async def acompletion_with_mcp(
|
|||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
_skip_mcp_handler=True,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import re
|
||||
import traceback
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload
|
||||
|
||||
|
|
@ -11,6 +11,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
iter_known_server_prefixes,
|
||||
logging_safe_mcp_headers,
|
||||
split_server_prefix_from_name,
|
||||
strip_known_server_prefix,
|
||||
|
|
@ -23,6 +24,7 @@ from litellm.types.llms.openai import (
|
|||
ResponsesAPIStreamingResponse,
|
||||
)
|
||||
from litellm.types.llms.openai import ToolParam as ResponsesToolParam
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
ChatCompletionMessageCustomToolCall,
|
||||
|
|
@ -45,6 +47,7 @@ else:
|
|||
|
||||
# NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling
|
||||
ToolParam: TypeAlias = Mapping[str, object]
|
||||
SplitTools: TypeAlias = tuple[list[ToolParam], list[Any]]
|
||||
|
||||
|
||||
class MCPToolResult(TypedDict):
|
||||
|
|
@ -56,14 +59,74 @@ class MCPToolResult(TypedDict):
|
|||
LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy"
|
||||
LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/"
|
||||
|
||||
# Matches any URL whose path ends with /mcp/<server_name> — covers both root-path
|
||||
# (http://host:port/mcp/name) and sub-path (http://host/base/mcp/name) proxy deployments.
|
||||
# A false-positive match (e.g. an external URL that happens to end with /mcp/<name>) results
|
||||
# in a "server not found" error from the internal gateway, not a silent failure or data leak,
|
||||
# so this broad pattern is intentional and preferred over anchoring to localhost only.
|
||||
_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$")
|
||||
|
||||
|
||||
def _mcp_server_url(tool: ToolParam) -> str | None:
|
||||
if not isinstance(tool, dict) or tool.get("type") != "mcp":
|
||||
return None
|
||||
server_url: Final = tool.get("server_url")
|
||||
return server_url if isinstance(server_url, str) else None
|
||||
|
||||
|
||||
def _names_gateway_explicitly(tool: ToolParam) -> bool:
|
||||
return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL)
|
||||
|
||||
|
||||
def _proxy_path_mcp_name(tool: ToolParam) -> str | None:
|
||||
server_url: Final = _mcp_server_url(tool)
|
||||
match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url)
|
||||
return None if match is None else match.group(1)
|
||||
|
||||
|
||||
def _registered_mcp_servers() -> Collection[MCPServer]:
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
return global_mcp_server_manager.get_registry().values()
|
||||
|
||||
|
||||
def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool:
|
||||
requested: Final = name.lower()
|
||||
return any(
|
||||
requested in (known.lower() for known in (*iter_known_server_prefixes(server), server.name))
|
||||
or name in (server.access_groups or ())
|
||||
for server in servers
|
||||
)
|
||||
|
||||
|
||||
async def _toolset_exists(name: str) -> bool:
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
return False
|
||||
return await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is not None
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, e)
|
||||
return False
|
||||
|
||||
|
||||
async def _gateway_served_names(
|
||||
names: Collection[str],
|
||||
servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers,
|
||||
toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists,
|
||||
) -> frozenset[str]:
|
||||
registered: Final = tuple(servers()) if names else ()
|
||||
return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)])
|
||||
|
||||
|
||||
async def _served_mcp_path_names(
|
||||
tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]]
|
||||
) -> frozenset[str]:
|
||||
names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None)
|
||||
return await served_names(names) if names else frozenset[str]()
|
||||
|
||||
|
||||
class LiteLLM_Proxy_MCP_Handler:
|
||||
"""
|
||||
Helper class with static methods for MCP integration with Responses API.
|
||||
|
|
@ -87,57 +150,41 @@ class LiteLLM_Proxy_MCP_Handler:
|
|||
|
||||
@staticmethod
|
||||
def _should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool:
|
||||
"""
|
||||
Returns True if any MCP tool should be handled via the litellm proxy MCP gateway.
|
||||
This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/<name>.
|
||||
"""
|
||||
if tools:
|
||||
for tool in tools:
|
||||
if isinstance(tool, dict) and tool.get("type") == "mcp":
|
||||
server_url = tool.get("server_url", "")
|
||||
if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL):
|
||||
return True
|
||||
if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(server_url):
|
||||
return True
|
||||
return False
|
||||
"""True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in
|
||||
/mcp/<name>. `_split_mcp_tools` then settles which of the latter the gateway actually serves."""
|
||||
return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ())
|
||||
|
||||
@staticmethod
|
||||
def _parse_mcp_tools(
|
||||
def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools:
|
||||
items: Final = tuple(tools or ())
|
||||
gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)]
|
||||
other_tools: Final[list[Any]] = [tool for tool in items if not _names_gateway_explicitly(tool)]
|
||||
return gateway_tools, other_tools
|
||||
|
||||
@staticmethod
|
||||
async def _split_mcp_tools(
|
||||
tools: Iterable[Mapping[str, object]] | None,
|
||||
) -> tuple[list[ToolParam], list[Any]]:
|
||||
"""
|
||||
Parse tools and separate MCP tools with litellm_proxy from other tools.
|
||||
served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names,
|
||||
) -> SplitTools:
|
||||
items: Final = tuple(tools or ())
|
||||
served: Final = await _served_mcp_path_names(items, served_names)
|
||||
return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(
|
||||
[
|
||||
{**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"}
|
||||
if (name := _proxy_path_mcp_name(tool)) in served
|
||||
else tool
|
||||
for tool in items
|
||||
]
|
||||
)
|
||||
|
||||
Returns:
|
||||
Tuple of (mcp_tools_with_litellm_proxy, other_tools)
|
||||
"""
|
||||
mcp_tools_with_litellm_proxy: Final[list[ToolParam]] = []
|
||||
other_tools: Final[list[Any]] = []
|
||||
|
||||
if tools:
|
||||
for tool in tools:
|
||||
if isinstance(tool, dict) and tool.get("type") == "mcp":
|
||||
server_url = tool.get("server_url", "")
|
||||
if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL):
|
||||
mcp_tools_with_litellm_proxy.append(tool)
|
||||
elif isinstance(server_url, str):
|
||||
# Also intercept URLs like http://localhost:4000/mcp/atlassian_test
|
||||
# by rewriting them to the internal litellm_proxy format.
|
||||
m = _PROXY_MCP_PATH_RE.match(server_url)
|
||||
if m:
|
||||
rewritten = {
|
||||
**tool,
|
||||
"server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}",
|
||||
}
|
||||
mcp_tools_with_litellm_proxy.append(rewritten)
|
||||
else:
|
||||
other_tools.append(tool)
|
||||
else:
|
||||
other_tools.append(tool)
|
||||
else:
|
||||
other_tools.append(tool)
|
||||
|
||||
return mcp_tools_with_litellm_proxy, other_tools
|
||||
@staticmethod
|
||||
async def routes_through_gateway(
|
||||
tools: Iterable[Mapping[str, object]] | None,
|
||||
served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names,
|
||||
) -> tuple[bool, ...]:
|
||||
items: Final = tuple(tools or ())
|
||||
served: Final = await _served_mcp_path_names(items, served_names)
|
||||
return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items)
|
||||
|
||||
@staticmethod
|
||||
async def _apply_toolset_permissions(
|
||||
|
|
|
|||
|
|
@ -195,6 +195,30 @@ model_list:
|
|||
session_affinity_ttl_seconds: 300
|
||||
```
|
||||
|
||||
## Custom dimensions
|
||||
|
||||
Add `custom_dimensions` under `complexity_router_config` to give domain keywords or regex patterns their own weighted signal
|
||||
|
||||
```yaml
|
||||
custom_dimensions:
|
||||
- name: internalFrameworks
|
||||
weight: 0.9
|
||||
keywords: [orbitmesh, fluxgate]
|
||||
- name: sqlMigration
|
||||
weight: 0.7
|
||||
patterns: ['\b(create|alter|drop)\s{1,4}table\b']
|
||||
```
|
||||
|
||||
Each dimension contributes its weight once when any matcher hits the current ask. Repeated matches do not increase it. The built-in score and tier boundaries are unchanged, and the total score is not renormalized. Keywords use the existing case-insensitive word-boundary and CJK rules. Regexes search the first 2048 characters case-insensitively and compile during configuration validation and router initialization, never per request
|
||||
|
||||
Only `heuristic`, `heuristic_first` and `hybrid` accept custom dimensions. Each name must be a unique ASCII identifier starting with a letter, at most 64 characters, and cannot reuse a built-in dimension name or a key in `dimension_weights`. Set its weight inline, greater than zero and at most one
|
||||
|
||||
Patterns are checked at configuration time against a grammar whose worst case stays a few milliseconds on 2048 characters. Every quantifier needs an explicit upper bound of at most 64 and must repeat a single character or character class, so `\s{1,4}` is accepted while `\s+`, `(a|aa){0,12}` and `(?:ab){0,64}` are refused. Backreferences, lookarounds, atomic groups and possessive quantifiers are refused as well. Each pattern is then costed: alternation branches and repeat lengths multiply the ways the engine can retry, and every later piece of the pattern is charged once per path that can reach it, so `a?a?a?a?a?a?a?a?` followed by a long fixed tail is refused even though each quantifier is small. The budget is 2048 work units per pattern and 8192 across the router. An invalid or over-budget pattern fails the write with a message naming the pattern and the rule it broke
|
||||
|
||||
Limits are 16 dimensions, 32 combined keywords/patterns per dimension, 256 characters per matcher and 4096 matcher characters per dimension. Matching runs inline on the request path with no timeout and no worker thread, because the grammar is what bounds the cost. These are routing hints, not security enforcement rules
|
||||
|
||||
The existing heuristic-v1 tuning quota covers custom dimensions and their weights: one changed router without an auto-router license, unlimited with the entitlement. Omitting `custom_dimensions` preserves existing scoring. Routing decisions and spend logs include signals such as `custom (sqlMigration)` without recording the configured pattern or matched text. The field is configured through YAML or the model API; this change adds no dashboard editor
|
||||
|
||||
## Usage
|
||||
|
||||
Once configured, use the model name like any other:
|
||||
|
|
|
|||
|
|
@ -67,6 +67,7 @@ from litellm.types.utils import (
|
|||
from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section
|
||||
from .config import (
|
||||
CALIBRATION_EXAMPLES_HEADING,
|
||||
CUSTOM_PATTERN_SCAN_CHARS,
|
||||
DEFAULT_CLASSIFICATION_RUBRIC,
|
||||
DEFAULT_CODE_KEYWORDS,
|
||||
DEFAULT_ESCALATION_KEYWORDS,
|
||||
|
|
@ -1119,6 +1120,10 @@ class ComplexityRouter(CustomLogger):
|
|||
self.config.custom_technical_keywords,
|
||||
)
|
||||
self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS
|
||||
self._custom_dimensions = tuple(
|
||||
(dimension, tuple(re.compile(pattern, re.IGNORECASE) for pattern in dimension.patterns))
|
||||
for dimension in self.config.custom_dimensions
|
||||
)
|
||||
if self.config.has_custom_tiers:
|
||||
self.escalation_keywords: tuple[str, ...] = ()
|
||||
elif self.config.escalation_keywords is not None:
|
||||
|
|
@ -1320,6 +1325,17 @@ class ComplexityRouter(CustomLogger):
|
|||
score: Final = score_high if match_count >= high_threshold else score_low
|
||||
return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count
|
||||
|
||||
def _score_custom_dimensions(self, prompt: str, user_text: str) -> tuple[tuple[DimensionScore, float], ...]:
|
||||
if not self._custom_dimensions:
|
||||
return ()
|
||||
scanned: Final = prompt[:CUSTOM_PATTERN_SCAN_CHARS]
|
||||
return tuple(
|
||||
(DimensionScore(dimension.name, 1.0, f"custom ({dimension.name})"), dimension.weight)
|
||||
for dimension, patterns in self._custom_dimensions
|
||||
if any(self._keyword_matches(user_text, keyword) for keyword in dimension.keywords)
|
||||
or any(pattern.search(scanned) is not None for pattern in patterns)
|
||||
)
|
||||
|
||||
def _score_multi_step(self, text: str) -> DimensionScore:
|
||||
"""Score based on multi-step patterns."""
|
||||
hits: Final = sum(1 for p in self._multi_step_patterns if p.search(text))
|
||||
|
|
@ -1415,12 +1431,13 @@ class ComplexityRouter(CustomLogger):
|
|||
self._score_question_complexity(prompt),
|
||||
]
|
||||
|
||||
# Collect signals
|
||||
signals: Final = [d.signal for d in dimensions if d.signal is not None]
|
||||
custom_dimensions: Final = self._score_custom_dimensions(prompt, user_text)
|
||||
signals: Final = [d.signal for d in (*dimensions, *(d for d, _ in custom_dimensions)) if d.signal is not None]
|
||||
|
||||
# Compute weighted score
|
||||
weights: Final = self.config.dimension_weights
|
||||
weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions)
|
||||
weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) + sum(
|
||||
dimension.score * weight for dimension, weight in custom_dimensions
|
||||
)
|
||||
|
||||
boundaries: Final = self._effective_tier_boundaries()
|
||||
clears_override_floor: Final = weighted_score >= self._effective_reasoning_override_min_score()
|
||||
|
|
|
|||
|
|
@ -5,13 +5,21 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas
|
|||
All values are configurable via proxy config.yaml.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
import math
|
||||
import re
|
||||
import warnings
|
||||
from collections.abc import Iterable, Mapping
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, Literal
|
||||
from typing import Annotated, Final, Literal, NamedTuple
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", DeprecationWarning)
|
||||
import sre_constants
|
||||
import sre_parse
|
||||
|
||||
from litellm.types.llms.openai import REASONING_EFFORT
|
||||
from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin
|
||||
|
||||
|
|
@ -569,6 +577,117 @@ class ClassifierLLMConfig(BaseModel):
|
|||
return self
|
||||
|
||||
|
||||
MAX_CUSTOM_PATTERN_REPEAT: Final[int] = 64
|
||||
MAX_CUSTOM_PATTERN_WORK: Final[int] = 2048
|
||||
MAX_CUSTOM_DIMENSIONS_WORK: Final[int] = 8192
|
||||
MAX_CUSTOM_PATTERN_DEPTH: Final[int] = 16
|
||||
CUSTOM_PATTERN_SCAN_CHARS: Final[int] = 2048
|
||||
|
||||
_ATOM_OPCODES: Final = frozenset(
|
||||
{sre_constants.LITERAL, sre_constants.NOT_LITERAL, sre_constants.ANY, sre_constants.IN, sre_constants.CATEGORY}
|
||||
)
|
||||
_REPEAT_OPCODES: Final = frozenset({sre_constants.MAX_REPEAT, sre_constants.MIN_REPEAT})
|
||||
|
||||
|
||||
class _PatternCost(NamedTuple):
|
||||
paths: int
|
||||
steps: int
|
||||
|
||||
|
||||
def _atom_steps(node: object) -> int:
|
||||
if isinstance(node, tuple) and len(node) == 2 and node[0] is sre_constants.IN:
|
||||
return 1 + len(node[1])
|
||||
return 1
|
||||
|
||||
|
||||
def _repeat_cost(argument: object) -> _PatternCost | str:
|
||||
if not isinstance(argument, tuple) or len(argument) != 3:
|
||||
return "unsupported repeat structure"
|
||||
low, high, body = argument
|
||||
if high > MAX_CUSTOM_PATTERN_REPEAT or len(body) != 1 or body[0][0] not in _ATOM_OPCODES:
|
||||
return "requires a single character or class repeated at most 64 times; use {n,m} instead of *, + or {n,}"
|
||||
choices: Final = high - low + 1
|
||||
return _PatternCost(choices, 1 + high * _atom_steps(body[0]) + choices)
|
||||
|
||||
|
||||
def _node_cost(node: object, depth: int) -> _PatternCost | str:
|
||||
if not isinstance(node, tuple) or len(node) != 2:
|
||||
return "unsupported regex structure"
|
||||
opcode, argument = node
|
||||
if opcode in _ATOM_OPCODES or opcode is sre_constants.AT:
|
||||
return _PatternCost(1, _atom_steps(node))
|
||||
if opcode is sre_constants.SUBPATTERN:
|
||||
return _sequence_cost(argument[-1], depth + 1)
|
||||
if opcode is sre_constants.BRANCH:
|
||||
costs: Final = tuple(_sequence_cost(branch, depth + 1) for branch in argument[1])
|
||||
refused: Final = next((cost for cost in costs if isinstance(cost, str)), None)
|
||||
if refused is not None:
|
||||
return refused
|
||||
return _PatternCost(
|
||||
sum(cost.paths for cost in costs if isinstance(cost, _PatternCost)),
|
||||
len(costs) + sum(cost.steps for cost in costs if isinstance(cost, _PatternCost)),
|
||||
)
|
||||
if opcode in _REPEAT_OPCODES:
|
||||
return _repeat_cost(argument)
|
||||
return "contains an unsupported regex construct"
|
||||
|
||||
|
||||
def _sequence_cost(nodes: Iterable[object], depth: int) -> _PatternCost | str:
|
||||
if depth > MAX_CUSTOM_PATTERN_DEPTH:
|
||||
return "nests deeper than 16 levels"
|
||||
costs: Final = tuple(_node_cost(node, depth) for node in nodes)
|
||||
refused: Final = next((cost for cost in costs if isinstance(cost, str)), None)
|
||||
if refused is not None:
|
||||
return refused
|
||||
valid: Final = tuple(cost for cost in costs if isinstance(cost, _PatternCost))
|
||||
# Choices multiply across a sequence; every continuation can execute once per preceding path.
|
||||
total: Final = _PatternCost(
|
||||
math.prod(cost.paths for cost in valid),
|
||||
1 + sum(cost.steps * math.prod(prior.paths for prior in valid[:index]) for index, cost in enumerate(valid)),
|
||||
)
|
||||
if total.steps > MAX_CUSTOM_PATTERN_WORK:
|
||||
return "exceeds the per-pattern regex work budget"
|
||||
return total
|
||||
|
||||
|
||||
def custom_pattern_work(pattern: str) -> int | str:
|
||||
try:
|
||||
re.compile(pattern, re.IGNORECASE)
|
||||
parsed: Final = sre_parse.parse(pattern, re.IGNORECASE)
|
||||
except (re.error, RecursionError, OverflowError):
|
||||
return "is not a valid regex"
|
||||
cost: Final = _sequence_cost(tuple(parsed), 0)
|
||||
return cost if isinstance(cost, str) else cost.steps
|
||||
|
||||
|
||||
class CustomDimension(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
name: str = Field(min_length=1, max_length=64, pattern=r"^[A-Za-z][A-Za-z0-9_]*$")
|
||||
weight: float = Field(gt=0, le=1, allow_inf_nan=False)
|
||||
keywords: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
|
||||
patterns: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_matchers(self) -> "CustomDimension":
|
||||
matchers: Final = (*self.keywords, *self.patterns)
|
||||
if not matchers or any(not matcher.strip() for matcher in matchers):
|
||||
raise ValueError("custom dimensions require nonblank keywords and/or patterns")
|
||||
if len(matchers) > 32 or sum(map(len, matchers)) > 4096:
|
||||
raise ValueError("custom dimensions allow at most 32 matchers and 4096 matcher characters each")
|
||||
costs: Final = tuple((pattern, custom_pattern_work(pattern)) for pattern in self.patterns)
|
||||
rejected: Final = tuple(f"pattern {pattern!r} {work}" for pattern, work in costs if isinstance(work, str))
|
||||
if rejected:
|
||||
raise ValueError("custom dimension " + "; ".join(rejected))
|
||||
return self
|
||||
|
||||
def pattern_work(self) -> int:
|
||||
"""Combined work estimate of the validated patterns."""
|
||||
return sum(
|
||||
work for work in (custom_pattern_work(pattern) for pattern in self.patterns) if isinstance(work, int)
|
||||
)
|
||||
|
||||
|
||||
class ComplexityRouterConfig(BaseModel):
|
||||
"""Configuration for the ComplexityRouter."""
|
||||
|
||||
|
|
@ -671,6 +790,19 @@ class ComplexityRouterConfig(BaseModel):
|
|||
description="Weights for each scoring dimension",
|
||||
)
|
||||
|
||||
custom_dimensions: tuple[CustomDimension, ...] = Field(
|
||||
default=(),
|
||||
max_length=16,
|
||||
description=(
|
||||
"Named binary dimensions added to the heuristic-v1 score. Each contributes its inline weight once "
|
||||
"when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters. "
|
||||
"Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, "
|
||||
"backreferences and lookarounds are rejected. Conservative work limits include alternation paths, "
|
||||
"repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. "
|
||||
"Only heuristic, heuristic_first and hybrid accept this field. Uses the existing heuristic tuning quota."
|
||||
),
|
||||
)
|
||||
|
||||
# Keyword lists (overridable)
|
||||
code_keywords: list[str] | None = Field(
|
||||
default=None,
|
||||
|
|
@ -1245,6 +1377,27 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_custom_dimensions(self) -> "ComplexityRouterConfig":
|
||||
if not self.custom_dimensions:
|
||||
return self
|
||||
if self.classifier_type not in ("heuristic", "heuristic_first", "hybrid"):
|
||||
raise ValueError("custom_dimensions requires classifier_type heuristic, heuristic_first or hybrid")
|
||||
names: Final = tuple(dimension.name.casefold() for dimension in self.custom_dimensions)
|
||||
reserved: Final = frozenset(name.casefold() for name in DEFAULT_DIMENSION_WEIGHTS)
|
||||
weighted: Final = frozenset(name.casefold() for name in self.dimension_weights)
|
||||
if len(frozenset(names)) != len(names) or frozenset(names) & reserved:
|
||||
raise ValueError("custom dimension names must be unique and must not shadow built-in dimensions")
|
||||
if frozenset(names) & weighted:
|
||||
raise ValueError("custom dimension weights must be inline, not in dimension_weights")
|
||||
work: Final = sum(dimension.pattern_work() for dimension in self.custom_dimensions)
|
||||
if work > MAX_CUSTOM_DIMENSIONS_WORK:
|
||||
raise ValueError(
|
||||
f"custom_dimensions regex work estimate is {work}; the limit across the router is "
|
||||
f"{MAX_CUSTOM_DIMENSIONS_WORK}"
|
||||
)
|
||||
return self
|
||||
|
||||
@field_validator("heuristic_first_max_tier", mode="before")
|
||||
@classmethod
|
||||
def _coerce_heuristic_first_max_tier(cls, value: object) -> object:
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ HEURISTIC_V1_TUNING_FIELDS: Final = (
|
|||
"reasoning_override_min_score",
|
||||
"token_thresholds",
|
||||
"dimension_weights",
|
||||
"custom_dimensions",
|
||||
"code_keywords",
|
||||
"reasoning_keywords",
|
||||
"technical_keywords",
|
||||
|
|
|
|||
|
|
@ -90,6 +90,7 @@ class PromptCachingDeploymentCheck(CustomLogger):
|
|||
enable_prompt_caching=(
|
||||
request_kwargs.get("enable_prompt_caching") is True if request_kwargs is not None else None
|
||||
),
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
model_id_dict: Final = await prompt_cache.async_get_model_id(
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
from pydantic import BaseModel, PrivateAttr, StrictInt
|
||||
from typing_extensions import Required, TypedDict
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
from litellm.types.llms.base import LiteLLMPydanticObjectBase
|
||||
|
||||
|
|
@ -172,6 +172,7 @@ class AugmentedAgentCard(AgentCard):
|
|||
class AgentObjectPermission(TypedDict, total=False):
|
||||
mcp_servers: list[str] | None
|
||||
mcp_access_groups: list[str] | None
|
||||
mcp_toolsets: ReadOnly[Sequence[str] | None]
|
||||
mcp_tool_permissions: dict[str, list[str]] | None
|
||||
models: list[str] | None
|
||||
agents: list[str] | None
|
||||
|
|
|
|||
|
|
@ -340,6 +340,7 @@ markers = [
|
|||
"asyncio: mark test as an asyncio test",
|
||||
"limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')",
|
||||
"no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests",
|
||||
"requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension",
|
||||
]
|
||||
filterwarnings = [
|
||||
# Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests)
|
||||
|
|
|
|||
|
|
@ -1,14 +1,12 @@
|
|||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
from unittest.mock import Mock, patch
|
||||
import pytest
|
||||
import base64
|
||||
import httpx
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}
|
||||
|
||||
|
|
@ -394,8 +392,6 @@ def test_bedrock_embedding_uses_correct_region_when_specified():
|
|||
os.environ["AWS_REGION_NAME"] = original_region_name
|
||||
else:
|
||||
os.environ.pop("AWS_REGION_NAME", None)
|
||||
|
||||
|
||||
def test_bedrock_embedding_region_bug_reproduction():
|
||||
"""
|
||||
Reproduces the bug where aws_region_name is ignored when passed explicitly.
|
||||
|
|
@ -458,13 +454,3 @@ def test_bedrock_embedding_region_bug_reproduction():
|
|||
os.environ["AWS_REGION_NAME"] = original_region_name
|
||||
else:
|
||||
os.environ.pop("AWS_REGION_NAME", None)
|
||||
|
||||
|
||||
def test_bedrock_titan_g1_text_02_model_info():
|
||||
"""Test that amazon.titan-embed-g1-text-02 has correct pricing metadata"""
|
||||
model_info = litellm.get_model_info("amazon.titan-embed-g1-text-02")
|
||||
assert model_info is not None, "Model info should not be None"
|
||||
assert model_info["litellm_provider"] == "bedrock"
|
||||
assert model_info["mode"] == "embedding"
|
||||
assert model_info["input_cost_per_token"] == 1e-07
|
||||
assert model_info["max_input_tokens"] == 8192
|
||||
|
|
|
|||
|
|
@ -1,34 +0,0 @@
|
|||
"""
|
||||
Tests for AWS Bedrock embedding model pricing in the model cost map.
|
||||
|
||||
Regression test for the Amazon Titan Text Embeddings V2 commercial price,
|
||||
which was previously set 10x too high (2e-07 instead of 2e-08).
|
||||
AWS lists Titan Text Embeddings V2 at $0.02 per 1M input tokens
|
||||
(= $0.00002 per 1K tokens = 2e-08 per token).
|
||||
"""
|
||||
|
||||
import importlib
|
||||
|
||||
|
||||
class TestBedrockEmbeddingPricing:
|
||||
"""Test suite for Bedrock embedding model pricing in the cost map."""
|
||||
|
||||
def test_titan_embed_v2_commercial_input_cost(self, monkeypatch):
|
||||
"""Titan Text Embeddings V2 should be priced at $0.02 / 1M tokens (2e-08)."""
|
||||
# Scope the local-cost-map flag to this test only, so it does not leak
|
||||
# into sibling tests. monkeypatch restores the environment on teardown.
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
|
||||
import litellm.litellm_core_utils.get_model_cost_map
|
||||
import litellm
|
||||
|
||||
# Reload so the cost map is re-read from the local file with the flag set.
|
||||
importlib.reload(litellm.litellm_core_utils.get_model_cost_map)
|
||||
importlib.reload(litellm)
|
||||
|
||||
model = litellm.model_cost["amazon.titan-embed-text-v2:0"]
|
||||
|
||||
assert model["input_cost_per_token"] == 2e-08
|
||||
assert model["output_cost_per_token"] == 0.0
|
||||
assert model["litellm_provider"] == "bedrock"
|
||||
assert model["mode"] == "embedding"
|
||||
|
|
@ -40,38 +40,6 @@ class TestBedrockGovCloudSupport:
|
|||
assert "us-gov-east-1" in all_regions
|
||||
assert "us-gov-west-1" in all_regions
|
||||
|
||||
def test_govcloud_models_in_model_cost(self):
|
||||
"""Test that GovCloud models are present in model cost configuration"""
|
||||
from litellm import model_cost
|
||||
|
||||
# Test Claude models in GovCloud
|
||||
assert (
|
||||
"bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
in model_cost
|
||||
)
|
||||
assert (
|
||||
"bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
in model_cost
|
||||
)
|
||||
assert (
|
||||
"bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost
|
||||
)
|
||||
assert (
|
||||
"bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost
|
||||
)
|
||||
assert "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0" in model_cost
|
||||
assert "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0" in model_cost
|
||||
|
||||
# Test Llama models in GovCloud
|
||||
assert "bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0" in model_cost
|
||||
assert "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0" in model_cost
|
||||
assert "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0" in model_cost
|
||||
assert "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0" in model_cost
|
||||
|
||||
# Test Titan models in GovCloud
|
||||
assert "bedrock/us-gov-east-1/amazon.titan-text-lite-v1" in model_cost
|
||||
assert "bedrock/us-gov-west-1/amazon.titan-text-lite-v1" in model_cost
|
||||
|
||||
def test_govcloud_model_routing(self):
|
||||
"""Test that GovCloud models are routed correctly"""
|
||||
# Test Claude model routing
|
||||
|
|
@ -148,135 +116,6 @@ class TestBedrockGovCloudSupport:
|
|||
assert not any("us-gov-east-1" in model for model in litellm.bedrock_models)
|
||||
assert not any("us-gov-west-1" in model for model in litellm.bedrock_models)
|
||||
|
||||
def test_govcloud_model_cost_properties(self):
|
||||
"""Test that GovCloud models have proper cost configuration"""
|
||||
from litellm import model_cost
|
||||
|
||||
# Check a specific GovCloud model has all required properties
|
||||
govcloud_model = model_cost[
|
||||
"bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
]
|
||||
|
||||
assert "max_tokens" in govcloud_model
|
||||
assert "max_input_tokens" in govcloud_model
|
||||
assert "max_output_tokens" in govcloud_model
|
||||
assert "input_cost_per_token" in govcloud_model
|
||||
assert "output_cost_per_token" in govcloud_model
|
||||
assert govcloud_model["litellm_provider"] == "bedrock"
|
||||
assert govcloud_model["mode"] == "chat"
|
||||
|
||||
def test_govcloud_model_pricing_verification(self):
|
||||
"""Test that GovCloud models have correct pricing that differs from base models"""
|
||||
from litellm import model_cost
|
||||
|
||||
# Claude Haiku 4.5 commercial list pricing is under the us.* inference profile id
|
||||
base_model = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
gov_east_model = (
|
||||
"bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
)
|
||||
gov_west_model = (
|
||||
"bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
)
|
||||
|
||||
# Verify base model pricing (us.* inference profile: $1.10/$5.50 per MTok)
|
||||
base_pricing = model_cost[base_model]
|
||||
assert base_pricing["input_cost_per_token"] == 1.1e-06
|
||||
assert base_pricing["output_cost_per_token"] == 5.5e-06
|
||||
|
||||
# Verify GovCloud models have different (higher) pricing
|
||||
gov_east_pricing = model_cost[gov_east_model]
|
||||
gov_west_pricing = model_cost[gov_west_model]
|
||||
|
||||
# GovCloud models should have ~20% higher pricing than base models
|
||||
assert gov_east_pricing["input_cost_per_token"] == 1.2e-06
|
||||
assert gov_east_pricing["output_cost_per_token"] == 6e-06
|
||||
assert gov_west_pricing["input_cost_per_token"] == 1.2e-06
|
||||
assert gov_west_pricing["output_cost_per_token"] == 6e-06
|
||||
|
||||
# Verify the pricing difference is approximately 20%
|
||||
assert (
|
||||
abs(
|
||||
gov_east_pricing["input_cost_per_token"]
|
||||
/ base_pricing["input_cost_per_token"]
|
||||
- 1.2
|
||||
)
|
||||
< 0.15
|
||||
)
|
||||
assert (
|
||||
abs(
|
||||
gov_east_pricing["output_cost_per_token"]
|
||||
/ base_pricing["output_cost_per_token"]
|
||||
- 1.2
|
||||
)
|
||||
< 0.15
|
||||
)
|
||||
assert (
|
||||
abs(
|
||||
gov_west_pricing["input_cost_per_token"]
|
||||
/ base_pricing["input_cost_per_token"]
|
||||
- 1.2
|
||||
)
|
||||
< 0.15
|
||||
)
|
||||
assert (
|
||||
abs(
|
||||
gov_west_pricing["output_cost_per_token"]
|
||||
/ base_pricing["output_cost_per_token"]
|
||||
- 1.2
|
||||
)
|
||||
< 0.15
|
||||
)
|
||||
|
||||
# Test Claude 3 Haiku pricing
|
||||
base_haiku_model = "anthropic.claude-3-haiku-20240307-v1:0"
|
||||
gov_east_haiku_model = (
|
||||
"bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0"
|
||||
)
|
||||
gov_west_haiku_model = (
|
||||
"bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0"
|
||||
)
|
||||
|
||||
# Verify base Haiku model pricing
|
||||
base_haiku_pricing = model_cost[base_haiku_model]
|
||||
assert base_haiku_pricing["input_cost_per_token"] == 2.5e-07 # 0.00000025
|
||||
assert base_haiku_pricing["output_cost_per_token"] == 1.25e-06 # 0.00000125
|
||||
|
||||
# Verify GovCloud Haiku models have different (higher) pricing
|
||||
gov_east_haiku_pricing = model_cost[gov_east_haiku_model]
|
||||
gov_west_haiku_pricing = model_cost[gov_west_haiku_model]
|
||||
|
||||
# GovCloud Haiku models should have 20% higher pricing than base models
|
||||
assert (
|
||||
gov_east_haiku_pricing["input_cost_per_token"] == 3e-07
|
||||
) # 0.0000003 (20% higher)
|
||||
assert (
|
||||
gov_east_haiku_pricing["output_cost_per_token"] == 1.5e-06
|
||||
) # 0.0000015 (20% higher)
|
||||
assert (
|
||||
gov_west_haiku_pricing["input_cost_per_token"] == 3e-07
|
||||
) # 0.0000003 (20% higher)
|
||||
assert (
|
||||
gov_west_haiku_pricing["output_cost_per_token"] == 1.5e-06
|
||||
) # 0.0000015 (20% higher)
|
||||
|
||||
# Verify the pricing difference is exactly 20%
|
||||
assert (
|
||||
gov_east_haiku_pricing["input_cost_per_token"]
|
||||
== base_haiku_pricing["input_cost_per_token"] * 1.2
|
||||
)
|
||||
assert (
|
||||
gov_east_haiku_pricing["output_cost_per_token"]
|
||||
== base_haiku_pricing["output_cost_per_token"] * 1.2
|
||||
)
|
||||
assert (
|
||||
gov_west_haiku_pricing["input_cost_per_token"]
|
||||
== base_haiku_pricing["input_cost_per_token"] * 1.2
|
||||
)
|
||||
assert (
|
||||
gov_west_haiku_pricing["output_cost_per_token"]
|
||||
== base_haiku_pricing["output_cost_per_token"] * 1.2
|
||||
)
|
||||
|
||||
@patch("litellm.completion")
|
||||
def test_govcloud_completion_cost_calculation(self, mock_completion):
|
||||
"""Test that completion requests use correct pricing for GovCloud models"""
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ Tests for Crusoe provider integration
|
|||
import os
|
||||
from unittest import mock
|
||||
|
||||
import litellm
|
||||
|
||||
CRUSOE_API_BASE = "https://managed-inference-api-proxy.crusoecloud.com/v1"
|
||||
|
||||
|
|
@ -71,38 +70,3 @@ def test_get_llm_provider_crusoe():
|
|||
)
|
||||
assert model == "meta-llama/Llama-3.3-70B-Instruct"
|
||||
assert provider == "crusoe"
|
||||
|
||||
|
||||
def test_crusoe_models_configuration():
|
||||
"""Test that Crusoe models are configured correctly"""
|
||||
from litellm import get_model_info
|
||||
|
||||
original_model_cost = litellm.model_cost
|
||||
original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP")
|
||||
try:
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
|
||||
crusoe_models = [
|
||||
"crusoe/meta-llama/Llama-3.3-70B-Instruct",
|
||||
"crusoe/deepseek-ai/DeepSeek-R1-0528",
|
||||
"crusoe/deepseek-ai/DeepSeek-V3-0324",
|
||||
"crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507",
|
||||
"crusoe/moonshotai/Kimi-K2-Thinking",
|
||||
"crusoe/openai/gpt-oss-120b",
|
||||
"crusoe/google/gemma-3-12b-it",
|
||||
]
|
||||
|
||||
for model in crusoe_models:
|
||||
model_info = get_model_info(model)
|
||||
assert model_info is not None, f"Model info not found for {model}"
|
||||
assert model_info.get("litellm_provider") == "crusoe", (
|
||||
f"{model} should have crusoe as provider"
|
||||
)
|
||||
assert model_info.get("mode") == "chat", f"{model} should be in chat mode"
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
if original_env is None:
|
||||
os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None)
|
||||
else:
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env
|
||||
|
|
|
|||
|
|
@ -1,8 +1,3 @@
|
|||
import os
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
|
|
@ -69,34 +64,6 @@ def test_hyperbolic_in_provider_lists():
|
|||
assert "https://api.hyperbolic.xyz/v1" in openai_compatible_endpoints
|
||||
|
||||
|
||||
def test_hyperbolic_models_configuration():
|
||||
"""Test that Hyperbolic models are properly configured"""
|
||||
import json
|
||||
|
||||
# Load model configuration directly from the JSON file
|
||||
json_path = os.path.join(
|
||||
os.path.dirname(__file__), "../../model_prices_and_context_window.json"
|
||||
)
|
||||
with open(json_path, "r") as f:
|
||||
model_data = json.load(f)
|
||||
|
||||
# Test a few key models
|
||||
test_models = [
|
||||
"hyperbolic/deepseek-ai/DeepSeek-V3",
|
||||
"hyperbolic/Qwen/Qwen2.5-Coder-32B-Instruct",
|
||||
"hyperbolic/deepseek-ai/DeepSeek-R1",
|
||||
]
|
||||
|
||||
for model in test_models:
|
||||
assert model in model_data
|
||||
model_info = model_data[model]
|
||||
assert model_info["litellm_provider"] == "hyperbolic"
|
||||
assert model_info["mode"] == "chat"
|
||||
assert "max_tokens" in model_info
|
||||
assert "input_cost_per_token" in model_info
|
||||
assert "output_cost_per_token" in model_info
|
||||
|
||||
|
||||
def test_hyperbolic_supported_params():
|
||||
"""Test that supported OpenAI parameters are correctly configured"""
|
||||
from litellm.llms.hyperbolic.chat.transformation import HyperbolicChatConfig
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ from unittest import mock
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import completion
|
||||
from litellm.llms.lambda_ai.chat.transformation import LambdaAIChatConfig
|
||||
|
||||
|
||||
|
|
@ -103,48 +102,6 @@ async def test_lambda_ai_completion_call():
|
|||
raise
|
||||
|
||||
|
||||
def test_lambda_ai_models_configuration():
|
||||
"""Test that Lambda AI models are configured correctly"""
|
||||
from litellm import get_model_info
|
||||
|
||||
# Reload model cost map to pick up local changes
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
|
||||
# Clear and repopulate lambda_ai_models list after reloading model_cost
|
||||
litellm.lambda_ai_models = set()
|
||||
litellm.add_known_models()
|
||||
|
||||
# Some Lambda AI models to test
|
||||
lambda_ai_models = [
|
||||
"lambda_ai/deepseek-llama3.3-70b",
|
||||
"lambda_ai/hermes3-8b",
|
||||
"lambda_ai/llama3.1-8b-instruct",
|
||||
"lambda_ai/llama3.2-11b-vision-instruct",
|
||||
"lambda_ai/qwen25-coder-32b-instruct",
|
||||
]
|
||||
|
||||
for model in lambda_ai_models:
|
||||
model_info = get_model_info(model)
|
||||
assert model_info is not None, f"Model info not found for {model}"
|
||||
assert (
|
||||
model_info.get("litellm_provider") == "lambda_ai"
|
||||
), f"{model} should have lambda_ai as provider"
|
||||
assert model_info.get("mode") == "chat", f"{model} should be in chat mode"
|
||||
assert (
|
||||
model_info.get("supports_function_calling") is True
|
||||
), f"{model} should support function calling"
|
||||
assert (
|
||||
model_info.get("supports_system_messages") is True
|
||||
), f"{model} should support system messages"
|
||||
|
||||
# Check vision support for vision models
|
||||
if "vision" in model:
|
||||
assert (
|
||||
model_info.get("supports_vision") is True
|
||||
), f"{model} should support vision"
|
||||
|
||||
|
||||
def test_lambda_ai_model_list_populated():
|
||||
"""Test that lambda_ai_models list is populated correctly"""
|
||||
# Ensure we're using local model cost map and repopulate models
|
||||
|
|
|
|||
|
|
@ -68,24 +68,6 @@ def test_morph_in_provider_lists():
|
|||
)
|
||||
|
||||
|
||||
def test_morph_model_info():
|
||||
"""Test that morph models have correct configuration."""
|
||||
import litellm
|
||||
|
||||
model_info = litellm.get_model_info("morph/morph-v3-large")
|
||||
|
||||
assert model_info["litellm_provider"] == "morph"
|
||||
assert model_info["mode"] == "chat"
|
||||
assert model_info["max_tokens"] == 16000
|
||||
assert model_info["max_input_tokens"] == 16000
|
||||
assert model_info["max_output_tokens"] == 16000
|
||||
assert model_info["input_cost_per_token"] == 9e-07 # $0.9/1M tokens
|
||||
assert model_info["output_cost_per_token"] == 1.9e-06 # $1.9/1M tokens
|
||||
assert model_info["supports_function_calling"] is False
|
||||
assert model_info["supports_vision"] is False
|
||||
assert model_info["supports_system_messages"] is True
|
||||
|
||||
|
||||
def test_morph_supported_params():
|
||||
"""Test that MorphChatConfig returns correct supported parameters."""
|
||||
config = MorphChatConfig()
|
||||
|
|
|
|||
|
|
@ -1,15 +1,11 @@
|
|||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, patch, MagicMock
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import Choices, Message, ModelResponse
|
||||
from litellm import ModelResponse
|
||||
from base_llm_unit_tests import BaseLLMChatTest, BaseOSeriesModelsTest
|
||||
|
||||
|
||||
|
|
@ -74,7 +70,6 @@ async def test_o1_handle_tool_calling_optional_params(
|
|||
- max_tokens is translated to 'max_completion_tokens'
|
||||
- role 'system' is translated to 'user'
|
||||
"""
|
||||
from openai import AsyncOpenAI
|
||||
from litellm.utils import ProviderConfigManager
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
|
|
@ -186,15 +181,6 @@ class TestOpenAIO3(BaseOSeriesModelsTest, BaseLLMChatTest):
|
|||
pass
|
||||
|
||||
|
||||
def test_o1_supports_vision():
|
||||
"""Test that o1 supports vision"""
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
for k, v in litellm.model_cost.items():
|
||||
if k.startswith("o1") and v.get("litellm_provider") == "openai":
|
||||
assert v.get("supports_vision") is True, f"{k} does not support vision"
|
||||
|
||||
|
||||
def test_o3_reasoning_effort():
|
||||
resp = litellm.completion(
|
||||
model="o3-mini",
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ from unittest import mock
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import completion
|
||||
from litellm.llms.v0.chat.transformation import V0ChatConfig
|
||||
|
||||
|
||||
|
|
@ -111,33 +110,3 @@ def test_v0_supported_params():
|
|||
]
|
||||
|
||||
assert set(supported_params) == set(expected_params)
|
||||
|
||||
|
||||
def test_v0_models_configuration():
|
||||
"""Test that v0 models are configured correctly"""
|
||||
from litellm import get_model_info
|
||||
|
||||
# Reload model cost map to pick up local changes
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
|
||||
# All v0 models
|
||||
v0_models = ["v0/v0-1.0-md", "v0/v0-1.5-md", "v0/v0-1.5-lg"]
|
||||
|
||||
for model in v0_models:
|
||||
model_info = get_model_info(model)
|
||||
assert model_info is not None, f"Model info not found for {model}"
|
||||
# All v0 models support vision (multimodal)
|
||||
assert (
|
||||
model_info.get("supports_vision") is True
|
||||
), f"{model} should support vision"
|
||||
assert (
|
||||
model_info.get("litellm_provider") == "v0"
|
||||
), f"{model} should have v0 as provider"
|
||||
assert model_info.get("mode") == "chat", f"{model} should be in chat mode"
|
||||
assert (
|
||||
model_info.get("supports_function_calling") is True
|
||||
), f"{model} should support function calling"
|
||||
assert (
|
||||
model_info.get("supports_system_messages") is True
|
||||
), f"{model} should support system messages"
|
||||
|
|
|
|||
|
|
@ -1,8 +1,6 @@
|
|||
# What is this?
|
||||
## Unit testing for the 'get_model_info()' function
|
||||
import os
|
||||
import traceback
|
||||
import json
|
||||
|
||||
|
||||
from typing import List, Dict, Any
|
||||
|
|
@ -11,7 +9,7 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm import get_model_info
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
def test_get_model_info_simple_model_name():
|
||||
|
|
@ -49,34 +47,6 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch):
|
|||
assert model_info["input_cost_per_token"] == 0.0
|
||||
|
||||
|
||||
def test_get_model_info_shows_correct_supports_vision():
|
||||
info = litellm.get_model_info("gemini/gemini-2.0-flash")
|
||||
print("info", info)
|
||||
assert info["supports_vision"] is True
|
||||
|
||||
|
||||
def test_get_model_info_shows_assistant_prefill():
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
info = litellm.get_model_info("deepseek/deepseek-chat")
|
||||
print("info", info)
|
||||
assert info.get("supports_assistant_prefill") is True
|
||||
|
||||
|
||||
def test_get_model_info_shows_supports_prompt_caching():
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
info = litellm.get_model_info("deepseek/deepseek-chat")
|
||||
print("info", info)
|
||||
assert info.get("supports_prompt_caching") is True
|
||||
|
||||
|
||||
def test_get_model_info_finetuned_models():
|
||||
info = litellm.get_model_info("ft:gpt-3.5-turbo:my-org:custom_suffix:id")
|
||||
print("info", info)
|
||||
assert info["input_cost_per_token"] == 0.000003
|
||||
|
||||
|
||||
def test_get_model_info_gemini_pro():
|
||||
info = litellm.get_model_info("gemini-2.0-flash")
|
||||
print("info", info)
|
||||
|
|
@ -219,7 +189,7 @@ def test_model_info_bedrock_converse_enforcement(monkeypatch):
|
|||
def test_get_model_info_custom_provider():
|
||||
# Custom provider example copied from https://docs.litellm.ai/docs/providers/custom_llm_server:
|
||||
import litellm
|
||||
from litellm import CustomLLM, completion, get_llm_provider
|
||||
from litellm import CustomLLM, completion
|
||||
|
||||
class MyCustomLLM(CustomLLM):
|
||||
def completion(self, *args, **kwargs) -> litellm.ModelResponse:
|
||||
|
|
|
|||
|
|
@ -1777,6 +1777,224 @@ class TestEnableAnthropicPromptCaching:
|
|||
assert messages == before
|
||||
|
||||
|
||||
class TestClaudeCodeOneShotAutoCaching:
|
||||
BILLING_TEXT = "x-anthropic-billing-header: cc_version=2.1.263; cc_entrypoint=cli; cc_is_subagent=true;"
|
||||
BILLING_SYSTEM = [{"type": "text", "text": BILLING_TEXT}]
|
||||
MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "unique fetched document"}]}]
|
||||
|
||||
@staticmethod
|
||||
def _kwargs(configured=None):
|
||||
kwargs = {
|
||||
"litellm_metadata": {},
|
||||
"proxy_server_request": {
|
||||
"headers": {
|
||||
"user-agent": "claude-cli/2.1.263 (external, cli)",
|
||||
"x-app": "cli-bg",
|
||||
}
|
||||
},
|
||||
}
|
||||
if configured is not None:
|
||||
kwargs["cache_control_injection_points"] = configured
|
||||
return kwargs
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"system",
|
||||
[
|
||||
BILLING_TEXT,
|
||||
BILLING_SYSTEM,
|
||||
[*BILLING_SYSTEM, {"type": "text", "text": " "}],
|
||||
[
|
||||
*BILLING_SYSTEM,
|
||||
{"type": "text", "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_entrypoint=cli;"},
|
||||
],
|
||||
],
|
||||
ids=["string", "text_block", "whitespace_block", "multiple_billing_blocks"],
|
||||
)
|
||||
@pytest.mark.parametrize("tools", [None, []], ids=["absent_tools", "empty_tools"])
|
||||
def test_skips_defaults_and_attribution_for_one_shot_subagent(self, monkeypatch, system, tools):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
messages = copy.deepcopy(self.MESSAGES)
|
||||
kwargs = self._kwargs()
|
||||
|
||||
result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control(
|
||||
messages,
|
||||
copy.deepcopy(system),
|
||||
kwargs,
|
||||
model="claude-sonnet-4-5",
|
||||
custom_llm_provider="anthropic",
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
assert result_messages == self.MESSAGES
|
||||
assert result_system == system
|
||||
assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"]
|
||||
|
||||
def test_user_agent_header_lookup_is_case_insensitive(self, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
kwargs = self._kwargs()
|
||||
user_agent = kwargs["proxy_server_request"]["headers"].pop("user-agent")
|
||||
kwargs["proxy_server_request"]["headers"]["User-Agent"] = user_agent
|
||||
|
||||
result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control(
|
||||
copy.deepcopy(self.MESSAGES),
|
||||
copy.deepcopy(self.BILLING_SYSTEM),
|
||||
kwargs,
|
||||
model="claude-sonnet-4-5",
|
||||
custom_llm_provider="anthropic",
|
||||
)
|
||||
|
||||
assert result_messages == self.MESSAGES
|
||||
assert result_system == self.BILLING_SYSTEM
|
||||
|
||||
def test_router_affinity_skips_string_billing_system(self, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
messages = copy.deepcopy(self.MESSAGES)
|
||||
kwargs = self._kwargs()
|
||||
kwargs["system"] = self.BILLING_TEXT
|
||||
|
||||
result = AnthropicCacheControlHook.messages_with_default_injections(
|
||||
messages=messages,
|
||||
models=("claude-sonnet-4-5",),
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
|
||||
assert result == messages
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"headers,system",
|
||||
[
|
||||
("not-a-mapping", BILLING_SYSTEM),
|
||||
(
|
||||
{"user-agent": "claude-cli/2.1.263 (external, cli)"},
|
||||
[{"type": "text", "text": "x-anthropic-billing-header: malformed"}],
|
||||
),
|
||||
({"user-agent": "claude-cli/2.1.263 (external, cli)"}, None),
|
||||
({"user-agent": "claude-cli/2.1.263 (external, cli)"}, ["not-a-mapping"]),
|
||||
(
|
||||
{"user-agent": "claude-cli/2.1.263 (external, cli)"},
|
||||
[{"type": "image", "text": BILLING_TEXT}],
|
||||
),
|
||||
],
|
||||
ids=["malformed_headers", "malformed_billing", "missing_system", "malformed_block", "non_text_block"],
|
||||
)
|
||||
def test_malformed_untrusted_context_keeps_defaults(self, monkeypatch, headers, system):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
|
||||
points = AnthropicCacheControlHook.get_default_injection_points(
|
||||
messages=copy.deepcopy(self.MESSAGES),
|
||||
system=copy.deepcopy(system),
|
||||
model="claude-sonnet-4-5",
|
||||
custom_llm_provider="anthropic",
|
||||
request_kwargs={"proxy_server_request": {"headers": headers}},
|
||||
)
|
||||
|
||||
assert len(points) == 2
|
||||
|
||||
def test_message_without_role_keeps_defaults(self, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
|
||||
points = AnthropicCacheControlHook.get_default_injection_points(
|
||||
messages=[{"content": "missing role"}],
|
||||
system=copy.deepcopy(self.BILLING_SYSTEM),
|
||||
model="claude-sonnet-4-5",
|
||||
custom_llm_provider="anthropic",
|
||||
request_kwargs=self._kwargs(),
|
||||
)
|
||||
|
||||
assert len(points) == 2
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"messages,system,tools",
|
||||
[
|
||||
(
|
||||
MESSAGES,
|
||||
BILLING_SYSTEM,
|
||||
[{"name": "WebFetch", "description": "fetch", "input_schema": {"type": "object"}}],
|
||||
),
|
||||
(MESSAGES, [*BILLING_SYSTEM, {"type": "text", "text": "Explore the repository"}], None),
|
||||
(
|
||||
[
|
||||
{"role": "user", "content": "first turn"},
|
||||
{"role": "assistant", "content": "reply"},
|
||||
*MESSAGES,
|
||||
],
|
||||
BILLING_SYSTEM,
|
||||
None,
|
||||
),
|
||||
],
|
||||
ids=["tools", "real_system", "history"],
|
||||
)
|
||||
def test_keeps_defaults_for_reusable_subagents(self, monkeypatch, messages, system, tools):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
kwargs = self._kwargs()
|
||||
|
||||
result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control(
|
||||
copy.deepcopy(messages),
|
||||
copy.deepcopy(system),
|
||||
kwargs,
|
||||
model="claude-sonnet-4-5",
|
||||
custom_llm_provider="anthropic",
|
||||
tools=copy.deepcopy(tools),
|
||||
)
|
||||
|
||||
assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == 2
|
||||
assert kwargs["litellm_metadata"]["litellm_gateway_injected_cache"] == ""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"user_agent,system",
|
||||
[
|
||||
("anthropic-sdk-python/0.75.0", BILLING_SYSTEM),
|
||||
(
|
||||
"claude-cli/2.1.263 (external, cli)",
|
||||
[
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"{BILLING_TEXT}\nadditional system instructions",
|
||||
}
|
||||
],
|
||||
),
|
||||
(
|
||||
"claude-cli/2.1.263 (external, cli)",
|
||||
[
|
||||
{
|
||||
"type": "text",
|
||||
"text": "x-anthropic-billing-header: cc_version=2.1.263; cc_is_subagent=false;",
|
||||
}
|
||||
],
|
||||
),
|
||||
],
|
||||
ids=["different_client", "appended_instructions", "not_a_subagent"],
|
||||
)
|
||||
def test_ambiguous_or_unmatched_signals_fail_open(self, monkeypatch, user_agent, system):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
kwargs = self._kwargs()
|
||||
kwargs["proxy_server_request"]["headers"]["user-agent"] = user_agent
|
||||
|
||||
result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control(
|
||||
copy.deepcopy(self.MESSAGES),
|
||||
copy.deepcopy(system),
|
||||
kwargs,
|
||||
model="claude-sonnet-4-5",
|
||||
custom_llm_provider="anthropic",
|
||||
)
|
||||
|
||||
assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == 2
|
||||
|
||||
def test_explicit_injection_points_remain_authoritative(self, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
kwargs = self._kwargs([{"location": "message", "role": "user"}])
|
||||
|
||||
result_messages, _ = AnthropicCacheControlHook.maybe_inject_cache_control(
|
||||
copy.deepcopy(self.MESSAGES),
|
||||
copy.deepcopy(self.BILLING_SYSTEM),
|
||||
kwargs,
|
||||
model="claude-sonnet-4-5",
|
||||
custom_llm_provider="anthropic",
|
||||
)
|
||||
|
||||
assert result_messages[0]["content"][-1]["cache_control"] == {"type": "ephemeral"}
|
||||
|
||||
|
||||
class TestPerKeyEnablePromptCaching:
|
||||
"""Per-request enable_prompt_caching override (stamped from key metadata) with the global flag off."""
|
||||
|
||||
|
|
@ -1977,6 +2195,25 @@ class TestConfiguredInjectionPointsStandDown:
|
|||
_, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs)
|
||||
assert result_sys == [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"configured",
|
||||
[None, CONFIGURED],
|
||||
ids=["automatic_defaults", "configured_points"],
|
||||
)
|
||||
def test_v1_messages_stands_down_for_root_cache_control(self, monkeypatch, configured):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
root_cache_control = {"type": "ephemeral"}
|
||||
kwargs = {"cache_control": root_cache_control, "litellm_metadata": {}}
|
||||
if configured is not None:
|
||||
kwargs["cache_control_injection_points"] = copy.deepcopy(configured)
|
||||
|
||||
result_messages, result_system = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs)
|
||||
|
||||
assert result_messages == self.V1_MESSAGES
|
||||
assert result_system == "sys"
|
||||
assert kwargs["cache_control"] is root_cache_control
|
||||
assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"]
|
||||
|
||||
def test_v1_messages_reentry_flow_preserves_tool_config_remainder(self):
|
||||
"""The advisor interceptor re-enters anthropic_messages() with the outer
|
||||
request's kwargs and post-injection messages. The first pass applies the
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
|
||||
|
||||
import litellm
|
||||
from litellm import LlmProviders
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import (
|
||||
|
|
@ -46,14 +45,6 @@ def test_xai_openai_compatible_provider_info():
|
|||
assert dynamic_api_key == "api-key"
|
||||
|
||||
|
||||
def test_xai_get_model_info_uses_xai_pricing_metadata():
|
||||
model_info = litellm.get_model_info("xai/grok-3-mini")
|
||||
|
||||
assert model_info["litellm_provider"] == "xai"
|
||||
assert model_info["key"] == "xai/grok-3-mini"
|
||||
assert model_info["mode"] == "chat"
|
||||
|
||||
|
||||
def test_xai_validate_environment_reads_api_key(monkeypatch):
|
||||
monkeypatch.setenv("XAI_API_KEY", "api-key")
|
||||
|
||||
|
|
|
|||
|
|
@ -26,6 +26,30 @@ FAKE_REGULAR_KEY = "sk-ant-api03-regular-key-for-testing-123456789"
|
|||
FAKE_AUTH_TOKEN = "sk-ant-aut01-fake-auth-token-for-testing-123456789"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"messages,system,expected",
|
||||
[
|
||||
([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: cc_is_subagent=true;", True),
|
||||
([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: =junk; cc_is_subagent=true;", False),
|
||||
([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: cc_version=; cc_is_subagent=true;", False),
|
||||
([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: malformed", False),
|
||||
([{"content": "missing role"}], "x-anthropic-billing-header: cc_is_subagent=true;", False),
|
||||
(["not-a-mapping"], "x-anthropic-billing-header: cc_is_subagent=true;", False),
|
||||
([{"role": "user", "content": "hi"}], ["not-a-mapping"], False),
|
||||
([{"role": "user", "content": "hi"}], None, False),
|
||||
],
|
||||
)
|
||||
def test_is_claude_code_one_shot_subagent_request(messages, system, expected):
|
||||
from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request
|
||||
|
||||
assert is_claude_code_one_shot_subagent_request(
|
||||
messages=messages,
|
||||
system=system,
|
||||
tools=None,
|
||||
user_agent="claude-cli/2.1.263 (external, cli)",
|
||||
) is expected
|
||||
|
||||
|
||||
class TestOptionallyHandleAnthropicOAuth:
|
||||
"""Tests for optionally_handle_anthropic_oauth function."""
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
import os
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
|
|
@ -38,25 +37,6 @@ class TestAzureMAIImageGeneration:
|
|||
assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro")
|
||||
assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1")
|
||||
|
||||
def test_mai_flash_and_2e_model_pricing_in_cost_map(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
|
||||
flash_info = litellm.get_model_info(
|
||||
model="azure_ai/MAI-Image-2.5-Flash",
|
||||
custom_llm_provider="azure_ai",
|
||||
)
|
||||
assert flash_info["input_cost_per_token"] == 1.75e-06
|
||||
assert flash_info["input_cost_per_image_token"] == 1.75e-06
|
||||
assert flash_info["output_cost_per_image_token"] == 3.3e-05
|
||||
|
||||
image_2e_info = litellm.get_model_info(
|
||||
model="azure_ai/MAI-Image-2e",
|
||||
custom_llm_provider="azure_ai",
|
||||
)
|
||||
assert image_2e_info["input_cost_per_token"] == 5e-06
|
||||
assert image_2e_info["output_cost_per_image_token"] == 1.95e-05
|
||||
|
||||
def test_get_mai_image_generation_url(self):
|
||||
url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url(
|
||||
api_base="https://my-resource.services.ai.azure.com",
|
||||
|
|
|
|||
|
|
@ -12,112 +12,6 @@ from importlib.resources import files
|
|||
|
||||
import pytest
|
||||
|
||||
FW_MODELS = {
|
||||
"azure_ai/FW-Kimi-K2.5": {
|
||||
"input_cost_per_token": 6.6e-07,
|
||||
"output_cost_per_token": 3.3e-06,
|
||||
"cache_read_input_token_cost": 1.1e-07,
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"supports_vision": True,
|
||||
},
|
||||
"azure_ai/FW-Kimi-K2.6": {
|
||||
"input_cost_per_token": 1.045e-06,
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"cache_read_input_token_cost": 1.76e-07,
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"supports_vision": True,
|
||||
},
|
||||
"azure_ai/FW-Kimi-K2.7-Code": {
|
||||
"input_cost_per_token": 1.05e-06,
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"cache_read_input_token_cost": 2.1e-07,
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"supports_vision": True,
|
||||
},
|
||||
"azure_ai/FW-Kimi-K3": {
|
||||
"input_cost_per_token": 3.3e-06,
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
"supports_vision": True,
|
||||
},
|
||||
"azure_ai/FW-Inkling": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
"output_cost_per_token": 4.05e-06,
|
||||
"cache_read_input_token_cost": 1.7e-07,
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
},
|
||||
"azure_ai/FW-DeepSeek-V3.2": {
|
||||
"input_cost_per_token": 6.2e-07,
|
||||
"output_cost_per_token": 1.85e-06,
|
||||
"cache_read_input_token_cost": 3.1e-07,
|
||||
"max_input_tokens": 163840,
|
||||
"max_output_tokens": 163840,
|
||||
},
|
||||
"azure_ai/FW-DeepSeek-V4-Pro": {
|
||||
"input_cost_per_token": 1.925e-06,
|
||||
"output_cost_per_token": 3.828e-06,
|
||||
"cache_read_input_token_cost": 1.65e-07,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 384000,
|
||||
},
|
||||
"azure_ai/FW-MiniMax-M3": {
|
||||
"input_cost_per_token": 3.3e-07,
|
||||
"output_cost_per_token": 1.32e-06,
|
||||
"cache_read_input_token_cost": 6.6e-08,
|
||||
"max_input_tokens": 512000,
|
||||
"max_output_tokens": 512000,
|
||||
"supports_vision": True,
|
||||
},
|
||||
"azure_ai/FW-MiniMax-M2.5": {
|
||||
"input_cost_per_token": 3.3e-07,
|
||||
"output_cost_per_token": 1.32e-06,
|
||||
"cache_read_input_token_cost": 3.3e-08,
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 1000000,
|
||||
},
|
||||
"azure_ai/FW-Nemotron-3-Ultra-NVFP4": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"output_cost_per_token": 2.4e-06,
|
||||
"cache_read_input_token_cost": 1.19e-07,
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
},
|
||||
"azure_ai/FW-GLM-5.2-Fast": {
|
||||
"input_cost_per_token": 2.1e-06,
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
"cache_read_input_token_cost": 2.1e-07,
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
},
|
||||
"azure_ai/FW-GLM-5.2": {
|
||||
"input_cost_per_token": 1.54e-06,
|
||||
"output_cost_per_token": 4.84e-06,
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
},
|
||||
"azure_ai/FW-GLM-5.1": {
|
||||
"input_cost_per_token": 1.54e-06,
|
||||
"output_cost_per_token": 4.84e-06,
|
||||
"cache_read_input_token_cost": 2.86e-07,
|
||||
"max_input_tokens": 202800,
|
||||
"max_output_tokens": 131072,
|
||||
},
|
||||
"azure_ai/FW-GLM-5": {
|
||||
"input_cost_per_token": 1.1e-06,
|
||||
"output_cost_per_token": 3.52e-06,
|
||||
"cache_read_input_token_cost": 2.2e-07,
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 128000,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def use_local_model_cost_map():
|
||||
|
|
@ -144,28 +38,6 @@ def use_local_model_cost_map():
|
|||
monkeypatch.undo()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_key,expected", list(FW_MODELS.items()))
|
||||
def test_azure_ai_fw_model_info(use_local_model_cost_map, model_key, expected):
|
||||
model_info = use_local_model_cost_map.get_model_info(model=model_key)
|
||||
|
||||
assert model_info["litellm_provider"] == "azure_ai"
|
||||
assert model_info["mode"] == "chat"
|
||||
assert model_info["input_cost_per_token"] == pytest.approx(expected["input_cost_per_token"])
|
||||
assert model_info["output_cost_per_token"] == pytest.approx(expected["output_cost_per_token"])
|
||||
assert model_info["cache_read_input_token_cost"] == pytest.approx(
|
||||
expected["cache_read_input_token_cost"]
|
||||
)
|
||||
assert model_info["max_input_tokens"] == expected["max_input_tokens"]
|
||||
assert model_info["max_output_tokens"] == expected["max_output_tokens"]
|
||||
assert model_info["max_tokens"] == expected["max_output_tokens"]
|
||||
assert model_info["supports_function_calling"] is True
|
||||
assert model_info["supports_reasoning"] is True
|
||||
assert model_info["supports_tool_choice"] is True
|
||||
assert model_info["supports_prompt_caching"] is True
|
||||
if expected.get("supports_vision"):
|
||||
assert model_info["supports_vision"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name,expected_prompt,expected_completion",
|
||||
[
|
||||
|
|
@ -197,22 +69,6 @@ def test_azure_ai_fw_cost_per_token(
|
|||
assert completion_cost == pytest.approx(expected_completion)
|
||||
|
||||
|
||||
def test_azure_ai_fw_nemotron_lightning_model_info(use_local_model_cost_map):
|
||||
model_info = use_local_model_cost_map.get_model_info(model="azure_ai/FW-Nemotron-Lightning-3.5-30B-A3B")
|
||||
|
||||
assert model_info["litellm_provider"] == "azure_ai"
|
||||
assert model_info["mode"] == "chat"
|
||||
assert model_info["input_cost_per_token"] == pytest.approx(6e-08)
|
||||
assert model_info["output_cost_per_token"] == pytest.approx(2.2e-07)
|
||||
assert model_info["cache_read_input_token_cost"] == pytest.approx(1e-08)
|
||||
assert model_info["max_input_tokens"] == 262144
|
||||
assert model_info["supports_function_calling"] is True
|
||||
assert model_info["supports_reasoning"] is True
|
||||
assert model_info["supports_tool_choice"] is True
|
||||
assert model_info["supports_prompt_caching"] is True
|
||||
assert model_info["supports_vision"] is False
|
||||
|
||||
|
||||
def test_azure_ai_fw_nemotron_lightning_supports_tool_choice(use_local_model_cost_map):
|
||||
from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig
|
||||
|
||||
|
|
|
|||
|
|
@ -33,33 +33,6 @@ def use_local_model_cost_map():
|
|||
monkeypatch.undo()
|
||||
|
||||
|
||||
def test_azure_ai_kimi_k26_model_info(use_local_model_cost_map):
|
||||
model_info = use_local_model_cost_map.get_model_info(model="azure_ai/kimi-k2.6")
|
||||
|
||||
assert model_info["litellm_provider"] == "azure_ai"
|
||||
assert model_info["mode"] == "chat"
|
||||
assert model_info["max_input_tokens"] == 262144
|
||||
assert model_info["max_output_tokens"] == 262144
|
||||
assert model_info["max_tokens"] == 262144
|
||||
assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07)
|
||||
assert model_info["output_cost_per_token"] == pytest.approx(4e-06)
|
||||
assert model_info["supports_function_calling"] is True
|
||||
assert model_info["supports_reasoning"] is True
|
||||
assert model_info["supports_tool_choice"] is True
|
||||
assert model_info["supports_vision"] is True
|
||||
|
||||
|
||||
def test_azure_ai_kimi_k26_raw_model_cost_entry(use_local_model_cost_map):
|
||||
model_info = use_local_model_cost_map.model_cost["azure_ai/kimi-k2.6"]
|
||||
|
||||
assert model_info["supported_modalities"] == ["text", "image"]
|
||||
assert model_info["supported_output_modalities"] == ["text"]
|
||||
assert model_info["supports_function_calling"] is True
|
||||
assert model_info["supports_reasoning"] is True
|
||||
assert model_info["supports_tool_choice"] is True
|
||||
assert model_info["supports_vision"] is True
|
||||
|
||||
|
||||
def test_azure_ai_kimi_k26_cost_per_token(use_local_model_cost_map):
|
||||
from litellm.llms.azure_ai.cost_calculator import cost_per_token
|
||||
from litellm.types.utils import Usage
|
||||
|
|
|
|||
|
|
@ -1,8 +1,5 @@
|
|||
"""Test Bedrock cross-region inference profile model mapping"""
|
||||
|
||||
import json
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
import pytest
|
||||
|
|
@ -102,13 +99,6 @@ GPT_5_6_PROFILES = [
|
|||
]
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _packaged_cost_map():
|
||||
"""The map litellm actually resolves against, for fields ModelInfoBase drops."""
|
||||
path = Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json"
|
||||
return json.loads(path.read_text())
|
||||
|
||||
|
||||
def _bedrock_response(model, usage):
|
||||
return ModelResponse(
|
||||
id="test",
|
||||
|
|
@ -126,17 +116,6 @@ def _bedrock_response(model, usage):
|
|||
)
|
||||
|
||||
|
||||
def test_bedrock_cross_region_inference_profile_mapping():
|
||||
"""Test that bedrock cross-region inference profile model is mapped"""
|
||||
model = "bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0"
|
||||
|
||||
model_info = _get_model_info_helper(model=model, custom_llm_provider="bedrock")
|
||||
|
||||
assert model_info is not None
|
||||
assert model_info["litellm_provider"] == "bedrock"
|
||||
assert model_info["input_cost_per_token"] == 8e-07
|
||||
|
||||
|
||||
def test_proxy_cost_calculation_scenario():
|
||||
"""Test exact GitHub issue scenario: proxy cost calculation"""
|
||||
model = "litellm_proxy/bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0"
|
||||
|
|
@ -176,38 +155,6 @@ def test_bedrock_gpt_5_6_profiles_route_to_converse(profile, local_model_cost_ma
|
|||
assert BedrockModelInfo.get_bedrock_route(f"bedrock/{profile.model_id}") == "converse"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id)
|
||||
def test_bedrock_gpt_5_6_published_rates(profile, local_model_cost_map):
|
||||
"""Geo and Global profiles carry their own published rates, per context tier."""
|
||||
model_info = _get_model_info_helper(
|
||||
model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock"
|
||||
)
|
||||
|
||||
assert model_info["litellm_provider"] == "bedrock_converse"
|
||||
assert model_info["mode"] == "chat"
|
||||
assert model_info["max_input_tokens"] == 1000000
|
||||
assert model_info["input_cost_per_token"] == profile.input_cost
|
||||
assert (
|
||||
model_info["input_cost_per_token_above_272k_tokens"]
|
||||
== profile.input_cost_above_272k
|
||||
)
|
||||
assert model_info["output_cost_per_token"] == profile.output_cost
|
||||
assert (
|
||||
model_info["output_cost_per_token_above_272k_tokens"]
|
||||
== profile.output_cost_above_272k
|
||||
)
|
||||
assert model_info["cache_creation_input_token_cost"] == profile.cache_write
|
||||
assert (
|
||||
model_info["cache_creation_input_token_cost_above_272k_tokens"]
|
||||
== profile.cache_write_above_272k
|
||||
)
|
||||
assert model_info["cache_read_input_token_cost"] == profile.cache_read
|
||||
assert (
|
||||
model_info["cache_read_input_token_cost_above_272k_tokens"]
|
||||
== profile.cache_read_above_272k
|
||||
)
|
||||
|
||||
|
||||
def test_bedrock_gpt_5_6_above_272k_tier_applies_to_cost(local_model_cost_map):
|
||||
"""A prompt over 272K tokens is billed at the long-context rate, not the base rate."""
|
||||
response = _bedrock_response(
|
||||
|
|
@ -267,31 +214,6 @@ def test_bedrock_gpt_5_6_bills_cache_write_tokens(local_model_cost_map):
|
|||
assert cost == pytest.approx(expected, rel=1e-9)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id)
|
||||
def test_bedrock_gpt_5_6_advertises_only_converse_supported_features(
|
||||
profile, local_model_cost_map
|
||||
):
|
||||
model_info = _get_model_info_helper(
|
||||
model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock"
|
||||
)
|
||||
|
||||
assert model_info["supports_function_calling"] is True
|
||||
assert model_info["supports_tool_choice"] is True
|
||||
assert model_info["supports_vision"] is True
|
||||
|
||||
# Bedrock rejects an explicit cachePoint block for these models, so the flag that
|
||||
# offers caller-driven caching stays off even though the cache rates are declared.
|
||||
assert not model_info.get("supports_prompt_caching")
|
||||
|
||||
# ModelInfoBase drops these two, so they are read from the map litellm resolves.
|
||||
raw = _packaged_cost_map()[profile.model_id]
|
||||
assert raw["supported_modalities"] == ["text", "image"]
|
||||
assert raw["supported_output_modalities"] == ["text"]
|
||||
# No bedrock_converse entry declares supported_endpoints; these models are reachable
|
||||
# on chat completions and on the Responses API without it.
|
||||
assert "supported_endpoints" not in raw
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id)
|
||||
def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(profile, local_model_cost_map):
|
||||
"""GPT-5.x on Converse maps reasoning_effort to reasoning.effort, so reasoning_effort
|
||||
|
|
|
|||
|
|
@ -8,9 +8,7 @@ gate, the URL construction for both paths, and the shared Bearer auth.
|
|||
"""
|
||||
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from botocore.exceptions import (
|
||||
|
|
@ -20,7 +18,6 @@ from botocore.exceptions import (
|
|||
)
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock_mantle.common_utils import mantle_base_segment
|
||||
from litellm.llms.bedrock_mantle.responses.transformation import (
|
||||
BedrockMantleResponsesAPIConfig,
|
||||
)
|
||||
|
|
@ -160,7 +157,6 @@ class TestBedrockMantleResponsesURL:
|
|||
assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses"
|
||||
assert url.count("/responses") == 1
|
||||
|
||||
|
||||
def test_url_aws_region_name_overrides_stale_api_base(self, monkeypatch):
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
|
||||
|
|
@ -1778,66 +1774,6 @@ class TestBedrockMantleResponsesSigV4:
|
|||
|
||||
|
||||
class TestBedrockMantleResponsesPricing:
|
||||
def test_gpt_5_5_pricing_and_mode(self, local_cost_map):
|
||||
info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5")
|
||||
assert info["mode"] == "responses"
|
||||
assert info["input_cost_per_token"] == pytest.approx(5.5e-06)
|
||||
assert info["output_cost_per_token"] == pytest.approx(3.3e-05)
|
||||
assert info["cache_read_input_token_cost"] == pytest.approx(5.5e-07)
|
||||
assert info["max_input_tokens"] == 1050000
|
||||
|
||||
def test_gpt_5_4_pricing_and_mode(self, local_cost_map):
|
||||
info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.4")
|
||||
assert info["mode"] == "responses"
|
||||
assert info["input_cost_per_token"] == pytest.approx(2.75e-06)
|
||||
assert info["output_cost_per_token"] == pytest.approx(1.65e-05)
|
||||
assert info["cache_read_input_token_cost"] == pytest.approx(2.75e-07)
|
||||
assert info["max_input_tokens"] == 1050000
|
||||
|
||||
def test_gpt_5_6_cyber_pricing_and_mode(self, local_cost_map):
|
||||
info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.6-cyber")
|
||||
assert info["mode"] == "responses"
|
||||
assert info["input_cost_per_token"] == pytest.approx(1.375e-05)
|
||||
assert info["cache_creation_input_token_cost"] == pytest.approx(1.71875e-05)
|
||||
assert info["cache_read_input_token_cost"] == pytest.approx(1.375e-06)
|
||||
assert info["output_cost_per_token"] == pytest.approx(8.25e-05)
|
||||
assert info["max_input_tokens"] == 272000
|
||||
|
||||
def test_gpt_daybreak_blue_pricing_and_route(self, local_cost_map):
|
||||
info = litellm.get_model_info("bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol")
|
||||
assert info["mode"] == "responses"
|
||||
assert info["input_cost_per_token"] == pytest.approx(5.5e-06)
|
||||
assert info["cache_creation_input_token_cost"] == pytest.approx(6.875e-06)
|
||||
assert info["cache_read_input_token_cost"] == pytest.approx(5.5e-07)
|
||||
assert info["output_cost_per_token"] == pytest.approx(3.3e-05)
|
||||
assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(1.1e-05)
|
||||
assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(4.95e-05)
|
||||
assert info["max_input_tokens"] == 1050000
|
||||
assert info["supported_endpoints"] == ["/v1/responses"]
|
||||
assert mantle_base_segment("openai.gpt-daybreak-blue-5.6-sol", litellm.model_cost) == "openai/v1"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, input_cost, cache_creation_cost, cache_read_cost, output_cost",
|
||||
[
|
||||
("openai.gpt-5.6-sol", 5.5e-06, 6.875e-06, 5.5e-07, 3.3e-05),
|
||||
("openai.gpt-5.6-terra", 2.2e-06, 2.75e-06, 2.2e-07, 1.32e-05),
|
||||
("openai.gpt-5.6-luna", 2.2e-07, 2.75e-07, 2.2e-08, 1.32e-06),
|
||||
],
|
||||
)
|
||||
def test_gpt_5_6_pricing_and_mode(
|
||||
self, local_cost_map, model, input_cost, cache_creation_cost, cache_read_cost, output_cost
|
||||
):
|
||||
info = litellm.get_model_info(f"bedrock_mantle/{model}")
|
||||
assert info["mode"] == "responses"
|
||||
assert info["input_cost_per_token"] == pytest.approx(input_cost)
|
||||
assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost)
|
||||
assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost)
|
||||
assert info["output_cost_per_token"] == pytest.approx(output_cost)
|
||||
assert info["max_input_tokens"] == 1050000
|
||||
assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2)
|
||||
assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2)
|
||||
assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2)
|
||||
assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.5)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, input_cost, output_cost",
|
||||
|
|
@ -1875,58 +1811,3 @@ class TestBedrockMantleResponsesPricing:
|
|||
def test_models_registered(self, local_cost_map):
|
||||
assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
|
||||
assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models
|
||||
|
||||
|
||||
def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]:
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
paths = {
|
||||
"root": repo_root / "model_prices_and_context_window.json",
|
||||
"bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json",
|
||||
}
|
||||
return json.loads(paths[map_name].read_text())
|
||||
|
||||
|
||||
class TestMantleGptRegistryEntries:
|
||||
"""Locks the OpenAI GPT entries to Bedrock Mantle's live behavior.
|
||||
|
||||
Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna
|
||||
and for gpt-5.5 and gpt-5.4 (oversize requests 400 with "prompt tokens (N)
|
||||
exceed model maximum (1050000)", and a 1,030,590-token request completes
|
||||
on every one of them), while the AWS model cards still quote 272K for
|
||||
gpt-5.5 and gpt-5.4. mode must stay "responses": Mantle's native
|
||||
/v1/chat/completions rejects function tools unless reasoning_effort is
|
||||
"none", so chat traffic has to keep bridging to the Responses API
|
||||
(see the responses_api_bridge tests above).
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize("map_name", ("root", "bundled_backup"))
|
||||
@pytest.mark.parametrize(
|
||||
"key",
|
||||
(
|
||||
"bedrock_mantle/openai.gpt-5.6-sol",
|
||||
"bedrock_mantle/openai.gpt-5.6-terra",
|
||||
"bedrock_mantle/openai.gpt-5.6-luna",
|
||||
),
|
||||
)
|
||||
def test_entry_matches_mantle_enforced_limits(self, map_name, key):
|
||||
entry = _repo_cost_map(map_name)[key]
|
||||
assert entry["max_input_tokens"] == 1050000
|
||||
assert entry["max_output_tokens"] == 128000
|
||||
assert entry["mode"] == "responses"
|
||||
assert entry["use_openai_responses_path"] is True
|
||||
assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"]
|
||||
|
||||
@pytest.mark.parametrize("map_name", ("root", "bundled_backup"))
|
||||
@pytest.mark.parametrize(
|
||||
"key",
|
||||
(
|
||||
"bedrock_mantle/openai.gpt-5.5",
|
||||
"bedrock_mantle/openai.gpt-5.4",
|
||||
),
|
||||
)
|
||||
def test_gpt_55_and_54_entries_match_mantle_enforced_limits(self, map_name, key):
|
||||
entry = _repo_cost_map(map_name)[key]
|
||||
assert entry["max_input_tokens"] == 1050000
|
||||
assert entry["max_output_tokens"] == 128000
|
||||
assert entry["mode"] == "responses"
|
||||
assert entry["use_openai_responses_path"] is True
|
||||
|
|
|
|||
|
|
@ -684,40 +684,6 @@ class TestBedrockMantleProviderResolution:
|
|||
class TestBedrockMantlePricing:
|
||||
"""Tests that verify Bedrock Mantle uses correct AWS Bedrock pricing, not OpenAI pricing."""
|
||||
|
||||
def test_gpt_oss_120b_pricing(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
|
||||
litellm.add_known_models()
|
||||
info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b")
|
||||
# Bedrock pricing: $0.15/M input, $0.60/M output
|
||||
assert info["input_cost_per_token"] == pytest.approx(1.5e-7)
|
||||
assert info["output_cost_per_token"] == pytest.approx(6e-7)
|
||||
|
||||
def test_gpt_oss_20b_pricing(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
|
||||
litellm.add_known_models()
|
||||
info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-20b")
|
||||
# Bedrock pricing: $0.075/M input, $0.30/M output
|
||||
assert info["input_cost_per_token"] == pytest.approx(7.5e-8)
|
||||
assert info["output_cost_per_token"] == pytest.approx(3e-7)
|
||||
|
||||
def test_pricing_significantly_cheaper_than_openai_native(self, monkeypatch):
|
||||
"""
|
||||
Verify Bedrock Mantle pricing is cheaper than OpenAI's direct API pricing.
|
||||
This is the core issue the provider addition fixes — previously users were being
|
||||
billed at OpenAI rates instead of the cheaper Bedrock rates.
|
||||
"""
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
|
||||
litellm.add_known_models()
|
||||
bedrock_info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b")
|
||||
# OpenAI direct pricing for gpt-oss-120b is ~$0.039/M input, $0.190/M output
|
||||
# Bedrock should be cheaper at $0.15/M input and $0.60/M output... wait
|
||||
# Actually, Bedrock ADDS value not reduces cost vs OpenAI direct for these models.
|
||||
# The key fix is that we now use Bedrock-specific prices instead of mapping to
|
||||
# some unrelated OpenAI model (like gpt-4) pricing.
|
||||
# Just validate the pricing is as expected from AWS docs.
|
||||
assert bedrock_info["input_cost_per_token"] == pytest.approx(1.5e-7)
|
||||
assert bedrock_info["output_cost_per_token"] == pytest.approx(6e-7)
|
||||
|
||||
def test_safeguard_models_have_larger_output_tokens(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
|
||||
litellm.add_known_models()
|
||||
|
|
@ -727,49 +693,6 @@ class TestBedrockMantlePricing:
|
|||
)
|
||||
assert info_safeguard["max_output_tokens"] > info_120b["max_output_tokens"]
|
||||
|
||||
def test_reasoning_support(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
|
||||
litellm.add_known_models()
|
||||
info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b")
|
||||
assert info.get("supports_reasoning") is True
|
||||
|
||||
def test_context_window(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
|
||||
litellm.add_known_models()
|
||||
info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b")
|
||||
assert info["max_input_tokens"] == 131072
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id,input_cost,output_cost,max_tokens",
|
||||
[
|
||||
("google.gemma-4-31b", 1.4e-07, 4e-07, 256000),
|
||||
("google.gemma-4-26b-a4b", 1.3e-07, 4e-07, 256000),
|
||||
("google.gemma-4-e2b", 4e-08, 8e-08, 128000),
|
||||
],
|
||||
)
|
||||
def test_gemma_4_bedrock_mantle_model_metadata(
|
||||
local_cost_map, model_id, input_cost, output_cost, max_tokens
|
||||
):
|
||||
full_model_name = f"bedrock_mantle/{model_id}"
|
||||
info = litellm.get_model_info(full_model_name)
|
||||
|
||||
assert info["mode"] == "chat"
|
||||
assert info["input_cost_per_token"] == pytest.approx(input_cost)
|
||||
assert info["output_cost_per_token"] == pytest.approx(output_cost)
|
||||
assert info["max_input_tokens"] == max_tokens
|
||||
assert info["max_output_tokens"] == max_tokens
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert (
|
||||
litellm.supports_parallel_function_calling(
|
||||
model=full_model_name, custom_llm_provider="bedrock_mantle"
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id",
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import get_model_info, supports_reasoning, supports_vision
|
||||
from litellm import supports_reasoning, supports_vision
|
||||
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
|
||||
from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
|
||||
from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
|
||||
|
|
@ -16,17 +16,6 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def force_local_model_cost(monkeypatch):
|
||||
"""Force local model cost map usage for all tests in this file."""
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
# Refresh model_cost from local map
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
|
||||
|
||||
litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url)
|
||||
|
||||
|
||||
def test_validate_environment_sets_session_affinity_from_litellm_session_id():
|
||||
config = FireworksAIConfig()
|
||||
|
||||
|
|
@ -404,15 +393,6 @@ def test_get_supported_openai_params_parallel_tool_calls_without_tool_choice(
|
|||
assert "tool_choice" not in supported_params
|
||||
|
||||
|
||||
def test_get_model_info_respects_explicit_fireworks_capabilities():
|
||||
"""Test that get_model_info preserves explicit capability flags from the model map."""
|
||||
model_info = get_model_info("fireworks_ai/accounts/fireworks/models/glm-5p1")
|
||||
|
||||
assert model_info["supports_function_calling"] is True
|
||||
assert model_info["supports_reasoning"] is True
|
||||
assert model_info["supports_tool_choice"] is True
|
||||
|
||||
|
||||
def test_get_provider_info_omits_false_supports_reasoning(monkeypatch):
|
||||
"""Test that Fireworks only overrides supports_reasoning for supported models."""
|
||||
config = FireworksAIConfig()
|
||||
|
|
|
|||
|
|
@ -56,17 +56,6 @@ def use_local_model_cost_map():
|
|||
monkeypatch.undo()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("alias", KIMI_ALIASES)
|
||||
def test_fireworks_kimi_raw_cost_entry_limits(use_local_model_cost_map, alias):
|
||||
entry = use_local_model_cost_map.model_cost[alias]
|
||||
|
||||
assert entry["litellm_provider"] == "fireworks_ai"
|
||||
assert entry["max_input_tokens"] == CONTEXT_WINDOW
|
||||
assert entry["max_output_tokens"] == OUTPUT_LIMIT
|
||||
assert entry["max_tokens"] == OUTPUT_LIMIT
|
||||
assert entry["max_output_tokens"] < entry["max_input_tokens"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("alias", KIMI_ALIASES)
|
||||
def test_fireworks_kimi_get_model_info_limits(use_local_model_cost_map, alias):
|
||||
model_info = use_local_model_cost_map.get_model_info(model=alias)
|
||||
|
|
|
|||
|
|
@ -1,13 +1,11 @@
|
|||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents
|
||||
|
||||
|
||||
def test_gemini_realtime_transformation_session_created():
|
||||
|
|
@ -308,20 +306,6 @@ def test_gemini_realtime_transformation_generation_complete():
|
|||
assert contains_audio_done_event, "Expected audio done event"
|
||||
|
||||
|
||||
def test_gemini_3_1_flash_live_preview_model_cost_map_entry():
|
||||
for key in (
|
||||
"gemini-3.1-flash-live-preview",
|
||||
"gemini/gemini-3.1-flash-live-preview",
|
||||
):
|
||||
assert key in litellm.model_cost
|
||||
info = litellm.model_cost[key]
|
||||
assert "/v1/realtime" in info.get("supported_endpoints", [])
|
||||
assert info.get("max_input_tokens") == 131072
|
||||
assert info.get("max_output_tokens") == 65536
|
||||
assert "video" in info.get("supported_modalities", [])
|
||||
assert info.get("supports_function_calling") is True
|
||||
|
||||
|
||||
def test_gemini_realtime_tool_call_transformation():
|
||||
"""Test transformation of Gemini toolCall to OpenAI function_call_arguments.done format."""
|
||||
config = GeminiRealtimeConfig()
|
||||
|
|
@ -1845,19 +1829,6 @@ def test_is_audio_only_live_model_uses_cost_map(model, expected, patch_gemini_au
|
|||
assert GeminiRealtimeConfig._is_audio_only_live_model(model) == expected
|
||||
|
||||
|
||||
def test_gemini_live_native_audio_entry_is_vertex_only():
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
catalog_path: Final = Path(__file__).parents[5] / "model_prices_and_context_window.json"
|
||||
catalog: Final = json.loads(catalog_path.read_text())
|
||||
vertex_key: Final = "gemini-live-2.5-flash-native-audio"
|
||||
assert catalog[vertex_key]["litellm_provider"] == "vertex_ai-language-models"
|
||||
assert catalog[vertex_key].get("gemini_native_audio") is True
|
||||
assert "gemini/gemini-live-2.5-flash-native-audio" not in catalog, "the Gemini API does not serve this model"
|
||||
|
||||
|
||||
def test_is_setup_message_and_is_content_message():
|
||||
config = GeminiRealtimeConfig()
|
||||
assert config.is_setup_message({"setup": {}}) is True
|
||||
|
|
|
|||
|
|
@ -231,26 +231,6 @@ def test_inception_in_provider_lists():
|
|||
assert "https://api.inceptionlabs.ai/v1" in litellm.openai_compatible_endpoints
|
||||
|
||||
|
||||
def test_inception_model_configuration(monkeypatch):
|
||||
from litellm import get_model_info
|
||||
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
litellm.inception_models = set()
|
||||
litellm.add_known_models()
|
||||
|
||||
info = get_model_info("inception/mercury-2")
|
||||
assert info.get("litellm_provider") == "inception"
|
||||
assert info.get("mode") == "chat"
|
||||
assert info.get("max_input_tokens") == 128000
|
||||
assert info.get("input_cost_per_token") == 2.5e-07
|
||||
assert info.get("output_cost_per_token") == 7.5e-07
|
||||
assert info.get("cache_read_input_token_cost") == 2.5e-08
|
||||
assert info.get("supports_function_calling") is True
|
||||
assert info.get("supports_tool_choice") is True
|
||||
assert info.get("supports_response_schema") is True
|
||||
|
||||
|
||||
def test_inception_model_list_populated(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
|
|
|
|||
|
|
@ -143,24 +143,6 @@ async def test_inception_fim_async():
|
|||
assert r.choices[0].text == "a + b"
|
||||
|
||||
|
||||
def test_inception_fim_model_configuration(monkeypatch):
|
||||
from litellm import get_model_info
|
||||
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
litellm.text_completion_inception_models = set()
|
||||
litellm.add_known_models()
|
||||
|
||||
assert (
|
||||
"text-completion-inception/mercury-edit-2"
|
||||
in litellm.text_completion_inception_models
|
||||
)
|
||||
info = get_model_info("text-completion-inception/mercury-edit-2")
|
||||
assert info.get("litellm_provider") == "text-completion-inception"
|
||||
assert info.get("mode") == "completion"
|
||||
assert info.get("max_input_tokens") == 32000
|
||||
|
||||
|
||||
def test_inception_fim_targets_fim_endpoint():
|
||||
"""
|
||||
End-to-end: a FIM request must hit `/v1/fim/completions` (NOT
|
||||
|
|
|
|||
|
|
@ -708,38 +708,6 @@ class TestKimiK26ModelRegistry:
|
|||
"""Load directly from the bundled backup so tests don't depend on remote fetch."""
|
||||
return GetModelCostMap.load_local_model_cost_map()
|
||||
|
||||
def test_kimi_k26_in_model_cost_map(self, model_cost_map):
|
||||
"""kimi-k2.6 should be present in the model cost map."""
|
||||
assert "moonshot/kimi-k2.6" in model_cost_map, "moonshot/kimi-k2.6 not found in model_cost"
|
||||
|
||||
def test_kimi_k26_pricing(self, model_cost_map):
|
||||
"""kimi-k2.6 pricing should match official Kimi API rates."""
|
||||
model_info = model_cost_map["moonshot/kimi-k2.6"]
|
||||
assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07)
|
||||
assert model_info["output_cost_per_token"] == pytest.approx(4e-06)
|
||||
assert model_info["cache_read_input_token_cost"] == pytest.approx(1.6e-07)
|
||||
|
||||
def test_kimi_k26_context_window(self, model_cost_map):
|
||||
"""kimi-k2.6 should have a 256K (262144 token) context window."""
|
||||
model_info = model_cost_map["moonshot/kimi-k2.6"]
|
||||
assert model_info["max_input_tokens"] == 262144
|
||||
assert model_info["max_output_tokens"] == 262144
|
||||
assert model_info["max_tokens"] == 262144
|
||||
|
||||
def test_kimi_k26_capabilities(self, model_cost_map):
|
||||
"""kimi-k2.6 should support function calling, vision, video input, tool choice, and reasoning."""
|
||||
model_info = model_cost_map["moonshot/kimi-k2.6"]
|
||||
assert model_info.get("supports_function_calling") is True
|
||||
assert model_info.get("supports_tool_choice") is True
|
||||
assert model_info.get("supports_vision") is True
|
||||
assert model_info.get("supports_video_input") is True
|
||||
assert model_info.get("supports_reasoning") is True
|
||||
|
||||
def test_kimi_k26_provider(self, model_cost_map):
|
||||
"""kimi-k2.6 should be assigned to the moonshot provider."""
|
||||
model_info = model_cost_map["moonshot/kimi-k2.6"]
|
||||
assert model_info["litellm_provider"] == "moonshot"
|
||||
|
||||
|
||||
class TestMoonshotResponseSchemaSupport:
|
||||
"""Every model currently live on api.moonshot.ai supports json_schema
|
||||
|
|
@ -762,10 +730,6 @@ class TestMoonshotResponseSchemaSupport:
|
|||
def model_cost_map(self):
|
||||
return GetModelCostMap.load_local_model_cost_map()
|
||||
|
||||
@pytest.mark.parametrize("model", LIVE_MODELS)
|
||||
def test_live_model_supports_response_schema(self, model, model_cost_map):
|
||||
assert model_cost_map[model].get("supports_response_schema") is True
|
||||
|
||||
def test_supports_response_schema_utility_reports_true(self, model_cost_map, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", model_cost_map)
|
||||
assert litellm.utils.supports_response_schema(model="moonshot/kimi-k2.5") is True
|
||||
|
|
|
|||
|
|
@ -110,22 +110,6 @@ class TestCognitionProviderIdentity:
|
|||
|
||||
|
||||
class TestCognitionCostTracking:
|
||||
@pytest.mark.parametrize(
|
||||
"model, input_cost, output_cost, cache_read_cost",
|
||||
[
|
||||
("cognition/swe-1.6", 5e-07, 2.5e-06, 2e-07),
|
||||
("cognition/swe-1.7", 5e-07, 2.5e-06, 2e-07),
|
||||
("cognition/swe-1.7-lightning", 2.5e-06, 1.25e-05, 1e-06),
|
||||
],
|
||||
)
|
||||
def test_cost_map_entries(self, model: str, input_cost: float, output_cost: float, cache_read_cost: float):
|
||||
info = litellm.get_model_info(model=model)
|
||||
|
||||
assert info["litellm_provider"] == "cognition"
|
||||
assert info["mode"] == "chat"
|
||||
assert info["input_cost_per_token"] == input_cost
|
||||
assert info["output_cost_per_token"] == output_cost
|
||||
assert info["cache_read_input_token_cost"] == cache_read_cost
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, expected_prompt_cost, expected_completion_cost",
|
||||
|
|
|
|||
|
|
@ -2,10 +2,9 @@
|
|||
Tests for JSON-based provider configuration system.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import patch
|
||||
|
||||
try:
|
||||
import pytest
|
||||
|
|
@ -318,25 +317,6 @@ class TestDarkbloom:
|
|||
assert config is not None
|
||||
assert config.custom_llm_provider == "darkbloom"
|
||||
|
||||
def test_darkbloom_model_cost_map(self):
|
||||
with open(
|
||||
os.path.join(workspace_path, "model_prices_and_context_window.json")
|
||||
) as f:
|
||||
model_cost = json.load(f)
|
||||
|
||||
expected_models = {
|
||||
"darkbloom/gemma-4-26b": (3e-08, 1.65e-07),
|
||||
"darkbloom/gpt-oss-20b": (1.45e-08, 7e-08),
|
||||
}
|
||||
for model, (input_cost, output_cost) in expected_models.items():
|
||||
assert model in model_cost
|
||||
assert model_cost[model]["litellm_provider"] == "darkbloom"
|
||||
assert model_cost[model]["max_output_tokens"] == 32768
|
||||
assert model_cost[model]["supports_function_calling"] is True
|
||||
assert model_cost[model]["supports_tool_choice"] is True
|
||||
assert model_cost[model]["input_cost_per_token"] == input_cost
|
||||
assert model_cost[model]["output_cost_per_token"] == output_cost
|
||||
|
||||
|
||||
class TestPublicAIIntegration:
|
||||
"""Integration tests for PublicAI provider"""
|
||||
|
|
|
|||
|
|
@ -59,23 +59,6 @@ class TestLibertAIProviderConfig:
|
|||
assert api_base == "https://custom.example.com/v1"
|
||||
assert api_key == "sk-test"
|
||||
|
||||
def test_libertai_model_cost_map(self):
|
||||
"""Test that libertai models are present in the model cost map"""
|
||||
model_cost = litellm.model_cost
|
||||
|
||||
assert "libertai/qwen3.6-27b" in model_cost
|
||||
info = model_cost["libertai/qwen3.6-27b"]
|
||||
assert info["litellm_provider"] == "libertai"
|
||||
assert info["mode"] == "chat"
|
||||
assert info["max_input_tokens"] == 262144
|
||||
assert info["max_output_tokens"] == 262144
|
||||
|
||||
# thinking variants are marked as reasoning models
|
||||
assert (
|
||||
model_cost["libertai/qwen3.6-27b-thinking"].get("supports_reasoning")
|
||||
is True
|
||||
)
|
||||
|
||||
def test_libertai_router_config(self):
|
||||
"""Test that libertai can be used in Router configuration"""
|
||||
from litellm import Router
|
||||
|
|
@ -95,20 +78,6 @@ class TestLibertAIProviderConfig:
|
|||
assert len(router.model_list) == 1
|
||||
assert router.model_list[0]["model_name"] == "libertai-chat"
|
||||
|
||||
def test_libertai_model_modes(self):
|
||||
"""Chat models carry mode 'chat'; the embedding model carries mode 'embedding'."""
|
||||
model_cost = litellm.model_cost
|
||||
|
||||
# chat model
|
||||
assert model_cost["libertai/qwen3.6-27b"]["mode"] == "chat"
|
||||
|
||||
# embedding model (bge-m3) must be normalized to mode 'embedding' so
|
||||
# /embeddings routing and the supported-endpoints matrix stay consistent
|
||||
assert "libertai/bge-m3" in model_cost
|
||||
bge = model_cost["libertai/bge-m3"]
|
||||
assert bge["litellm_provider"] == "libertai"
|
||||
assert bge["mode"] == "embedding"
|
||||
|
||||
def test_libertai_supported_endpoints_matrix(self):
|
||||
"""The runtime-served backup matrix (GET /public/supported_endpoints) lists libertai."""
|
||||
import json
|
||||
|
|
|
|||
|
|
@ -193,19 +193,6 @@ class TestMetaAnthropicMessages:
|
|||
|
||||
|
||||
class TestMuseSparkModelInfo:
|
||||
def test_muse_spark_pricing_and_capabilities(self):
|
||||
info = litellm.get_model_info("meta/muse-spark-1.1")
|
||||
|
||||
assert info["litellm_provider"] == "meta"
|
||||
assert info["input_cost_per_token"] == 1.25e-06
|
||||
assert info["output_cost_per_token"] == 4.25e-06
|
||||
assert info["cache_read_input_token_cost"] == 1.5e-07
|
||||
assert info["max_input_tokens"] == 1048576
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_web_search"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
|
||||
def test_muse_spark_cost_calculation(self):
|
||||
from litellm import completion_cost
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ Unit tests for Perplexity embedding transformation logic.
|
|||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import struct
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
|
@ -298,25 +297,3 @@ class TestPerplexityEmbeddingProviderConfig:
|
|||
)
|
||||
assert config is not None
|
||||
assert isinstance(config, PerplexityEmbeddingConfig)
|
||||
|
||||
|
||||
class TestPerplexityEmbeddingModelInfo:
|
||||
"""Test that Perplexity embedding models are in model_prices_and_context_window."""
|
||||
|
||||
def test_model_info_available(self):
|
||||
import litellm
|
||||
|
||||
info = litellm.get_model_info("perplexity/pplx-embed-v1-0.6b")
|
||||
assert info is not None
|
||||
assert info["mode"] == "embedding"
|
||||
assert info["max_input_tokens"] == 32768
|
||||
assert info["output_vector_size"] == 1024
|
||||
|
||||
def test_model_info_4b_available(self):
|
||||
import litellm
|
||||
|
||||
info = litellm.get_model_info("perplexity/pplx-embed-v1-4b")
|
||||
assert info is not None
|
||||
assert info["mode"] == "embedding"
|
||||
assert info["max_input_tokens"] == 32768
|
||||
assert info["output_vector_size"] == 2560
|
||||
|
|
|
|||
|
|
@ -26,7 +26,6 @@ from litellm.types.utils import (
|
|||
Usage,
|
||||
PromptTokensDetailsWrapper,
|
||||
)
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
|
||||
class TestPerplexityCostCalculator:
|
||||
|
|
@ -317,21 +316,6 @@ class TestPerplexityCostCalculator:
|
|||
|
||||
assert math.isclose(total_cost, expected_total, rel_tol=1e-6)
|
||||
|
||||
def test_model_info_access(self):
|
||||
"""Test that model info correctly returns the new cost fields."""
|
||||
model_info = get_model_info(
|
||||
model="sonar-deep-research", custom_llm_provider="perplexity"
|
||||
)
|
||||
|
||||
# Check that the new fields are accessible
|
||||
assert "citation_cost_per_token" in model_info
|
||||
assert model_info["citation_cost_per_token"] == 2e-6
|
||||
assert model_info["search_context_cost_per_query"] == {
|
||||
"search_context_size_low": 0.005,
|
||||
"search_context_size_medium": 0.005,
|
||||
"search_context_size_high": 0.005,
|
||||
}
|
||||
|
||||
@pytest.mark.parametrize("citation_tokens", [0, 10, 25, 100])
|
||||
@pytest.mark.parametrize("search_queries", [0, 1, 5, 10])
|
||||
@pytest.mark.parametrize("reasoning_tokens", [0, 15, 30])
|
||||
|
|
@ -477,37 +461,6 @@ class TestPerplexityCostCalculator:
|
|||
assert math.isclose(prompt_cost, expected_prompt, rel_tol=1e-9)
|
||||
assert math.isclose(completion_cost, expected_completion, rel_tol=1e-9)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id, usd_per_1m_input, usd_per_1m_output, usd_per_1m_cache_read",
|
||||
[
|
||||
("deepseek-v4-flash-0731", 0.13, 0.26, 0.028),
|
||||
("glm-5.2", 1.4, 4.4, 0.14),
|
||||
("kimi-k3", 3.0, 15.0, 0.3),
|
||||
("kimi-k2.7-code", 0.95, 4.0, 0.19),
|
||||
],
|
||||
)
|
||||
def test_agent_api_entries_carry_perplexity_published_rates(
|
||||
self, model_id, usd_per_1m_input, usd_per_1m_output, usd_per_1m_cache_read
|
||||
):
|
||||
"""The Agent API third-party models are priced from Perplexity's own catalog
|
||||
(GET https://api.perplexity.ai/v1/models, `pricing` in usd_per_1m_tokens).
|
||||
Perplexity's model id already starts with `perplexity/`, so the cost-map key
|
||||
doubles the prefix. Regression: glm-5.2 shipped glm-5.3's 0.26 cache-read rate,
|
||||
copied from the neighbouring catalog row, an 86% overcharge on cached input.
|
||||
"""
|
||||
info = get_model_info(
|
||||
model=f"perplexity/{model_id}", custom_llm_provider="perplexity"
|
||||
)
|
||||
|
||||
assert info["key"] == f"perplexity/perplexity/{model_id}"
|
||||
assert info["litellm_provider"] == "perplexity"
|
||||
assert info["mode"] == "responses"
|
||||
assert math.isclose(info["input_cost_per_token"], usd_per_1m_input / 1e6, rel_tol=1e-9)
|
||||
assert math.isclose(info["output_cost_per_token"], usd_per_1m_output / 1e6, rel_tol=1e-9)
|
||||
assert math.isclose(
|
||||
info["cache_read_input_token_cost"], usd_per_1m_cache_read / 1e6, rel_tol=1e-9
|
||||
)
|
||||
|
||||
def test_agent_api_fallback_rates_price_a_response_without_metered_cost(self):
|
||||
"""Perplexity meters cost on the response, but when `usage.cost` is absent the
|
||||
calculator falls back to the mapped per-token rates. Regression: that fallback
|
||||
|
|
|
|||
|
|
@ -136,19 +136,6 @@ class TestVertexAIVideoConfig:
|
|||
# Should NOT include endpoint
|
||||
assert not url.endswith(":predictLongRunning")
|
||||
|
||||
def test_veo_31_lite_model_cost_entries_match_pricing(self):
|
||||
for path in (ROOT_MODEL_COST_PATH, BACKUP_MODEL_COST_PATH):
|
||||
model_cost = _load_model_cost_map(path)
|
||||
info = model_cost.get(VEO_31_LITE_VERTEX_MODEL)
|
||||
|
||||
assert info is not None, f"{VEO_31_LITE_VERTEX_MODEL} missing from {path}"
|
||||
assert info["litellm_provider"] == "vertex_ai-video-models"
|
||||
assert info["mode"] == "video_generation"
|
||||
assert info["max_input_tokens"] == 1024
|
||||
assert info["output_cost_per_second"] == 0.05
|
||||
assert info["output_cost_per_second_1080p"] == 0.08
|
||||
assert info["supported_modalities"] == ["text", "image"]
|
||||
|
||||
def test_veo_31_lite_provider_routing_from_local_model_map(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
|
|
|
|||
|
|
@ -63,11 +63,6 @@ TIER_COST_FIELDS = (
|
|||
"output_cost_per_token_above_200k_tokens",
|
||||
"cache_read_input_token_cost_above_200k_tokens",
|
||||
)
|
||||
STALE_TIER_FIELDS = (
|
||||
"input_cost_per_token_above_128k_tokens",
|
||||
"output_cost_per_token_above_128k_tokens",
|
||||
"cache_read_input_token_cost_above_128k_tokens",
|
||||
)
|
||||
|
||||
|
||||
def expected_retirement_date(slug: str) -> str:
|
||||
|
|
@ -102,11 +97,10 @@ def test_redirected_slug_keeps_its_retirement_date(cost_map: dict, slug: str):
|
|||
assert cost_map[slug]["deprecation_date"] == expected_retirement_date(slug)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("slug", REDIRECTED_SLUGS)
|
||||
def test_no_slug_keeps_the_superseded_128k_tier(cost_map: dict, slug: str):
|
||||
"""The 128k tier belonged to the retired model; grok-4.3 tiers at 200k."""
|
||||
for field in STALE_TIER_FIELDS:
|
||||
assert field not in cost_map[slug], field
|
||||
def test_a_live_xai_model_is_untouched(cost_map: dict):
|
||||
"""Guard against the repricing leaking onto models xAI still serves directly."""
|
||||
assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"]
|
||||
assert "deprecation_date" not in cost_map["xai/grok-4.6"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("slug", REDIRECTED_SLUGS)
|
||||
|
|
@ -116,12 +110,7 @@ def test_redirected_slug_carries_the_target_tier_rates(cost_map: dict, slug: str
|
|||
entry = cost_map[slug]
|
||||
for field in TIER_COST_FIELDS:
|
||||
assert entry[field] == target[field], field
|
||||
|
||||
|
||||
def test_a_live_xai_model_is_untouched(cost_map: dict):
|
||||
"""Guard against the repricing leaking onto models xAI still serves directly."""
|
||||
assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"]
|
||||
assert "deprecation_date" not in cost_map["xai/grok-4.6"]
|
||||
assert {k for k in entry if "_above_" in k} == {k for k in target if "_above_" in k}
|
||||
|
||||
|
||||
def test_both_cost_maps_agree_on_the_redirected_slugs():
|
||||
|
|
|
|||
|
|
@ -2,11 +2,9 @@
|
|||
Tests for Z.AI (Zhipu AI) provider - GLM models
|
||||
"""
|
||||
|
||||
import json
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm import completion
|
||||
|
|
@ -57,32 +55,9 @@ def test_zai_in_provider_lists():
|
|||
assert "zai" in litellm.provider_list
|
||||
|
||||
|
||||
def test_zai_models_in_model_cost(local_model_cost_map):
|
||||
"""Test that ZAI models are in the model cost map"""
|
||||
|
||||
zai_models = [
|
||||
"zai/glm-4.7",
|
||||
"zai/glm-4.6",
|
||||
"zai/glm-4.5",
|
||||
"zai/glm-4.5v",
|
||||
"zai/glm-4.5-x",
|
||||
"zai/glm-4.5-air",
|
||||
"zai/glm-4.5-airx",
|
||||
"zai/glm-4-32b-0414-128k",
|
||||
"zai/glm-4.5-flash",
|
||||
]
|
||||
|
||||
for model in zai_models:
|
||||
assert model in litellm.model_cost, f"Model {model} not found in model_cost"
|
||||
assert litellm.model_cost[model]["litellm_provider"] == "zai"
|
||||
|
||||
|
||||
def test_zai_glm46_cost_calculation(local_model_cost_map):
|
||||
"""Test the cost calculation for glm-4.6"""
|
||||
|
||||
key = "zai/glm-4.6"
|
||||
info = litellm.model_cost[key]
|
||||
|
||||
prompt_cost, completion_cost = cost_per_token(
|
||||
model="zai/glm-4.6",
|
||||
prompt_tokens=1000000, # 1M tokens
|
||||
|
|
@ -94,26 +69,6 @@ def test_zai_glm46_cost_calculation(local_model_cost_map):
|
|||
assert math.isclose(completion_cost, 2.2, rel_tol=1e-6)
|
||||
|
||||
|
||||
def test_zai_flash_model_is_free(local_model_cost_map):
|
||||
"""Test that glm-4.5-flash has zero cost"""
|
||||
|
||||
key = "zai/glm-4.5-flash"
|
||||
info = litellm.model_cost[key]
|
||||
|
||||
assert info["input_cost_per_token"] == 0
|
||||
assert info["output_cost_per_token"] == 0
|
||||
|
||||
|
||||
def test_glm47_supports_reasoning(local_model_cost_map):
|
||||
"""Test that GLM-4.7 supports reasoning"""
|
||||
|
||||
key = "zai/glm-4.7"
|
||||
assert key in litellm.model_cost, f"Model {key} not found in model_cost"
|
||||
|
||||
info = litellm.model_cost[key]
|
||||
assert info["supports_reasoning"] is True
|
||||
|
||||
|
||||
def test_glm47_cost_calculation(local_model_cost_map):
|
||||
"""Test cost calculation for GLM-4.7"""
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from starlette.datastructures import Headers
|
|||
|
||||
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
MCPRequestHandler,
|
||||
UnloadableEntitlementError,
|
||||
_is_mcp_admitted_user_subject,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
|
|
@ -4396,6 +4397,147 @@ class TestAgentMCPPermissions:
|
|||
)
|
||||
assert sorted(result) == ["tool_a", "tool_b"]
|
||||
|
||||
def _agent_object_permission(self, *, toolset_ids, servers=(), tool_permissions=None):
|
||||
agent_object_permission = MagicMock()
|
||||
agent_object_permission.mcp_servers = list(servers)
|
||||
agent_object_permission.mcp_access_groups = []
|
||||
agent_object_permission.mcp_tool_permissions = tool_permissions
|
||||
agent_object_permission.mcp_toolsets = list(toolset_ids)
|
||||
return agent_object_permission
|
||||
|
||||
def _mock_manager_with_toolsets(self, toolset_perms):
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.expand_permission_list = MagicMock(side_effect=lambda servers: list(servers))
|
||||
mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {})
|
||||
mock_manager.resolve_toolset_tool_permissions = AsyncMock(return_value=toolset_perms)
|
||||
return mock_manager
|
||||
|
||||
def _agent_toolset_patches(self, agent_object_permission, mock_manager):
|
||||
return (
|
||||
patch.object( # test-quality-ok: stub the agent perm loader; the resolver reads module globals with no injection seam
|
||||
MCPRequestHandler, "_get_agent_object_permission", AsyncMock(return_value=agent_object_permission)
|
||||
),
|
||||
patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling toolset tests
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
mock_manager,
|
||||
),
|
||||
patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here
|
||||
MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])
|
||||
),
|
||||
)
|
||||
|
||||
async def test_get_allowed_mcp_servers_for_agent_includes_toolset_servers(self):
|
||||
"""An agent granted only mcp_toolsets reaches the toolset's servers, exactly as a
|
||||
key, team, or org granted only toolsets does"""
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets")
|
||||
agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"], servers=["server-direct"])
|
||||
mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]})
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
|
||||
stack.enter_context(patcher)
|
||||
result = await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth)
|
||||
|
||||
assert sorted(result) == ["server-a", "server-direct"]
|
||||
mock_manager.resolve_toolset_tool_permissions.assert_awaited_once_with(toolset_ids=["toolset-1"])
|
||||
|
||||
async def test_get_allowed_mcp_servers_toolset_only_agent_caps_key_servers(self):
|
||||
"""Regression: an agent whose only grant is a toolset used to resolve to [] and place
|
||||
no ceiling at all, so a key bound to it kept every server the key itself granted"""
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets")
|
||||
agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"])
|
||||
mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]})
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
|
||||
stack.enter_context(patcher)
|
||||
stack.enter_context(
|
||||
patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"])
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch.object( # test-quality-ok: team resolution has its own tests; pin it empty here
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[])
|
||||
)
|
||||
)
|
||||
result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
|
||||
|
||||
assert result == ["server-a"]
|
||||
|
||||
async def test_get_allowed_mcp_servers_agent_dangling_toolset_denies(self):
|
||||
"""An agent toolset that resolves to nothing is a known restriction with unknown
|
||||
contents: deny, never fall through to the key's own servers"""
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets")
|
||||
agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-gone"])
|
||||
mock_manager = self._mock_manager_with_toolsets({})
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
|
||||
stack.enter_context(patcher)
|
||||
with pytest.raises(UnloadableEntitlementError):
|
||||
await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth)
|
||||
stack.enter_context(
|
||||
patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"])
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch.object( # test-quality-ok: team resolution has its own tests; pin it empty here
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[])
|
||||
)
|
||||
)
|
||||
result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
|
||||
|
||||
assert result == []
|
||||
|
||||
async def test_get_agent_tool_permissions_for_server_unions_direct_and_toolset_tools(self):
|
||||
"""The agent's tool ceiling on a server is its direct tool grants plus the tools its
|
||||
toolsets grant there, and None only when neither names the server"""
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets")
|
||||
agent_object_permission = self._agent_object_permission(
|
||||
toolset_ids=["toolset-1"], tool_permissions={"server-a": ["tool_direct"]}
|
||||
)
|
||||
mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_via_toolset"], "server-b": ["tool_b"]})
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
|
||||
stack.enter_context(patcher)
|
||||
server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-a", user_api_key_auth)
|
||||
server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-b", user_api_key_auth)
|
||||
server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-c", user_api_key_auth)
|
||||
|
||||
assert sorted(server_a_tools) == ["tool_direct", "tool_via_toolset"]
|
||||
assert server_b_tools == ["tool_b"]
|
||||
assert server_c_tools is None
|
||||
|
||||
async def test_get_allowed_tools_for_server_toolset_only_agent_caps_key_tools(self):
|
||||
"""Regression: a key allowing [tool_a, tool_b] bound to an agent whose toolset grants
|
||||
only tool_a on the server ends with [tool_a]; the toolset used to be ignored"""
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets")
|
||||
agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"])
|
||||
mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_a"]})
|
||||
key_perm = MagicMock()
|
||||
key_perm.mcp_tool_permissions = {"server-a": ["tool_a", "tool_b"]}
|
||||
key_perm.mcp_toolsets = []
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
|
||||
stack.enter_context(patcher)
|
||||
stack.enter_context(
|
||||
patch.object( # test-quality-ok: stub the key perm loader; the resolver reads module globals with no injection seam
|
||||
MCPRequestHandler, "_get_key_object_permission", return_value=key_perm
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch.object( # test-quality-ok: team resolution has its own tests; pin it absent here
|
||||
MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)
|
||||
)
|
||||
)
|
||||
result = await MCPRequestHandler.get_allowed_tools_for_server("server-a", user_api_key_auth)
|
||||
|
||||
assert result == ["tool_a"]
|
||||
|
||||
async def test_get_agent_object_permission_uses_shared_helper(self):
|
||||
"""``_get_agent_object_permission`` must resolve the agent's
|
||||
``object_permission_id`` and then defer to the shared
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import contextvars
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -1368,6 +1369,295 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing():
|
|||
mock_logger.info.assert_any_call("Successfully fetched %s tools total from all MCP servers", 0)
|
||||
|
||||
|
||||
def _denied_scope_manager(known_server_names_to_ids: dict[str, str]) -> MagicMock:
|
||||
servers = {name: MagicMock(server_id=server_id) for name, server_id in known_server_names_to_ids.items()}
|
||||
manager = MagicMock()
|
||||
manager.get_mcp_server_by_name = lambda name, client_ip=None: servers.get(name)
|
||||
return manager
|
||||
|
||||
|
||||
def _scope_resolver(resolved_without_agent: dict[str, str], access_groups: tuple[str, ...] = ()) -> AsyncMock:
|
||||
async def resolve(user_api_key_auth, mcp_servers, client_ip=None):
|
||||
if user_api_key_auth is not None and user_api_key_auth.agent_id:
|
||||
return []
|
||||
return [
|
||||
SimpleNamespace(
|
||||
server_id=server_id,
|
||||
server_name=server_name,
|
||||
alias=None,
|
||||
short_prefix=None,
|
||||
access_groups=list(access_groups),
|
||||
)
|
||||
for server_name, server_id in resolved_without_agent.items()
|
||||
]
|
||||
|
||||
return AsyncMock(side_effect=resolve)
|
||||
|
||||
|
||||
async def _denied_scoped_list(
|
||||
user_api_key_auth: UserAPIKeyAuth,
|
||||
mcp_servers: list[str],
|
||||
mock_manager: MagicMock,
|
||||
resolver: AsyncMock,
|
||||
) -> HTTPException:
|
||||
from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
|
||||
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
|
||||
resolver,
|
||||
),
|
||||
patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
|
||||
mock_manager,
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _get_tools_from_mcp_servers(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
mcp_auth_header=None,
|
||||
mcp_servers=mcp_servers,
|
||||
)
|
||||
return exc_info.value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_list_denied_by_agent_binding_raises_403_naming_agent():
|
||||
"""The agent-binding veto must raise a 403 naming the agent, never a silent 200 with no tools."""
|
||||
pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
|
||||
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
|
||||
resolver = _scope_resolver(resolved_without_agent={"github": "srv-github"})
|
||||
|
||||
denial = await _denied_scoped_list(
|
||||
user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver
|
||||
)
|
||||
|
||||
assert denial.status_code == 403
|
||||
message = denial.detail["error"]
|
||||
assert "MCP server 'github'" in message
|
||||
assert "agent 'agent-123'" in message
|
||||
assert "mcp_servers" in message
|
||||
rerun_auth = resolver.await_args_list[1].kwargs["user_api_key_auth"]
|
||||
assert rerun_auth.agent_id is None
|
||||
assert rerun_auth.user_id == "test_user"
|
||||
assert resolver.await_args_list[1].kwargs["mcp_servers"] == ["github"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_scope_lists_nothing_instead_of_raising_a_nameless_denial():
|
||||
"""An empty ``x-mcp-servers`` header scopes to no servers; that is an empty listing, not a 403."""
|
||||
pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
|
||||
from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
|
||||
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
|
||||
resolver = AsyncMock(return_value=[])
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
|
||||
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
|
||||
resolver,
|
||||
),
|
||||
patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
|
||||
_denied_scope_manager({"github": "srv-github"}),
|
||||
),
|
||||
):
|
||||
listing = await _get_tools_from_mcp_servers(
|
||||
user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=[]
|
||||
)
|
||||
|
||||
assert listing.tools == []
|
||||
resolver.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_list_denied_for_non_agent_key_raises_generic_403():
|
||||
"""A denial for a key with no agent binding stays generic and skips the agent-stripped rerun."""
|
||||
pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
|
||||
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
|
||||
resolver = AsyncMock(return_value=[])
|
||||
|
||||
denial = await _denied_scoped_list(
|
||||
user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver
|
||||
)
|
||||
|
||||
assert denial.status_code == 403
|
||||
message = denial.detail["error"]
|
||||
assert "github" in message
|
||||
assert "agent" not in message
|
||||
resolver.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_list_unknown_name_raises_same_generic_403_as_unauthorized():
|
||||
"""Unknown and registered-but-unauthorized names raise byte-identical generic 403s, so a
|
||||
caller cannot probe which server names exist."""
|
||||
pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
|
||||
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
|
||||
|
||||
unknown = await _denied_scoped_list(
|
||||
user_api_key_auth, ["github"], _denied_scope_manager({}), _scope_resolver(resolved_without_agent={})
|
||||
)
|
||||
unauthorized = await _denied_scoped_list(
|
||||
user_api_key_auth,
|
||||
["github"],
|
||||
_denied_scope_manager({"github": "srv-github"}),
|
||||
_scope_resolver(resolved_without_agent={}),
|
||||
)
|
||||
|
||||
assert unknown.status_code == unauthorized.status_code == 403
|
||||
assert unknown.detail["error"] == unauthorized.detail["error"]
|
||||
assert "github" in unknown.detail["error"]
|
||||
assert "agent" not in unknown.detail["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_list_access_group_vetoed_by_agent_names_agent_and_group():
|
||||
"""An access-group scope vetoed by the agent binding raises the 403 naming the agent and the
|
||||
group instead of the silent empty list."""
|
||||
pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
|
||||
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
|
||||
|
||||
denial = await _denied_scoped_list(
|
||||
user_api_key_auth,
|
||||
["prod-group"],
|
||||
_denied_scope_manager({}),
|
||||
_scope_resolver(resolved_without_agent={"github": "srv-github"}, access_groups=("prod-group",)),
|
||||
)
|
||||
|
||||
assert denial.status_code == 403
|
||||
message = denial.detail["error"]
|
||||
assert "access group 'prod-group'" in message
|
||||
assert "agent 'agent-123'" in message
|
||||
assert "mcp_access_groups" in message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_list_mixed_unknown_and_vetoed_group_names_the_group_that_resolved():
|
||||
"""With an unknown name ahead of the agent-vetoed group in the scope, the 403 must name the group
|
||||
whose servers the key can reach, never the unknown name, or the admin is told to grant a group
|
||||
that does not exist."""
|
||||
pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
|
||||
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
|
||||
|
||||
denial = await _denied_scoped_list(
|
||||
user_api_key_auth,
|
||||
["no-such-group", "prod-group"],
|
||||
_denied_scope_manager({}),
|
||||
_scope_resolver(resolved_without_agent={"github": "srv-github"}, access_groups=("prod-group",)),
|
||||
)
|
||||
|
||||
assert denial.status_code == 403
|
||||
message = denial.detail["error"]
|
||||
assert "access group 'prod-group'" in message
|
||||
assert "no-such-group" not in message
|
||||
assert "agent 'agent-123'" in message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_list_agent_key_denied_by_key_grants_raises_generic_403():
|
||||
"""When the agent-stripped rerun still resolves nothing, the 403 stays generic instead of
|
||||
blaming the agent binding."""
|
||||
pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
|
||||
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
|
||||
resolver = _scope_resolver(resolved_without_agent={})
|
||||
|
||||
denial = await _denied_scoped_list(
|
||||
user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver
|
||||
)
|
||||
|
||||
assert denial.status_code == 403
|
||||
message = denial.detail["error"]
|
||||
assert "github" in message
|
||||
assert "agent" not in message
|
||||
assert resolver.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_list_agent_veto_attributed_for_differently_cased_server_name():
|
||||
"""The scope filter matches `/mcp/GitHub` to a server named `github` case-insensitively, so the
|
||||
agent-attributed 403 must match the same way instead of falling back to the generic denial."""
|
||||
pytest.importorskip("litellm.proxy._experimental.mcp_server.server")
|
||||
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123")
|
||||
|
||||
denial = await _denied_scoped_list(
|
||||
user_api_key_auth,
|
||||
["GitHub"],
|
||||
_denied_scope_manager({"github": "srv-github"}),
|
||||
_scope_resolver(resolved_without_agent={"github": "srv-github"}),
|
||||
)
|
||||
|
||||
assert denial.status_code == 403
|
||||
message = denial.detail["error"]
|
||||
assert "MCP server 'GitHub'" in message
|
||||
assert "agent 'agent-123'" in message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error():
|
||||
"""The MCP protocol handler surfaces a permission HTTPException as a clean JSON-RPC error
|
||||
(McpError, INVALID_REQUEST) carrying the denial message, instead of a raw 500."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import handle_list_tools
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.types import INVALID_REQUEST
|
||||
|
||||
denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'"
|
||||
denial = HTTPException(status_code=403, detail={"error": denial_message})
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the protocol handler reads auth from module context; no injection seam
|
||||
"litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context",
|
||||
new=AsyncMock(return_value=(None, None, None, None, None, None, None)),
|
||||
),
|
||||
patch( # test-quality-ok: the listing helper is the handler's only collaborator; the suite's seam
|
||||
"litellm.proxy._experimental.mcp_server.server._list_mcp_tools",
|
||||
new=AsyncMock(side_effect=denial),
|
||||
),
|
||||
):
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await handle_list_tools()
|
||||
|
||||
assert exc_info.value.error.code == INVALID_REQUEST
|
||||
assert exc_info.value.error.message == denial_message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict():
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import mcp_server_tool_call
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'"
|
||||
denial = HTTPException(status_code=403, detail={"error": denial_message})
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the protocol handler reads auth from module context; no injection seam
|
||||
"litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context",
|
||||
new=AsyncMock(return_value=(None, None, None, None, None, None, None)),
|
||||
),
|
||||
patch( # test-quality-ok: the tool-call helper is the handler's only collaborator; the suite's seam
|
||||
"litellm.proxy._experimental.mcp_server.server.call_mcp_tool",
|
||||
new=AsyncMock(side_effect=denial),
|
||||
),
|
||||
):
|
||||
result = await mcp_server_tool_call("github-search_issues", {})
|
||||
|
||||
assert result.isError is True
|
||||
assert result.content[0].text == f"Error: {denial_message}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_server_tool_call_body_with_none_arguments():
|
||||
"""Test that proxy_server_request body handles None arguments correctly"""
|
||||
|
|
@ -3518,6 +3808,35 @@ async def test_call_mcp_tool_user_unauthorized_access():
|
|||
assert "User not allowed to call this tool" in exc_info.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_mcp_tool_scoped_denial_names_the_binding_agent():
|
||||
from litellm.proxy._experimental.mcp_server.server import call_mcp_tool
|
||||
|
||||
agent_bound_key = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-123")
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_allowed_mcp_servers",
|
||||
AsyncMock(return_value=[]),
|
||||
),
|
||||
patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
|
||||
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
|
||||
_scope_resolver({"github": "srv-github"}),
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await call_mcp_tool(
|
||||
name="github-search_issues",
|
||||
arguments={},
|
||||
user_api_key_auth=agent_bound_key,
|
||||
mcp_servers=["github"],
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "MCP server 'github'" in exc_info.value.detail["error"]
|
||||
assert "agent 'agent-123'" in exc_info.value.detail["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_mcp_tool_unauthorized_403_does_not_leak_server_credentials():
|
||||
"""Regression for LIT-4703 / GH #29936.
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ Covers:
|
|||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -1271,3 +1272,37 @@ class TestMcpServerToolCallErrorHandling:
|
|||
|
||||
assert result.isError is True
|
||||
assert "User not allowed to call this tool" in result.content[0].text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_mcp_tool_call_scoped_denial_names_the_binding_agent() -> None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.tool_search import handle_mcp_tool_call
|
||||
|
||||
agent_bound_key = UserAPIKeyAuth(api_key="test_key", agent_id="agent-123")
|
||||
|
||||
async def resolve(user_api_key_auth, mcp_servers, client_ip=None):
|
||||
if user_api_key_auth.agent_id:
|
||||
return []
|
||||
return [
|
||||
SimpleNamespace(
|
||||
server_id="srv-github", server_name="github", alias=None, short_prefix=None, access_groups=[]
|
||||
)
|
||||
]
|
||||
|
||||
with patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
|
||||
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
|
||||
new=AsyncMock(side_effect=resolve),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await handle_mcp_tool_call(
|
||||
tool_name="github-create_issue",
|
||||
arguments={},
|
||||
user_api_key_dict=agent_bound_key,
|
||||
mcp_servers=["github"],
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "MCP server 'github'" in exc_info.value.detail["error"]
|
||||
assert "agent 'agent-123'" in exc_info.value.detail["error"]
|
||||
|
|
|
|||
|
|
@ -2216,3 +2216,26 @@ async def test_top_k_above_router_default_is_respected():
|
|||
|
||||
assert len(filtered) == 6
|
||||
print("✅ Configured top_k above the semantic-router default of 5 is honored")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves():
|
||||
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
|
||||
|
||||
gateway_reference = {"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"}
|
||||
external_tool = {
|
||||
"type": "mcp",
|
||||
"server_label": "zapier",
|
||||
"server_url": "https://mcp.zapier.com/api/mcp/mcp",
|
||||
"allowed_tools": ["zapier_send_email"],
|
||||
}
|
||||
|
||||
async def served_names(names):
|
||||
assert names == {"mcp"}
|
||||
return frozenset()
|
||||
|
||||
narrowed = await SemanticToolFilterHook._narrow_mcp_references(
|
||||
[gateway_reference, external_tool], ["srv-tool_1"], served_names=served_names
|
||||
)
|
||||
|
||||
assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_tool]
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext
|
|||
"""
|
||||
|
||||
import base64
|
||||
import logging
|
||||
from typing import Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -15,6 +16,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
abbreviate_api_key,
|
||||
check_complete_credentials,
|
||||
custom_auth_common_checks_warning,
|
||||
log_once_if_budget_reservation_disabled,
|
||||
warn_once_if_custom_auth_skips_common_checks,
|
||||
get_end_user_id_from_request_body,
|
||||
get_key_mcp_rpm_limit,
|
||||
|
|
@ -101,6 +103,41 @@ class TestWarnOnceIfCustomAuthSkipsCommonChecks:
|
|||
assert logger.warning.call_count == 0
|
||||
|
||||
|
||||
class TestLogOnceIfBudgetReservationDisabled:
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_sentinel(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"litellm.constants.budget_reservation_disabled_info_emitted",
|
||||
False,
|
||||
)
|
||||
|
||||
def test_logs_info_only_once_when_enabled(self, caplog):
|
||||
with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
|
||||
log_once_if_budget_reservation_disabled(disabled=False)
|
||||
assert not any(
|
||||
"disable_budget_reservation is enabled" in record.message
|
||||
for record in caplog.records
|
||||
)
|
||||
for _ in range(3):
|
||||
log_once_if_budget_reservation_disabled(disabled=True)
|
||||
|
||||
records = [
|
||||
record
|
||||
for record in caplog.records
|
||||
if "disable_budget_reservation is enabled" in record.message
|
||||
]
|
||||
assert len(records) == 1
|
||||
assert records[0].levelno == logging.INFO
|
||||
|
||||
def test_logs_to_injected_logger_only_once(self):
|
||||
logger = MagicMock()
|
||||
log_once_if_budget_reservation_disabled(disabled=False, logger=logger)
|
||||
for _ in range(3):
|
||||
log_once_if_budget_reservation_disabled(disabled=True, logger=logger)
|
||||
assert logger.info.call_count == 1
|
||||
assert "disable_budget_reservation is enabled" in logger.info.call_args[0][0]
|
||||
|
||||
|
||||
class TestGetKeyModelRpmLimit:
|
||||
"""Tests for get_key_model_rpm_limit function."""
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
|
|
@ -146,6 +147,35 @@ async def test_disable_budget_reservation_skips_reservation():
|
|||
assert user_api_key_auth_obj.budget_reservation is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disable_budget_reservation_does_not_log_per_request(caplog):
|
||||
user_api_key_auth_obj = UserAPIKeyAuth(token="test_token")
|
||||
|
||||
with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
|
||||
for _ in range(3):
|
||||
await _reserve_budget_after_common_checks(
|
||||
user_api_key_auth_obj=user_api_key_auth_obj,
|
||||
request_data={"model": "gpt-4o"},
|
||||
route="/v1/chat/completions",
|
||||
llm_router=None,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
skip_budget_checks=False,
|
||||
general_settings={"disable_budget_reservation": True},
|
||||
)
|
||||
|
||||
records = [
|
||||
record
|
||||
for record in caplog.records
|
||||
if "disable_budget_reservation is enabled" in record.message
|
||||
]
|
||||
assert records == []
|
||||
assert user_api_key_auth_obj.budget_reservation is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_reservation_runs_when_not_disabled():
|
||||
"""Control for #27639: with the flag absent, the reservation still runs and is stored."""
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ Pins covered:
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
|
|
@ -1633,6 +1634,32 @@ async def test_ProxyConfig_load_config_minimal_yaml(tmp_path, monkeypatch):
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("setting", ["true", "false", "null", "'true'", None])
|
||||
async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monkeypatch, caplog, setting):
|
||||
config_file = tmp_path / "budget.yaml"
|
||||
flag = f" disable_budget_reservation: {setting}\n" if setting is not None else ""
|
||||
config_file.write_text(
|
||||
"model_list: []\nlitellm_settings: {}\ngeneral_settings:\n"
|
||||
" master_key: null\n" + flag
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
|
||||
monkeypatch.setattr("litellm.constants.budget_reservation_disabled_info_emitted", False)
|
||||
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
|
||||
config = ProxyConfig()
|
||||
|
||||
with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
|
||||
for _ in range(3):
|
||||
await config.load_config(router=None, config_file_path=str(config_file))
|
||||
|
||||
records = [
|
||||
record for record in caplog.records
|
||||
if "disable_budget_reservation is enabled" in record.message
|
||||
]
|
||||
assert [record.levelno for record in records] == ([logging.INFO] if setting == "true" else [])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_load_config_resolves_router_settings_plugins(tmp_path, monkeypatch):
|
||||
"""Regression: router_settings.plugins dotted-path strings must be resolved to
|
||||
|
|
|
|||
|
|
@ -1,6 +1,13 @@
|
|||
import json
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
from httpx import Response
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
from litellm.responses.mcp import chat_completions_handler
|
||||
|
|
@ -1344,3 +1351,39 @@ async def test_acompletion_with_mcp_streaming_drains_inner_stream_after_exhausti
|
|||
|
||||
assert len(all_chunks) == 3
|
||||
assert initial_stream.drained_after_exhaustion is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_acompletion_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch):
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
|
||||
|
||||
zapier_tool = {"type": "mcp", "server_label": "zapier", "server_url": "https://mcp.zapier.com/api/mcp/mcp"}
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None))
|
||||
monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {})
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
provider = respx.post("https://api.openai.com/v1/chat/completions").mock(
|
||||
return_value=Response(
|
||||
200,
|
||||
json={
|
||||
"id": "chatcmpl-zapier",
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": "gpt-4.1",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
result = await acompletion_with_mcp(
|
||||
model="openai/gpt-4.1",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=[zapier_tool],
|
||||
api_key="sk-test",
|
||||
acompletion=True,
|
||||
)
|
||||
|
||||
assert isinstance(result, ModelResponse)
|
||||
assert result.id == "chatcmpl-zapier"
|
||||
assert json.loads(provider.calls.last.request.content)["tools"] == [zapier_tool]
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from openai.types.responses.tool_param import Mcp
|
||||
import importlib
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing
|
||||
|
|
@ -724,6 +725,178 @@ def test_extract_tool_call_details_still_prefers_openai_arguments():
|
|||
assert arguments == '{"city": "Paris"}'
|
||||
|
||||
|
||||
def _registered(
|
||||
server_id: str,
|
||||
name: str,
|
||||
alias: str | None = None,
|
||||
server_name: str | None = None,
|
||||
access_groups: list[str] | None = None,
|
||||
):
|
||||
from litellm.types.mcp import MCPTransport
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
return MCPServer(
|
||||
server_id=server_id,
|
||||
name=name,
|
||||
alias=alias,
|
||||
server_name=server_name,
|
||||
transport=MCPTransport.http,
|
||||
access_groups=access_groups,
|
||||
)
|
||||
|
||||
|
||||
async def _no_toolset(_: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
ZAPIER_TOOL: Mcp = {
|
||||
"type": "mcp",
|
||||
"server_label": "zapier",
|
||||
"server_url": "https://mcp.zapier.com/api/mcp/mcp",
|
||||
"require_approval": "never",
|
||||
}
|
||||
EXPLICIT_GATEWAY_TOOL = {"type": "mcp", "server_label": "github", "server_url": "litellm_proxy/mcp/github"}
|
||||
FUNCTION_TOOL = {"type": "function", "name": "get_weather", "parameters": {}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gateway_served_names_matches_alias_server_name_name_access_group_and_toolset():
|
||||
from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names
|
||||
|
||||
servers = (
|
||||
_registered("id-1", "github-name", alias="github", server_name="github-server", access_groups=["prod-group"]),
|
||||
_registered("id-2", "deepwiki"),
|
||||
)
|
||||
|
||||
async def toolset_exists(name: str) -> bool:
|
||||
return name == "my-toolset"
|
||||
|
||||
served = await _gateway_served_names(
|
||||
{"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset", "mcp", "nope"},
|
||||
servers=lambda: servers,
|
||||
toolset_exists=toolset_exists,
|
||||
)
|
||||
|
||||
assert served == {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gateway_served_names_matches_server_id_short_prefix_and_alias_case_like_the_gateway():
|
||||
from litellm.proxy._experimental.mcp_server.utils import compute_short_server_prefix
|
||||
from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names
|
||||
|
||||
server_id = "0b9ae4ca-1bd2-4faa-b183-7dd812597e3b"
|
||||
short_prefix = compute_short_server_prefix(server_id)
|
||||
servers = (_registered(server_id, "github-name", alias="github", access_groups=["prod-group"]),)
|
||||
|
||||
served = await _gateway_served_names(
|
||||
{server_id, short_prefix, "GitHub", "PROD-GROUP", "nope"}, servers=lambda: servers, toolset_exists=_no_toolset
|
||||
)
|
||||
|
||||
assert served == {server_id, short_prefix, "GitHub"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routes_through_gateway_flags_explicit_and_served_tools_only():
|
||||
served_tool = {"type": "mcp", "server_label": "github", "server_url": "http://localhost:4000/mcp/github"}
|
||||
|
||||
async def served_names(names):
|
||||
assert names == {"github", "mcp"}
|
||||
return frozenset({"github"})
|
||||
|
||||
flags = await LiteLLM_Proxy_MCP_Handler.routes_through_gateway(
|
||||
[ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, served_tool, FUNCTION_TOOL], served_names=served_names
|
||||
)
|
||||
|
||||
assert flags == (False, True, True, False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_split_mcp_tools_leaves_external_mcp_path_urls_for_the_provider():
|
||||
|
||||
async def served_names(names):
|
||||
assert names == {"mcp"}
|
||||
return frozenset()
|
||||
|
||||
gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(
|
||||
[ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names
|
||||
)
|
||||
|
||||
assert gateway_tools == [EXPLICIT_GATEWAY_TOOL]
|
||||
assert other_tools == [ZAPIER_TOOL, FUNCTION_TOOL]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_split_mcp_tools_repoints_served_proxy_urls_at_the_gateway():
|
||||
served_tool = {
|
||||
"type": "mcp",
|
||||
"server_label": "toolset",
|
||||
"server_url": "http://localhost:4000/mcp/my-toolset",
|
||||
"require_approval": "never",
|
||||
"allowed_tools": ["get_me"],
|
||||
}
|
||||
unserved_tool = {"type": "mcp", "server_label": "typo", "server_url": "http://localhost:4000/mcp/githb"}
|
||||
|
||||
async def served_names(names):
|
||||
return frozenset({"my-toolset"})
|
||||
|
||||
gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(
|
||||
[served_tool, unserved_tool], served_names=served_names
|
||||
)
|
||||
|
||||
assert gateway_tools == [{**served_tool, "server_url": "litellm_proxy/mcp/my-toolset"}]
|
||||
assert other_tools == [unserved_tool]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_split_mcp_tools_skips_resolution_when_nothing_points_at_the_proxy():
|
||||
async def served_names(names):
|
||||
raise AssertionError("no lookup expected")
|
||||
|
||||
gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(
|
||||
[EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names
|
||||
)
|
||||
|
||||
assert gateway_tools == [EXPLICIT_GATEWAY_TOOL]
|
||||
assert other_tools == [FUNCTION_TOOL]
|
||||
|
||||
|
||||
def test_should_use_gateway_still_triggers_on_http_mcp_path():
|
||||
assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([ZAPIER_TOOL]) is True
|
||||
assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([EXPLICIT_GATEWAY_TOOL]) is True
|
||||
assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([FUNCTION_TOOL]) is False
|
||||
assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(None) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aresponses_api_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch):
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
|
||||
from litellm.responses import main as responses_main
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None))
|
||||
monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {})
|
||||
provider_tools: list[object] = []
|
||||
|
||||
def fake_provider(**kwargs: object) -> object:
|
||||
request_params = cast(dict[str, object], kwargs["response_api_optional_request_params"])
|
||||
provider_tools.append(request_params.get("tools"))
|
||||
|
||||
async def respond() -> ResponsesAPIResponse:
|
||||
return ResponsesAPIResponse(id="resp_zapier", created_at=0, output=[])
|
||||
|
||||
return respond()
|
||||
|
||||
monkeypatch.setattr(responses_main.base_llm_http_handler, "response_api_handler", fake_provider)
|
||||
|
||||
response = await responses_main.aresponses_api_with_mcp(
|
||||
input="Reply with the single word ok.", model="openai/gpt-4.1", tools=[ZAPIER_TOOL]
|
||||
)
|
||||
|
||||
assert isinstance(response, ResponsesAPIResponse)
|
||||
assert provider_tools == [[ZAPIER_TOOL]]
|
||||
|
||||
|
||||
def _response_with_reasoning_and_tool_call() -> Any:
|
||||
"""A first-turn response as a reasoning model returns it: reasoning item, then a function call."""
|
||||
return ResponsesAPIResponse(
|
||||
|
|
|
|||
|
|
@ -7,7 +7,8 @@ Tests the rule-based complexity scoring and tier assignment logic.
|
|||
import asyncio
|
||||
import logging
|
||||
import sys
|
||||
from typing import Dict, List
|
||||
import time
|
||||
from typing import Dict, Final, List
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -47,6 +48,7 @@ from litellm.router_strategy.complexity_router.config import (
|
|||
ClassifierLLMConfig,
|
||||
ComplexityRouterConfig,
|
||||
ComplexityTier,
|
||||
custom_pattern_work,
|
||||
)
|
||||
from litellm.router_strategy.complexity_router.tier_predictor import (
|
||||
TierGlobalStatistic,
|
||||
|
|
@ -756,6 +758,205 @@ class TestCustomTechnicalKeywords:
|
|||
assert custom_score > baseline_score
|
||||
|
||||
|
||||
class TestCustomDimensions:
|
||||
@pytest.mark.parametrize(
|
||||
"matchers,prompt",
|
||||
[
|
||||
pytest.param(
|
||||
{"keywords": ["orbitmesh", "fluxgate"]},
|
||||
"Connect ORBITMESH and fluxgate for the requested change",
|
||||
id="keywords",
|
||||
),
|
||||
pytest.param(
|
||||
{"patterns": [r"\bCREATE\s{1,4}TABLE\b", r"\bALTER\s{1,4}TABLE\b"]},
|
||||
"create table widgets (id integer); ALTER TABLE widgets ADD label text;",
|
||||
id="regex",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_custom_dimension_changes_only_matching_requests(
|
||||
self, mock_router_instance: MagicMock, matchers: dict[str, object], prompt: str
|
||||
) -> None:
|
||||
baseline: Final = ComplexityRouter("test-router", mock_router_instance)
|
||||
configured: Final = ComplexityRouter(
|
||||
"test-router",
|
||||
mock_router_instance,
|
||||
{"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, **matchers}]},
|
||||
)
|
||||
baseline_tier, baseline_score, baseline_signals = baseline.classify(prompt)
|
||||
tier, score, signals = configured.classify(prompt)
|
||||
assert baseline_tier == ComplexityTier.SIMPLE
|
||||
assert tier != ComplexityTier.SIMPLE
|
||||
assert score == pytest.approx(baseline_score + 0.7)
|
||||
assert signals == [*baseline_signals, "custom (internalFrameworks)"]
|
||||
plain: Final = "Hello!"
|
||||
assert configured.classify(plain) == baseline.classify(plain)
|
||||
assert configured.classify(plain)[0] == ComplexityTier.SIMPLE
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dimension_overrides,config_overrides",
|
||||
[
|
||||
pytest.param({"keywords": []}, {}, id="missing-matchers"),
|
||||
pytest.param({"keywords": [" "]}, {}, id="blank-keyword"),
|
||||
pytest.param({"patterns": ["\t"]}, {}, id="blank-pattern"),
|
||||
pytest.param({"patterns": ["("]}, {}, id="invalid-regex"),
|
||||
pytest.param({"patterns": [r"a*b"]}, {}, id="unbounded-star"),
|
||||
pytest.param({"patterns": [r"a{2,}b"]}, {}, id="unbounded-brace"),
|
||||
pytest.param({"patterns": [r"a{0,65}b"]}, {}, id="repeat-over-64"),
|
||||
pytest.param({"patterns": [r"(a{0,8}){0,8}b"]}, {}, id="nested-repeat"),
|
||||
pytest.param({"patterns": [r"(a|aa){0,12}b"]}, {}, id="alternation-in-repeat"),
|
||||
pytest.param({"patterns": [r"(?:ab){0,64}c"]}, {}, id="group-repeat"),
|
||||
pytest.param({"patterns": ["a?" * 9 + "b"]}, {}, id="pattern-work-over-budget"),
|
||||
pytest.param({"patterns": ["(?:a|aa)" * 9 + "z"]}, {}, id="ambiguous-alternation-chain"),
|
||||
pytest.param({"patterns": ["a?" * 8 + "a{64}" * 10 + "z"]}, {}, id="cheap-prefix-expensive-tail"),
|
||||
pytest.param({"patterns": [r"(a)\1"]}, {}, id="backreference"),
|
||||
pytest.param({"patterns": [r"(?=x)y"]}, {}, id="lookahead"),
|
||||
pytest.param({"patterns": [r"(?>ab)"]}, {}, id="atomic-group"),
|
||||
pytest.param({"patterns": [r"a*+b"]}, {}, id="possessive"),
|
||||
pytest.param({"name": "CODEPRESENCE"}, {"dimension_weights": {"tokenCount": 0.1}}, id="reserved-name"),
|
||||
pytest.param({}, {"dimension_weights": {"INTERNALFRAMEWORKS": 0.7}}, id="weight-in-map"),
|
||||
pytest.param({"weight": 0}, {}, id="zero-weight"),
|
||||
pytest.param({"weight": 1.1}, {}, id="excess-weight"),
|
||||
pytest.param({"weight": float("nan")}, {}, id="nan-weight"),
|
||||
pytest.param({"weight": float("inf")}, {}, id="infinite-weight"),
|
||||
pytest.param({"name": "bad-name"}, {}, id="invalid-name"),
|
||||
pytest.param({"name": "x" * 65}, {}, id="long-name"),
|
||||
pytest.param({"keywords": [""]}, {}, id="empty-matcher"),
|
||||
pytest.param({"keywords": ["x" * 257]}, {}, id="long-matcher"),
|
||||
pytest.param({"keywords": ["x"] * 32, "patterns": ["y"]}, {}, id="combined-matcher-count"),
|
||||
pytest.param({"keywords": ["x" * 256] * 17}, {}, id="matcher-character-budget"),
|
||||
pytest.param({"unknown": True}, {}, id="extra-field"),
|
||||
],
|
||||
)
|
||||
def test_custom_dimension_invalid_configuration_rejected(
|
||||
self, dimension_overrides: dict[str, object], config_overrides: dict[str, object]
|
||||
) -> None:
|
||||
with pytest.raises(ValidationError, match=r"custom_dimensions|custom dimension"):
|
||||
ComplexityRouterConfig.model_validate(
|
||||
{
|
||||
"custom_dimensions": [
|
||||
{
|
||||
"name": "internalFrameworks",
|
||||
"weight": 0.7,
|
||||
"keywords": ["orbitmesh"],
|
||||
**dimension_overrides,
|
||||
}
|
||||
],
|
||||
**config_overrides,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"names",
|
||||
[
|
||||
pytest.param(("internalFrameworks", "INTERNALFRAMEWORKS"), id="duplicate-casefolded-name"),
|
||||
pytest.param(tuple(f"dimension{i}" for i in range(17)), id="dimension-count"),
|
||||
],
|
||||
)
|
||||
def test_custom_dimension_names_and_count_are_bounded(self, names: tuple[str, ...]) -> None:
|
||||
with pytest.raises(ValidationError, match=r"custom_dimensions|custom dimension"):
|
||||
ComplexityRouterConfig.model_validate(
|
||||
{"custom_dimensions": [{"name": name, "weight": 0.7, "keywords": ["orbitmesh"]} for name in names]}
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("classifier_type", ("heuristic_v2", "llm", "custom"))
|
||||
def test_custom_dimensions_reject_classifiers_outside_the_tuning_gate(self, classifier_type: str) -> None:
|
||||
classifier_config: Final = (
|
||||
{"classifier_plugin": _FixedTierClassifier("SIMPLE")}
|
||||
if classifier_type == "custom"
|
||||
else {"classifier_llm_config": {"model": "judge"}}
|
||||
if classifier_type == "llm"
|
||||
else {}
|
||||
)
|
||||
with pytest.raises(ValidationError, match="custom_dimensions requires classifier_type"):
|
||||
ComplexityRouterConfig.model_validate(
|
||||
{
|
||||
"classifier_type": classifier_type,
|
||||
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}],
|
||||
**classifier_config,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh"))
|
||||
async def test_custom_dimensions_public_hook_scores_only_current_ask(
|
||||
self, mock_router_instance: MagicMock, current_ask: str
|
||||
) -> None:
|
||||
router: Final = ComplexityRouter(
|
||||
"test-router",
|
||||
mock_router_instance,
|
||||
{
|
||||
"tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"},
|
||||
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}],
|
||||
},
|
||||
)
|
||||
result: Final = await router.async_pre_routing_hook(
|
||||
model="test-router",
|
||||
request_kwargs={},
|
||||
messages=[
|
||||
{"role": "system", "content": "orbitmesh"},
|
||||
{"role": "user", "content": "orbitmesh"},
|
||||
{"role": "assistant", "content": "orbitmesh is ready"},
|
||||
{"role": "user", "content": current_ask},
|
||||
{"role": "tool", "tool_call_id": "previous", "content": "orbitmesh"},
|
||||
],
|
||||
)
|
||||
assert result is not None
|
||||
assert result.routing_decision is not None
|
||||
assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (current_ask == "orbitmesh")
|
||||
assert result.model == ("top" if current_ask == "orbitmesh" else "cheap")
|
||||
assert "orbitmesh" not in " ".join(result.routing_decision["signals"])
|
||||
|
||||
def test_custom_patterns_scan_only_the_first_2048_characters(self, mock_router_instance: MagicMock) -> None:
|
||||
router: Final = ComplexityRouter(
|
||||
"test-router",
|
||||
mock_router_instance,
|
||||
{"custom_dimensions": [{"name": "late", "weight": 0.7, "patterns": [r"zzz{1,3}"]}]},
|
||||
)
|
||||
assert "custom (late)" in router.classify("a" * 2040 + " zzz")[2]
|
||||
assert "custom (late)" not in router.classify("a" * 2048 + " zzz")[2]
|
||||
|
||||
def test_custom_dimensions_router_wide_regex_work_is_capped(self) -> None:
|
||||
heavy: Final = {"weight": 0.5, "patterns": ["a?" * 8 + "z"]}
|
||||
ComplexityRouterConfig.model_validate({"custom_dimensions": [{"name": f"d{i}", **heavy} for i in range(6)]})
|
||||
with pytest.raises(ValidationError, match="regex work estimate is 8939"):
|
||||
ComplexityRouterConfig.model_validate({"custom_dimensions": [{"name": f"d{i}", **heavy} for i in range(7)]})
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pattern,work",
|
||||
[
|
||||
pytest.param(r"\b(create|alter|drop)\s{1,4}table\b", 135, id="sql-ddl"),
|
||||
pytest.param("a?" * 8 + "z", 1277, id="optional-chain-near-cap"),
|
||||
pytest.param(r"a{0,15}a{0,15}z", 801, id="adjacent-bounded-near-cap"),
|
||||
pytest.param(r"[a-z0-9_]{3,63}\.(com|net|io)", 1291, id="class-repeat-plus-alternation"),
|
||||
pytest.param("(?:a|aa)" * 8 + "z", 1787, id="ambiguous-alternation-near-cap"),
|
||||
pytest.param("a{64}" * 10 + "z", 662, id="long-deterministic-tail"),
|
||||
],
|
||||
)
|
||||
def test_custom_pattern_work_stays_cheap_on_adversarial_text(
|
||||
self, mock_router_instance: MagicMock, pattern: str, work: int
|
||||
) -> None:
|
||||
assert custom_pattern_work(pattern) == work
|
||||
router: Final = ComplexityRouter(
|
||||
"test-router",
|
||||
mock_router_instance,
|
||||
{
|
||||
"custom_dimensions": [
|
||||
{"name": "bounded", "weight": 0.7, "patterns": [pattern]},
|
||||
{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]},
|
||||
]
|
||||
},
|
||||
)
|
||||
adversarial: Final = "orbitmesh " + "a" * 4000
|
||||
started: Final = time.perf_counter()
|
||||
tier, score, signals = router.classify(adversarial)
|
||||
elapsed: Final = time.perf_counter() - started
|
||||
assert signals == ["long (1002 tokens)", "custom (internalFrameworks)"]
|
||||
assert score == pytest.approx(0.8)
|
||||
assert tier == ComplexityTier.REASONING
|
||||
assert elapsed < 0.1
|
||||
|
||||
|
||||
class TestAsyncPreRoutingHookEdgeCases:
|
||||
"""Test edge cases for async_pre_routing_hook method."""
|
||||
|
||||
|
|
|
|||
|
|
@ -332,6 +332,67 @@ async def test_per_request_enable_prompt_caching_reaches_the_affinity_key(monkey
|
|||
assert filtered == [deployments[1]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claude_code_one_shot_subagent_does_not_reuse_an_auto_injected_affinity_key(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
cache = DualCache()
|
||||
check = PromptCachingDeploymentCheck(cache=cache)
|
||||
deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)
|
||||
messages = cast(List[AllMessageValues], [{"role": "user", "content": "unique " * 3000}])
|
||||
request_kwargs = {
|
||||
"system": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "x-anthropic-billing-header: cc_version=2.1.263; cc_is_subagent=true;",
|
||||
}
|
||||
],
|
||||
"proxy_server_request": {"headers": {"user-agent": "claude-cli/2.1.263 (external, cli)"}},
|
||||
}
|
||||
auto_injected_messages = AnthropicCacheControlHook.messages_with_default_injections(
|
||||
messages=messages,
|
||||
models=(AUTO_CACHING_MODEL,),
|
||||
)
|
||||
assert auto_injected_messages != messages
|
||||
await PromptCachingCache(cache=cache).async_add_model_id(
|
||||
model_id="dep-2", messages=auto_injected_messages, tools=None
|
||||
)
|
||||
|
||||
filtered = await check.async_filter_deployments(
|
||||
model=MODEL_GROUP_ALIAS,
|
||||
healthy_deployments=deployments,
|
||||
messages=messages,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
assert filtered == deployments
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_root_cache_control_does_not_reuse_an_auto_injected_affinity_key(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
cache = DualCache()
|
||||
check = PromptCachingDeploymentCheck(cache=cache)
|
||||
deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)
|
||||
messages = _auto_caching_messages()
|
||||
auto_injected_messages = AnthropicCacheControlHook.messages_with_default_injections(
|
||||
messages=messages,
|
||||
models=(AUTO_CACHING_MODEL,),
|
||||
)
|
||||
assert auto_injected_messages != messages
|
||||
await PromptCachingCache(cache=cache).async_add_model_id(
|
||||
model_id="dep-2", messages=auto_injected_messages, tools=None
|
||||
)
|
||||
|
||||
filtered = await check.async_filter_deployments(
|
||||
model=MODEL_GROUP_ALIAS,
|
||||
healthy_deployments=deployments,
|
||||
messages=messages,
|
||||
request_kwargs={"cache_control": {"type": "ephemeral"}},
|
||||
)
|
||||
|
||||
assert filtered == deployments
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_marked_cache_control_keeps_routing_off_another_requests_prefix(monkeypatch, local_model_cost_map):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -58,6 +59,7 @@ class TestTuningFingerprint:
|
|||
"reasoning_override_min_score": 0.05,
|
||||
"token_thresholds": {"simple": 20, "complex": 500},
|
||||
"dimension_weights": {"codePresence": 0.9},
|
||||
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}],
|
||||
"code_keywords": ["orionflow"],
|
||||
"reasoning_keywords": ["deduce"],
|
||||
"technical_keywords": ["ledgerkit"],
|
||||
|
|
@ -216,6 +218,28 @@ class TestQuota:
|
|||
is None
|
||||
)
|
||||
|
||||
def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self) -> None:
|
||||
baselines: Final = snapshot_tuning_baselines(())
|
||||
original: Final = _router("a", {})
|
||||
config: Final = {
|
||||
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}]
|
||||
}
|
||||
edited_config: Final = {
|
||||
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.9, "keywords": ["orbitmesh"]}]
|
||||
}
|
||||
added: Final = _router("a", config)
|
||||
edited: Final = _router("a", edited_config)
|
||||
second: Final = _router("b", config)
|
||||
|
||||
assert tuning_fingerprint(config) != tuning_fingerprint(edited_config)
|
||||
assert mutable_tuned_identities((added,), baselines) == {router_identity(original)}
|
||||
assert tuning_quota_violation(candidate=added, others=(original,), baselines=baselines, limit=1) is None
|
||||
assert tuning_quota_violation(candidate=edited, others=(added,), baselines=baselines, limit=1) is None
|
||||
assert tuning_quota_violation(candidate=second, others=(edited,), baselines=baselines, limit=1) is not None
|
||||
assert tuning_quota_violation(candidate=original, others=(edited,), baselines=baselines, limit=1) is None
|
||||
assert mutable_tuned_identities((original,), baselines) == frozenset()
|
||||
assert tuning_quota_violation(candidate=second, others=(original,), baselines=baselines, limit=1) is None
|
||||
|
||||
def test_violation_message_names_the_limit_and_remedy(self) -> None:
|
||||
message = tuning_limit_violation(held=2, limit=1)
|
||||
assert message is not None
|
||||
|
|
|
|||
|
|
@ -1,170 +0,0 @@
|
|||
"""
|
||||
Test suite for AWS Bedrock extended beta model support
|
||||
Tests model configuration, pricing, and regional availability for:
|
||||
- DeepSeek V3.2
|
||||
- Minimax M2.1
|
||||
- Moonshot AI Kimi K2.5
|
||||
- Qwen3 Coder Next
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Set env var to use local model cost map instead of fetching from remote
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "true"
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm import get_model_info
|
||||
|
||||
# Model configurations: (model_name, regions, max_input, max_output)
|
||||
MODEL_CONFIGS = [
|
||||
(
|
||||
"deepseek.v3.2",
|
||||
[
|
||||
"ap-northeast-1",
|
||||
"ap-south-1",
|
||||
"ap-southeast-3",
|
||||
"eu-north-1",
|
||||
"sa-east-1",
|
||||
"us-east-1",
|
||||
"us-east-2",
|
||||
"us-west-2",
|
||||
],
|
||||
163840,
|
||||
163840,
|
||||
),
|
||||
(
|
||||
"minimax.minimax-m2.1",
|
||||
[
|
||||
"ap-northeast-1",
|
||||
"ap-south-1",
|
||||
"ap-southeast-3",
|
||||
"eu-central-1",
|
||||
"eu-north-1",
|
||||
"eu-south-1",
|
||||
"eu-west-1",
|
||||
"eu-west-2",
|
||||
"sa-east-1",
|
||||
"us-east-1",
|
||||
"us-east-2",
|
||||
"us-west-2",
|
||||
],
|
||||
196000,
|
||||
8192,
|
||||
),
|
||||
(
|
||||
"moonshotai.kimi-k2.5",
|
||||
[
|
||||
"ap-northeast-1",
|
||||
"ap-south-1",
|
||||
"ap-southeast-3",
|
||||
"eu-north-1",
|
||||
"sa-east-1",
|
||||
"us-east-1",
|
||||
"us-east-2",
|
||||
"us-west-2",
|
||||
],
|
||||
262144,
|
||||
262144,
|
||||
),
|
||||
(
|
||||
"qwen.qwen3-coder-next",
|
||||
[
|
||||
"ap-northeast-1",
|
||||
"ap-south-1",
|
||||
"ap-southeast-3",
|
||||
"eu-central-1",
|
||||
"eu-south-1",
|
||||
"eu-west-1",
|
||||
"eu-west-2",
|
||||
"sa-east-1",
|
||||
"us-east-1",
|
||||
"us-east-2",
|
||||
"us-west-2",
|
||||
],
|
||||
262144,
|
||||
8192,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class TestBedrockNewModels:
|
||||
"""Unified test suite for all new Bedrock models"""
|
||||
|
||||
@pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS)
|
||||
def test_model_info_primary_region(
|
||||
self, model_name, regions, max_input, max_output
|
||||
):
|
||||
"""Test model configuration in primary region (us-east-1)"""
|
||||
model = f"bedrock/us-east-1/{model_name}"
|
||||
model_info = get_model_info(model)
|
||||
|
||||
assert model_info is not None, f"Model {model_name} not found"
|
||||
assert model_info["max_input_tokens"] == max_input
|
||||
assert model_info["max_output_tokens"] == max_output
|
||||
assert model_info["litellm_provider"] == "bedrock"
|
||||
assert model_info["mode"] == "chat"
|
||||
assert model_info["supports_function_calling"] is True
|
||||
|
||||
@pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS)
|
||||
def test_pricing_configured(self, model_name, regions, max_input, max_output):
|
||||
"""Verify pricing is set for all models"""
|
||||
model = f"bedrock/us-east-1/{model_name}"
|
||||
model_info = get_model_info(model)
|
||||
|
||||
assert (
|
||||
model_info["input_cost_per_token"] > 0
|
||||
), f"Missing input cost for {model_name}"
|
||||
assert (
|
||||
model_info["output_cost_per_token"] > 0
|
||||
), f"Missing output cost for {model_name}"
|
||||
|
||||
@pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS)
|
||||
def test_region_count(self, model_name, regions, max_input, max_output):
|
||||
"""Verify each bedrock/{region}/{model_name} resolves via get_model_info"""
|
||||
for region in regions:
|
||||
model = f"bedrock/{region}/{model_name}"
|
||||
model_info = get_model_info(model)
|
||||
assert model_info is not None, f"Model {model_name} not found in {region}"
|
||||
assert model_info["max_input_tokens"] == max_input
|
||||
assert model_info["max_output_tokens"] == max_output
|
||||
|
||||
@pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS)
|
||||
def test_sample_regional_variants(self, model_name, regions, max_input, max_output):
|
||||
"""Test sample regional variants (us-east-1, eu-west-1, ap-northeast-1)"""
|
||||
for region in ["us-east-1", "ap-northeast-1"]:
|
||||
if region in regions:
|
||||
model = f"bedrock/{region}/{model_name}"
|
||||
model_info = get_model_info(model)
|
||||
assert (
|
||||
model_info is not None
|
||||
), f"Model {model_name} not found in {region}"
|
||||
assert model_info["max_input_tokens"] == max_input
|
||||
assert model_info["litellm_provider"] == "bedrock"
|
||||
|
||||
|
||||
class TestModelSpecificFeatures:
|
||||
"""Model-specific capability tests"""
|
||||
|
||||
def test_deepseek_v3_2_context_window(self):
|
||||
"""DeepSeek V3.2 has 163K context window"""
|
||||
model_info = get_model_info("bedrock/us-east-1/deepseek.v3.2")
|
||||
assert model_info["max_input_tokens"] == 163840
|
||||
|
||||
def test_minimax_m2_1_context_window(self):
|
||||
"""Minimax M2.1 has 196K input, 8K output"""
|
||||
model_info = get_model_info("bedrock/us-east-1/minimax.minimax-m2.1")
|
||||
assert model_info["max_input_tokens"] == 196000
|
||||
assert model_info["max_output_tokens"] == 8192
|
||||
|
||||
def test_moonshotai_kimi_k2_5_context_window(self):
|
||||
"""Moonshot AI Kimi K2.5 has 256K context window"""
|
||||
model_info = get_model_info("bedrock/us-east-1/moonshotai.kimi-k2.5")
|
||||
assert model_info["max_input_tokens"] == 262144
|
||||
assert model_info["max_output_tokens"] == 262144
|
||||
|
||||
def test_qwen3_coder_next_context_window(self):
|
||||
"""Qwen3 Coder Next has 256K input, 8K output"""
|
||||
model_info = get_model_info("bedrock/us-east-1/qwen.qwen3-coder-next")
|
||||
assert model_info["max_input_tokens"] == 262144
|
||||
assert model_info["max_output_tokens"] == 8192
|
||||
|
|
@ -1,51 +0,0 @@
|
|||
"""
|
||||
Test suite for NVIDIA Nemotron Super 3 120B on AWS Bedrock
|
||||
Verifies model configuration, pricing, and regional availability.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "true"
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm import get_model_info
|
||||
|
||||
|
||||
MODEL_NAME = "nvidia.nemotron-super-3-120b"
|
||||
|
||||
|
||||
class TestNemotronSuper3120B:
|
||||
"""Test model definition for nvidia.nemotron-super-3-120b"""
|
||||
|
||||
def test_model_info_primary_region(self):
|
||||
"""Test model resolves in us-east-1"""
|
||||
model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}")
|
||||
|
||||
assert model_info is not None, f"Model {MODEL_NAME} not found"
|
||||
assert model_info["max_input_tokens"] == 256000
|
||||
assert model_info["max_output_tokens"] == 32768
|
||||
assert model_info["litellm_provider"] == "bedrock_converse"
|
||||
assert model_info["mode"] == "chat"
|
||||
assert model_info["supports_function_calling"] is True
|
||||
|
||||
def test_pricing_configured(self):
|
||||
"""Verify pricing matches AWS Bedrock rates"""
|
||||
model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}")
|
||||
|
||||
assert model_info["input_cost_per_token"] == 1.5e-07
|
||||
assert model_info["output_cost_per_token"] == 6.5e-07
|
||||
|
||||
def test_context_window(self):
|
||||
"""Nemotron Super 3 120B has 256K input, 32K output on Bedrock"""
|
||||
model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}")
|
||||
|
||||
assert model_info["max_input_tokens"] == 256000
|
||||
assert model_info["max_output_tokens"] == 32768
|
||||
|
||||
def test_resolves_without_region(self):
|
||||
"""Test model resolves with just bedrock/ prefix"""
|
||||
model_info = get_model_info(f"bedrock/{MODEL_NAME}")
|
||||
|
||||
assert model_info is not None, f"Model {MODEL_NAME} not found without region"
|
||||
assert model_info["max_input_tokens"] == 256000
|
||||
|
|
@ -1,47 +0,0 @@
|
|||
"""
|
||||
Validate that AWS GovCloud (Bedrock us-gov-*) Haiku 4.5 entries carry
|
||||
the 1-hour cache write tier.
|
||||
|
||||
AWS Bedrock GovCloud pricing applies a +20% premium over global
|
||||
Anthropic rates. Global Haiku 4.5 1h cache write is $2.00/MTok; us-gov
|
||||
is therefore $2.40/MTok — exactly 1.6x the 5-minute rate of $1.50/MTok.
|
||||
|
||||
Source: https://aws.amazon.com/bedrock/pricing/
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model_data():
|
||||
json_path = os.path.join(
|
||||
os.path.dirname(__file__), "../../model_prices_and_context_window.json"
|
||||
)
|
||||
with open(json_path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
HAIKU_USGOV_KEYS = [
|
||||
"bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_key", HAIKU_USGOV_KEYS)
|
||||
def test_usgov_haiku_4_5_1hr_cache_write(model_data, model_key):
|
||||
assert model_key in model_data, f"Missing model entry: {model_key}"
|
||||
info = model_data[model_key]
|
||||
assert (
|
||||
info["cache_creation_input_token_cost"] == 1.5e-06
|
||||
), f"{model_key}: 5m cache write should be $1.50/MTok"
|
||||
assert (
|
||||
info["cache_creation_input_token_cost_above_1hr"] == 2.4e-06
|
||||
), f"{model_key}: 1h cache write should be $2.40/MTok"
|
||||
ratio = (
|
||||
info["cache_creation_input_token_cost_above_1hr"]
|
||||
/ info["cache_creation_input_token_cost"]
|
||||
)
|
||||
assert abs(ratio - 1.6) < 1e-9, f"{model_key}: 1h/5m ratio is {ratio}, expected 1.6"
|
||||
|
|
@ -31,34 +31,6 @@ def model_data():
|
|||
return json.load(f)
|
||||
|
||||
|
||||
SONNET_4_5_USGOV_KEYS = [
|
||||
"bedrock/us-gov-east-1/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0",
|
||||
"bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0",
|
||||
"us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_key", SONNET_4_5_USGOV_KEYS)
|
||||
def test_usgov_sonnet_4_5_pricing(model_data, model_key):
|
||||
"""Each us-gov sonnet-4-5 entry must carry the +20%-over-global rates
|
||||
that AWS publishes on the GovCloud pricing page.
|
||||
"""
|
||||
assert model_key in model_data, f"Missing model entry: {model_key}"
|
||||
info = model_data[model_key]
|
||||
|
||||
assert info["input_cost_per_token"] == 3.6e-06, (
|
||||
f"{model_key}: input_cost_per_token should be $3.60/MTok (got {info['input_cost_per_token']})"
|
||||
)
|
||||
assert info["output_cost_per_token"] == 1.8e-05, f"{model_key}: output_cost_per_token should be $18.00/MTok"
|
||||
assert info["cache_creation_input_token_cost"] == 4.5e-06, f"{model_key}: 5m cache write should be $4.50/MTok"
|
||||
assert info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06, (
|
||||
f"{model_key}: 1h cache write should be $7.20/MTok"
|
||||
)
|
||||
assert info["cache_read_input_token_cost"] == 3.6e-07, f"{model_key}: cache read should be $0.36/MTok"
|
||||
|
||||
|
||||
def test_usgov_carries_20_percent_premium_over_global(model_data):
|
||||
"""The us-gov rates must equal 1.2x the global anthropic.* rates,
|
||||
matching AWS's documented GovCloud uplift.
|
||||
|
|
@ -92,18 +64,6 @@ EXPECTED_USGOV_ABOVE_200K = {
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field,expected", EXPECTED_USGOV_ABOVE_200K.items())
|
||||
def test_usgov_cross_region_above_200k_carries_gov_premium(model_data, field, expected):
|
||||
"""The `_above_200k_tokens` tier on the us-gov cross-region inference
|
||||
profile must also carry the +20% GovCloud uplift. The original PR
|
||||
corrected the base rates but left the 200k-tier fields at the +10%
|
||||
commercial-US rates, undercharging long-context requests.
|
||||
"""
|
||||
info = model_data[USGOV_CROSS_REGION_KEY]
|
||||
assert field in info, f"{USGOV_CROSS_REGION_KEY}: missing field {field}"
|
||||
assert info[field] == expected, f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})"
|
||||
|
||||
|
||||
def test_usgov_cross_region_above_200k_ratio_to_global(model_data):
|
||||
"""Cross-check via the property-based invariant: every `_above_200k_tokens`
|
||||
field on the us-gov cross-region profile must equal 1.2x the global
|
||||
|
|
@ -117,167 +77,6 @@ def test_usgov_cross_region_above_200k_ratio_to_global(model_data):
|
|||
assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2"
|
||||
|
||||
|
||||
CLAUDE_GOV_EXPECTED = {
|
||||
"anthropic.claude-sonnet-5": {
|
||||
"input_cost_per_token": 2.4e-06,
|
||||
"output_cost_per_token": 1.2e-05,
|
||||
"cache_creation_input_token_cost": 3e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 4.8e-06,
|
||||
"cache_read_input_token_cost": 2.4e-07,
|
||||
},
|
||||
"anthropic.claude-opus-4-8": {
|
||||
"input_cost_per_token": 6e-06,
|
||||
"output_cost_per_token": 3e-05,
|
||||
"cache_creation_input_token_cost": 7.5e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1.2e-05,
|
||||
"cache_read_input_token_cost": 6e-07,
|
||||
},
|
||||
"anthropic.claude-opus-5": {
|
||||
"input_cost_per_token": 6e-06,
|
||||
"output_cost_per_token": 3e-05,
|
||||
"cache_creation_input_token_cost": 7.5e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1.2e-05,
|
||||
"cache_read_input_token_cost": 6e-07,
|
||||
},
|
||||
"anthropic.claude-fable-5-1": {
|
||||
"input_cost_per_token": 1.2e-05,
|
||||
"output_cost_per_token": 6e-05,
|
||||
"cache_creation_input_token_cost": 1.5e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2.4e-05,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
USGOV_CLAUDE_KEY_TEMPLATES = {
|
||||
"bedrock/us-gov-east-1/{base_key}": "bedrock",
|
||||
"bedrock/us-gov-west-1/{base_key}": "bedrock",
|
||||
"us-gov.{base_key}": "bedrock_converse",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("base_key", CLAUDE_GOV_EXPECTED)
|
||||
@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items())
|
||||
def test_usgov_claude_pricing(model_data, key_template, expected_provider, base_key):
|
||||
"""Sonnet 5, Opus 4.8, Opus 5, and Fable 5.1 gov entries, both in-region keys
|
||||
and the us-gov. geo inference profile the model cards list for GovCloud, must
|
||||
carry the 1.2x GovCloud premium over the global anthropic.* rates. No public
|
||||
AWS source (offer files, pricing page) lists Claude GovCloud rows; the premium
|
||||
is the one AWS quotes for Opus 4.8 in GovCloud ($6/$30 per million).
|
||||
"""
|
||||
gov_key = key_template.format(base_key=base_key)
|
||||
assert gov_key in model_data, f"Missing model entry: {gov_key}"
|
||||
info = model_data[gov_key]
|
||||
assert info["litellm_provider"] == expected_provider
|
||||
assert "search_context_cost_per_query" not in info
|
||||
for field, expected in CLAUDE_GOV_EXPECTED[base_key].items():
|
||||
assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})"
|
||||
ratio = info[field] / model_data[base_key][field]
|
||||
assert abs(ratio - 1.2) < 1e-9, f"{gov_key}: {field} gov/global ratio is {ratio}, expected 1.2"
|
||||
|
||||
|
||||
CONVERSE_GOV_EXPECTED = {
|
||||
"nvidia.nemotron-nano-3-30b": (7.2e-08, 2.88e-07),
|
||||
"nvidia.nemotron-nano-9b-v2": (7.2e-08, 2.76e-07),
|
||||
"nvidia.nemotron-nano-12b-v2": (2.4e-07, 7.2e-07),
|
||||
"nvidia.nemotron-super-3-120b": (1.8e-07, 7.8e-07),
|
||||
"openai.gpt-oss-20b-1:0": (8.4e-08, 3.6e-07),
|
||||
"openai.gpt-oss-120b-1:0": (1.8e-07, 7.2e-07),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("base_key", CONVERSE_GOV_EXPECTED)
|
||||
@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items())
|
||||
def test_usgov_converse_model_pricing(model_data, key_template, expected_provider, base_key):
|
||||
"""Nemotron and gpt-oss gov entries, in-region and the us-gov. geo inference
|
||||
profile both GovCloud regions list as ACTIVE, must match the AWS Bedrock
|
||||
offer file, which prices both regions identically at 1.2x commercial.
|
||||
"""
|
||||
gov_key = key_template.format(base_key=base_key)
|
||||
assert gov_key in model_data, f"Missing model entry: {gov_key}"
|
||||
info = model_data[gov_key]
|
||||
expected_input, expected_output = CONVERSE_GOV_EXPECTED[base_key]
|
||||
assert info["input_cost_per_token"] == expected_input
|
||||
assert info["output_cost_per_token"] == expected_output
|
||||
assert info["litellm_provider"] == expected_provider
|
||||
base = model_data[base_key]
|
||||
assert abs(info["input_cost_per_token"] / base["input_cost_per_token"] - 1.2) < 1e-9
|
||||
assert abs(info["output_cost_per_token"] / base["output_cost_per_token"] - 1.2) < 1e-9
|
||||
|
||||
|
||||
def test_usgov_west_llama3_8b_output_price_fixed(model_data):
|
||||
"""The us-gov-west-1 llama3-8b entry carried the 70B output rate ($2.65/MTok);
|
||||
the AWS Bedrock offer file prices output at $0.60/MTok. AWS lists the model
|
||||
in us-gov-west-1 only, so there is no east entry to check.
|
||||
"""
|
||||
info = model_data["bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0"]
|
||||
assert info["input_cost_per_token"] == 3e-07
|
||||
assert info["output_cost_per_token"] == 6e-07
|
||||
|
||||
|
||||
MANTLE_GOV_TIERED_EXPECTED = {
|
||||
"openai.gpt-5.6-luna": {
|
||||
"input_cost_per_token": 2.64e-07,
|
||||
"input_cost_per_token_above_272k_tokens": 5.28e-07,
|
||||
"cache_creation_input_token_cost": 3.3e-07,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 6.6e-07,
|
||||
"cache_read_input_token_cost": 2.64e-08,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 5.28e-08,
|
||||
"output_cost_per_token": 1.584e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 2.376e-06,
|
||||
},
|
||||
"openai.gpt-5.6-terra": {
|
||||
"input_cost_per_token": 2.64e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 5.28e-06,
|
||||
"cache_creation_input_token_cost": 3.3e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 6.6e-06,
|
||||
"cache_read_input_token_cost": 2.64e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 5.28e-07,
|
||||
"output_cost_per_token": 1.584e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 2.376e-05,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MANTLE_GOV_TIERED_EXPECTED)
|
||||
def test_usgov_west_mantle_terra_luna_pricing(model_data, model):
|
||||
"""Terra and Luna carry 1.2x commercial across every tier in the
|
||||
us-gov-west-1 offer file; the us-gov-east-1 offer file has no SKUs for them.
|
||||
"""
|
||||
gov_key = f"bedrock_mantle/us-gov-west-1/{model}"
|
||||
assert gov_key in model_data, f"Missing model entry: {gov_key}"
|
||||
info = model_data[gov_key]
|
||||
for field, expected in MANTLE_GOV_TIERED_EXPECTED[model].items():
|
||||
assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})"
|
||||
assert info["litellm_provider"] == "bedrock_mantle"
|
||||
assert f"bedrock_mantle/us-gov-east-1/{model}" not in model_data
|
||||
|
||||
|
||||
@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"])
|
||||
def test_usgov_mantle_gpt_5_4_pricing_has_no_long_context_tier(model_data, region):
|
||||
"""gpt-5.4 gov rates come from the offer file, which publishes only the
|
||||
standard tier in GovCloud: no long-context SKUs exist there, unlike commercial.
|
||||
"""
|
||||
gov_key = f"bedrock_mantle/{region}/openai.gpt-5.4"
|
||||
assert gov_key in model_data, f"Missing model entry: {gov_key}"
|
||||
info = model_data[gov_key]
|
||||
assert info["input_cost_per_token"] == 3.3e-06
|
||||
assert info["cache_read_input_token_cost"] == 3.3e-07
|
||||
assert info["output_cost_per_token"] == 1.98e-05
|
||||
assert not any(field.endswith("_above_272k_tokens") for field in info)
|
||||
|
||||
|
||||
def test_usgov_mantle_grok_4_3_west_only(model_data):
|
||||
"""grok-4.3 is priced in the us-gov-west-1 offer file only; the east offer
|
||||
file carries grok-4.6 instead.
|
||||
"""
|
||||
info = model_data["bedrock_mantle/us-gov-west-1/xai.grok-4.3"]
|
||||
assert info["input_cost_per_token"] == 1.5e-06
|
||||
assert info["output_cost_per_token"] == 3e-06
|
||||
assert info["cache_read_input_token_cost"] == 2.4e-07
|
||||
assert "bedrock_mantle/us-gov-east-1/xai.grok-4.3" not in model_data
|
||||
|
||||
|
||||
def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data):
|
||||
"""us-gov-east-1 serves claude-3-haiku through the us-gov. inference profile
|
||||
only, so the profile row must bill exactly like the in-region gov row.
|
||||
|
|
@ -290,96 +89,6 @@ def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data):
|
|||
}
|
||||
|
||||
|
||||
GROK_4_6_GOV_KEYS = {
|
||||
"us-gov.xai.grok-4.6": ("us.xai.grok-4.6", "bedrock_converse"),
|
||||
"bedrock_mantle/us-gov-west-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"),
|
||||
"bedrock_mantle/us-gov-east-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gov_key", GROK_4_6_GOV_KEYS)
|
||||
def test_usgov_grok_4_6_pricing(model_data, gov_key):
|
||||
"""Both GovCloud regions serve grok-4.6 through the us-gov. profile only, and
|
||||
both offer files price its standard SKU at 1.2x the commercial US rate.
|
||||
"""
|
||||
base_key, expected_provider = GROK_4_6_GOV_KEYS[gov_key]
|
||||
assert gov_key in model_data, f"Missing model entry: {gov_key}"
|
||||
info = model_data[gov_key]
|
||||
assert info["litellm_provider"] == expected_provider
|
||||
assert info["input_cost_per_token"] == 2.64e-06
|
||||
assert info["output_cost_per_token"] == 7.92e-06
|
||||
assert info["cache_read_input_token_cost"] == 6.6e-07
|
||||
for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"):
|
||||
assert abs(info[field] / model_data[base_key][field] - 1.2) < 1e-9
|
||||
|
||||
|
||||
NOVA_GOV_WEST_EXPECTED = {
|
||||
"amazon.nova-lite-v1:0": (7.2e-08, 2.88e-07),
|
||||
"amazon.nova-micro-v1:0": (4.2e-08, 1.68e-07),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("base_key", NOVA_GOV_WEST_EXPECTED)
|
||||
def test_usgov_west_nova_lite_micro_pricing(model_data, base_key):
|
||||
"""Nova Lite and Micro are on-demand in us-gov-west-1 only; the offer file
|
||||
prices them at 1.2x commercial, like the Nova Pro row that was already there.
|
||||
"""
|
||||
gov_key = f"bedrock/us-gov-west-1/{base_key}"
|
||||
assert gov_key in model_data, f"Missing model entry: {gov_key}"
|
||||
info = model_data[gov_key]
|
||||
expected_input, expected_output = NOVA_GOV_WEST_EXPECTED[base_key]
|
||||
assert info["litellm_provider"] == "bedrock"
|
||||
assert info["input_cost_per_token"] == expected_input
|
||||
assert info["output_cost_per_token"] == expected_output
|
||||
assert abs(info["input_cost_per_token"] / model_data[base_key]["input_cost_per_token"] - 1.2) < 1e-9
|
||||
assert abs(info["output_cost_per_token"] / model_data[base_key]["output_cost_per_token"] - 1.2) < 1e-9
|
||||
assert f"bedrock/us-gov-east-1/{base_key}" not in model_data
|
||||
|
||||
|
||||
def test_usgov_west_nova_2_multimodal_embeddings_pricing(model_data):
|
||||
"""Every meter of the multimodal embedding model (tokens, images, audio and
|
||||
video seconds) carries the 1.2x uplift the us-gov-west-1 offer file lists.
|
||||
"""
|
||||
gov_key = "bedrock/us-gov-west-1/amazon.nova-2-multimodal-embeddings-v1:0"
|
||||
assert gov_key in model_data, f"Missing model entry: {gov_key}"
|
||||
info = model_data[gov_key]
|
||||
assert info["litellm_provider"] == "bedrock"
|
||||
assert info["mode"] == "embedding"
|
||||
assert info["input_cost_per_token"] == 1.62e-07
|
||||
assert info["input_cost_per_image"] == 7.2e-05
|
||||
assert info["input_cost_per_audio_per_second"] == 0.000168
|
||||
assert info["input_cost_per_video_per_second"] == 0.00084
|
||||
assert "bedrock/us-gov-east-1/amazon.nova-2-multimodal-embeddings-v1:0" not in model_data
|
||||
|
||||
|
||||
MANTLE_GOV_FLAT_EXPECTED = {
|
||||
"google.gemma-4-e2b": (4.8e-08, 9.6e-08, ("us-gov-west-1",)),
|
||||
"google.gemma-4-26b-a4b": (1.56e-07, 4.8e-07, ("us-gov-west-1",)),
|
||||
"google.gemma-4-31b": (1.68e-07, 4.8e-07, ("us-gov-west-1",)),
|
||||
"openai.gpt-oss-20b": (8.4e-08, 3.6e-07, ("us-gov-west-1", "us-gov-east-1")),
|
||||
"openai.gpt-oss-120b": (1.8e-07, 7.2e-07, ("us-gov-west-1", "us-gov-east-1")),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MANTLE_GOV_FLAT_EXPECTED)
|
||||
def test_usgov_mantle_gemma_and_gpt_oss_pricing(model_data, model):
|
||||
"""Gemma 4 is priced in the us-gov-west-1 offer file only and gpt-oss in both;
|
||||
each Mantle gov row carries the offer file's standard SKU, and no row exists
|
||||
for a region whose offer file has no SKU.
|
||||
"""
|
||||
expected_input, expected_output, regions = MANTLE_GOV_FLAT_EXPECTED[model]
|
||||
for region in ("us-gov-west-1", "us-gov-east-1"):
|
||||
gov_key = f"bedrock_mantle/{region}/{model}"
|
||||
if region not in regions:
|
||||
assert gov_key not in model_data
|
||||
continue
|
||||
assert gov_key in model_data, f"Missing model entry: {gov_key}"
|
||||
info = model_data[gov_key]
|
||||
assert info["litellm_provider"] == "bedrock_mantle"
|
||||
assert info["input_cost_per_token"] == expected_input
|
||||
assert info["output_cost_per_token"] == expected_output
|
||||
|
||||
|
||||
GOV_ROW_SOURCES = {
|
||||
"us-gov.anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1",
|
||||
"bedrock/us-gov-west-1/anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1",
|
||||
|
|
@ -417,33 +126,3 @@ def test_usgov_rows_keep_commercial_limits_and_capabilities(model_data, gov_key)
|
|||
assert _non_pricing_fields(gov) == _non_pricing_fields(model_data[GOV_ROW_SOURCES[gov_key]])
|
||||
assert "search_context_cost_per_query" not in gov
|
||||
assert "source" not in gov
|
||||
|
||||
|
||||
AZURE_GOV_EXPECTED = {
|
||||
"azure/us-gov/gpt-5.1": {
|
||||
"input_cost_per_token": 1.71875e-06,
|
||||
"cache_read_input_token_cost": 1.71875e-07,
|
||||
"output_cost_per_token": 1.375e-05,
|
||||
},
|
||||
"azure/us-gov/o3-mini": {
|
||||
"input_cost_per_token": 1.513e-06,
|
||||
"cache_read_input_token_cost": 7.57e-07,
|
||||
"output_cost_per_token": 6.05e-06,
|
||||
},
|
||||
"azure/us-gov/text-embedding-3-large": {"input_cost_per_token": 1.63e-07},
|
||||
"azure/us-gov/text-embedding-3-small": {"input_cost_per_token": 2.5e-08},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gov_key", AZURE_GOV_EXPECTED)
|
||||
def test_azure_usgov_pricing(model_data, gov_key):
|
||||
"""Azure Government meters from the Azure retail prices API
|
||||
(usgovvirginia/usgovarizona, serviceName 'Foundry Models'). No Government
|
||||
retirement schedule is published, so these entries carry no deprecation_date.
|
||||
"""
|
||||
assert gov_key in model_data, f"Missing model entry: {gov_key}"
|
||||
info = model_data[gov_key]
|
||||
for field, expected in AZURE_GOV_EXPECTED[gov_key].items():
|
||||
assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})"
|
||||
assert info["litellm_provider"] == "azure"
|
||||
assert "deprecation_date" not in info
|
||||
|
|
|
|||
|
|
@ -14,7 +14,6 @@ import os
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import BEDROCK_CONVERSE_MODELS
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
|
||||
|
|
@ -27,89 +26,6 @@ def _load_root_cost_map() -> dict:
|
|||
return json.load(f)
|
||||
|
||||
|
||||
|
||||
def test_fable_5_model_pricing_and_capabilities():
|
||||
model_data = _load_root_cost_map()
|
||||
|
||||
expected_models = [
|
||||
("claude-fable-5", "anthropic"),
|
||||
("anthropic.claude-fable-5", "bedrock_converse"),
|
||||
("vertex_ai/claude-fable-5", "vertex_ai-anthropic_models"),
|
||||
# Unlike Opus 4.8 (200k on Foundry), Fable 5 has the full 1M context
|
||||
# window on Microsoft Foundry.
|
||||
("azure_ai/claude-fable-5", "azure_ai"),
|
||||
]
|
||||
|
||||
for model_name, provider in expected_models:
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
|
||||
assert info["litellm_provider"] == provider
|
||||
assert info["mode"] == "chat"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["max_tokens"] == 128000
|
||||
|
||||
# $10 / $50 per MTok (2x Opus 4.8), with the standard 1.25x 5m
|
||||
# cache-write, 2x 1h cache-write, and 0.1x cache-read multipliers.
|
||||
assert info["input_cost_per_token"] == 1e-05
|
||||
assert info["output_cost_per_token"] == 5e-05
|
||||
assert info["cache_creation_input_token_cost"] == 1.25e-05
|
||||
assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05
|
||||
assert info["cache_read_input_token_cost"] == 1e-06
|
||||
|
||||
# Flat-rate across the full 1M context window.
|
||||
assert "input_cost_per_token_above_200k_tokens" not in info
|
||||
assert "output_cost_per_token_above_200k_tokens" not in info
|
||||
|
||||
assert info["supports_assistant_prefill"] is False
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_xhigh_reasoning_effort"] is True
|
||||
assert info["supports_max_reasoning_effort"] is True
|
||||
|
||||
|
||||
def test_fable_5_bedrock_regional_model_pricing():
|
||||
model_data = _load_root_cost_map()
|
||||
|
||||
# Fable 5 launched with us/eu geo inference profiles plus a global profile
|
||||
# (no au/apac/jp). Global uses base pricing; geo profiles carry the
|
||||
# standard 10% regional premium.
|
||||
expected_models = {
|
||||
"global.anthropic.claude-fable-5": {
|
||||
"input_cost_per_token": 1e-05,
|
||||
"output_cost_per_token": 5e-05,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_read_input_token_cost": 1e-06,
|
||||
},
|
||||
"us.anthropic.claude-fable-5": {
|
||||
"input_cost_per_token": 1.1e-05,
|
||||
"output_cost_per_token": 5.5e-05,
|
||||
"cache_creation_input_token_cost": 1.375e-05,
|
||||
"cache_read_input_token_cost": 1.1e-06,
|
||||
},
|
||||
"eu.anthropic.claude-fable-5": {
|
||||
"input_cost_per_token": 1.1e-05,
|
||||
"output_cost_per_token": 5.5e-05,
|
||||
"cache_creation_input_token_cost": 1.375e-05,
|
||||
"cache_read_input_token_cost": 1.1e-06,
|
||||
},
|
||||
}
|
||||
|
||||
for model_name, expected in expected_models.items():
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
assert info["litellm_provider"] == "bedrock_converse"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["bedrock_output_config_effort_ceiling"] == "xhigh"
|
||||
for key, value in expected.items():
|
||||
assert info[key] == value
|
||||
|
||||
|
||||
def test_fable_5_geo_multiplier_without_fast_mode():
|
||||
"""First-party ``inference_geo='us'`` carries the 1.1x premium, but unlike
|
||||
the Opus line there is no fast-mode variant for Fable 5; a ``fast`` key
|
||||
|
|
@ -144,13 +60,6 @@ def test_fable_5_registered_for_bedrock_converse():
|
|||
assert "anthropic.claude-fable-5" in BEDROCK_CONVERSE_MODELS
|
||||
|
||||
|
||||
def test_fable_5_provider_resolves_via_model_info(local_model_cost_map):
|
||||
info = litellm.get_model_info(model="claude-fable-5")
|
||||
assert info["litellm_provider"] == "anthropic"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cost_map",
|
||||
[_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
|
||||
|
|
@ -222,46 +131,6 @@ FABLE_5_1_VARIANTS = (
|
|||
)
|
||||
|
||||
|
||||
def test_fable_5_1_model_pricing_and_capabilities():
|
||||
model_data = _load_root_cost_map()
|
||||
|
||||
expected_models = [
|
||||
("claude-fable-5-1", "anthropic"),
|
||||
("anthropic.claude-fable-5-1", "bedrock_converse"),
|
||||
("vertex_ai/claude-fable-5-1", "vertex_ai-anthropic_models"),
|
||||
("azure_ai/claude-fable-5-1", "azure_ai"),
|
||||
]
|
||||
|
||||
for model_name, provider in expected_models:
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
|
||||
assert info["litellm_provider"] == provider
|
||||
assert info["mode"] == "chat"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["max_tokens"] == 128000
|
||||
|
||||
assert info["input_cost_per_token"] == 1e-05
|
||||
assert info["output_cost_per_token"] == 5e-05
|
||||
assert info["cache_creation_input_token_cost"] == 1.25e-05
|
||||
assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05
|
||||
|
||||
assert "input_cost_per_token_above_200k_tokens" not in info
|
||||
assert "output_cost_per_token_above_200k_tokens" not in info
|
||||
|
||||
assert info["supports_assistant_prefill"] is False
|
||||
assert info["supports_forced_tool_use"] is False
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_xhigh_reasoning_effort"] is True
|
||||
assert info["supports_max_reasoning_effort"] is True
|
||||
assert info["prompt_cache_min_tokens"] == 512
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cost_map",
|
||||
[_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
|
||||
|
|
@ -280,48 +149,6 @@ def test_fable_5_1_cache_reads_cost_a_quarter_of_fable_5(cost_map):
|
|||
), model_name
|
||||
|
||||
|
||||
def test_fable_5_1_bedrock_regional_model_pricing():
|
||||
model_data = _load_root_cost_map()
|
||||
|
||||
expected_models = {
|
||||
"global.anthropic.claude-fable-5-1": {
|
||||
"input_cost_per_token": 1e-05,
|
||||
"output_cost_per_token": 5e-05,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
},
|
||||
"us.anthropic.claude-fable-5-1": {
|
||||
"input_cost_per_token": 1.1e-05,
|
||||
"output_cost_per_token": 5.5e-05,
|
||||
"cache_creation_input_token_cost": 1.375e-05,
|
||||
"cache_read_input_token_cost": 2.75e-07,
|
||||
},
|
||||
"eu.anthropic.claude-fable-5-1": {
|
||||
"input_cost_per_token": 1.1e-05,
|
||||
"output_cost_per_token": 5.5e-05,
|
||||
"cache_creation_input_token_cost": 1.375e-05,
|
||||
"cache_read_input_token_cost": 2.75e-07,
|
||||
},
|
||||
}
|
||||
|
||||
for model_name, expected in expected_models.items():
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
assert info["litellm_provider"] == "bedrock_converse"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["bedrock_output_config_effort_ceiling"] == "xhigh"
|
||||
for key, value in expected.items():
|
||||
assert info[key] == value
|
||||
|
||||
|
||||
def test_fable_5_1_geo_multiplier_without_fast_mode():
|
||||
"""Fable 5.1 has no fast mode, so a ``fast`` key here would misprice
|
||||
``speed='fast'`` requests."""
|
||||
model_data = _load_root_cost_map()
|
||||
assert model_data["claude-fable-5-1"]["provider_specific_entry"] == {"us": 1.1}
|
||||
|
||||
|
||||
def test_fable_5_1_present_in_bundled_backup():
|
||||
backup = GetModelCostMap.load_local_model_cost_map()
|
||||
root = _load_root_cost_map()
|
||||
|
|
@ -334,13 +161,6 @@ def test_fable_5_1_registered_for_bedrock_converse():
|
|||
assert "anthropic.claude-fable-5-1" in BEDROCK_CONVERSE_MODELS
|
||||
|
||||
|
||||
def test_fable_5_1_provider_resolves_via_model_info(local_model_cost_map):
|
||||
info = litellm.get_model_info(model="claude-fable-5-1")
|
||||
assert info["litellm_provider"] == "anthropic"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -7,57 +7,6 @@ import json
|
|||
import os
|
||||
|
||||
|
||||
def test_bedrock_haiku_4_5_configuration():
|
||||
"""Test that all Bedrock Claude Haiku 4.5 models use bedrock_converse provider"""
|
||||
# Load model configuration
|
||||
json_path = os.path.join(
|
||||
os.path.dirname(__file__), "../../model_prices_and_context_window.json"
|
||||
)
|
||||
with open(json_path) as f:
|
||||
model_data = json.load(f)
|
||||
|
||||
# All Bedrock Haiku 4.5 variants that should use bedrock_converse
|
||||
bedrock_haiku_models = [
|
||||
"anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"anthropic.claude-haiku-4-5@20251001",
|
||||
"us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"apac.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"jp.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"global.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"au.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
]
|
||||
|
||||
for model in bedrock_haiku_models:
|
||||
assert model in model_data, f"Model {model} not found in config"
|
||||
model_info = model_data[model]
|
||||
|
||||
# Verify uses bedrock_converse (not legacy bedrock provider)
|
||||
assert (
|
||||
model_info["litellm_provider"] == "bedrock_converse"
|
||||
), f"{model} should use bedrock_converse provider, got {model_info['litellm_provider']}"
|
||||
|
||||
# Verify supports vision (key missing capability)
|
||||
assert (
|
||||
model_info.get("supports_vision") is True
|
||||
), f"{model} should support vision"
|
||||
|
||||
# Verify core capabilities
|
||||
assert model_info.get("supports_computer_use") is True
|
||||
assert model_info.get("supports_function_calling") is True
|
||||
assert model_info.get("supports_tool_choice") is True
|
||||
assert model_info.get("supports_prompt_caching") is True
|
||||
assert model_info.get("supports_response_schema") is True
|
||||
assert model_info.get("supports_pdf_input") is True
|
||||
assert model_info.get("supports_assistant_prefill") is True
|
||||
assert model_info.get("supports_reasoning") is True
|
||||
|
||||
# Verify token limits
|
||||
assert model_info["max_input_tokens"] == 200000
|
||||
assert model_info["max_output_tokens"] == 64000
|
||||
assert model_info["mode"] == "chat"
|
||||
|
||||
|
||||
def test_bedrock_haiku_4_5_matches_sonnet_capabilities():
|
||||
"""
|
||||
Test that Haiku 4.5 has same capabilities as Sonnet 4.5
|
||||
|
|
@ -97,36 +46,3 @@ def test_bedrock_haiku_4_5_matches_sonnet_capabilities():
|
|||
assert haiku_info.get(capability) == sonnet_info.get(
|
||||
capability
|
||||
), f"Capability {capability} mismatch: Haiku={haiku_info.get(capability)}, Sonnet={sonnet_info.get(capability)}"
|
||||
|
||||
|
||||
def test_anthropic_api_haiku_4_5_configuration():
|
||||
"""Test that Anthropic API Claude Haiku 4.5 has correct configuration"""
|
||||
# Load model configuration
|
||||
json_path = os.path.join(
|
||||
os.path.dirname(__file__), "../../model_prices_and_context_window.json"
|
||||
)
|
||||
with open(json_path) as f:
|
||||
model_data = json.load(f)
|
||||
|
||||
# Anthropic API models (not Bedrock)
|
||||
anthropic_models = [
|
||||
"claude-haiku-4-5-20251001",
|
||||
"claude-haiku-4-5",
|
||||
]
|
||||
|
||||
for model in anthropic_models:
|
||||
assert model in model_data, f"Model {model} not found in config"
|
||||
model_info = model_data[model]
|
||||
|
||||
# Should use anthropic provider (not bedrock)
|
||||
assert (
|
||||
model_info["litellm_provider"] == "anthropic"
|
||||
), f"{model} should use anthropic provider"
|
||||
|
||||
# Should support vision
|
||||
assert (
|
||||
model_info.get("supports_vision") is True
|
||||
), f"{model} should support vision"
|
||||
|
||||
# Should have larger output token limit (64K for Anthropic API)
|
||||
assert model_info["max_output_tokens"] == 64000
|
||||
|
|
|
|||
|
|
@ -71,125 +71,6 @@ def test_claude_4_6_australia_region_uses_au_prefix_not_apac():
|
|||
), "apac.anthropic.claude-sonnet-4-6 should not be in bedrock_converse_models"
|
||||
|
||||
|
||||
def test_opus_4_6_model_pricing_and_capabilities():
|
||||
json_path = os.path.join(
|
||||
os.path.dirname(__file__), "../../model_prices_and_context_window.json"
|
||||
)
|
||||
with open(json_path) as f:
|
||||
model_data = json.load(f)
|
||||
|
||||
expected_models = {
|
||||
"claude-opus-4-6": {
|
||||
"provider": "anthropic",
|
||||
"has_long_context_pricing": False,
|
||||
"max_input_tokens": 1000000,
|
||||
},
|
||||
"claude-opus-4-6-20260205": {
|
||||
"provider": "anthropic",
|
||||
"has_long_context_pricing": False,
|
||||
"max_input_tokens": 1000000,
|
||||
},
|
||||
"anthropic.claude-opus-4-6-v1": {
|
||||
"provider": "bedrock_converse",
|
||||
"has_long_context_pricing": False,
|
||||
"max_input_tokens": 1000000,
|
||||
},
|
||||
"vertex_ai/claude-opus-4-6": {
|
||||
"provider": "vertex_ai-anthropic_models",
|
||||
"has_long_context_pricing": False,
|
||||
"max_input_tokens": 1000000,
|
||||
},
|
||||
"azure_ai/claude-opus-4-6": {
|
||||
"provider": "azure_ai",
|
||||
"has_long_context_pricing": False,
|
||||
"max_input_tokens": 1000000,
|
||||
},
|
||||
}
|
||||
|
||||
for model_name, config in expected_models.items():
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
|
||||
assert info["litellm_provider"] == config["provider"]
|
||||
assert info["mode"] == "chat"
|
||||
assert info["max_input_tokens"] == config["max_input_tokens"]
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["max_tokens"] == 128000
|
||||
|
||||
assert info["input_cost_per_token"] == 5e-06
|
||||
assert info["output_cost_per_token"] == 2.5e-05
|
||||
assert info["cache_creation_input_token_cost"] == 6.25e-06
|
||||
assert info["cache_read_input_token_cost"] == 5e-07
|
||||
|
||||
if config["has_long_context_pricing"]:
|
||||
assert info["input_cost_per_token_above_200k_tokens"] == 1e-05
|
||||
assert info["output_cost_per_token_above_200k_tokens"] == 3.75e-05
|
||||
assert info["cache_creation_input_token_cost_above_200k_tokens"] == 1.25e-05
|
||||
assert info["cache_read_input_token_cost_above_200k_tokens"] == 1e-06
|
||||
else:
|
||||
assert "input_cost_per_token_above_200k_tokens" not in info
|
||||
assert "output_cost_per_token_above_200k_tokens" not in info
|
||||
assert "cache_creation_input_token_cost_above_200k_tokens" not in info
|
||||
assert "cache_read_input_token_cost_above_200k_tokens" not in info
|
||||
|
||||
assert info["supports_assistant_prefill"] is False
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
|
||||
|
||||
def test_opus_4_6_bedrock_regional_model_pricing():
|
||||
json_path = os.path.join(
|
||||
os.path.dirname(__file__), "../../model_prices_and_context_window.json"
|
||||
)
|
||||
with open(json_path) as f:
|
||||
model_data = json.load(f)
|
||||
|
||||
expected_models = {
|
||||
"global.anthropic.claude-opus-4-6-v1": {
|
||||
"input_cost_per_token": 5e-06,
|
||||
"output_cost_per_token": 2.5e-05,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
},
|
||||
"us.anthropic.claude-opus-4-6-v1": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
},
|
||||
"eu.anthropic.claude-opus-4-6-v1": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
},
|
||||
"au.anthropic.claude-opus-4-6-v1": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
},
|
||||
}
|
||||
|
||||
for model_name, expected in expected_models.items():
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
assert info["litellm_provider"] == "bedrock_converse"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["max_tokens"] == 128000
|
||||
assert info["supports_assistant_prefill"] is False
|
||||
assert "input_cost_per_token_above_200k_tokens" not in info
|
||||
assert "output_cost_per_token_above_200k_tokens" not in info
|
||||
assert "cache_creation_input_token_cost_above_200k_tokens" not in info
|
||||
assert "cache_read_input_token_cost_above_200k_tokens" not in info
|
||||
for key, value in expected.items():
|
||||
assert info[key] == value
|
||||
|
||||
|
||||
def test_opus_4_6_alias_and_dated_metadata_match():
|
||||
json_path = os.path.join(
|
||||
os.path.dirname(__file__), "../../model_prices_and_context_window.json"
|
||||
|
|
|
|||
|
|
@ -16,7 +16,6 @@ import os
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import BEDROCK_CONVERSE_MODELS
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
|
||||
|
|
@ -29,102 +28,6 @@ def _load_root_cost_map() -> dict:
|
|||
return json.load(f)
|
||||
|
||||
|
||||
|
||||
def test_opus_4_8_model_pricing_and_capabilities():
|
||||
model_data = _load_root_cost_map()
|
||||
|
||||
expected_models = {
|
||||
"claude-opus-4-8": {
|
||||
"provider": "anthropic",
|
||||
"max_input_tokens": 1000000,
|
||||
},
|
||||
"anthropic.claude-opus-4-8": {
|
||||
"provider": "bedrock_converse",
|
||||
"max_input_tokens": 1000000,
|
||||
},
|
||||
"vertex_ai/claude-opus-4-8": {
|
||||
"provider": "vertex_ai-anthropic_models",
|
||||
"max_input_tokens": 1000000,
|
||||
},
|
||||
"azure_ai/claude-opus-4-8": {
|
||||
"provider": "azure_ai",
|
||||
"max_input_tokens": 1000000,
|
||||
},
|
||||
}
|
||||
|
||||
for model_name, config in expected_models.items():
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
|
||||
assert info["litellm_provider"] == config["provider"]
|
||||
assert info["mode"] == "chat"
|
||||
assert info["max_input_tokens"] == config["max_input_tokens"]
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["max_tokens"] == 128000
|
||||
|
||||
# Base pricing matches Opus 4.7: $5 / $25 per MTok, with the standard
|
||||
# 1.25x cache-write and 0.1x cache-read multipliers.
|
||||
assert info["input_cost_per_token"] == 5e-06
|
||||
assert info["output_cost_per_token"] == 2.5e-05
|
||||
assert info["cache_creation_input_token_cost"] == 6.25e-06
|
||||
assert info["cache_read_input_token_cost"] == 5e-07
|
||||
|
||||
# Opus 4.x flagships are flat-rate across the full context window.
|
||||
assert "input_cost_per_token_above_200k_tokens" not in info
|
||||
assert "output_cost_per_token_above_200k_tokens" not in info
|
||||
|
||||
assert info["supports_assistant_prefill"] is False
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
|
||||
assert model_data["claude-opus-4-8"]["supports_native_structured_output"] is True
|
||||
|
||||
|
||||
def test_opus_4_8_bedrock_regional_model_pricing():
|
||||
model_data = _load_root_cost_map()
|
||||
|
||||
# Global endpoints use base pricing; regional endpoints carry a 10% premium.
|
||||
expected_models = {
|
||||
"global.anthropic.claude-opus-4-8": {
|
||||
"input_cost_per_token": 5e-06,
|
||||
"output_cost_per_token": 2.5e-05,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
},
|
||||
"us.anthropic.claude-opus-4-8": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
},
|
||||
"eu.anthropic.claude-opus-4-8": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
},
|
||||
"au.anthropic.claude-opus-4-8": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
},
|
||||
}
|
||||
|
||||
for model_name, expected in expected_models.items():
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
assert info["litellm_provider"] == "bedrock_converse"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["bedrock_output_config_effort_ceiling"] == "xhigh"
|
||||
for key, value in expected.items():
|
||||
assert info[key] == value
|
||||
|
||||
|
||||
def test_opus_4_8_fast_mode_multiplier():
|
||||
"""Opus 4.8 dropped fast-mode pricing to 2x base ($10/$50 per MTok);
|
||||
Opus 4.7 was 6x ($30/$150)."""
|
||||
|
|
@ -134,44 +37,10 @@ def test_opus_4_8_fast_mode_multiplier():
|
|||
assert entry["fast"] == 2.0
|
||||
|
||||
|
||||
def test_opus_4_8_present_in_bundled_backup():
|
||||
"""The bundled backup is the runtime fallback (and what tests load with
|
||||
``LITELLM_LOCAL_MODEL_COST_MAP=True``) — it must carry the same entries as
|
||||
the root cost map, otherwise the model resolves on one path but not the
|
||||
other."""
|
||||
backup = GetModelCostMap.load_local_model_cost_map()
|
||||
for model_name in (
|
||||
"claude-opus-4-8",
|
||||
"anthropic.claude-opus-4-8",
|
||||
"global.anthropic.claude-opus-4-8",
|
||||
"us.anthropic.claude-opus-4-8",
|
||||
"eu.anthropic.claude-opus-4-8",
|
||||
"au.anthropic.claude-opus-4-8",
|
||||
"vertex_ai/claude-opus-4-8",
|
||||
"vertex_ai/claude-opus-4-8@default",
|
||||
"azure_ai/claude-opus-4-8",
|
||||
):
|
||||
assert model_name in backup, f"Missing from backup cost map: {model_name}"
|
||||
assert backup["claude-opus-4-8"]["supports_native_structured_output"] is True
|
||||
|
||||
|
||||
def test_opus_4_8_registered_for_bedrock_converse():
|
||||
assert "anthropic.claude-opus-4-8" in BEDROCK_CONVERSE_MODELS
|
||||
|
||||
|
||||
def test_opus_4_8_provider_resolves_via_model_info(local_model_cost_map):
|
||||
"""Regression: ``claude-opus-4-8`` must resolve to provider ``anthropic``.
|
||||
|
||||
Before the cost-map entry existed, the model was unknown to LiteLLM, so it
|
||||
could not be tied to the ``anthropic`` provider and an ``anthropic/*``
|
||||
wildcard deployment would not match it.
|
||||
"""
|
||||
info = litellm.get_model_info(model="claude-opus-4-8")
|
||||
assert info["litellm_provider"] == "anthropic"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cost_map",
|
||||
[_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
|
||||
|
|
|
|||
|
|
@ -17,7 +17,6 @@ import os
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import BEDROCK_CONVERSE_MODELS
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
|
||||
|
|
@ -52,91 +51,6 @@ def _load_root_cost_map() -> dict:
|
|||
return json.load(f)
|
||||
|
||||
|
||||
|
||||
def test_opus_5_pricing_and_capabilities():
|
||||
model_data = _load_root_cost_map()
|
||||
|
||||
expected_providers = {
|
||||
"claude-opus-5": "anthropic",
|
||||
"anthropic.claude-opus-5": "bedrock_converse",
|
||||
"vertex_ai/claude-opus-5": "vertex_ai-anthropic_models",
|
||||
"azure_ai/claude-opus-5": "azure_ai",
|
||||
}
|
||||
|
||||
for model_name, provider in expected_providers.items():
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
|
||||
assert info["litellm_provider"] == provider
|
||||
assert info["mode"] == "chat"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["max_tokens"] == 128000
|
||||
|
||||
# Opus 5 ships at Opus 4.8's rates: $5 / $25 per MTok, with the standard
|
||||
# 1.25x cache-write, 2x 1-hour cache-write, and 0.1x cache-read multipliers.
|
||||
assert info["input_cost_per_token"] == 5e-06
|
||||
assert info["output_cost_per_token"] == 2.5e-05
|
||||
assert info["cache_creation_input_token_cost"] == 6.25e-06
|
||||
assert info["cache_creation_input_token_cost_above_1hr"] == 1e-05
|
||||
assert info["cache_read_input_token_cost"] == 5e-07
|
||||
|
||||
# Flat rate across the full 1M window, no long-context premium.
|
||||
assert "input_cost_per_token_above_200k_tokens" not in info
|
||||
assert "output_cost_per_token_above_200k_tokens" not in info
|
||||
|
||||
# gen-5 adaptive-thinking profile: effort-driven, no sampling params, no
|
||||
# assistant prefill.
|
||||
assert info["supports_adaptive_thinking"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_sampling_params"] is False
|
||||
assert info["supports_assistant_prefill"] is False
|
||||
assert info["supports_xhigh_reasoning_effort"] is True
|
||||
assert info["supports_max_reasoning_effort"] is True
|
||||
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
|
||||
|
||||
def test_opus_5_bedrock_regional_pricing():
|
||||
"""Global/base endpoints use base pricing; the us./eu./au./jp. regional
|
||||
cross-region inference profiles carry a 10% premium."""
|
||||
model_data = _load_root_cost_map()
|
||||
|
||||
base_pricing = {
|
||||
"input_cost_per_token": 5e-06,
|
||||
"output_cost_per_token": 2.5e-05,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
}
|
||||
regional_pricing = {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
}
|
||||
|
||||
expected = {
|
||||
"anthropic.claude-opus-5": base_pricing,
|
||||
"global.anthropic.claude-opus-5": base_pricing,
|
||||
"us.anthropic.claude-opus-5": regional_pricing,
|
||||
"eu.anthropic.claude-opus-5": regional_pricing,
|
||||
"au.anthropic.claude-opus-5": regional_pricing,
|
||||
"jp.anthropic.claude-opus-5": regional_pricing,
|
||||
}
|
||||
|
||||
for model_name, pricing in expected.items():
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
assert info["litellm_provider"] == "bedrock_converse"
|
||||
for key, value in pricing.items():
|
||||
assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS)
|
||||
def test_opus_5_bedrock_entries_declare_no_effort_ceiling(model_name):
|
||||
"""Bedrock accepts every effort level for Opus 5, so no clamp belongs here.
|
||||
|
|
@ -216,18 +130,6 @@ def test_opus_5_registered_for_bedrock_converse():
|
|||
assert "anthropic.claude-opus-5" in BEDROCK_CONVERSE_MODELS
|
||||
|
||||
|
||||
def test_opus_5_provider_resolves_via_model_info(local_model_cost_map):
|
||||
"""Regression: ``claude-opus-5`` must resolve to provider ``anthropic``.
|
||||
|
||||
Without the cost-map entry the model is unknown to LiteLLM, so it cannot be
|
||||
tied to the ``anthropic`` provider and an ``anthropic/*`` wildcard deployment
|
||||
would not match it."""
|
||||
info = litellm.get_model_info(model="claude-opus-5")
|
||||
assert info["litellm_provider"] == "anthropic"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cost_map",
|
||||
[_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
|
||||
|
|
|
|||
|
|
@ -15,7 +15,6 @@ import os
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import BEDROCK_CONVERSE_MODELS
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
|
||||
|
|
@ -41,96 +40,6 @@ def _load_root_cost_map() -> dict:
|
|||
return json.load(f)
|
||||
|
||||
|
||||
|
||||
def test_sonnet_5_pricing_and_capabilities():
|
||||
model_data = _load_root_cost_map()
|
||||
|
||||
expected_providers = {
|
||||
"claude-sonnet-5": "anthropic",
|
||||
"anthropic.claude-sonnet-5": "bedrock_converse",
|
||||
"vertex_ai/claude-sonnet-5": "vertex_ai-anthropic_models",
|
||||
"azure_ai/claude-sonnet-5": "azure_ai",
|
||||
}
|
||||
|
||||
for model_name, provider in expected_providers.items():
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
|
||||
assert info["litellm_provider"] == provider
|
||||
assert info["mode"] == "chat"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["max_tokens"] == 128000
|
||||
|
||||
# Introductory Sonnet 5 pricing through 2026-08-31: $2 / $10 per MTok,
|
||||
# with the 1.25x cache-write and 0.1x cache-read multipliers. On
|
||||
# 2026-09-01 flip these five fields back to the sticker rate, here and
|
||||
# in both cost-map JSON files (all ten claude-sonnet-5 entries):
|
||||
# input_cost_per_token: 3e-06
|
||||
# output_cost_per_token: 1.5e-05
|
||||
# cache_creation_input_token_cost: 3.75e-06
|
||||
# cache_creation_input_token_cost_above_1hr: 6e-06
|
||||
# cache_read_input_token_cost: 3e-07
|
||||
# Regional Bedrock profiles (us./eu./au./jp.) stay at 1.1x those values:
|
||||
# 3.3e-06 / 1.65e-05 / 4.125e-06 / 6.6e-06 / 3.3e-07 (see
|
||||
# test_sonnet_5_bedrock_regional_pricing below).
|
||||
assert info["input_cost_per_token"] == 2e-06
|
||||
assert info["output_cost_per_token"] == 1e-05
|
||||
assert info["cache_creation_input_token_cost"] == 2.5e-06
|
||||
assert info["cache_creation_input_token_cost_above_1hr"] == 4e-06
|
||||
assert info["cache_read_input_token_cost"] == 2e-07
|
||||
|
||||
# gen-5 adaptive-thinking profile: effort-driven, no sampling params, no
|
||||
# assistant prefill.
|
||||
assert info["supports_adaptive_thinking"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_sampling_params"] is False
|
||||
assert info["supports_assistant_prefill"] is False
|
||||
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
|
||||
|
||||
def test_sonnet_5_bedrock_regional_pricing():
|
||||
"""Global/base endpoints use base pricing; the us./eu./au./jp. regional
|
||||
cross-region inference profiles carry a 10% premium."""
|
||||
model_data = _load_root_cost_map()
|
||||
|
||||
base_pricing = {
|
||||
"input_cost_per_token": 2e-06,
|
||||
"output_cost_per_token": 1e-05,
|
||||
"cache_creation_input_token_cost": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 4e-06,
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
}
|
||||
regional_pricing = {
|
||||
"input_cost_per_token": 2.2e-06,
|
||||
"output_cost_per_token": 1.1e-05,
|
||||
"cache_creation_input_token_cost": 2.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 4.4e-06,
|
||||
"cache_read_input_token_cost": 2.2e-07,
|
||||
}
|
||||
|
||||
expected = {
|
||||
"anthropic.claude-sonnet-5": base_pricing,
|
||||
"global.anthropic.claude-sonnet-5": base_pricing,
|
||||
"us.anthropic.claude-sonnet-5": regional_pricing,
|
||||
"eu.anthropic.claude-sonnet-5": regional_pricing,
|
||||
"au.anthropic.claude-sonnet-5": regional_pricing,
|
||||
"jp.anthropic.claude-sonnet-5": regional_pricing,
|
||||
}
|
||||
|
||||
for model_name, pricing in expected.items():
|
||||
assert model_name in model_data, f"Missing model entry: {model_name}"
|
||||
info = model_data[model_name]
|
||||
assert info["litellm_provider"] == "bedrock_converse"
|
||||
assert info["bedrock_output_config_effort_ceiling"] == "xhigh"
|
||||
for key, value in pricing.items():
|
||||
assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}"
|
||||
|
||||
|
||||
def test_sonnet_5_present_in_bundled_backup():
|
||||
"""The bundled backup is the runtime fallback (and what tests load with
|
||||
``LITELLM_LOCAL_MODEL_COST_MAP=True``); it must carry the same entries as the
|
||||
|
|
@ -144,18 +53,6 @@ def test_sonnet_5_registered_for_bedrock_converse():
|
|||
assert "anthropic.claude-sonnet-5" in BEDROCK_CONVERSE_MODELS
|
||||
|
||||
|
||||
def test_sonnet_5_provider_resolves_via_model_info(local_model_cost_map):
|
||||
"""Regression: ``claude-sonnet-5`` must resolve to provider ``anthropic``.
|
||||
|
||||
Before the cost-map entry existed, the model was unknown to LiteLLM, so it
|
||||
could not be tied to the ``anthropic`` provider and an ``anthropic/*``
|
||||
wildcard deployment would not match it."""
|
||||
info = litellm.get_model_info(model="claude-sonnet-5")
|
||||
assert info["litellm_provider"] == "anthropic"
|
||||
assert info["max_input_tokens"] == 1000000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cost_map",
|
||||
[_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
|
||||
|
|
|
|||
|
|
@ -27,17 +27,6 @@ BACKUP_MAP = os.path.join(
|
|||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _use_local_model_cost_map(monkeypatch):
|
||||
original_model_cost = litellm.model_cost
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
|
||||
|
||||
def _load(path: str) -> dict:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
|
@ -47,50 +36,6 @@ def _cloudflare_keys(data: dict) -> set:
|
|||
return {k for k in data if k.startswith("cloudflare/")}
|
||||
|
||||
|
||||
def test_glm_5_2_entry_is_present_and_well_formed():
|
||||
entry = litellm.model_cost["cloudflare/@cf/zai-org/glm-5.2"]
|
||||
assert entry["litellm_provider"] == "cloudflare"
|
||||
assert entry["mode"] == "chat"
|
||||
assert entry["supports_function_calling"] is True
|
||||
assert entry["input_cost_per_token"] > 0
|
||||
assert entry["output_cost_per_token"] > 0
|
||||
|
||||
|
||||
def test_vision_model_is_flagged_supports_vision():
|
||||
entry = litellm.model_cost["cloudflare/@cf/meta/llama-3.2-11b-vision-instruct"]
|
||||
assert entry["litellm_provider"] == "cloudflare"
|
||||
assert entry.get("supports_vision") is True
|
||||
|
||||
|
||||
def test_additional_current_models_are_present():
|
||||
for key in (
|
||||
"cloudflare/@cf/openai/gpt-oss-120b",
|
||||
"cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast",
|
||||
):
|
||||
entry = litellm.model_cost[key]
|
||||
assert entry["litellm_provider"] == "cloudflare"
|
||||
assert entry["mode"] == "chat"
|
||||
assert entry["supports_function_calling"] is True
|
||||
assert entry["input_cost_per_token"] > 0
|
||||
assert entry["output_cost_per_token"] > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"key, published_price_per_audio_minute",
|
||||
[
|
||||
("cloudflare/@cf/openai/whisper", 0.00045),
|
||||
("cloudflare/@cf/openai/whisper-large-v3-turbo", 0.00051),
|
||||
],
|
||||
)
|
||||
def test_whisper_transcription_pricing_is_stored_per_second(key, published_price_per_audio_minute):
|
||||
entry = litellm.model_cost[key]
|
||||
assert entry["litellm_provider"] == "cloudflare"
|
||||
assert entry["mode"] == "audio_transcription"
|
||||
assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"]
|
||||
assert entry["output_cost_per_second"] == 0.0
|
||||
assert entry["input_cost_per_second"] == pytest.approx(published_price_per_audio_minute / 60)
|
||||
|
||||
|
||||
def test_root_and_backup_have_identical_cloudflare_keys():
|
||||
if not os.path.exists(ROOT_MAP):
|
||||
pytest.skip("root cost map only ships in source checkouts")
|
||||
|
|
|
|||
|
|
@ -32,22 +32,6 @@ def _load(path):
|
|||
return json.load(f)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", DAYBREAK_MODELS)
|
||||
def test_daybreak_capability_contract(model):
|
||||
info = _load(MAIN_PATH).get(model)
|
||||
assert info is not None, f"{model} missing from model_prices_and_context_window.json"
|
||||
|
||||
assert info["litellm_provider"] == "openai"
|
||||
assert info["mode"] == "chat"
|
||||
assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"]
|
||||
|
||||
assert info["supports_computer_use"] is True
|
||||
assert info["supports_parallel_function_calling"] is True
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_vision"] is True
|
||||
|
||||
|
||||
def test_blue_alias_matches_its_snapshot_computer_use():
|
||||
cost_map = _load(MAIN_PATH)
|
||||
|
||||
|
|
|
|||
|
|
@ -39,26 +39,6 @@ class TestDeepSeekModelCostEntries:
|
|||
"""Verify that provider-prefixed DeepSeek entries contain the same
|
||||
capability flags as their bare-name counterparts in the JSON files."""
|
||||
|
||||
def test_deepseek_chat_supports_response_schema_in_backup(self):
|
||||
data = _load_backup_json()
|
||||
entry = data.get("deepseek/deepseek-chat", {})
|
||||
assert entry.get("supports_response_schema") is True
|
||||
|
||||
def test_deepseek_reasoner_supports_response_schema_in_backup(self):
|
||||
data = _load_backup_json()
|
||||
entry = data.get("deepseek/deepseek-reasoner", {})
|
||||
assert entry.get("supports_response_schema") is True
|
||||
|
||||
def test_deepseek_chat_supports_system_messages_in_backup(self):
|
||||
data = _load_backup_json()
|
||||
entry = data.get("deepseek/deepseek-chat", {})
|
||||
assert entry.get("supports_system_messages") is True
|
||||
|
||||
def test_deepseek_reasoner_supports_system_messages_in_backup(self):
|
||||
data = _load_backup_json()
|
||||
entry = data.get("deepseek/deepseek-reasoner", {})
|
||||
assert entry.get("supports_system_messages") is True
|
||||
|
||||
def test_deepseek_chat_max_input_tokens_matches_bare_in_backup(self):
|
||||
data = _load_backup_json()
|
||||
bare = data.get("deepseek-chat", {})
|
||||
|
|
@ -71,26 +51,6 @@ class TestDeepSeekModelCostEntries:
|
|||
prefixed = data.get("deepseek/deepseek-reasoner", {})
|
||||
assert prefixed.get("max_output_tokens") == bare.get("max_output_tokens")
|
||||
|
||||
def test_main_json_deepseek_chat_supports_response_schema(self):
|
||||
main_path = os.path.join(
|
||||
os.path.dirname(os.path.dirname(litellm.__file__)),
|
||||
"model_prices_and_context_window.json",
|
||||
)
|
||||
with open(main_path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
entry = data.get("deepseek/deepseek-chat", {})
|
||||
assert entry.get("supports_response_schema") is True
|
||||
|
||||
def test_main_json_deepseek_reasoner_supports_response_schema(self):
|
||||
main_path = os.path.join(
|
||||
os.path.dirname(os.path.dirname(litellm.__file__)),
|
||||
"model_prices_and_context_window.json",
|
||||
)
|
||||
with open(main_path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
entry = data.get("deepseek/deepseek-reasoner", {})
|
||||
assert entry.get("supports_response_schema") is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# API-level tests – verify supports_response_schema returns True
|
||||
|
|
|
|||
|
|
@ -14,26 +14,9 @@ import os
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def _local_model_cost_map():
|
||||
"""
|
||||
Point litellm at the bundled cost map for the duration of this module
|
||||
only. ``mp.undo()`` restores both the environment variable and
|
||||
``litellm.model_cost`` so nothing leaks into later tests.
|
||||
"""
|
||||
mp = pytest.MonkeyPatch()
|
||||
mp.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
mp.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
get_model_info.cache_clear()
|
||||
yield
|
||||
mp.undo()
|
||||
get_model_info.cache_clear()
|
||||
|
||||
|
||||
NEW_ENTRIES = {
|
||||
"fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": {
|
||||
"input_cost_per_token": 1.32e-06,
|
||||
|
|
@ -54,19 +37,6 @@ def model_data():
|
|||
return json.load(f)
|
||||
|
||||
|
||||
def test_fireworks_serverless_entries_exist(model_data):
|
||||
"""The new prefixed entry carries the pricing and metadata from #37274."""
|
||||
for key, expected in NEW_ENTRIES.items():
|
||||
assert key in model_data, f"{key} is missing from model_prices_and_context_window.json"
|
||||
entry = model_data[key]
|
||||
for field, value in expected.items():
|
||||
assert entry[field] == pytest.approx(value), f"{key}.{field}"
|
||||
assert entry["litellm_provider"] == "fireworks_ai"
|
||||
assert entry["mode"] == "chat"
|
||||
assert entry["supports_function_calling"] is True
|
||||
assert entry["supports_vision"] is False
|
||||
|
||||
|
||||
def test_bare_fireworks_ids_resolve_through_prefixed_entries():
|
||||
"""Bare IDs from #37274 resolve via the provider-prefix lookup path."""
|
||||
for bare_id, prefixed_key in [
|
||||
|
|
|
|||
|
|
@ -1,54 +1,6 @@
|
|||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["azure_ai/gpt-5.5", "azure_ai/gpt-5.5-2026-04-23"])
|
||||
def test_azure_ai_gpt_5_5_model_info(model):
|
||||
json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
|
||||
with open(json_path) as f:
|
||||
model_cost = json.load(f)
|
||||
|
||||
info = model_cost.get(model)
|
||||
assert (
|
||||
info is not None
|
||||
), f"{model} not found in model_prices_and_context_window.json"
|
||||
|
||||
assert info["litellm_provider"] == "azure_ai"
|
||||
assert info["mode"] == "chat"
|
||||
|
||||
assert info["input_cost_per_token"] == 5e-06
|
||||
assert info["output_cost_per_token"] == 3e-05
|
||||
assert info["cache_read_input_token_cost"] == 5e-07
|
||||
|
||||
assert info["input_cost_per_token_above_272k_tokens"] == 1e-05
|
||||
assert info["output_cost_per_token_above_272k_tokens"] == 4.5e-05
|
||||
assert info["cache_read_input_token_cost_above_272k_tokens"] == 1e-06
|
||||
|
||||
assert info["input_cost_per_token_priority"] == 1e-05
|
||||
assert info["output_cost_per_token_priority"] == 6e-05
|
||||
|
||||
assert info["max_input_tokens"] == 1050000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["max_tokens"] == 128000
|
||||
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_response_schema"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_web_search"] is True
|
||||
# gpt-5.5 dropped minimal reasoning effort support (true on gpt-5.4)
|
||||
assert info["supports_minimal_reasoning_effort"] is False
|
||||
|
||||
routed_model, provider, _, _ = get_llm_provider(model=model)
|
||||
assert routed_model == model.split("/", 1)[1]
|
||||
assert provider == "azure_ai"
|
||||
|
||||
|
||||
def test_azure_ai_gpt_5_5_backup_matches_main():
|
||||
"""Ensure the bundled model cost map stays in sync with the canonical file."""
|
||||
|
|
|
|||
|
|
@ -1,10 +1,8 @@
|
|||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from typing_extensions import get_args, get_type_hints
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import ModelInfoBase
|
||||
|
||||
REALTIME_ONLY_GPT_MODELS = (
|
||||
|
|
@ -43,44 +41,11 @@ REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS = (
|
|||
ALL_REALTIME_ONLY_GPT_MODELS = REALTIME_ONLY_GPT_MODELS + REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS
|
||||
|
||||
|
||||
def _load_cost_map() -> dict:
|
||||
json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
|
||||
with open(json_path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def test_realtime_is_a_valid_mode_literal():
|
||||
hints = get_type_hints(ModelInfoBase, include_extras=False)
|
||||
assert "realtime" in get_args(hints["mode"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS)
|
||||
def test_realtime_only_gpt_models_are_mode_realtime(model):
|
||||
"""These models only serve /v1/realtime and are rejected by /v1/chat/completions
|
||||
("This is not a chat model ..."), so they must not be tagged mode=chat."""
|
||||
info = _load_cost_map()[model]
|
||||
assert info["supported_endpoints"] == ["/v1/realtime"]
|
||||
assert info["mode"] == "realtime"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS)
|
||||
def test_realtime_only_gpt_4o_models_are_mode_realtime(model):
|
||||
"""gpt-4o(-mini)-realtime-preview are realtime-only and must not be mode=chat."""
|
||||
assert _load_cost_map()[model]["mode"] == "realtime"
|
||||
|
||||
|
||||
def test_get_model_info_reports_realtime_mode(monkeypatch):
|
||||
"""get_model_info must resolve the retag against the bundled cost map, not the
|
||||
hosted map fetched from main, which lags this repo until the next promotion."""
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
litellm.get_model_info.cache_clear()
|
||||
try:
|
||||
assert litellm.get_model_info("gpt-realtime-mini")["mode"] == "realtime"
|
||||
finally:
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
|
||||
def test_backup_matches_main_for_realtime_models():
|
||||
repo_root = Path(__file__).parents[2]
|
||||
with open(repo_root / "model_prices_and_context_window.json") as f:
|
||||
|
|
|
|||
|
|
@ -3,8 +3,6 @@ from pathlib import Path
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[2]
|
||||
MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json"
|
||||
|
|
@ -27,56 +25,6 @@ def _load(path):
|
|||
return json.load(f)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MEDIUM_3_5_MODELS)
|
||||
def test_medium_3_5_specs(model):
|
||||
info = _load(MAIN_PATH).get(model)
|
||||
assert info is not None, f"{model} missing from model_prices_and_context_window.json"
|
||||
|
||||
assert info["litellm_provider"] == "mistral"
|
||||
assert info["mode"] == "chat"
|
||||
|
||||
assert info["input_cost_per_token"] == 1.5e-06
|
||||
assert info["output_cost_per_token"] == 7.5e-06
|
||||
|
||||
assert info["max_input_tokens"] == 262144
|
||||
assert info["max_output_tokens"] == 262144
|
||||
assert info["max_tokens"] == 262144
|
||||
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_response_schema"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_assistant_prefill"] is True
|
||||
|
||||
routed_model, provider, _, _ = get_llm_provider(model=model)
|
||||
assert routed_model == model.split("/", 1)[1]
|
||||
assert provider == "mistral"
|
||||
|
||||
|
||||
def test_mistral_medium_latest_resolves_to_medium_3_5(local_model_cost_map):
|
||||
"""LIT-3883: the -latest alias was retargeted to Medium 3.5; get_model_info must
|
||||
return the 3.5 pricing/context/reasoning, not the stale Medium 3.1 values."""
|
||||
info = litellm.get_model_info(model="mistral/mistral-medium-latest")
|
||||
|
||||
assert info["input_cost_per_token"] == 1.5e-06
|
||||
assert info["output_cost_per_token"] == 7.5e-06
|
||||
assert info["max_input_tokens"] == 262144
|
||||
assert info["supports_reasoning"] is True
|
||||
|
||||
|
||||
def test_mistral_medium_2508_keeps_medium_3_1_specs():
|
||||
"""The date-pinned 2508 alias is Medium 3.1 and must not inherit 3.5 pricing."""
|
||||
info = _load(MAIN_PATH).get("mistral/mistral-medium-2508")
|
||||
assert info is not None, "mistral/mistral-medium-2508 missing from cost map"
|
||||
|
||||
assert info["input_cost_per_token"] == 4e-07
|
||||
assert info["output_cost_per_token"] == 2e-06
|
||||
assert info["max_input_tokens"] == 131072
|
||||
assert info.get("supports_reasoning") is not True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", SYNCED_MODELS)
|
||||
def test_backup_matches_main(model):
|
||||
"""Ensure the bundled (backup) cost map stays in sync with the canonical file."""
|
||||
|
|
|
|||
|
|
@ -18,29 +18,6 @@ def _load(path):
|
|||
return json.load(f)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", SMALL_4_0_MODELS)
|
||||
def test_small_4_0_specs(model):
|
||||
info = _load(MAIN_PATH).get(model)
|
||||
assert info is not None, f"{model} missing from model_prices_and_context_window.json"
|
||||
|
||||
assert info["litellm_provider"] == "mistral"
|
||||
assert info["mode"] == "chat"
|
||||
|
||||
assert info["input_cost_per_token"] == 1.5e-07
|
||||
assert info["output_cost_per_token"] == 6e-07
|
||||
|
||||
assert info["max_input_tokens"] == 262144
|
||||
assert info["max_output_tokens"] == 262144
|
||||
assert info["max_tokens"] == 262144
|
||||
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_response_schema"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_assistant_prefill"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", SMALL_4_0_MODELS)
|
||||
def test_backup_matches_main(model):
|
||||
main_cost = _load(MAIN_PATH)
|
||||
|
|
|
|||
|
|
@ -23,46 +23,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di
|
|||
return json.load(f)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING)
|
||||
def test_muse_spark_1_2_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float):
|
||||
info = _load_cost_map().get(model)
|
||||
assert info is not None, f"{model} not found in model_prices_and_context_window.json"
|
||||
|
||||
assert info["litellm_provider"] == "meta"
|
||||
assert info["mode"] == "chat"
|
||||
|
||||
assert info["input_cost_per_token"] == input_cost
|
||||
assert info["output_cost_per_token"] == output_cost
|
||||
assert info["cache_read_input_token_cost"] == cached_cost
|
||||
|
||||
assert info["max_input_tokens"] == 1048576
|
||||
assert info["max_output_tokens"] == 131072
|
||||
assert info["max_tokens"] == 131072
|
||||
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_parallel_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_response_schema"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_pdf_input"] is True
|
||||
assert info["supports_web_search"] is True
|
||||
assert info["supports_minimal_reasoning_effort"] is True
|
||||
assert info["supports_xhigh_reasoning_effort"] is True
|
||||
|
||||
assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"]
|
||||
assert info["supported_modalities"] == ["text", "image", "video"]
|
||||
assert info["supported_output_modalities"] == ["text"]
|
||||
|
||||
assert info["search_context_cost_per_query"] == {
|
||||
"search_context_size_high": WEB_SEARCH_COST_PER_QUERY,
|
||||
"search_context_size_low": WEB_SEARCH_COST_PER_QUERY,
|
||||
"search_context_size_medium": WEB_SEARCH_COST_PER_QUERY,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING)
|
||||
def test_muse_spark_1_2_cost_per_token(
|
||||
local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float
|
||||
|
|
|
|||
|
|
@ -23,46 +23,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di
|
|||
return json.load(f)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING)
|
||||
def test_muse_spark_1_3_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float):
|
||||
info = _load_cost_map().get(model)
|
||||
assert info is not None, f"{model} not found in model_prices_and_context_window.json"
|
||||
|
||||
assert info["litellm_provider"] == "meta"
|
||||
assert info["mode"] == "chat"
|
||||
|
||||
assert info["input_cost_per_token"] == input_cost
|
||||
assert info["output_cost_per_token"] == output_cost
|
||||
assert info["cache_read_input_token_cost"] == cached_cost
|
||||
|
||||
assert info["max_input_tokens"] == 1048576
|
||||
assert info["max_output_tokens"] == 131072
|
||||
assert info["max_tokens"] == 131072
|
||||
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_parallel_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_response_schema"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_pdf_input"] is True
|
||||
assert info["supports_web_search"] is True
|
||||
assert info["supports_minimal_reasoning_effort"] is True
|
||||
assert info["supports_xhigh_reasoning_effort"] is True
|
||||
|
||||
assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"]
|
||||
assert info["supported_modalities"] == ["text", "image", "video"]
|
||||
assert info["supported_output_modalities"] == ["text"]
|
||||
|
||||
assert info["search_context_cost_per_query"] == {
|
||||
"search_context_size_high": WEB_SEARCH_COST_PER_QUERY,
|
||||
"search_context_size_low": WEB_SEARCH_COST_PER_QUERY,
|
||||
"search_context_size_medium": WEB_SEARCH_COST_PER_QUERY,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING)
|
||||
def test_muse_spark_1_3_cost_per_token(
|
||||
local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float
|
||||
|
|
|
|||
|
|
@ -20,14 +20,6 @@ def test_replicate_models_have_valid_key_prefix(model_cost: dict[str, Any]) -> N
|
|||
)
|
||||
|
||||
|
||||
def test_replicate_openai_gpt_oss_20b_key_exists(model_cost: dict[str, Any]) -> None:
|
||||
assert "replicate/openai/gpt-oss-20b" in model_cost
|
||||
info = model_cost["replicate/openai/gpt-oss-20b"]
|
||||
assert info["litellm_provider"] == "replicate"
|
||||
assert info["mode"] == "chat"
|
||||
assert info["supports_function_calling"] is True
|
||||
|
||||
|
||||
def test_replicate_backup_matches_main() -> None:
|
||||
repo_root = Path(__file__).parents[2]
|
||||
main_path = repo_root / "model_prices_and_context_window.json"
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ from typing import Final
|
|||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
REPO_ROOT: Final = Path(__file__).parents[2]
|
||||
|
||||
|
|
@ -77,59 +76,6 @@ def cost_map() -> CostMap:
|
|||
return COST_MAP_ADAPTER.validate_python(json.load(f))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", SERVERLESS_CHAT_MODELS)
|
||||
def test_together_serverless_chat_model_is_mapped(cost_map: CostMap, model: str):
|
||||
info = cost_map.get(model)
|
||||
assert info is not None, f"{model} missing from model_prices_and_context_window.json"
|
||||
assert info["litellm_provider"] == "together_ai"
|
||||
assert info["mode"] == "chat"
|
||||
assert info["input_cost_per_token"] >= 0
|
||||
assert info["output_cost_per_token"] >= info["input_cost_per_token"]
|
||||
assert "deprecation_date" not in info
|
||||
|
||||
routed_model, provider, _, _ = get_llm_provider(model=model)
|
||||
assert routed_model == model.removeprefix("together_ai/")
|
||||
assert provider == "together_ai"
|
||||
|
||||
|
||||
def test_together_kimi_k3_pricing_and_capabilities(cost_map: CostMap):
|
||||
info = cost_map["together_ai/moonshotai/Kimi-K3"]
|
||||
assert info["input_cost_per_token"] == 3e-06
|
||||
assert info["output_cost_per_token"] == 1.5e-05
|
||||
assert info["max_input_tokens"] == 1048576
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_response_schema"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
|
||||
|
||||
def test_together_glm_52_pricing(cost_map: CostMap):
|
||||
info = cost_map["together_ai/zai-org/GLM-5.2"]
|
||||
assert info["input_cost_per_token"] == 1.4e-06
|
||||
assert info["output_cost_per_token"] == 4.4e-06
|
||||
assert info["max_input_tokens"] == 1048575
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
|
||||
|
||||
def test_together_glm_53_flash_pricing_and_capabilities(cost_map: CostMap):
|
||||
info = cost_map["together_ai/zai-org/GLM-5.3-Flash"]
|
||||
assert info["input_cost_per_token"] == 1.5e-07
|
||||
assert info["output_cost_per_token"] == 5e-07
|
||||
assert info["cache_read_input_token_cost"] == 3e-08
|
||||
assert info["max_input_tokens"] == 1048575
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_parallel_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_response_schema"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
|
||||
|
||||
def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost_map: CostMap):
|
||||
inflated = sorted(
|
||||
model
|
||||
|
|
@ -142,21 +88,6 @@ def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost
|
|||
assert inflated == []
|
||||
|
||||
|
||||
def test_together_multilingual_e5_embedding_entry(cost_map: CostMap):
|
||||
info = cost_map["together_ai/intfloat/multilingual-e5-large-instruct"]
|
||||
assert info["mode"] == "embedding"
|
||||
assert info["input_cost_per_token"] == 2e-08
|
||||
assert info["max_input_tokens"] == 514
|
||||
assert info["output_vector_size"] == 1024
|
||||
|
||||
|
||||
def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: CostMap):
|
||||
info = cost_map["together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo"]
|
||||
assert info["input_cost_per_token"] == 1.04e-06
|
||||
assert info["output_cost_per_token"] == 1.04e-06
|
||||
assert info["max_input_tokens"] == 131072
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS))
|
||||
def test_together_deprecated_model_carries_deprecation_date(cost_map: CostMap, model: str):
|
||||
info = cost_map.get(model)
|
||||
|
|
@ -210,32 +141,7 @@ CACHED_INPUT_MODELS: Final = (
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", CACHED_INPUT_MODELS)
|
||||
def test_together_cached_input_model_carries_cache_read_pricing(cost_map: CostMap, model: str):
|
||||
info = cost_map.get(model)
|
||||
assert info is not None, f"{model} missing from model_prices_and_context_window.json"
|
||||
assert info.get("supports_prompt_caching") is True
|
||||
cache_read = info.get("cache_read_input_token_cost")
|
||||
assert isinstance(cache_read, float)
|
||||
assert 0 < cache_read < info["input_cost_per_token"]
|
||||
assert "cache_creation_input_token_cost" not in info
|
||||
|
||||
|
||||
def test_together_prompt_caching_flag_implies_cache_read_rate(cost_map: CostMap):
|
||||
for model, info in cost_map.items():
|
||||
if model.startswith("together_ai/") and info.get("supports_prompt_caching"):
|
||||
assert "cache_read_input_token_cost" in info, f"{model} flags caching without a cache read rate"
|
||||
|
||||
|
||||
def test_together_deepseek_v4_flash_cache_read_rate(cost_map: CostMap):
|
||||
info = cost_map["together_ai/deepseek-ai/DeepSeek-V4-Flash-0731"]
|
||||
assert info["input_cost_per_token"] == 1.4e-07
|
||||
assert info["cache_read_input_token_cost"] == 3e-08
|
||||
assert info["output_cost_per_token"] == 2.8e-07
|
||||
|
||||
|
||||
def test_together_qwen_37_max_repriced_to_current_together_rate(cost_map: CostMap):
|
||||
info = cost_map["together_ai/Qwen/Qwen3.7-Max"]
|
||||
assert info["input_cost_per_token"] == 2.5e-06
|
||||
assert info["output_cost_per_token"] == 7.5e-06
|
||||
assert info["cache_read_input_token_cost"] == 5e-07
|
||||
|
|
|
|||
24
tests/test_litellm_rust/conftest.py
Normal file
24
tests/test_litellm_rust/conftest.py
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(items):
|
||||
rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in {
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
}
|
||||
if not rust_enabled:
|
||||
skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension")
|
||||
for item in items:
|
||||
item.add_marker(skip)
|
||||
return
|
||||
|
||||
try:
|
||||
from litellm.rust_bridge import _native # noqa: F401 # validates the installed extension
|
||||
except ImportError as error:
|
||||
raise pytest.UsageError(
|
||||
"LITELLM_RUST=1 requires a compiled litellm.rust_bridge._native extension"
|
||||
) from error
|
||||
72
tests/test_litellm_rust/test_ocr.py
Normal file
72
tests/test_litellm_rust/test_ocr.py
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
import json
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ocr_server():
|
||||
requests = []
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self):
|
||||
requests.append(
|
||||
{
|
||||
"headers": {name.lower(): value for name, value in self.headers.items()},
|
||||
"body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))),
|
||||
}
|
||||
)
|
||||
if self.headers.get("User-Agent", "").startswith("python-httpx"):
|
||||
self.send_response(418)
|
||||
self.end_headers()
|
||||
return
|
||||
response = json.dumps(
|
||||
{
|
||||
"pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}],
|
||||
"model": "mistral-ocr-latest",
|
||||
"usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
|
||||
}
|
||||
).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(response)))
|
||||
self.end_headers()
|
||||
self.wfile.write(response)
|
||||
|
||||
def log_message(self, format, *args):
|
||||
pass
|
||||
|
||||
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
thread = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield server, requests
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join()
|
||||
|
||||
|
||||
def test_ocr_with_rust_extension(ocr_server):
|
||||
server, requests = ocr_server
|
||||
host, port = server.server_address
|
||||
|
||||
response = litellm.ocr(
|
||||
model="mistral/mistral-ocr-latest",
|
||||
document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
|
||||
api_key="test-key",
|
||||
api_base=f"http://{host}:{port}",
|
||||
)
|
||||
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert len(requests) == 1
|
||||
assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx")
|
||||
assert requests[0]["body"] == {
|
||||
"model": "mistral-ocr-latest",
|
||||
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
|
||||
}
|
||||
|
|
@ -313,6 +313,26 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => {
|
|||
return agentData;
|
||||
};
|
||||
|
||||
export const parseMcpPermissionsForForm = (agent: any) => ({
|
||||
allowed_mcp_servers_and_groups: {
|
||||
servers: agent.object_permission?.mcp_servers ?? [],
|
||||
accessGroups: agent.object_permission?.mcp_access_groups ?? [],
|
||||
toolsets: agent.object_permission?.mcp_toolsets ?? [],
|
||||
},
|
||||
mcp_tool_permissions: agent.object_permission?.mcp_tool_permissions ?? {},
|
||||
});
|
||||
|
||||
/**
|
||||
* Always includes every MCP key (empty when cleared) so removals persist;
|
||||
* the proxy merges object_permission per key, leaving non-MCP grants untouched.
|
||||
*/
|
||||
export const buildMcpObjectPermission = (values: any) => ({
|
||||
mcp_servers: values.allowed_mcp_servers_and_groups?.servers ?? [],
|
||||
mcp_access_groups: values.allowed_mcp_servers_and_groups?.accessGroups ?? [],
|
||||
mcp_toolsets: values.allowed_mcp_servers_and_groups?.toolsets ?? [],
|
||||
mcp_tool_permissions: values.mcp_tool_permissions ?? {},
|
||||
});
|
||||
|
||||
/**
|
||||
* Parse agent data for form fields
|
||||
*/
|
||||
|
|
@ -356,5 +376,6 @@ export const parseAgentForForm = (agent: any) => {
|
|||
: [],
|
||||
// extra_headers: already an array of strings
|
||||
extra_headers: agent.extra_headers ?? [],
|
||||
...parseMcpPermissionsForForm(agent),
|
||||
};
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import React from "react";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
|
||||
import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event";
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
|
|
@ -10,6 +11,12 @@ vi.mock("@/components/networking", () => ({
|
|||
getAgentInfo: vi.fn(),
|
||||
patchAgentCall: vi.fn(),
|
||||
getAgentCreateMetadata: vi.fn(),
|
||||
getProxyBaseUrl: vi.fn(() => ""),
|
||||
getUiConfig: vi.fn(async () => ({})),
|
||||
fetchMCPServers: vi.fn(async () => []),
|
||||
fetchMCPAccessGroups: vi.fn(async () => []),
|
||||
fetchMCPToolsets: vi.fn(async () => []),
|
||||
listMCPTools: vi.fn(async () => ({ tools: [] })),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({
|
||||
|
|
@ -111,7 +118,14 @@ const bedrockAgentcoreInfo: AgentCreateInfo = {
|
|||
|
||||
const setup = () => userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
|
||||
|
||||
const renderView = () => render(<AgentInfoView agentId="agent-1" onClose={vi.fn()} accessToken="tok" isAdmin={true} />);
|
||||
const renderView = () => {
|
||||
const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
|
||||
return render(
|
||||
<QueryClientProvider client={queryClient}>
|
||||
<AgentInfoView agentId="agent-1" onClose={vi.fn()} accessToken="tok" isAdmin={true} />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
};
|
||||
|
||||
const openEditor = async (user: ReturnType<typeof setup>) => {
|
||||
await user.click(await screen.findByRole("tab", { name: "Settings" }));
|
||||
|
|
@ -161,6 +175,7 @@ describe("AgentInfoView update payload", () => {
|
|||
rpm_limit: 222,
|
||||
session_tpm_limit: 333,
|
||||
session_rpm_limit: 444,
|
||||
object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} },
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -201,6 +216,7 @@ describe("AgentInfoView update payload", () => {
|
|||
rpm_limit: 222,
|
||||
session_tpm_limit: 333,
|
||||
session_rpm_limit: 444,
|
||||
object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} },
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -278,9 +294,30 @@ describe("AgentInfoView update payload", () => {
|
|||
api_base: "https://other.example.com",
|
||||
model: "langgraph/asst_1",
|
||||
},
|
||||
object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} },
|
||||
});
|
||||
});
|
||||
|
||||
it("keeps the agent's existing MCP grants in the update payload", async () => {
|
||||
const existingMcpGrants = {
|
||||
mcp_servers: ["srv-1"],
|
||||
mcp_access_groups: ["grp-a"],
|
||||
mcp_toolsets: ["toolset-1"],
|
||||
mcp_tool_permissions: { "srv-1": ["tool_x"] },
|
||||
};
|
||||
vi.mocked(networking.getAgentInfo).mockResolvedValue({
|
||||
...A2A_AGENT,
|
||||
object_permission: existingMcpGrants,
|
||||
} as never);
|
||||
const user = setup();
|
||||
renderView();
|
||||
await openEditor(user);
|
||||
|
||||
await save(user);
|
||||
|
||||
expect(patchedPayload().object_permission).toEqual(existingMcpGrants);
|
||||
});
|
||||
|
||||
it("preserves the full AgentCore runtime ARN (including the resource id after runtime/) across an unedited save", async () => {
|
||||
vi.mocked(networking.getAgentCreateMetadata).mockResolvedValue([bedrockAgentcoreInfo]);
|
||||
vi.mocked(networking.getAgentInfo).mockResolvedValue(BEDROCK_AGENTCORE_AGENT as never);
|
||||
|
|
|
|||
|
|
@ -24,6 +24,18 @@ vi.mock("./agent_form_fields", () => ({
|
|||
unmountedA2AFieldNames: () => [],
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({
|
||||
useMCPServers: () => ({ data: [{ server_id: "srv-1", server_name: "github" }] }),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({
|
||||
default: () => <div data-testid="mcp-server-selector" />,
|
||||
}));
|
||||
|
||||
vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({
|
||||
default: () => <div data-testid="mcp-tool-permissions" />,
|
||||
}));
|
||||
|
||||
const agent = {
|
||||
agent_id: "agent-1",
|
||||
agent_name: "support-agent",
|
||||
|
|
@ -62,5 +74,18 @@ describe("AgentInfoView settings", () => {
|
|||
expect(token).toBe("sk-test");
|
||||
expect(agentId).toBe("agent-1");
|
||||
expect(payload.tpm_limit).toBe(42);
|
||||
const clearedMcpGrants = { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} };
|
||||
expect(payload.object_permission).toEqual(clearedMcpGrants);
|
||||
});
|
||||
|
||||
it("shows MCP grants with server names on the overview tab", async () => {
|
||||
vi.mocked(networking.getAgentInfo).mockResolvedValue({
|
||||
...agent,
|
||||
object_permission: { mcp_servers: ["srv-1"] },
|
||||
} as unknown as Agent);
|
||||
|
||||
render(<AgentInfoView agentId="agent-1" onClose={vi.fn()} accessToken="sk-test" isAdmin={true} />);
|
||||
|
||||
expect(await screen.findByText("github (srv-1)")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -15,16 +15,27 @@ import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo }
|
|||
import { Agent } from "@/components/agents/types";
|
||||
import { KeyResponse } from "@/components/key_team_helpers/key_list";
|
||||
import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
|
||||
import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
|
||||
import KeyInfoView from "@/components/templates/key_info_view";
|
||||
import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector";
|
||||
import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions";
|
||||
import AgentVirtualKeys from "./agent_virtual_keys";
|
||||
import AgentFormFields, { unmountedA2AFieldNames } from "./agent_form_fields";
|
||||
import DynamicAgentFormFields, { buildDynamicAgentData, unmountedDynamicFieldNames } from "./dynamic_agent_form_fields";
|
||||
import { AGENT_FORM_CONFIG, buildAgentDataFromForm, parseAgentForForm } from "./agent_config";
|
||||
import {
|
||||
AGENT_FORM_CONFIG,
|
||||
buildAgentDataFromForm,
|
||||
buildMcpObjectPermission,
|
||||
parseAgentForForm,
|
||||
parseMcpPermissionsForForm,
|
||||
} from "./agent_config";
|
||||
import {
|
||||
AgentFormField,
|
||||
AgentFormValues,
|
||||
AgentNumberInput,
|
||||
AgentRequestPayload,
|
||||
McpServerSelection,
|
||||
labelWithHint,
|
||||
omitFieldValues,
|
||||
useCollapsiblePanels,
|
||||
} from "./AgentFormKit";
|
||||
|
|
@ -111,7 +122,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
} else {
|
||||
const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType);
|
||||
if (typeInfo) {
|
||||
form.reset(parseDynamicAgentForForm(data, typeInfo));
|
||||
form.reset({ ...parseDynamicAgentForForm(data, typeInfo), ...parseMcpPermissionsForForm(data) });
|
||||
} else {
|
||||
form.reset(parseAgentForForm(data));
|
||||
}
|
||||
|
|
@ -131,7 +142,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
if (agentType !== "a2a") {
|
||||
const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType);
|
||||
if (typeInfo) {
|
||||
form.reset(parseDynamicAgentForForm(agent, typeInfo));
|
||||
form.reset({ ...parseDynamicAgentForForm(agent, typeInfo), ...parseMcpPermissionsForForm(agent) });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -139,6 +150,14 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
|
||||
const selectedAgentTypeInfo = agentTypeMetadata.find((t) => t.agent_type === detectedAgentType);
|
||||
const watchedFormValues = useWatch({ control: form.control });
|
||||
const mcpSelection = useWatch({ control: form.control, name: "allowed_mcp_servers_and_groups" });
|
||||
const mcpToolPermissions = useWatch({ control: form.control, name: "mcp_tool_permissions" });
|
||||
const { data: mcpServers = [] } = useMCPServers();
|
||||
|
||||
const mcpServerLabel = (serverId: string) => {
|
||||
const server = mcpServers.find((s) => s.server_id === serverId);
|
||||
return server?.server_name ? `${server.server_name} (${serverId})` : serverId;
|
||||
};
|
||||
|
||||
const discoveryRequest = useMemo(
|
||||
() => buildDiscoveryRequest(detectedAgentType, watchedFormValues || {}, selectedAgentTypeInfo),
|
||||
|
|
@ -199,7 +218,10 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
? overlayDiscoveredCardParams(built, appliedDiscoveredSelection.selected_card)
|
||||
: built;
|
||||
|
||||
await patchAgentCall(accessToken, agentId, updateData);
|
||||
await patchAgentCall(accessToken, agentId, {
|
||||
...updateData,
|
||||
object_permission: buildMcpObjectPermission(values),
|
||||
});
|
||||
toast.success("Agent updated successfully");
|
||||
setIsEditing(false);
|
||||
fetchAgentInfo();
|
||||
|
|
@ -337,13 +359,20 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
{agent.object_permission &&
|
||||
(agent.object_permission.mcp_servers?.length ||
|
||||
agent.object_permission.mcp_access_groups?.length ||
|
||||
agent.object_permission.mcp_toolsets?.length ||
|
||||
(agent.object_permission.mcp_tool_permissions &&
|
||||
Object.keys(agent.object_permission.mcp_tool_permissions).length > 0)) && (
|
||||
<div style={{ marginTop: 24 }}>
|
||||
<h3 className="text-lg font-medium">MCP Tool Permissions</h3>
|
||||
<DetailList className="mt-4">
|
||||
{agent.object_permission.mcp_servers && agent.object_permission.mcp_servers.length > 0 && (
|
||||
<DetailItem label="MCP Servers">{agent.object_permission.mcp_servers.join(", ")}</DetailItem>
|
||||
<DetailItem label="MCP Servers">
|
||||
<div className="space-y-1">
|
||||
{agent.object_permission.mcp_servers.map((serverId) => (
|
||||
<div key={serverId}>{mcpServerLabel(serverId)}</div>
|
||||
))}
|
||||
</div>
|
||||
</DetailItem>
|
||||
)}
|
||||
{agent.object_permission.mcp_access_groups &&
|
||||
agent.object_permission.mcp_access_groups.length > 0 && (
|
||||
|
|
@ -351,13 +380,16 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
{agent.object_permission.mcp_access_groups.join(", ")}
|
||||
</DetailItem>
|
||||
)}
|
||||
{agent.object_permission.mcp_toolsets && agent.object_permission.mcp_toolsets.length > 0 && (
|
||||
<DetailItem label="MCP Toolsets">{agent.object_permission.mcp_toolsets.join(", ")}</DetailItem>
|
||||
)}
|
||||
{agent.object_permission.mcp_tool_permissions &&
|
||||
Object.keys(agent.object_permission.mcp_tool_permissions).length > 0 && (
|
||||
<DetailItem label="Tool permissions per server">
|
||||
<div className="space-y-1">
|
||||
{Object.entries(agent.object_permission.mcp_tool_permissions).map(([serverId, tools]) => (
|
||||
<div key={serverId}>
|
||||
<span className="font-medium">{serverId}:</span>{" "}
|
||||
<span className="font-medium">{mcpServerLabel(serverId)}:</span>{" "}
|
||||
{Array.isArray(tools) ? tools.join(", ") : String(tools)}
|
||||
</div>
|
||||
))}
|
||||
|
|
@ -457,6 +489,41 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
{rateLimitField("session_rpm_limit", "Session RPM Limit")}
|
||||
</div>
|
||||
|
||||
<Separator className="my-6" />
|
||||
<h3 className="text-lg font-medium mb-4">MCP Servers</h3>
|
||||
<FieldGroup>
|
||||
<AgentFormField
|
||||
name="allowed_mcp_servers_and_groups"
|
||||
label={labelWithHint(
|
||||
"Allowed MCP Servers",
|
||||
"Select which MCP servers or access groups this agent can access. Keys bound to this agent can only reach servers granted here.",
|
||||
)}
|
||||
>
|
||||
{({ value, onChange }) => (
|
||||
<MCPServerSelector
|
||||
onChange={onChange}
|
||||
value={{
|
||||
servers: (value as McpServerSelection | undefined)?.servers ?? [],
|
||||
accessGroups: (value as McpServerSelection | undefined)?.accessGroups ?? [],
|
||||
toolsets: (value as McpServerSelection | undefined)?.toolsets ?? [],
|
||||
}}
|
||||
accessToken={accessToken ?? ""}
|
||||
placeholder="Select MCP servers or access groups (optional)"
|
||||
/>
|
||||
)}
|
||||
</AgentFormField>
|
||||
</FieldGroup>
|
||||
<div className="mt-4">
|
||||
<MCPToolPermissions
|
||||
accessToken={accessToken ?? ""}
|
||||
selectedServers={mcpSelection?.servers ?? []}
|
||||
toolPermissions={mcpToolPermissions ?? {}}
|
||||
onChange={(toolPerms: Record<string, string[]>) =>
|
||||
form.setValue("mcp_tool_permissions", toolPerms)
|
||||
}
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div className="mt-6 flex justify-end gap-2">
|
||||
<Button
|
||||
type="button"
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue