From d76ec3a47e42405ed684bc0e14cd2ab873770f43 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 21:57:53 +0000 Subject: [PATCH] fix(proxy)!: answer 400 for malformed tools on count_tokens instead of crashing (#46583) * fix(proxy): answer 400 for malformed tools on count_tokens instead of crashing A non-list tools value crashed the vector store lookup during auth and the built-in tool cost lookup in the failure logging path, so the proxy answered an unhandled 500 and closed the keep-alive connection. count_tokens now answers a 400 invalid_request_error for a malformed body, drops malformed tools when drop_params is on, and both lookups skip values of the wrong shape. * fix(proxy): skip non-list tools in the tool allowlist check Also validate the count_tokens body without recursion, which the recursion detector rejects. * test(integration): restore the int tools case on count_tokens The proxy now answers a non-list tools value with a 400, so the case is deterministic again. * fix(proxy): reject malformed tools at the LLM route entry and keep vector store checks on malformed ids * fix(proxy): apply drop_params like core when dropping malformed tools, and reject malformed tool vector store ids Check tools after request setup so the Claude Code default and the logged call type apply, read the request, global, and model group drop_params the way core does, drop tool_choice and friends with tools, cover Gemini generateContent, and answer 400 for a malformed tools[].vector_store_ids with the strict flag off * fix(proxy): keep the request's call type when logging a gate rejection after request setup The body-shape fallback relabeled a /v1/messages rejection as acompletion and a list-input /v1/responses rejection as aembedding, so it now only applies when the logging object carries no known call type * fix(proxy): share the malformed tools check across tool routes and honor drop_params per entry * fix(proxy): accept null vector_store_ids, read every named store id in auth, check the responses websocket * refactor(proxy): leave the responses websocket out of the request-level tool check * fix(proxy): let a request's additional_drop_params replace the deployment list and keep compact's call type on rejections * fix(proxy): read router default drop lists and deployment ids in the malformed tools check * fix(proxy): leave malformed tools to the router when it can pick another deployment Fallbacks, auto-routers, model strings, comma batches and team model names let the router pick a deployment after the proxy checks the request, so the proxy now forwards a malformed tools field there as main does, and only rejects it when the deployment it checks is the one that serves. Token counting never reroutes, so count_tokens and input_tokens keep reading the group they count with. additional_drop_params names are read the way core reads them, so a list with a non-string entry still names its strings. * test(integration): cover malformed tools on reroutable models, router defaults and deployment ids * fix(proxy): resolve malformed tools at the deployment the router picks * fix(proxy): keep provider-repaired tool schemas, honor fallback_on_unsupported_params, skip cooldown on malformed tools Empty schemas and JSON-string schemas pass the tools check the way litellm's provider paths already read them. The rejection is a BadRequestError subclass carrying the deployment's provider, so fallback_on_unsupported_params: false keeps the 400, and the router never cools a deployment down for a request litellm rejected before calling it. Adds the audit's pre_call guardrail and cross-provider bridge cells. * test(integration): expect a 400 for an empty-string tools field on /v1/messages Real Anthropic rejects tools "" with a 400 on main too, so the case joins the malformed tools test * fix(router): keep web search routing for tools passed as a tuple * fix(router): keep a deployment's own tools when drop_params drops a malformed request tools * test(router): list _merged_tools as covered through the deployment tools merge tests * fix(router): drop the request's tool settings with its malformed tools when deployment tools replace them * fix(router): skip a tool whose type is not a string in the web search filter instead of raising TypeError * fix(prometheus): keep deployment health metrics clean on a malformed tools rejection * test: type the new malformed tools tests * test: mark new test locals Final and type kwargs dicts --- litellm/__init__.py | 1 + litellm/exceptions.py | 4 + litellm/integrations/prometheus.py | 7 +- .../llm_cost_calc/tool_call_cost_tracking.py | 13 +- litellm/litellm_core_utils/tool_fields.py | 164 + .../chat/guardrail_translation/handler.py | 5 +- .../base_llm/guardrail_translation/utils.py | 13 +- .../chat/guardrail_translation/handler.py | 11 +- .../guardrail_translation/handler.py | 5 +- .../proxy/anthropic_endpoints/endpoints.py | 57 +- litellm/proxy/common_request_processing.py | 59 + .../proxy/guardrails/tool_name_extraction.py | 10 +- .../proxy/response_api_endpoints/endpoints.py | 34 +- litellm/proxy/utils.py | 5 +- litellm/router.py | 34 +- litellm/router_utils/common_utils.py | 18 +- litellm/router_utils/cooldown_handlers.py | 6 + litellm/utils.py | 31 +- .../vector_stores/vector_store_registry.py | 39 +- .../router_code_coverage.py | 2 + ...st_anthropic_messages_caller_tools_wire.py | 10 +- .../test_bedrock_mantle_count_tokens_wire.py | 17 +- .../providers/test_malformed_tools_chaos.py | 619 ++++ .../providers/test_malformed_tools_wire.py | 2731 +++++++++++++++++ ...ometheus_deployment_state_proxy_rejects.py | 16 +- .../test_tool_call_cost_tracking.py | 23 +- .../litellm_core_utils/test_tool_fields.py | 187 ++ .../anthropic_endpoints/test_endpoints.py | 148 + .../response_api_endpoints/test_endpoints.py | 47 +- .../proxy/test_tools_allowlist_enforcement.py | 15 + .../test_post_call_failure_hook.py | 43 + .../test_router_utils_common_utils.py | 13 + tests/unit/test_router/test_router.py | 90 + .../test_vector_store_registry.py | 50 + 34 files changed, 4435 insertions(+), 92 deletions(-) create mode 100644 litellm/litellm_core_utils/tool_fields.py create mode 100644 tests/integration/providers/test_malformed_tools_chaos.py create mode 100644 tests/integration/providers/test_malformed_tools_wire.py create mode 100644 tests/unit/litellm_core_utils/test_tool_fields.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 88e56c14187..83e7642c39f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1416,6 +1416,7 @@ from .exceptions import ( InvalidRequestError, BadRequestError, ImageFetchError, + MalformedRequestFieldError as MalformedRequestFieldError, VectorStoreSearchError, NotFoundError, PermissionDeniedError, diff --git a/litellm/exceptions.py b/litellm/exceptions.py index b3a0802db13..733f1aa72d5 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -313,6 +313,10 @@ class VectorStoreSearchError(BadRequestError): ) +class MalformedRequestFieldError(BadRequestError): + pass + + class UnprocessableEntityError(openai.UnprocessableEntityError): def __init__( self, diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index fa5496004e3..97a1073d7c0 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -27,6 +27,7 @@ from litellm.constants import ( PROXY_REJECTED_BEFORE_ROUTING_KEY, ) from litellm.exceptions import ( + MalformedRequestFieldError, validate_rate_limit_category, validate_rate_limit_type, ) @@ -3104,7 +3105,11 @@ class PrometheusLogger(CustomLogger): # On LiteLLM-side rejects (no deployment picked), route request_kwargs["model"] # into requested_model and leave deployment-scoped labels empty. - deployment_selected: Final = bool(model_id) and not _litellm_params.get(PROXY_REJECTED_BEFORE_ROUTING_KEY) + deployment_selected: Final = ( + bool(model_id) + and not _litellm_params.get(PROXY_REJECTED_BEFORE_ROUTING_KEY) + and not isinstance(exception, MalformedRequestFieldError) + ) if deployment_selected: label_litellm_model_name = litellm_model_name label_model_id = model_id diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index ce5dceb2e5c..04510c20db0 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -858,9 +858,9 @@ class StandardBuiltInToolCostTracking: @staticmethod def get_web_search_options(kwargs: Mapping[str, object]) -> WebSearchOptions | None: - if "web_search_options" in kwargs: - web_search_options: Final = cast(WebSearchOptions, kwargs.get("web_search_options", {})) - return WebSearchOptions(**web_search_options) + web_search_options: Final = kwargs.get("web_search_options") + if isinstance(web_search_options, dict): + return WebSearchOptions(**cast(WebSearchOptions, web_search_options)) tools: Final = StandardBuiltInToolCostTracking._get_tools_from_kwargs( kwargs=kwargs, tool_type="web_search_preview" @@ -876,10 +876,9 @@ class StandardBuiltInToolCostTracking: _get_web_search_options = get_web_search_options @staticmethod - def _get_tools_from_kwargs(kwargs: Mapping[str, object], tool_type: str) -> list[object] | None: - if "tools" in kwargs: - return cast(list[object], kwargs.get("tools", [])) - return None + def _get_tools_from_kwargs(kwargs: Mapping[str, object], tool_type: str) -> Sequence[object] | None: + tools: Final = kwargs.get("tools") + return cast(Sequence[object], tools) if isinstance(tools, (list, tuple)) else None @staticmethod def get_file_search_tool_call(kwargs: Mapping[str, object]) -> FileSearchTool | None: diff --git a/litellm/litellm_core_utils/tool_fields.py b/litellm/litellm_core_utils/tool_fields.py new file mode 100644 index 00000000000..f90c8f5d174 --- /dev/null +++ b/litellm/litellm_core_utils/tool_fields.py @@ -0,0 +1,164 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from functools import reduce +from typing import Annotated, Final, TypeAlias + +from pydantic import BeforeValidator, ConfigDict, Json, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +import litellm +from litellm.litellm_core_utils.core_helpers import normalize_drop_params + +ROUTES_WITH_TOOLS: Final = frozenset( + { + "acompletion", + "aresponses", + "acompact_responses", + "anthropic_messages", + "agenerate_content", + "agenerate_content_stream", + "acreate_interaction", + } +) +TOOL_FIELDS: Final = frozenset({"tools", "tool_choice", "parallel_tool_calls", "toolConfig", "tool_config"}) +_TOOLS: Final = TypeAdapter(list[object]) +_FIELD_LIST: Final = TypeAdapter(list[object], config=ConfigDict(strict=True)) +_JSON_OBJECT: Final = TypeAdapter(Json[dict[str, object]]) + + +def _schema_as_read(value: object) -> object: + if not value: + return None + if not isinstance(value, str): + return value + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return value + + +_Schema: TypeAlias = Annotated[dict[str, object] | None, BeforeValidator(_schema_as_read)] + + +class _ToolFunction(TypedDict, total=False): + parameters: ReadOnly[_Schema] + + +class _ToolEntry(TypedDict, total=False): + function: ReadOnly[_ToolFunction | None] + parameters: ReadOnly[_Schema] + input_schema: ReadOnly[_Schema] + vector_store_ids: ReadOnly[list[str] | None] + + +_TOOL_ENTRY: Final = TypeAdapter(_ToolEntry) +_WEB_SEARCH_OPTIONS: Final = TypeAdapter(dict[str, object] | None, config=ConfigDict(strict=True)) + + +@dataclass(frozen=True, slots=True) +class DropPolicy: + drop_params: bool + additional_drop_params: frozenset[str] + + def drops(self, field: str) -> bool: + return self.drop_params or field in self.additional_drop_params + + +@dataclass(frozen=True, slots=True) +class MalformedField: + field: str + message: str + dropped: Callable[[Mapping[str, object]], Mapping[str, object]] + + +def _field_names(value: object) -> frozenset[str]: + try: + return frozenset(name for name in _FIELD_LIST.validate_python(value) if isinstance(name, str)) + except ValidationError: + return frozenset() + + +def drop_policy(params: Mapping[str, object]) -> DropPolicy: + return DropPolicy( + drop_params=litellm.drop_params is True or normalize_drop_params(params.get("drop_params")) is True, + additional_drop_params=_field_names(params.get("additional_drop_params")), + ) + + +def first_validation_error(e: ValidationError, prefix: tuple[str | int, ...] = ()) -> str: + error: Final = e.errors(include_url=False)[0] + return f"{'.'.join(str(part) for part in (*prefix, *error['loc']))}: {error['msg']}" + + +def _tool_entry_error(index: int, entry: object) -> str | None: + try: + _TOOL_ENTRY.validate_python(entry, strict=True) + except ValidationError as e: + return first_validation_error(e, prefix=("tools", index)) + return None + + +def _without_tool_fields(data: Mapping[str, object]) -> Mapping[str, object]: + return {name: value for name, value in data.items() if name not in TOOL_FIELDS} + + +def _malformed_tools(data: Mapping[str, object]) -> MalformedField | None: + tools: Final = data.get("tools") + if tools is None: + return None + try: + entries: Final = _TOOLS.validate_python(tools) + except ValidationError as e: + return MalformedField("tools", first_validation_error(e, prefix=("tools",)), _without_tool_fields) + entry_errors: Final = tuple(_tool_entry_error(index, entry) for index, entry in enumerate(entries)) + first_error: Final = next((error for error in entry_errors if error is not None), None) + if first_error is None: + return None + kept: Final = [entry for entry, error in zip(entries, entry_errors) if error is None] + return MalformedField( + "tools", first_error, lambda fields: {**fields, "tools": kept} if kept else _without_tool_fields(fields) + ) + + +def _without_web_search_options(data: Mapping[str, object]) -> Mapping[str, object]: + return {name: value for name, value in data.items() if name != "web_search_options"} + + +def _malformed_web_search_options(data: Mapping[str, object]) -> MalformedField | None: + try: + _WEB_SEARCH_OPTIONS.validate_python(data.get("web_search_options")) + except ValidationError as e: + return MalformedField( + "web_search_options", + first_validation_error(e, prefix=("web_search_options",)), + _without_web_search_options, + ) + return None + + +def resolve_tool_fields( + data: Mapping[str, object], candidates: Callable[[], tuple[DropPolicy, ...]] +) -> Mapping[str, object] | MalformedField: + checked: Final = (_malformed_tools(data), _malformed_web_search_options(data)) + malformed: Final = tuple(field for field in checked if field is not None) + if not malformed: + return data + policies: Final = candidates() + rejected: Final = next( + (field for field in malformed if not any(policy.drops(field.field) for policy in policies)), None + ) + if rejected is not None: + return rejected + return reduce( + lambda resolved, field: field.dropped(resolved), + (field for field in malformed if all(policy.drops(field.field) for policy in policies)), + data, + ) + + +def resolve_deployment_tool_fields( + kwargs: Mapping[str, object], route: object +) -> Mapping[str, object] | MalformedField: + if "proxy_server_request" not in kwargs or route not in ROUTES_WITH_TOOLS: + return kwargs + return resolve_tool_fields(kwargs, lambda: (drop_policy(kwargs),)) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 6bc1c4ead64..6b6974becac 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -45,6 +45,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, + list_field, merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, @@ -924,10 +925,10 @@ class AnthropicMessagesHandler(BaseTranslation): images=(), ) - def extract_request_tool_names(self, data: dict) -> list[str]: + def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]: """Extract every tool name in an Anthropic messages request: tools[].name, plus tools[].function.name for OpenAI-format tools the bridge forwards verbatim.""" - return [name for tool in data.get("tools") or [] for name in anthropic_tool_names(tool)] + return [name for tool in list_field(data, "tools") for name in anthropic_tool_names(tool)] @classmethod def _extract_input_text_and_images( diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 39ff2fa0528..55a4b3da0b3 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -4,7 +4,7 @@ import json from collections.abc import Callable, Iterator, Mapping, Sequence from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -292,6 +292,17 @@ def openai_tool_name(tool: object) -> str | None: return flat_name if isinstance(flat_name, str) else None +_FIELDS: Final = TypeAdapter(Mapping[object, object], config=ConfigDict(strict=True)) +_ITEMS: Final = TypeAdapter(Sequence[object], config=ConfigDict(strict=True)) + + +def list_field(data: object, key: str) -> Sequence[object]: + try: + return _ITEMS.validate_python(_FIELDS.validate_python(data).get(key)) + except ValidationError: + return () + + def anthropic_tool_names(tool: object) -> tuple[str, ...]: """Every name a /v1/messages tool dict can act under: the flat Anthropic ``name`` plus ``function.name`` for OpenAI-format tools the bridge forwards verbatim. Allowlist checks diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index c168c3f88fc..3070bedd30d 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -39,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, + list_field, merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, openai_tool_name, @@ -156,7 +157,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs["structured_messages"] = [structured_messages[index] for index in scoped_message_indices] # Pass tools (function definitions) to the guardrail tools: Final = data.get("tools") - if tools and not scan_only_tool_results: + if isinstance(tools, list) and tools and not scan_only_tool_results: inputs["tools"] = tools # Include model information if available model: Final = data.get("model") @@ -258,19 +259,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return "no scannable content after message scoping" return None if images else "no scannable content" - def extract_request_tool_names(self, data: dict) -> list[str]: + def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]: """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name, and the tools a replayed tool_search_output loads).""" names: Final[list[str]] = [] - for tool in data.get("tools") or []: + for tool in list_field(data, "tools"): if isinstance(tool, dict) and tool.get("type") == "function": fn = tool.get("function") if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) - for fn in data.get("functions") or []: + for fn in list_field(data, "functions"): if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) - return [*names, *replayed_tool_search_tool_names(cast(object, data))] # cast-ok: upcast, read field by field + return [*names, *replayed_tool_search_tool_names(data)] def _extract_inputs( self, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 435c11c0602..84ff790ae48 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -55,6 +55,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, effective_skip_system_message_for_guardrail, + list_field, merge_guardrailed_scoped_messages, role_out_of_guardrail_scope, scoped_structured_message_indices, @@ -764,11 +765,11 @@ class OpenAIResponsesHandler(BaseTranslation): data.get("input"), data.get("instructions"), structured_messages, merged ) - def extract_request_tool_names(self, data: dict) -> list[str]: + def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]: """Extract tool names from Responses API request (tools[].name for function and custom, tools[].server_label for mcp).""" names: Final[list[str]] = [] - for tool in data.get("tools") or []: + for tool in list_field(data, "tools"): if not isinstance(tool, dict): continue if tool.get("type") in ("function", "custom") and tool.get("name"): diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index a11268ad432..1fa276d9ab2 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -6,9 +6,9 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final -from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import JSONResponse -from pydantic import JsonValue +from pydantic import JsonValue, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -18,6 +18,7 @@ from litellm.anthropic_interface.exceptions import ( AnthropicExceptionMapping, ) from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.anthropic.pass_through.context_management import ( AnthropicContextManagementError, ) @@ -32,6 +33,7 @@ from litellm.proxy.common_request_processing import ( log_llm_api_exception, proxy_exception_from_http_exception, proxy_exception_from_route_error, + resolve_counted_tool_fields, resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id @@ -47,6 +49,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( with_litellm_call_id, ) from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli +from litellm.router import Router router: Final = APIRouter() @@ -91,6 +94,31 @@ def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSO return JSONResponse(status_code=status_code, content=content, headers=exc.headers) +_TOKEN_COUNT_FIELDS: Final = ("model", "messages", "tools", "system") + + +def _token_count_request(fields: Mapping[str, object]) -> TokenCountRequest | ProxyException: + try: + return TokenCountRequest.model_validate({name: fields[name] for name in _TOKEN_COUNT_FIELDS if name in fields}) + except ValidationError as e: + return ProxyException( + message=first_validation_error(e), + type="invalid_request_error", + param=str(e.errors(include_url=False)[0]["loc"][0]), + code=status.HTTP_400_BAD_REQUEST, + ) + + +def _token_count_fields( + data: Mapping[str, object], auto_drop_params: bool, llm_router: Router | None, team_id: str | None +) -> Mapping[str, object] | ProxyException: + fields: Final = { + name: data[name] for name in (*_TOKEN_COUNT_FIELDS, "drop_params", "additional_drop_params") if name in data + } + resolving: Final = {**fields, "drop_params": True} if auto_drop_params else fields + return resolve_counted_tool_fields(resolving, llm_router, team_id) + + def _strip_total_tokens_from_anthropic_response(response: object) -> None: """Remove the OpenAI-flavored `usage.total_tokens` field that LiteLLM injects into Anthropic /v1/messages responses. @@ -351,7 +379,7 @@ async def count_tokens( Returns: {"input_tokens": } """ - from litellm.proxy.proxy_server import count_request_tokens, proxy_config + from litellm.proxy.proxy_server import count_request_tokens, llm_router, proxy_config litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: @@ -368,26 +396,21 @@ async def count_tokens( if not messages: raise HTTPException(status_code=400, detail="messages parameter is required") - # Create TokenCountRequest for the internal endpoint - from litellm.proxy._types import TokenCountRequest - - token_request: Final = TokenCountRequest( - model=model_name, - messages=messages, - tools=data.get("tools"), - system=data.get("system"), + auto_drop_params: Final = should_auto_drop_params_for_agentic_cli( + request.headers.get("user-agent", ""), "drop_params" in data, proxy_config ) + token_fields: Final = _token_count_fields(data, auto_drop_params, llm_router, user_api_key_dict.team_id) + token_request: Final = ( + token_fields if isinstance(token_fields, ProxyException) else _token_count_request(token_fields) + ) + if isinstance(token_request, ProxyException): + return _anthropic_error_json_response(token_request, request) token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, messages_format="anthropic", - **_count_tokens_body_options( - data, - should_auto_drop_params_for_agentic_cli( - request.headers.get("user-agent", ""), "drop_params" in data, proxy_config - ), - ), + **_count_tokens_body_options(data, auto_drop_params), ) except HTTPException as e: return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5f798fce6cf..2d387b43aa8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -80,6 +80,7 @@ from litellm.litellm_core_utils.served_output_texts import ( from litellm.litellm_core_utils.streaming_handler import ( backfill_missing_cache_usage_fields, ) +from litellm.litellm_core_utils.tool_fields import DropPolicy, MalformedField, drop_policy, resolve_tool_fields from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( can_key_call_resolved_model, @@ -129,6 +130,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports ) from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError @@ -633,6 +635,63 @@ def serialize_http_exception_detail( return str(detail), None +_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True)) +_UNSEEN_DEPLOYMENT: Final = DropPolicy(drop_params=True, additional_drop_params=frozenset()) + + +def _is_auto_router(params: Mapping[str, object]) -> bool: + model: Final = params.get("model") + return isinstance(model, str) and model.startswith(AUTO_ROUTER_MODEL_PREFIX) + + +def _named_deployments(llm_router: Router, model: str, team_id: str | None) -> tuple[Mapping[str, object], ...]: + by_id: Final = None if model in llm_router.model_names else llm_router.get_deployment(model_id=model) + if by_id is not None: + return (_REQUEST_FIELDS.validate_python(by_id.litellm_params.model_dump(exclude_none=True)),) + group: Final = llm_router.get_model_from_alias(model=model) or model + if group not in llm_router.model_names or llm_router.get_model_list_from_routing_groups(model_name=group): + return () + deployments: Final = tuple( + _REQUEST_FIELDS.validate_python(deployment["litellm_params"]) + for deployment in llm_router.get_all_deployments(model_name=group, team_id=team_id) + ) + return () if any(_is_auto_router(params) for params in deployments) else deployments + + +def _router_defaults(llm_router: Router | None) -> Mapping[str, object]: + configured: Final[object] = getattr(llm_router, "default_litellm_params", None) + if configured is None: + return {} + return {name: value for name, value in _REQUEST_FIELDS.validate_python(configured).items() if value is not None} + + +def _counted_drop_policies( + data: Mapping[str, object], llm_router: Router | None, team_id: str | None +) -> tuple[DropPolicy, ...]: + requested: Final = {**_router_defaults(llm_router), **data} + model: Final = data.get("model") + deployments: Final = ( + _named_deployments(llm_router, model, team_id) if llm_router is not None and isinstance(model, str) else () + ) + if not deployments: + return (drop_policy(requested), _UNSEEN_DEPLOYMENT) + return tuple(drop_policy({**deployment, **requested}) for deployment in deployments) + + +def resolve_counted_tool_fields( + data: Mapping[str, object], llm_router: Router | None, team_id: str | None +) -> Mapping[str, object] | ProxyException: + resolved: Final = resolve_tool_fields(data, lambda: _counted_drop_policies(data, llm_router, team_id)) + if isinstance(resolved, MalformedField): + return ProxyException( + message=resolved.message, + type="invalid_request_error", + param=resolved.field, + code=status.HTTP_400_BAD_REQUEST, + ) + return resolved + + def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException: raw_detail: Final = attribute_of(exc, "detail", str(exc)) message, structured_fields = serialize_http_exception_detail(raw_detail) diff --git a/litellm/proxy/guardrails/tool_name_extraction.py b/litellm/proxy/guardrails/tool_name_extraction.py index 72e85f35f7d..526f58dcd36 100644 --- a/litellm/proxy/guardrails/tool_name_extraction.py +++ b/litellm/proxy/guardrails/tool_name_extraction.py @@ -6,23 +6,23 @@ knowledge lives in one place. Uses guardrail translation handlers where availabl with standalone extractors for generate_content and MCP. """ +from collections.abc import Mapping from typing import Any, Final from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.llms import load_guardrail_translation_mappings +from litellm.llms.base_llm.guardrail_translation.utils import list_field from litellm.types.utils import CallTypes # Call types that have no guardrail translation handler; we use standalone extractors STANDALONE_EXTRACTORS: Final[dict[str, Any]] = {} -def _extract_generate_content_tool_names(data: dict) -> list[str]: +def _extract_generate_content_tool_names(data: Mapping[str, object]) -> list[str]: """Google generateContent: tools[].functionDeclarations[].name""" names: Final[list[str]] = [] - for tool in data.get("tools") or []: - if not isinstance(tool, dict): - continue - for decl in tool.get("functionDeclarations") or []: + for tool in list_field(data, "tools"): + for decl in list_field(tool, "functionDeclarations"): if isinstance(decl, dict) and decl.get("name"): names.append(str(decl["name"])) return names diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index a44f1d10fda..207e966ef28 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -34,6 +34,7 @@ from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup +from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, build_blocked_response, @@ -45,7 +46,12 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth, user_api_key_auth_websocket, ) -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, cost_headers, create_response +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + cost_headers, + create_response, + resolve_counted_tool_fields, +) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _safe_set_request_parsed_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -57,6 +63,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -69,7 +76,6 @@ from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.utils import ProxyLogging - from litellm.router import Router router: Final = APIRouter() _RESPONSES_WS_CONFIG_VALUE_ADAPTER: Final[TypeAdapter[object | None]] = TypeAdapter(object | None) @@ -1376,6 +1382,15 @@ def _proxy_token_counter() -> _TokenCounter: _token_counter_dep: Final = Depends(_proxy_token_counter) +def _proxy_llm_router() -> Router | None: + from litellm.proxy.proxy_server import llm_router + + return llm_router + + +_llm_router_dep: Final = Depends(_proxy_llm_router) + + def _responses_invalid_request_response(message: str, param: str | None, code: str | None) -> JSONResponse: body: Final[_ResponsesApiErrorBody] = { "error": { @@ -1432,6 +1447,8 @@ def _responses_input_as_token_count_messages( async def responses_input_tokens( request: Request, token_counter: _TokenCounter = _token_counter_dep, + llm_router: Router | None = _llm_router_dep, + user_api_key_dict: UserAPIKeyAuth = _user_api_key_auth_dep, ): """ Count the input tokens of a Responses API request without calling the model. @@ -1464,6 +1481,13 @@ async def responses_input_tokens( code="missing_required_parameter", ) + counted_fields: Final = { + name: data[name] for name in ("model", "tools", "drop_params", "additional_drop_params") if name in data + } + token_fields: Final = resolve_counted_tool_fields(counted_fields, llm_router, user_api_key_dict.team_id) + if isinstance(token_fields, ProxyException): + return _responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None) + try: payload: Final[_TokenCountPayload] = { "model": model_name, @@ -1471,9 +1495,13 @@ async def responses_input_tokens( input_value=input_value, instructions=data.get("instructions"), ), - "tools": data.get("tools"), + "tools": token_fields.get("tools"), } token_request: Final = TokenCountRequest.model_validate(payload) + except ValidationError as e: + return _responses_invalid_request_response( + message=first_validation_error(e), param=str(e.errors(include_url=False)[0]["loc"][0]), code=None + ) except Exception as e: return _responses_invalid_request_response( message=f"Invalid request for token counting: {e}", param=None, code=None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index dc6423c979b..36ace3c1697 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3493,6 +3493,7 @@ class ProxyLogging: Is triggered when self._is_proxy_only_error() returns True """ litellm_logging_obj: Logging | None = request_data.get("litellm_logging_obj", None) + set_up_by_request: Final = litellm_logging_obj is not None if litellm_logging_obj is None: from litellm._uuid import uuid @@ -3557,7 +3558,9 @@ class ProxyLogging: input = request_data["input"] litellm_logging_obj.model_call_details["input"] = input body_shape_call_type = CallTypes.aembedding.value - resolved_call_type: Final = _call_type_for_route(route) or body_shape_call_type + resolved_call_type: Final = _call_type_for_route(route) or ( + None if set_up_by_request else body_shape_call_type + ) if resolved_call_type is not None and litellm_logging_obj.call_type != CallTypes.pass_through.value: litellm_logging_obj.call_type = resolved_call_type litellm_logging_obj.model_call_details["call_type"] = resolved_call_type diff --git a/litellm/router.py b/litellm/router.py index 316c7f27ec0..d041f958e17 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -123,6 +123,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_sensitive_structure, ) from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy from litellm.llms.anthropic.common_utils import AnthropicModelInfo, anthropic_error_frame_exception from litellm.llms.base_llm.passthrough.transformation import replace_path_segment from litellm.llms.base_llm.vector_store.transformation import ( @@ -537,6 +538,7 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) +_NO_DEFAULT_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -573,6 +575,21 @@ def _deployment_that_just_failed(request_metadata: object) -> str | None: return model_id if isinstance(model_id, str) else None +def _drops_malformed_request_tools( + deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object] +) -> bool: + if isinstance(request.get("tools") or [], (list, tuple)) or not deployment_params.get("tools"): + return False + set_defaults: Final = {name: value for name, value in defaults.items() if value is not None} + return drop_policy({**deployment_params, **set_defaults, **request}).drops("tools") + + +def _merged_tools(deployment_tools: Sequence[object], request_tools: object) -> list[object] | None: + if not isinstance(request_tools, (list, tuple)): + return None + return [*deployment_tools, *request_tools] if deployment_tools or request_tools else None + + def _deployment_pick_attributes(model: str, request_kwargs: Mapping[str, object] | None) -> Mapping[str, str | int]: """Bounded attributes for one deployment pick; attempt is 1-based within the current model group.""" kwargs: Final = request_kwargs or {} @@ -4507,7 +4524,9 @@ class Router: return deployment_pydantic_obj @staticmethod - def _merge_tools_from_deployment(deployment: dict, kwargs: dict) -> None: + def _merge_tools_from_deployment( + deployment: dict, kwargs: dict, defaults: Mapping[str, object] = _NO_DEFAULT_PARAMS + ) -> None: """ Merge tools from deployment litellm_params with request kwargs. When both have tools, concatenate them (deployment tools first, then request tools). @@ -4518,11 +4537,12 @@ class Router: dep_params = dep_params_raw else: dep_params = dep_params_raw.model_dump(exclude_none=True) - dep_tools: Final = dep_params.get("tools") or [] - req_tools: Final = kwargs.get("tools") or [] - if dep_tools or req_tools: - merged: Final = list(dep_tools) + list(req_tools) - kwargs["tools"] = merged + if _drops_malformed_request_tools(dep_params, defaults, kwargs): + for name in TOOL_FIELDS: + kwargs.pop(name, None) + merged_tools: Final = _merged_tools(dep_params.get("tools") or [], kwargs.get("tools") or []) + if merged_tools is not None: + kwargs["tools"] = merged_tools if "tool_choice" not in kwargs and dep_params.get("tool_choice") is not None: kwargs["tool_choice"] = dep_params["tool_choice"] @@ -4543,7 +4563,7 @@ class Router: deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) ): kwargs.pop(key, None) - self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs) + self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=self.default_litellm_params) model_info = deployment.get("model_info", {}).copy() deployment_litellm_model_name = deployment["litellm_params"]["model"] diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index c4e4d9ccc0b..9b9a3d185cc 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -309,6 +309,9 @@ def filter_team_based_models( ] +_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "web_search_preview"}) + + def _deployment_supports_web_search(deployment: dict) -> bool: """ Check if a deployment supports web search. @@ -343,14 +346,13 @@ def filter_web_search_deployments( if isinstance(healthy_deployments, dict): return healthy_deployments - is_web_search_request = False - tools: Final = request_kwargs.get("tools") or [] - for tool in tools: - # These are the two websearch tools for OpenAI / Azure. - if tool.get("type") == "web_search" or tool.get("type") == "web_search_preview": - is_web_search_request = True - break - + tools: Final[object] = request_kwargs.get("tools") + is_web_search_request: Final = isinstance(tools, (list, tuple)) and any( + isinstance(tool, dict) + and isinstance(tool_type := tool.get("type"), str) + and tool_type in _WEB_SEARCH_TOOL_TYPES + for tool in tools + ) if not is_web_search_request: return healthy_deployments diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index c7e3dab0d1a..bb64d912c4c 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -329,6 +329,12 @@ def _should_run_cooldown_logic( ) return False + if isinstance(original_exception, litellm.MalformedRequestFieldError): + verbose_router_logger.debug( + "Should Not Run Cooldown Logic: litellm rejected the request before calling the deployment" + ) + return False + if request_kwargs is not None and _is_per_user_provider_request(request_kwargs): verbose_router_logger.debug( "Should Not Run Cooldown Logic: the failure came from one caller's per-user provider session" diff --git a/litellm/utils.py b/litellm/utils.py index 5a933d8f2b8..76ef50c3a2b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -92,6 +92,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens +from litellm.litellm_core_utils.tool_fields import MalformedField, resolve_deployment_tool_fields from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.rust_bridge.catalog import decision from litellm.rust_bridge.configuration import Decision @@ -474,6 +475,7 @@ from .exceptions import ( BudgetExceededError, ContentPolicyViolationError, ContextWindowExceededError, + MalformedRequestFieldError, ModelNotMappedError, NotFoundError, OpenAIError, @@ -1495,7 +1497,23 @@ def check_coroutine(value) -> bool: return get_coroutine_checker().is_async_callable(value) -async def async_pre_call_deployment_hook(kwargs: dict[str, Any], call_type: str): +def _deployment_provider(kwargs: Mapping[str, object]) -> str: + model: Final = kwargs.get("model") + custom_llm_provider: Final = kwargs.get("custom_llm_provider") + api_base: Final = kwargs.get("api_base") + if not isinstance(model, str): + return "" + try: + return litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + api_base=api_base if isinstance(api_base, str) else None, + )[1] + except BadRequestError: + return "" + + +async def async_pre_call_deployment_hook(kwargs: dict[str, Any], call_type: str) -> dict[str, Any]: """ Allow modifying the request just before it's sent to the deployment. @@ -1506,7 +1524,16 @@ async def async_pre_call_deployment_hook(kwargs: dict[str, Any], call_type: str) except ValueError: typed_call_type = None # unknown call type - modified_kwargs = kwargs.copy() + route: Final[object] = getattr(kwargs.get("litellm_logging_obj"), "call_type", None) + resolved: Final = resolve_deployment_tool_fields(kwargs, route) + if isinstance(resolved, MalformedField): + raise MalformedRequestFieldError( + message=resolved.message, + model=kwargs.get("model"), + llm_provider=_deployment_provider(kwargs), + body={"type": "invalid_request_error", "param": resolved.field}, + ) + modified_kwargs = dict(resolved) custom_logger_class: Final = _get_cached_custom_logger() for callback in litellm.callbacks: diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index e0b536d6c57..b93d9874dbc 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -1,6 +1,6 @@ # litellm/proxy/vector_stores/vector_store_registry.py import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from itertools import chain from typing import ( @@ -11,7 +11,7 @@ from typing import ( get_args, ) -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import remove_items_at_indices @@ -34,6 +34,26 @@ else: PrismaClient = Any _DB_ROW: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) +_ITEMS: Final = TypeAdapter(Sequence[object], config=ConfigDict(strict=True)) +_TOOL: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True)) +_STORE_ID_ITEMS: Final = TypeAdapter(list[object], config=ConfigDict(strict=True)) + + +def _named_store_ids(store_ids: object) -> tuple[str, ...]: + if isinstance(store_ids, str): + return (store_ids,) + try: + items: Final = _STORE_ID_ITEMS.validate_python(store_ids) + except ValidationError: + return () + return tuple(store_id for store_id in items if isinstance(store_id, str)) + + +def _tool_vector_store_ids(tool: object) -> tuple[str, ...]: + try: + return _named_store_ids(_TOOL.validate_python(tool).get("vector_store_ids")) + except ValidationError: + return () class VectorStoreIndexRegistry: @@ -131,7 +151,7 @@ class VectorStoreRegistry: return VectorStoreToolParams(**kwargs) - def get_vector_store_ids_to_run(self, non_default_params: dict, tools: list[dict] | None = None) -> list[str]: + def get_vector_store_ids_to_run(self, non_default_params: dict, tools: object = None) -> list[str]: """ Returns the vector store ids to run @@ -175,7 +195,7 @@ class VectorStoreRegistry: tools_to_remove: Final[list[int]] = [] for i, tool in enumerate(tools): - tool_vector_store_ids = tool.get("vector_store_ids", []) + tool_vector_store_ids = _tool_vector_store_ids(tool) if not tool_vector_store_ids: continue @@ -395,14 +415,15 @@ class VectorStoreRegistry: return vector_stores_to_run - def _get_vector_store_ids_from_tool_calls(self, tools: list[dict] | None, vector_store_ids: list[str]) -> list[str]: + def _get_vector_store_ids_from_tool_calls(self, tools: object, vector_store_ids: list[str]) -> list[str]: """ Returns the vector store ids from the tool calls """ - return [ - *vector_store_ids, - *chain.from_iterable(tool["vector_store_ids"] for tool in tools or () if "vector_store_ids" in tool), - ] + try: + validated_tools: Final = _ITEMS.validate_python(tools) + except ValidationError: + return vector_store_ids + return [*vector_store_ids, *chain.from_iterable(_tool_vector_store_ids(tool) for tool in validated_tools)] def load_vector_stores_from_config(self, vector_stores_config: list[dict]): """ diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 6762e4cd35c..e669107900d 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -112,6 +112,8 @@ ignored_function_names = [ "_get_wildcard_deployments", # Tested through the get_model_list_of_routed_group wildcard test in test_router.py "_is_fallback_hop", # Tested through the order fallback hop tests in test_router_order_fallback.py "_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py + "_merged_tools", # Tested through the _update_kwargs_with_deployment tools merge tests in test_router.py + "_drops_malformed_request_tools", # Tested through the _update_kwargs_with_deployment tools merge tests in test_router.py "_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py "_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py "_register_continuation_prefill_check", # Tested through the continuation tests in test_router.py diff --git a/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py b/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py index 2e8a9d0d7e4..15c0593f80d 100644 --- a/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py +++ b/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py @@ -364,7 +364,7 @@ def test_spend_log_keeps_the_tools_the_caller_sent(rig: Rig, model: str, guardra assert object_value(rows[0]["proxy_server_request"])["tools"] == [strict_tool()] -@pytest.mark.parametrize("tools", [7, "x" * 5000, ["lint"]], ids=["int", "five_kb", "list_of_str"]) +@pytest.mark.parametrize("tools", [7, "x" * 5000, ["lint"], ""], ids=["int", "five_kb", "list_of_str", "empty_string"]) def test_a_malformed_tools_field_is_answered_with_an_error_and_the_proxy_keeps_serving( rig: Rig, tools: JsonValue ) -> None: @@ -379,14 +379,6 @@ def test_a_malformed_tools_field_is_answered_with_an_error_and_the_proxy_keeps_s assert reply_text(control_status, control_raw, streaming=False) == f"linted {control_nonce.removeprefix('nonce-')}" -def test_an_empty_string_tools_field_reaches_claude_as_sent(rig: Rig) -> None: - request_nonce: Final = nonce() - status, raw = rig.send(messages_body(CLAUDE_MODEL, request_nonce, tools="")) - - assert reply_text(status, raw, streaming=False) == f"linted {request_nonce.removeprefix('nonce-')}" - assert sent_body(rig.received(request_nonce), CLAUDE_TARGET)["tools"] == "" - - @pytest.mark.parametrize( ("model", "target"), [(CLAUDE_MODEL, CLAUDE_TARGET), (OPENAI_MODEL, RESPONSES_TARGET)], diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 5179514b2b8..701d185b241 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -1147,8 +1147,8 @@ def test_messages_count_tokens_surfaces_the_mantle_error_for_a_non_text_system( @pytest.mark.parametrize( "tools", - ["", "x" * 5120, ["get_weather"]], - ids=["empty_string", "5kb_string", "list_of_strings"], + [5, "", "x" * 5120, ["get_weather"]], + ids=["int", "empty_string", "5kb_string", "list_of_strings"], ) def test_messages_count_tokens_rejects_malformed_tools_without_calling_either_peer( counting_proxy: OwnedProxy, mantle_port: int, tools: JsonValue @@ -1161,8 +1161,8 @@ def test_messages_count_tokens_rejects_malformed_tools_without_calling_either_pe ): model: Final = _deployment(scenario, runtime.url) refused: Final = _count(gateway, {**_bare(model), "tools": tools}) - assert 400 <= refused.status_code < 600, refused.text - assert "input_tokens" not in refused.text, refused.text + assert refused.status_code == 400, refused.text + assert object_value(_JSON_OBJECT.validate_json(refused.content)["error"])["type"] == "invalid_request_error" assert runtime.drain() == () and mantle.drain() == () assert _generated_then_counted(gateway.client, gateway.key, model, _bare(model)) == (200, 200, _MANTLE_COUNT) assert _mantle_bodies(mantle) == (_MANTLE_BARE,) @@ -1185,8 +1185,13 @@ def test_messages_count_tokens_without_messages_is_rejected( assert runtime.drain() == () and mantle.drain() == () +@pytest.mark.parametrize( + "fields", + [{}, {"tools": 5}, {"web_search_options": None}], + ids=["well_formed", "int_tools", "null_web_search_options"], +) def test_messages_count_tokens_unauthenticated_request_never_reaches_either_peer( - counting_proxy: OwnedProxy, mantle_port: int + counting_proxy: OwnedProxy, mantle_port: int, fields: dict[str, JsonValue] ) -> None: gateway: Final = counting_proxy.gateway with ( @@ -1196,7 +1201,7 @@ def test_messages_count_tokens_unauthenticated_request_never_reaches_either_peer ): model: Final = _deployment(scenario, runtime.url) response: Final = gateway.request( - "POST", "/v1/messages/count_tokens", _full(model), key="sk-not-a-key-this-proxy-issued" + "POST", "/v1/messages/count_tokens", {**_full(model), **fields}, key="sk-not-a-key-this-proxy-issued" ) assert response.status_code == 401, response.text assert "input_tokens" not in response.text, response.text diff --git a/tests/integration/providers/test_malformed_tools_chaos.py b/tests/integration/providers/test_malformed_tools_chaos.py new file mode 100644 index 00000000000..0ab50a5afbe --- /dev/null +++ b/tests/integration/providers/test_malformed_tools_chaos.py @@ -0,0 +1,619 @@ +from __future__ import annotations + +import json +import os +import re +import threading +import uuid +from collections.abc import Iterator +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from typing import Final, Literal, TypeAlias + +import httpx +import psutil +import pytest +from integration._support.anthropic_sse import error_body, message_json, message_stream +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.database import write_rows +from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.worker_readiness import wait_for_worker_readiness +from pydantic import JsonValue, TypeAdapter + +_OPENAI_BACKEND: Final = "gpt-4o-mini" +_CLAUDE_BACKEND: Final = "claude-haiku-4-5" +_GEMINI_BACKEND: Final = "gemini-2.5-flash" +_PROVIDER_KEY: Final = "scripted-malformed-tools-chaos-key" +_OWNED_CHAT: Final = "malformed-tools-owned-chat" +_OWNED_STORE: Final = "vs_malformed_tools_owned_store" +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 +_RESPAWN_SECONDS: Final = graceful_stop_seconds() + 60 +_PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1"))) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_NOT_A_LIST: Final = "tools: Input should be a valid list" +_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} +_CHAT_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "description": "Look up the weather", "parameters": _WEATHER_SCHEMA}, +} +_CLAUDE_TOOL: Final[dict[str, JsonValue]] = { + "name": "get_weather", + "description": "Look up the weather", + "input_schema": _WEATHER_SCHEMA, +} +_RESPONSES_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "name": "get_weather", + "description": "Look up the weather", + "parameters": _WEATHER_SCHEMA, +} +_GEMINI_TOOL: Final[dict[str, JsonValue]] = { + "functionDeclarations": [ + {"name": "get_weather", "description": "Look up the weather", "parameters": _WEATHER_SCHEMA} + ] +} +Kind: TypeAlias = Literal["valid", "rejected", "dropped"] + + +def _marker(label: str) -> str: + return f"{label} marker-{uuid.uuid4().hex}" + + +def _tag(text: str) -> str: + found: Final = _MARKER.search(text) + return found.group(1) if found else "none" + + +def _answer(tag: str) -> str: + return f"answer {tag}" + + +def _gemini_body(tag: str) -> bytes: + return json.dumps( + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": _answer(tag)}]}, "index": 0, "finishReason": "STOP"} + ], + "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 2, "totalTokenCount": 11}, + "modelVersion": _GEMINI_BACKEND, + } + ).encode() + + +def _scripted(request: Request) -> Reply: + text: Final = request.body.decode() + tag: Final = _tag(text) + path: Final = request.target.split("?")[0] + stream: Final = bool(text) and _JSON_OBJECT.validate_json(request.body).get("stream") is True + if path == "/v1/chat/completions": + return chat_reply(f"chatcmpl-{tag}", _OPENAI_BACKEND, _answer(tag), stream=stream) + if path == "/v1/responses": + return responses_reply(f"resp_{tag}", _OPENAI_BACKEND, _answer(tag), stream=stream) + if path == "/v1/messages" and _JSON_OBJECT.validate_json(request.body).get("tools") == 5: + return Reply(status=400, body=error_body(400, _NOT_A_LIST)) + if path == "/v1/messages" and stream: + return Reply( + content_type="text/event-stream", chunks=message_stream(f"msg_{tag}", _CLAUDE_BACKEND, _answer(tag)) + ) + if path == "/v1/messages": + return Reply(body=message_json(f"msg_{tag}", _CLAUDE_BACKEND, _answer(tag))) + if path.endswith(":streamGenerateContent"): + return Reply(content_type="text/event-stream", chunks=(b"data: " + _gemini_body(tag) + b"\r\n\r\n",)) + if path.endswith(":generateContent"): + return Reply(body=_gemini_body(tag)) + return Reply(status=404, body=json.dumps({"error": f"unscripted {request.target}"}).encode()) + + +@dataclass(frozen=True, slots=True) +class _Switchboard: + down: threading.Event + holding: threading.Event + release: threading.Event + held: SimpleQueue[str] + + def respond(self, request: Request) -> Reply: + if self.down.is_set(): + return Reply(drop_connection=True) + if self.holding.is_set(): + self.held.put(request.target) + assert self.release.wait(timeout=120), "Held request was never released" + return _scripted(request) + + def reset(self) -> None: + self.down.clear() + self.release.set() + self.holding.clear() + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + wire: Wire + switchboard: _Switchboard + chat: str + claude: str + gemini: str + + +@pytest.fixture(scope="module") +def rig() -> Iterator[_Rig]: + switchboard: Final = _Switchboard(threading.Event(), threading.Event(), threading.Event(), SimpleQueue()) + with ( + gateway_from_environment() as gateway, + wire_server(switchboard.respond) as wire, + gateway.scenario() as scenario, + ): + built: Final = _Rig( + gateway=gateway, + wire=wire, + switchboard=switchboard, + chat=scenario.model(api_base=f"{wire.url}/v1", api_key=_PROVIDER_KEY), + claude=scenario.model(model=f"anthropic/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_PROVIDER_KEY), + gemini=scenario.model(model=f"gemini/{_GEMINI_BACKEND}", api_base=wire.url, api_key=_PROVIDER_KEY), + ) + for group in (built.chat, built.claude, built.gemini): + wait_for_worker_readiness(gateway, group, 1, _PROXY_WORKERS) + try: + yield built + finally: + switchboard.reset() + + +@pytest.fixture(autouse=True) +def _reset(rig: _Rig) -> Iterator[None]: + rig.switchboard.reset() + rig.wire.drain() + yield + rig.switchboard.reset() + + +@dataclass(frozen=True, slots=True) +class _Call: + path: str + body: dict[str, JsonValue] + marker: str + kind: Kind + + +def _tools_for(kind: Kind, tool: dict[str, JsonValue]) -> dict[str, JsonValue]: + if kind == "rejected": + return {"tools": 5} + if kind == "dropped": + return {"tools": [tool, 5], "drop_params": True} + return {"tools": [tool]} + + +def _call(rig: _Rig, index: int, kind: Kind, label: str) -> _Call: + marker: Final = _marker(label) + stream: Final = (index // 4) % 2 == 1 + family: Final = index % 4 + if family == 0: + return _Call( + "/v1/chat/completions", + { + "model": rig.chat, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + **_tools_for(kind, _CHAT_TOOL), + }, + marker, + kind, + ) + if family == 1: + return _Call( + "/v1/messages", + { + "model": rig.claude, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + **_tools_for(kind, _CLAUDE_TOOL), + }, + marker, + kind, + ) + if family == 2: + return _Call( + "/v1/responses", + {"model": rig.chat, "input": marker, "stream": stream, **_tools_for(kind, _RESPONSES_TOOL)}, + marker, + kind, + ) + action: Final = "streamGenerateContent?alt=sse" if stream else "generateContent" + return _Call( + f"/v1beta/models/{rig.gemini}:{action}", + {"contents": [{"role": "user", "parts": [{"text": marker}]}], **_tools_for(kind, _GEMINI_TOOL)}, + marker, + kind, + ) + + +def _kind(index: int) -> Kind: + kinds: Final[tuple[Kind, ...]] = ("valid", "rejected", "dropped") + return kinds[index % 3] + + +def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=120, trust_env=False)) for _ in range(count) + ) + + +def _post(client: httpx.Client, key: str, call: _Call) -> httpx.Response: + return client.post(call.path, json=call.body, headers={"Authorization": f"Bearer {key}"}) + + +def _status_or_dropped(client: httpx.Client, key: str, call: _Call) -> int | None: + try: + return _post(client, key, call).status_code + except httpx.TransportError: + return None + + +def _deliveries(wire: Wire) -> tuple[str, ...]: + return tuple(request.body.decode() for request in wire.drain()) + + +def _seen(deliveries: tuple[str, ...], marker: str) -> int: + return sum(marker in body for body in deliveries) + + +def _delivered_tool_names(deliveries: tuple[str, ...], marker: str) -> tuple[str, ...]: + (body,) = tuple(body for body in deliveries if marker in body) + return tuple(re.findall(r'"name": ?"([a-z_]+)"', body)) + + +def _assert_wave(calls: tuple[_Call, ...], statuses: tuple[int | None, ...], deliveries: tuple[str, ...]) -> None: + expected_statuses: Final = tuple(400 if call.kind == "rejected" else 200 for call in calls) + assert statuses == expected_statuses, tuple(zip(statuses, (call.marker for call in calls), strict=True)) + seen: Final = tuple(_seen(deliveries, call.marker) for call in calls) + assert seen == tuple(0 if call.kind == "rejected" else 1 for call in calls), seen + names: Final = tuple(_delivered_tool_names(deliveries, call.marker) for call in calls if call.kind != "rejected") + assert all("get_weather" in found for found in names), names + + +def test_mixed_burst_rejects_malformed_tools_and_forwards_the_rest_once(rig: _Rig) -> None: + with ExitStack() as stack: + calls: Final = tuple(_call(rig, index, _kind(index), "burst") for index in range(40)) + clients: Final = _clients(stack, str(rig.gateway.client.base_url), len(calls)) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(calls))) + statuses: Final = tuple( + pool.map( + lambda pair: _status_or_dropped(pair[0], rig.gateway.key, pair[1]), zip(clients, calls, strict=True) + ) + ) + _assert_wave(calls, statuses, _deliveries(rig.wire)) + + +def test_upstream_outage_still_rejects_malformed_tools_and_recovers(rig: _Rig) -> None: + with ExitStack() as stack: + clients: Final = _clients(stack, str(rig.gateway.client.base_url), 12) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + + def wave(label: str) -> tuple[tuple[_Call, ...], tuple[int | None, ...]]: + calls: Final = tuple(_call(rig, index, _kind(index), label) for index in range(len(clients))) + statuses: Final = tuple( + pool.map( + lambda pair: _status_or_dropped(pair[0], rig.gateway.key, pair[1]), zip(clients, calls, strict=True) + ) + ) + return calls, statuses + + before_calls, before_statuses = wave("before outage") + _assert_wave(before_calls, before_statuses, _deliveries(rig.wire)) + rig.switchboard.down.set() + outage_calls, outage_statuses = wave("during outage") + for call, status in zip(outage_calls, outage_statuses, strict=True): + assert (status == 400) if call.kind == "rejected" else (status is not None and status >= 500), ( + call.kind, + status, + ) + assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 + rig.wire.drain() + rig.switchboard.down.clear() + after_calls, after_statuses = wave("after outage") + _assert_wave(after_calls, after_statuses, _deliveries(rig.wire)) + + +def test_held_upstream_does_not_delay_malformed_tools_rejections(rig: _Rig) -> None: + with ExitStack() as stack: + held_calls: Final = tuple(_call(rig, index * 4, "valid", "held") for index in range(8)) + rejected_calls: Final = tuple(_call(rig, index, "rejected", "rejected") for index in range(8)) + clients: Final = _clients(stack, str(rig.gateway.client.base_url), len(held_calls) + len(rejected_calls)) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + rig.switchboard.release.clear() + rig.switchboard.holding.set() + stack.callback(rig.switchboard.release.set) + pending: Final[tuple[Future[int | None], ...]] = tuple( + pool.submit(_status_or_dropped, client, rig.gateway.key, call) + for client, call in zip(clients[: len(held_calls)], held_calls, strict=True) + ) + eventually(rig.switchboard.held.qsize, lambda size: size == len(held_calls), seconds=60) + rejected: Final = tuple( + pool.map( + lambda pair: _status_or_dropped(pair[0], rig.gateway.key, pair[1]), + zip(clients[len(held_calls) :], rejected_calls, strict=True), + ) + ) + assert rejected == (400,) * len(rejected_calls), rejected + assert not any(future.done() for future in pending), "a held request finished before its release" + rig.switchboard.release.set() + assert tuple(future.result(timeout=120) for future in pending) == (200,) * len(held_calls) + deliveries: Final = _deliveries(rig.wire) + assert tuple(_seen(deliveries, call.marker) for call in held_calls) == (1,) * len(held_calls) + assert tuple(_seen(deliveries, call.marker) for call in rejected_calls) == (0,) * len(rejected_calls) + + +@dataclass(frozen=True, slots=True) +class _Owned: + proxy: OwnedProxy + wire: Wire + held: SimpleQueue[str] + holding: threading.Event + release: threading.Event + blocked_tool: str + denied_tool: str + + +def _owned_config(path: Path, upstream: str, denied_tool: str) -> Path: + path.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": _OWNED_CHAT, + "litellm_params": { + "model": f"openai/{_OPENAI_BACKEND}", + "api_base": f"{upstream}/v1", + "api_key": _PROVIDER_KEY, + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + "litellm_settings": {"drop_params": True}, + "router_settings": {"disable_cooldowns": True}, + "guardrails": [ + { + "guardrail_name": "malformed-tools-policy", + "litellm_params": {"guardrail": "tool_policy", "mode": "pre_call", "default_on": True}, + }, + { + "guardrail_name": "malformed-tools-permission", + "litellm_params": { + "guardrail": "tool_permission", + "mode": "pre_call", + "default_on": False, + "default_action": "allow", + "on_disallowed_action": "block", + "rules": [{"id": "deny_tool", "tool_name": f"^{denied_tool}$", "decision": "deny"}], + }, + }, + ], + "vector_store_registry": [ + { + "vector_store_name": "malformed-tools-owned-store", + "litellm_params": { + "vector_store_id": _OWNED_STORE, + "custom_llm_provider": "openai", + "api_base": upstream, + "api_key": _PROVIDER_KEY, + }, + } + ], + } + ) + ) + return path + + +@pytest.fixture(scope="module") +def owned(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + suffix: Final = uuid.uuid4().hex[:12] + blocked_tool: Final = f"wipe_disk_{suffix}" + denied_tool: Final = f"run_shell_{suffix}" + held: Final[SimpleQueue[str]] = SimpleQueue() + holding: Final = threading.Event() + release: Final = threading.Event() + + def respond(request: Request) -> Reply: + if holding.is_set(): + held.put(request.target) + assert release.wait(timeout=120), "Held request was never released" + return _scripted(request) + + directory: Final = tmp_path_factory.mktemp("malformed-tools-owned") + with ExitStack() as stack: + gateway: Final = stack.enter_context(gateway_from_environment()) + policy: Final = gateway.request( + "POST", "/v1/tool/policy", {"tool_name": blocked_tool, "input_policy": "blocked"} + ) + assert policy.status_code == 200, policy.text + stack.callback(write_rows, 'DELETE FROM "LiteLLM_ToolTable" WHERE tool_name = %s', (blocked_tool,)) + wire: Final = stack.enter_context(wire_server(respond)) + stack.callback(release.set) + proxy: Final = stack.enter_context( + owned_proxy_process( + gateway, + directory, + {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": _PROVIDER_KEY}, + config=_owned_config(directory / "malformed-tools.yaml", wire.url, denied_tool), + workers=2, + ) + ) + yield _Owned(proxy, wire, held, holding, release, blocked_tool, denied_tool) + + +def _owned_chat(owned: _Owned, marker: str, **fields: JsonValue) -> httpx.Response: + return owned.proxy.gateway.request( + "POST", + "/v1/chat/completions", + {"model": _OWNED_CHAT, "messages": [{"role": "user", "content": marker}], **fields}, + ) + + +def _named_tool(name: str) -> dict[str, JsonValue]: + return {"type": "function", "function": {"name": name, "parameters": _WEATHER_SCHEMA}} + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_global_drop_params_wins_over_an_explicit_request_false(owned: _Owned) -> None: + marker: Final = _marker("global drop") + response: Final = _owned_chat(owned, marker, tools=5, drop_params=False) + assert response.status_code == 200, response.text + assert _answer(_tag(marker)) in response.text, response.text + (sent,) = tuple(request for request in owned.wire.drain() if marker in request.body.decode()) + assert "tools" not in _JSON_OBJECT.validate_json(sent.body), sent.body + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_tool_policy_still_blocks_a_tool_kept_beside_a_dropped_entry(owned: _Owned) -> None: + marker: Final = _marker("tool policy") + response: Final = _owned_chat(owned, marker, tools=[_named_tool(owned.blocked_tool), 5]) + assert response.status_code == 400, response.text + assert "Violated tool policy" in response.text and owned.blocked_tool in response.text, response.text + assert tuple(request for request in owned.wire.drain() if marker in request.body.decode()) == () + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_tool_permission_still_denies_a_tool_kept_beside_a_dropped_entry(owned: _Owned) -> None: + marker: Final = _marker("tool permission") + response: Final = _owned_chat( + owned, marker, tools=[_named_tool(owned.denied_tool), "x"], guardrails=["malformed-tools-permission"] + ) + assert response.status_code == 400, response.text + assert "Violated guardrail policy" in response.text and owned.denied_tool in response.text, response.text + assert tuple(request for request in owned.wire.drain() if marker in request.body.decode()) == () + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_anthropic_pass_through_forwards_malformed_tools_unchanged(owned: _Owned) -> None: + marker: Final = _marker("pass-through") + body: Final[dict[str, JsonValue]] = { + "model": _CLAUDE_BACKEND, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + "tools": 5, + } + response: Final = owned.proxy.gateway.request("POST", "/anthropic/v1/messages", body) + assert response.status_code == 400, response.text + assert _JSON_OBJECT.validate_json(response.content) == _JSON_OBJECT.validate_json(error_body(400, _NOT_A_LIST)) + (sent,) = tuple(request for request in owned.wire.drain() if marker in request.body.decode()) + assert sent.target == "/v1/messages", sent.target + assert _JSON_OBJECT.validate_json(sent.body) == body, sent.body + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_anthropic_pass_through_streams_well_formed_tools(owned: _Owned) -> None: + marker: Final = _marker("pass-through stream") + body: Final[dict[str, JsonValue]] = { + "model": _CLAUDE_BACKEND, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + "tools": [_CLAUDE_TOOL], + "stream": True, + } + response: Final = owned.proxy.gateway.request("POST", "/anthropic/v1/messages", body) + assert response.status_code == 200, response.text + assert response.content == b"".join(message_stream(f"msg_{_tag(marker)}", _CLAUDE_BACKEND, _answer(_tag(marker)))) + (sent,) = tuple(request for request in owned.wire.drain() if marker in request.body.decode()) + assert _JSON_OBJECT.validate_json(sent.body) == body, sent.body + + +def _local_port(client: httpx.Client) -> int: + with client.stream("GET", "/health/liveliness") as response: + port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1]) + response.read() + assert response.status_code == 200, response.text + return port + + +def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: + return frozenset( + connection.raddr.port + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.raddr and connection.laddr.port == proxy_port + ) + + +def _owned_call(index: int, label: str) -> _Call: + marker: Final = _marker(label) + tools: Final[JsonValue] = [_CHAT_TOOL, 5] if index % 2 else [_CHAT_TOOL] + return _Call( + "/v1/chat/completions", + { + "model": _OWNED_CHAT, + "messages": [{"role": "user", "content": marker}], + "stream": index % 4 == 1, + "tools": tools, + }, + marker, + "dropped" if index % 2 else "valid", + ) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_worker_sigkill_mid_burst_leaves_the_replacement_dropping_bad_entries(owned: _Owned) -> None: + proxy_url: Final = owned.proxy.gateway.client.base_url + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.proxy.log.read_text())), + lambda pids: len(pids) >= 2, + seconds=30, + ) + started_before: Final = len(workers) + serving: Final = owned.proxy.worker_pids() + owned.wire.drain() + with ExitStack() as stack: + clients: Final = _clients(stack, str(proxy_url), 12) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + ports: Final = tuple(_local_port(client) for client in clients) + calls: Final = tuple(_owned_call(index, "worker kill") for index in range(len(clients))) + owned.release.clear() + owned.holding.set() + stack.callback(owned.holding.clear) + stack.callback(owned.release.set) + futures: Final = tuple( + pool.submit(_status_or_dropped, client, owned.proxy.gateway.key, call) + for client, call in zip(clients, calls, strict=True) + ) + eventually(owned.held.qsize, lambda size: size == len(clients), seconds=60) + shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in serving} + assert sum(map(len, shares.values())) == len(clients), shares + victim: Final = min((pid for pid in serving if shares[pid]), key=lambda pid: len(shares[pid])) + with owned.proxy.replacing_worker(victim) as kill: + kill() + owned.holding.clear() + owned.release.set() + results: Final = tuple(future.result(timeout=120) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else 200), (port, result, shares) + eventually( + lambda: len(_STARTED_WORKER.findall(owned.proxy.log.read_text())), + lambda started: started > started_before, + _RESPAWN_SECONDS, + ) + second_wave: Final = tuple(_owned_call(index * 2 + 1, "after kill") for index in range(6)) + (fresh,) = _clients(stack, str(proxy_url), 1) + assert tuple(_post(fresh, owned.proxy.gateway.key, call).status_code for call in second_wave) == (200,) * len( + second_wave + ) + deliveries: Final = _deliveries(owned.wire) + everything: Final = calls + second_wave + assert tuple(_seen(deliveries, call.marker) for call in everything) == (1,) * len(everything) + assert tuple(_delivered_tool_names(deliveries, call.marker) for call in everything) == ( + ("get_weather",), + ) * len(everything) + assert owned.proxy.process.poll() is None diff --git a/tests/integration/providers/test_malformed_tools_wire.py b/tests/integration/providers/test_malformed_tools_wire.py new file mode 100644 index 00000000000..25a86f7fc7c --- /dev/null +++ b/tests/integration/providers/test_malformed_tools_wire.py @@ -0,0 +1,2731 @@ +from __future__ import annotations + +import asyncio +import http.client +import json +import os +import re +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, replace +from itertools import chain, product +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from google import genai +from google.genai import types +from integration._support.anthropic_sse import error_body, message_json, message_stream +from integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, + held, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.openai_wire import chat_reply, openai_error, responses_reply +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.responses_vendor import same_response +from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.worker_readiness import wait_for_worker_readiness +from integration.authorization._guardrail_opt_out import upstream_observations +from pydantic import JsonValue, TypeAdapter + +from tests.integration.caching.response_cache_case import await_entry, redis_store + +_OPENAI_BACKEND: Final = "gpt-4o-mini" +_CLAUDE_BACKEND: Final = "claude-haiku-4-5" +_GEMINI_BACKEND: Final = "gemini-2.5-flash" +_PROVIDER_KEY: Final = "scripted-malformed-tools-key" +_COUNT: Final = 37 +_CONFIG_STORE: Final = "vs_integration_config_store" +_OTHER_STORE: Final = "vs_malformed_tools_other_store" +_CLAUDE_CODE_AGENT: Final = "claude-cli/2.1.0 (external, cli)" +_CODEX_AGENT: Final = "codex_cli_rs/0.50.0 (Linux 6.8; x86_64) xterm-256color" +_REJECT_MARK: Final = "upstream-401" +_DOWN: Final = "/down/" +_REFUSALS: Final = MappingProxyType( + { + "context_window_fallbacks": ( + "/context-full/", + { + "message": "This model's maximum context length is 8192 tokens.", + "type": "invalid_request_error", + "param": "messages", + "code": "context_length_exceeded", + }, + ), + "content_policy_fallbacks": ( + "/content-policy/", + { + "message": "Your request was rejected as a result of our safety system.", + "type": "invalid_request_error", + "param": None, + "code": "content_policy_violation", + }, + ), + } +) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1"))) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_SPEND_ROWS: Final = ( + 'SELECT call_type, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE request_id = %s OR litellm_call_id = %s' +) +_NOT_A_LIST: Final = "tools: Input should be a valid list" +_NOT_AN_OBJECT: Final = "tools.0: Input should be a valid dictionary" +_SECOND_NOT_AN_OBJECT: Final = "tools.1: Input should be a valid dictionary" +_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 +_TOOL_FIELDS: Final = ("tools", "tool_choice", "parallel_tool_calls", "toolConfig", "tool_config") +_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} + + +def _marker() -> str: + return f"marker-{uuid.uuid4().hex}" + + +def _tag(text: str) -> str: + found: Final = _MARKER.search(text) + return found.group(1) if found else "none" + + +def _answer(tag: str) -> str: + return f"answer {tag}" + + +def _gemini_body(tag: str) -> dict[str, JsonValue]: + return { + "responseId": f"gemini-{tag}-{uuid.uuid4().hex}", + "candidates": [ + {"content": {"role": "model", "parts": [{"text": _answer(tag)}]}, "index": 0, "finishReason": "STOP"} + ], + "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 2, "totalTokenCount": 11}, + "modelVersion": _GEMINI_BACKEND, + } + + +def _interaction_body(tag: str) -> dict[str, JsonValue]: + return { + "id": f"interaction-{tag}", + "object": "interaction", + "model": _GEMINI_BACKEND, + "status": "completed", + "created": "2026-10-15T00:00:00Z", + "steps": [{"type": "content", "content": {"type": "text", "text": _answer(tag)}}], + "usage": {"total_tokens": 11}, + } + + +def _rejected_by_provider(path: str) -> Reply: + if path.startswith("/v1/messages"): + return Reply(status=401, body=error_body(401, "scripted 401")) + if path == "/v1beta/interactions": + return Reply( + status=401, + body=json.dumps({"error": {"code": 401, "message": "scripted 401", "status": "UNAUTHENTICATED"}}).encode(), + ) + return openai_error(401) + + +def _unavailable(path: str) -> Reply: + if "/v1/messages" in path: + return Reply(status=503, body=error_body(503, "scripted outage")) + return openai_error(503) + + +def _respond(request: Request) -> Reply: + text: Final = request.body.decode() + tag: Final = _tag(text) + path: Final = request.target.split("?")[0] + if path.startswith(_DOWN): + return _unavailable(path) + refusal: Final = next((error for prefix, error in _REFUSALS.values() if path.startswith(prefix)), None) + if refusal is not None: + return Reply(status=400, body=json.dumps({"error": refusal}).encode()) + if _REJECT_MARK in text: + return _rejected_by_provider(path) + stream: Final = bool(text) and _JSON_OBJECT.validate_json(request.body).get("stream") is True + if path == "/v1/chat/completions": + return chat_reply(f"chatcmpl-{tag}", _OPENAI_BACKEND, _answer(tag), stream=stream) + if path == "/v1/responses": + return responses_reply(f"resp_{tag}", _OPENAI_BACKEND, _answer(tag), stream=stream) + if path == "/v1/messages/count_tokens": + return Reply(body=json.dumps({"input_tokens": _COUNT}).encode()) + if path == "/v1/messages" and stream: + return Reply( + content_type="text/event-stream", chunks=message_stream(f"msg_{tag}", _CLAUDE_BACKEND, _answer(tag)) + ) + if path == "/v1/messages": + return Reply(body=message_json(f"msg_{tag}", _CLAUDE_BACKEND, _answer(tag))) + if path.endswith(":streamGenerateContent"): + return Reply( + content_type="text/event-stream", chunks=(f"data: {json.dumps(_gemini_body(tag))}\r\n\r\n".encode(),) + ) + if path.endswith(":generateContent"): + return Reply(body=json.dumps(_gemini_body(tag)).encode()) + if path == "/v1beta/interactions": + return Reply(body=json.dumps(_interaction_body(tag)).encode()) + return Reply(status=404, body=json.dumps({"error": f"unscripted {request.target}"}).encode()) + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + wire: Wire + openai: str + anthropic: str + gemini: str + drop_group: str + gemini_drop_group: str + mixed_group: str + field_group: str + tools_drop_group: str + + +def _deployment_ids(scenario: Scenario, name: str, deployments: tuple[dict[str, JsonValue], ...]) -> tuple[str, ...]: + created: Final = tuple( + scenario.gateway.post("/model/new", {"model_name": name, "litellm_params": params, "model_info": {}}) + for params in deployments + ) + identities: Final = tuple(string_value(object_value(deployment["model_info"])["id"]) for deployment in created) + for identity in identities: + scenario.cleanups.callback(scenario.delete_model, identity) + return identities + + +def _same_name_deployments(scenario: Scenario, api_base: str, variants: tuple[Mapping[str, JsonValue], ...]) -> str: + name: Final = f"integration-mixed-{uuid.uuid4().hex}" + _deployment_ids( + scenario, + name, + tuple( + {"model": f"openai/{_OPENAI_BACKEND}", "api_key": _PROVIDER_KEY, "api_base": api_base, **variant} + for variant in variants + ), + ) + return name + + +@pytest.fixture(scope="module") +def rig() -> Iterator[_Rig]: + with gateway_from_environment() as gateway, wire_server(_respond) as wire, gateway.scenario() as scenario: + openai_base: Final = f"{wire.url}/v1" + built: Final = _Rig( + gateway=gateway, + wire=wire, + openai=scenario.model(api_base=openai_base, api_key=_PROVIDER_KEY), + anthropic=scenario.model(model=f"anthropic/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_PROVIDER_KEY), + gemini=scenario.model(model=f"gemini/{_GEMINI_BACKEND}", api_base=wire.url, api_key=_PROVIDER_KEY), + drop_group=scenario.model(api_base=openai_base, api_key=_PROVIDER_KEY, drop_params=True), + gemini_drop_group=scenario.model( + model=f"gemini/{_GEMINI_BACKEND}", api_base=wire.url, api_key=_PROVIDER_KEY, drop_params=True + ), + mixed_group=_same_name_deployments(scenario, openai_base, ({"drop_params": True}, {})), + field_group=scenario.model(api_base=openai_base, api_key=_PROVIDER_KEY, additional_drop_params=["tools"]), + tools_drop_group=scenario.model( + api_base=openai_base, api_key=_PROVIDER_KEY, drop_params=True, tools=[CHAT.tool] + ), + ) + for group, deployments in ( + (built.openai, 1), + (built.anthropic, 1), + (built.gemini, 1), + (built.drop_group, 1), + (built.gemini_drop_group, 1), + (built.mixed_group, 2), + (built.field_group, 1), + (built.tools_drop_group, 1), + ): + wait_for_worker_readiness(gateway, group, deployments, _PROXY_WORKERS) + yield built + + +@pytest.fixture(autouse=True) +def _drained(rig: _Rig) -> None: + rig.wire.drain() + + +@dataclass(frozen=True, slots=True) +class _FieldGroups: + anthropic: str + gemini: str + + +@pytest.fixture(scope="module") +def field_groups(rig: _Rig) -> Iterator[_FieldGroups]: + with rig.gateway.scenario() as scenario: + built: Final = _FieldGroups( + anthropic=scenario.model( + model=f"anthropic/{_CLAUDE_BACKEND}", + api_base=rig.wire.url, + api_key=_PROVIDER_KEY, + additional_drop_params=["tools"], + ), + gemini=scenario.model( + model=f"gemini/{_GEMINI_BACKEND}", + api_base=rig.wire.url, + api_key=_PROVIDER_KEY, + additional_drop_params=["tools"], + ), + ) + for group in (built.anthropic, built.gemini): + wait_for_worker_readiness(rig.gateway, group, 1, _PROXY_WORKERS) + yield built + + +@pytest.fixture(scope="module") +def limited_key(rig: _Rig) -> Iterator[str]: + with rig.gateway.scenario() as scenario: + yield scenario.key(rpm_limit=0) + + +@dataclass(frozen=True, slots=True) +class _Family: + name: str + tool: dict[str, JsonValue] + model: Callable[[_Rig], str] + path: Callable[[str, bool], str] + body: Callable[[str, str, bool], dict[str, JsonValue]] + target: Callable[[bool], str] + call_type: Callable[[bool], str] + upstream_id: Callable[[str], str | None] + tool_choice: tuple[str, JsonValue] + exact_tools: bool + + +def _chat_body(model: str, marker: str, stream: bool) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}], "stream": stream} + + +def _messages_body(model: str, marker: str, stream: bool) -> dict[str, JsonValue]: + return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}], "stream": stream} + + +def _responses_body(model: str, marker: str, stream: bool) -> dict[str, JsonValue]: + return {"model": model, "input": marker, "stream": stream} + + +def _generate_body(model: str, marker: str, stream: bool) -> dict[str, JsonValue]: + return {"contents": [{"role": "user", "parts": [{"text": marker}]}]} + + +def _generate_path(model: str, stream: bool) -> str: + return ( + f"/v1beta/models/{model}:streamGenerateContent?alt=sse" if stream else f"/v1beta/models/{model}:generateContent" + ) + + +CHAT: Final = _Family( + name="chat", + tool={ + "type": "function", + "function": {"name": "get_weather", "description": "Look up the weather", "parameters": _WEATHER_SCHEMA}, + }, + model=lambda rig: rig.openai, + path=lambda model, stream: "/v1/chat/completions", + body=_chat_body, + target=lambda stream: "/v1/chat/completions", + call_type=lambda stream: "acompletion", + upstream_id=lambda tag: f"chatcmpl-{tag}", + tool_choice=("tool_choice", {"type": "function", "function": {"name": "get_weather"}}), + exact_tools=True, +) +MESSAGES: Final = _Family( + name="messages", + tool={"name": "get_weather", "description": "Look up the weather", "input_schema": _WEATHER_SCHEMA}, + model=lambda rig: rig.anthropic, + path=lambda model, stream: "/v1/messages", + body=_messages_body, + target=lambda stream: "/v1/messages", + call_type=lambda stream: "anthropic_messages", + upstream_id=lambda tag: f"msg_{tag}", + tool_choice=("tool_choice", {"type": "tool", "name": "get_weather"}), + exact_tools=True, +) +RESPONSES: Final = _Family( + name="responses", + tool={ + "type": "function", + "name": "get_weather", + "description": "Look up the weather", + "parameters": _WEATHER_SCHEMA, + }, + model=lambda rig: rig.openai, + path=lambda model, stream: "/v1/responses", + body=_responses_body, + target=lambda stream: "/v1/responses", + call_type=lambda stream: "aresponses", + upstream_id=lambda tag: None, + tool_choice=("tool_choice", {"type": "function", "name": "get_weather"}), + exact_tools=True, +) +GENERATE: Final = _Family( + name="generate_content", + tool={ + "functionDeclarations": [ + {"name": "get_weather", "description": "Look up the weather", "parameters": _WEATHER_SCHEMA} + ] + }, + model=lambda rig: rig.gemini, + path=_generate_path, + body=_generate_body, + target=lambda stream: ":streamGenerateContent" if stream else ":generateContent", + call_type=lambda stream: "agenerate_content_stream" if stream else "agenerate_content", + upstream_id=lambda tag: None, + tool_choice=("toolConfig", {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["get_weather"]}}), + exact_tools=False, +) +_FAMILIES: Final = (CHAT, MESSAGES, RESPONSES, GENERATE) +_FAMILY_IDS: Final = tuple(family.name for family in _FAMILIES) +_STREAMS: Final = pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +_EACH_FAMILY: Final = pytest.mark.parametrize("family", _FAMILIES, ids=_FAMILY_IDS) + + +@dataclass(frozen=True, slots=True) +class _Route: + path: Callable[[str], str] + body: Callable[[str, str], dict[str, JsonValue]] + model: Callable[[_Rig], str] + + +_ROUTES: Final[Mapping[str, _Route]] = MappingProxyType( + { + "chat": _Route(lambda model: "/v1/chat/completions", lambda m, k: _chat_body(m, k, False), CHAT.model), + "chat_unversioned": _Route(lambda model: "/chat/completions", lambda m, k: _chat_body(m, k, False), CHAT.model), + "chat_azure_alias": _Route( + lambda model: f"/openai/deployments/{model}/chat/completions", + lambda m, k: _chat_body(m, k, False), + CHAT.model, + ), + "chat_stream": _Route(lambda model: "/v1/chat/completions", lambda m, k: _chat_body(m, k, True), CHAT.model), + "messages": _Route(lambda model: "/v1/messages", lambda m, k: _messages_body(m, k, False), MESSAGES.model), + "messages_stream": _Route( + lambda model: "/v1/messages", lambda m, k: _messages_body(m, k, True), MESSAGES.model + ), + "responses": _Route(lambda model: "/v1/responses", lambda m, k: _responses_body(m, k, False), RESPONSES.model), + "responses_stream": _Route( + lambda model: "/v1/responses", lambda m, k: _responses_body(m, k, True), RESPONSES.model + ), + "responses_compact": _Route( + lambda model: "/v1/responses/compact", lambda m, k: {"model": m, "input": k}, RESPONSES.model + ), + "generate_content": _Route( + lambda model: _generate_path(model, False), lambda m, k: _generate_body(m, k, False), GENERATE.model + ), + "stream_generate_content": _Route( + lambda model: _generate_path(model, True), lambda m, k: _generate_body(m, k, True), GENERATE.model + ), + "interactions": _Route( + lambda model: "/v1beta/interactions", lambda m, k: {"model": m, "input": k}, lambda rig: rig.gemini + ), + } +) +_MALFORMED_TOOLS: Final[Mapping[str, tuple[JsonValue, str]]] = MappingProxyType( + { + "int": (5, _NOT_A_LIST), + "empty_string": ("", _NOT_A_LIST), + "five_kb_string": ("x" * 5120, _NOT_A_LIST), + "list_of_names": (["get_weather"], _NOT_AN_OBJECT), + "empty_object": ({}, _NOT_A_LIST), + "false": (False, _NOT_A_LIST), + "zero": (0, _NOT_A_LIST), + } +) +_MALFORMED_ENTRIES: Final[Mapping[str, tuple[JsonValue, str]]] = MappingProxyType( + { + "int_entry": (5, _NOT_AN_OBJECT), + "string_function": ( + {"type": "function", "function": "get_weather"}, + "tools.0.function: Input should be a valid dictionary", + ), + "string_function_parameters": ( + {"type": "function", "function": {"name": "get_weather", "parameters": "bad"}}, + "tools.0.function.parameters: Input should be a valid dictionary", + ), + "list_input_schema": ( + {"name": "get_weather", "input_schema": [1]}, + "tools.0.input_schema: Input should be a valid dictionary", + ), + "int_parameters": ( + {"type": "function", "name": "get_weather", "parameters": 5}, + "tools.0.parameters: Input should be a valid dictionary", + ), + "string_vector_store_ids": ( + {"type": "file_search", "vector_store_ids": _OTHER_STORE}, + "tools.0.vector_store_ids: Input should be a valid list", + ), + } +) + + +def _send( + rig: _Rig, + family: _Family, + marker: str, + stream: bool, + *, + model: str | None = None, + key: str | None = None, + call_id: str | None = None, + user_agent: str | None = None, + **fields: JsonValue, +) -> httpx.Response: + chosen: Final = model or family.model(rig) + headers: Final = { + **({"x-litellm-call-id": call_id} if call_id else {}), + **({"user-agent": user_agent} if user_agent else {}), + } + return rig.gateway.request( + "POST", family.path(chosen, stream), {**family.body(chosen, marker, stream), **fields}, key=key, headers=headers + ) + + +def _delivered(rig: _Rig, marker: str) -> tuple[Request, ...]: + return tuple(request for request in rig.wire.drain() if marker in request.body.decode()) + + +def _one_delivery(rig: _Rig, marker: str, target: str) -> dict[str, JsonValue]: + (sent,) = _delivered(rig, marker) + assert sent.method == "POST" and target in sent.target, sent.target + return _JSON_OBJECT.validate_json(sent.body) + + +def _tool_names(family: _Family, body: Mapping[str, JsonValue]) -> tuple[str, ...]: + tools: Final = body.get("tools") + if not isinstance(tools, list): + return () + if family is GENERATE: + declarations: Final = chain.from_iterable(_declarations(object_value(tool)) for tool in tools) + return tuple(string_value(object_value(declaration)["name"]) for declaration in declarations) + if family is CHAT: + return tuple(string_value(object_value(object_value(tool)["function"])["name"]) for tool in tools) + return tuple(string_value(object_value(tool)["name"]) for tool in tools) + + +def _declarations(tool: Mapping[str, JsonValue]) -> list[JsonValue]: + declarations: Final = tool.get("functionDeclarations", tool.get("function_declarations")) + assert isinstance(declarations, list), tool + return declarations + + +def _assert_answered(family: _Family, response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + tag: Final = _tag(marker) + assert _answer(tag) in response.text, response.text + upstream_id: Final = family.upstream_id(tag) + if upstream_id is not None: + assert upstream_id in response.text, response.text + if family is RESPONSES and not response.headers.get("content-type", "").startswith("text/event-stream"): + assert same_response(string_value(_JSON_OBJECT.validate_json(response.content)["id"]), f"resp_{tag}") + + +def _rejection(response: httpx.Response) -> tuple[str, str, JsonValue]: + body: Final = _JSON_OBJECT.validate_json(response.content) + error: Final = object_value(body["error"]) + if body.get("type") == "error": + return string_value(error["type"]), string_value(error["message"]), None + return string_value(error["type"]), string_value(error["message"]), error.get("param") + + +def _assert_rejected(response: httpx.Response, message: str, param: str = "tools") -> None: + assert response.status_code == 400, response.text + kind, said, named = _rejection(response) + assert kind == "invalid_request_error", response.text + assert message in said, response.text + assert named in (param, None), response.text + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows(_SPEND_ROWS, (call_id, call_id)) + + +def _logged_once(call_id: str) -> dict[str, JsonValue]: + (row,) = held(lambda: _spend_rows(call_id), lambda rows: len(rows) == 1, holding=3, seconds=70) + return row + + +def _tool_values(sent: Mapping[str, JsonValue]) -> frozenset[str]: + return frozenset(name for name in _TOOL_FIELDS if sent.get(name) is not None) + + +def _call_id() -> str: + return f"malformed-tools-{uuid.uuid4().hex}" + + +@_STREAMS +@_EACH_FAMILY +def test_well_formed_tools_reach_the_provider_and_log_one_success_row(rig: _Rig, family: _Family, stream: bool) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = _send(rig, family, marker, stream, call_id=call_id, tools=[family.tool]) + _assert_answered(family, response, marker) + sent: Final = _one_delivery(rig, marker, family.target(stream)) + assert _tool_names(family, sent) == ("get_weather",), sent + if family.exact_tools: + assert sent["tools"] == [family.tool], sent + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (family.call_type(stream), "success"), row + + +def test_openai_sdk_chat_with_tools_answers_with_the_provider_reply(rig: _Rig) -> None: + marker: Final = _marker() + client: Final = openai.OpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False), + ) + completion: Final = client.chat.completions.create( + model=rig.openai, + messages=[{"role": "user", "content": marker}], + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Look up the weather", + "parameters": _WEATHER_SCHEMA, + }, + } + ], + ) + assert completion.id == f"chatcmpl-{_tag(marker)}", completion + assert completion.choices[0].message.content == _answer(_tag(marker)), completion + assert _one_delivery(rig, marker, "/v1/chat/completions")["tools"] == [CHAT.tool] + + +def test_openai_sdk_rejects_malformed_tools_with_a_400(rig: _Rig) -> None: + marker: Final = _marker() + client: Final = openai.OpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False), + ) + with pytest.raises(openai.BadRequestError) as raised: + client.chat.completions.create( + model=rig.openai, messages=[{"role": "user", "content": marker}], extra_body={"tools": 5} + ) + assert raised.value.status_code == 400, raised.value + assert object_value(raised.value.body)["param"] == "tools", raised.value.body + assert _delivered(rig, marker) == () + + +def test_async_openai_sdk_streams_responses_with_tools(rig: _Rig) -> None: + marker: Final = _marker() + + async def stream_text() -> str: + client: Final = openai.AsyncOpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False), + ) + events: Final = await client.responses.create( + model=rig.openai, + input=marker, + stream=True, + tools=[ + { + "type": "function", + "name": "get_weather", + "description": "Look up the weather", + "parameters": _WEATHER_SCHEMA, + "strict": False, + } + ], + ) + completed: Final = [event async for event in events if event.type == "response.completed"] + assert len(completed) == 1, completed + return completed[0].response.output_text + + assert asyncio.run(stream_text()) == _answer(_tag(marker)) + assert _tool_names(RESPONSES, _one_delivery(rig, marker, "/v1/responses")) == ("get_weather",) + + +def test_anthropic_sdk_streams_messages_with_tools(rig: _Rig) -> None: + marker: Final = _marker() + client: Final = anthropic.Anthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False), + ) + with client.messages.stream( + model=rig.anthropic, + max_tokens=16, + messages=[{"role": "user", "content": marker}], + tools=[{"name": "get_weather", "description": "Look up the weather", "input_schema": _WEATHER_SCHEMA}], + ) as stream: + final: Final = stream.get_final_message() + assert final.id == f"msg_{_tag(marker)}", final + assert _one_delivery(rig, marker, "/v1/messages")["tools"] == [MESSAGES.tool] + + +def test_google_genai_sdk_generates_with_tools(rig: _Rig) -> None: + marker: Final = _marker() + client: Final = genai.Client(api_key=rig.gateway.key, http_options={"base_url": str(rig.gateway.client.base_url)}) + response: Final = client.models.generate_content( + model=rig.gemini, + contents=marker, + config=types.GenerateContentConfig( + tools=[ + types.Tool( + function_declarations=[ + types.FunctionDeclaration( + name="get_weather", + description="Look up the weather", + parameters_json_schema=_WEATHER_SCHEMA, + ) + ] + ) + ], + automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True), + ), + ) + assert response.text == _answer(_tag(marker)), response + assert _tool_names(GENERATE, _one_delivery(rig, marker, ":generateContent")) == ("get_weather",) + + +@pytest.mark.parametrize("value", _MALFORMED_TOOLS.values(), ids=_MALFORMED_TOOLS.keys()) +@pytest.mark.parametrize("route", _ROUTES.values(), ids=_ROUTES.keys()) +def test_malformed_tools_answer_400_and_never_reach_the_provider( + rig: _Rig, route: _Route, value: tuple[JsonValue, str] +) -> None: + marker: Final = _marker() + model: Final = route.model(rig) + tools, message = value + response: Final = rig.gateway.request("POST", route.path(model), {**route.body(model, marker), "tools": tools}) + _assert_rejected(response, message) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize("entry", _MALFORMED_ENTRIES.values(), ids=_MALFORMED_ENTRIES.keys()) +@pytest.mark.parametrize("family", (CHAT, MESSAGES), ids=("chat", "messages")) +def test_a_malformed_tool_entry_answers_400_naming_its_index( + rig: _Rig, family: _Family, entry: tuple[JsonValue, str] +) -> None: + marker: Final = _marker() + tool, message = entry + _assert_rejected(_send(rig, family, marker, False, tools=[tool]), message) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize("value", [5, "x", []], ids=["int", "string", "list"]) +def test_malformed_web_search_options_answer_400(rig: _Rig, value: JsonValue) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, web_search_options=value) + expected: Final = "web_search_options: Input should be a valid dictionary" + _assert_rejected(response, expected, param="web_search_options") + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize("value", [5, "x", []], ids=["int", "string", "list"]) +def test_malformed_web_search_options_drop_under_drop_params(rig: _Rig, value: JsonValue) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, web_search_options=value, drop_params=True) + _assert_answered(CHAT, response, marker) + assert "web_search_options" not in _one_delivery(rig, marker, "/v1/chat/completions") + + +@_EACH_FAMILY +def test_wrong_key_with_malformed_tools_answers_401(rig: _Rig, family: _Family) -> None: + marker: Final = _marker() + response: Final = _send(rig, family, marker, False, key=f"sk-wrong-{uuid.uuid4().hex}", tools=5) + assert response.status_code == 401, response.text + assert _delivered(rig, marker) == () + + +def _count_body(rig: _Rig, marker: str, **fields: JsonValue) -> dict[str, JsonValue]: + return {"model": rig.anthropic, "messages": [{"role": "user", "content": marker}], **fields} + + +@pytest.mark.parametrize( + "fields", + [ + {"tools": 5}, + {"tools": ""}, + {"tools": "x" * 5120}, + {"tools": ["get_weather"]}, + {"messages": "x"}, + ], + ids=["int_tools", "empty_string_tools", "five_kb_tools", "list_of_names", "string_messages"], +) +def test_count_tokens_answers_an_anthropic_400_and_keeps_the_connection( + rig: _Rig, fields: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + port: Final = rig.gateway.client.base_url.port + assert port is not None + connection: Final = http.client.HTTPConnection(rig.gateway.client.base_url.host, port, timeout=30) + headers: Final = {"Authorization": f"Bearer {rig.gateway.key}", "Content-Type": "application/json"} + try: + connection.request( + "POST", "/v1/messages/count_tokens", body=json.dumps(_count_body(rig, marker, **fields)), headers=headers + ) + counted: Final = connection.getresponse() + counted_body: Final = _JSON_OBJECT.validate_json(counted.read()) + socket_after_count: Final = connection.sock + connection.request( + "POST", "/v1/messages", body=json.dumps(_messages_body(rig.anthropic, marker, False)), headers=headers + ) + followed: Final = connection.getresponse() + followed_body: Final = followed.read().decode() + socket_after_follow_up: Final = connection.sock + finally: + connection.close() + assert counted.status == 400, counted_body + assert counted_body["type"] == "error", counted_body + error: Final = object_value(counted_body["error"]) + assert error["type"] == "invalid_request_error", counted_body + assert string_value(error["message"]).startswith(next(iter(fields))), counted_body + assert followed.status == 200, followed_body + assert f"msg_{_tag(marker)}" in followed_body, followed_body + assert socket_after_count is not None and socket_after_follow_up is socket_after_count + (sent,) = _delivered(rig, marker) + assert sent.target == "/v1/messages", sent.target + + +def test_anthropic_sdk_count_tokens_rejects_malformed_tools(rig: _Rig) -> None: + marker: Final = _marker() + client: Final = anthropic.Anthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False), + ) + with pytest.raises(anthropic.BadRequestError) as raised: + client.messages.count_tokens( + model=rig.anthropic, messages=[{"role": "user", "content": marker}], extra_body={"tools": 5} + ) + assert raised.value.status_code == 400, raised.value + assert object_value(object_value(raised.value.body)["error"])["message"] == _NOT_A_LIST, raised.value.body + assert _delivered(rig, marker) == () + + +def test_count_tokens_forwards_well_formed_tools(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/messages/count_tokens", _count_body(rig, marker, tools=[MESSAGES.tool]) + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT}, response.text + assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool] + + +@pytest.mark.parametrize( + ("headers", "fields"), + [({}, {"drop_params": True}), ({"user-agent": _CLAUDE_CODE_AGENT}, {})], + ids=["drop_params", "claude_code_agent"], +) +def test_count_tokens_drops_only_the_bad_entry( + rig: _Rig, headers: dict[str, str], fields: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + _count_body(rig, marker, tools=[MESSAGES.tool, 5], **fields), + headers=headers, + ) + assert response.status_code == 200, response.text + assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool] + + +def test_responses_input_tokens_counts_well_formed_tools(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/responses/input_tokens", {"model": rig.anthropic, "input": marker, "tools": [RESPONSES.tool]} + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content) == { + "object": "response.input_tokens", + "input_tokens": _COUNT, + }, response.text + sent: Final = _one_delivery(rig, marker, "/v1/messages/count_tokens") + assert "get_weather" in json.dumps(sent["tools"]), sent + + +def test_responses_input_tokens_rejects_malformed_tools(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/responses/input_tokens", {"model": rig.anthropic, "input": marker, "tools": 5} + ) + _assert_rejected(response, _NOT_A_LIST) + assert _delivered(rig, marker) == () + + +def test_responses_input_tokens_drops_only_the_bad_entry(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": rig.anthropic, "input": marker, "tools": [RESPONSES.tool, 5], "drop_params": True}, + ) + assert response.status_code == 200, response.text + sent: Final = _one_delivery(rig, marker, "/v1/messages/count_tokens") + tools: Final = sent["tools"] + assert isinstance(tools, list) and len(tools) == 1 and "get_weather" in json.dumps(tools), sent + + +def _allowlisted_key(scenario: Scenario) -> str: + return scenario.key(metadata={"allowed_tools": ["get_weather"]}) + + +@_EACH_FAMILY +def test_allowlisted_key_keeps_its_allowlist_with_malformed_tools(rig: _Rig, family: _Family) -> None: + with rig.gateway.scenario() as scenario: + key: Final = _allowlisted_key(scenario) + allowed: Final = _marker() + _assert_answered(family, _send(rig, family, allowed, False, key=key, tools=[family.tool]), allowed) + assert len(_delivered(rig, allowed)) == 1 + malformed: Final = _marker() + _assert_rejected(_send(rig, family, malformed, False, key=key, tools=5), _NOT_A_LIST) + denied_tool: Final = json.loads(json.dumps(family.tool).replace("get_weather", "delete_files")) + denied: Final = _marker() + refused: Final = _send(rig, family, denied, False, key=key, tools=[denied_tool]) + assert refused.status_code == 403, refused.text + assert "delete_files" in refused.text, refused.text + assert _delivered(rig, malformed) + _delivered(rig, denied) == () + + +def _file_search(store_ids: JsonValue) -> dict[str, JsonValue]: + return {"type": "file_search", "vector_store_ids": store_ids} + + +def test_store_limited_key_searches_its_store_from_a_file_search_tool(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(object_permission={"vector_stores": [_CONFIG_STORE]}) + response: Final = _send(rig, CHAT, marker, False, key=key, tools=[_file_search([_CONFIG_STORE])]) + _assert_answered(CHAT, response, marker) + sent: Final = _one_delivery(rig, marker, "/v1/chat/completions") + assert "file_search" not in json.dumps(sent), sent + searches: Final = eventually( + lambda: tuple( + observation + for observation in upstream_observations(rig.gateway) + if observation["path"] == f"/vector_stores/{_CONFIG_STORE}/search" and marker in json.dumps(observation) + ), + lambda found: len(found) == 1, + ) + assert len(searches) == 1, searches + + +@pytest.mark.parametrize( + ("store_ids", "status", "named"), + [ + ([_OTHER_STORE], 401, _OTHER_STORE), + (_CONFIG_STORE, 400, "tools.0.vector_store_ids: Input should be a valid list"), + (_OTHER_STORE, 401, _OTHER_STORE), + ([5, _OTHER_STORE], 401, _OTHER_STORE), + ([5, _CONFIG_STORE], 400, "tools.0.vector_store_ids.0: Input should be a valid string"), + ], + ids=["other_store", "allowed_store_as_string", "other_store_as_string", "int_beside_other", "int_beside_allowed"], +) +def test_store_limited_key_reads_every_string_store_id( + rig: _Rig, store_ids: JsonValue, status: int, named: str +) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(object_permission={"vector_stores": [_CONFIG_STORE]}) + response: Final = _send(rig, CHAT, marker, False, key=key, tools=[_file_search(store_ids)]) + assert response.status_code == status, response.text + assert named in response.text, response.text + assert _delivered(rig, marker) == () + + +def test_store_limited_key_treats_null_store_ids_as_none(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(object_permission={"vector_stores": [_CONFIG_STORE]}) + response: Final = _send(rig, CHAT, marker, False, key=key, tools=[_file_search(None)]) + _assert_answered(CHAT, response, marker) + assert _one_delivery(rig, marker, "/v1/chat/completions")["tools"] == [_file_search(None)] + + +@pytest.mark.parametrize( + ("first", "last", "status"), + [("5", "WELL_FORMED", 200), ("WELL_FORMED", "5", 400)], + ids=["well_formed_last", "malformed_last"], +) +def test_duplicate_tools_keys_use_the_last_one(rig: _Rig, first: str, last: str, status: int) -> None: + marker: Final = _marker() + well_formed: Final = json.dumps([CHAT.tool]) + prefix: Final = json.dumps({"model": rig.openai, "messages": [{"role": "user", "content": marker}]})[:-1] + raw: Final = ( + f'{prefix}, "tools": {first.replace("WELL_FORMED", well_formed)}, ' + f'"tools": {last.replace("WELL_FORMED", well_formed)}}}' + ) + response: Final = rig.gateway.client.post( + "/v1/chat/completions", + content=raw.encode(), + headers={"Authorization": f"Bearer {rig.gateway.key}", "Content-Type": "application/json"}, + ) + assert response.status_code == status, response.text + delivered: Final = _delivered(rig, marker) + if status == 400: + _assert_rejected(response, _NOT_A_LIST) + assert delivered == () + return + (sent,) = delivered + assert _JSON_OBJECT.validate_json(sent.body)["tools"] == [CHAT.tool] + + +@pytest.mark.parametrize( + ("family", "route", "expected"), + [ + (CHAT, None, "acompletion"), + (MESSAGES, None, "anthropic_messages"), + (RESPONSES, None, "aresponses"), + (GENERATE, None, "agenerate_content"), + (None, "interactions", "acreate_interaction"), + ], + ids=["chat", "messages", "responses", "generate_content", "interactions"], +) +def test_a_malformed_tools_rejection_logs_one_failure_row_with_its_call_type( + rig: _Rig, family: _Family | None, route: str | None, expected: str +) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = ( + _send(rig, family, marker, False, call_id=call_id, tools=5) + if family is not None + else rig.gateway.request( + "POST", + "/v1beta/interactions", + {"model": rig.gemini, "input": [{"type": "text", "text": marker}], "tools": 5}, + headers={"x-litellm-call-id": call_id}, + ) + ) + _assert_rejected(response, _NOT_A_LIST) + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (expected, "failure"), row + + +@pytest.mark.parametrize( + ("path", "body", "expected"), + [ + ("/v1/chat/completions", "chat", "acompletion"), + ("/v1/messages", "messages", "anthropic_messages"), + ("/v1beta/interactions", "interactions", "acreate_interaction"), + ], + ids=["chat", "messages", "interactions_list_input"], +) +def test_a_provider_401_with_well_formed_tools_logs_its_endpoint_call_type( + rig: _Rig, path: str, body: str, expected: str +) -> None: + marker: Final = f"{_marker()} {_REJECT_MARK}" + call_id: Final = _call_id() + payloads: Final[Mapping[str, dict[str, JsonValue]]] = { + "chat": {**_chat_body(rig.openai, marker, False), "tools": [CHAT.tool]}, + "messages": {**_messages_body(rig.anthropic, marker, False), "tools": [MESSAGES.tool]}, + "interactions": { + "model": rig.gemini, + "input": [{"type": "text", "text": marker}], + "tools": [{"type": "function", "name": "get_weather", "parameters": _WEATHER_SCHEMA}], + }, + } + response: Final = rig.gateway.request("POST", path, payloads[body], headers={"x-litellm-call-id": call_id}) + assert response.status_code == 401, response.text + assert len(_delivered(rig, marker)) == 1 + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (expected, "failure"), row + + +@pytest.mark.parametrize("drop", [True, "true"], ids=["bool", "string"]) +@_EACH_FAMILY +def test_request_drop_params_drops_malformed_tools_and_their_companions( + rig: _Rig, family: _Family, drop: JsonValue +) -> None: + marker: Final = _marker() + companion, value = family.tool_choice + response: Final = _send(rig, family, marker, False, tools=5, drop_params=drop, **{companion: value}) + _assert_answered(family, response, marker) + sent: Final = _one_delivery(rig, marker, family.target(False)) + assert _tool_values(sent) == frozenset(), sent + + +@_EACH_FAMILY +def test_drop_params_keeps_the_good_entries_and_the_tool_choice(rig: _Rig, family: _Family) -> None: + marker: Final = _marker() + companion, value = family.tool_choice + response: Final = _send(rig, family, marker, False, tools=[family.tool, 5], drop_params=True, **{companion: value}) + _assert_answered(family, response, marker) + sent: Final = _one_delivery(rig, marker, family.target(False)) + assert _tool_names(family, sent) == ("get_weather",), sent + if family.exact_tools: + assert sent["tools"] == [family.tool], sent + assert companion in sent, sent + + +def test_additional_drop_params_naming_tools_drops_them(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, tools=5, additional_drop_params=["tools"]) + _assert_answered(CHAT, response, marker) + assert "tools" not in _one_delivery(rig, marker, "/v1/chat/completions") + + +def test_additional_drop_params_naming_another_field_still_rejects_tools(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, tools=5, additional_drop_params=["web_search_options"]) + _assert_rejected(response, _NOT_A_LIST) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize("family", (CHAT, GENERATE), ids=("chat", "generate_content")) +def test_a_drop_params_deployment_group_drops_malformed_tools(rig: _Rig, family: _Family) -> None: + marker: Final = _marker() + model: Final = rig.drop_group if family is CHAT else rig.gemini_drop_group + response: Final = _send(rig, family, marker, False, model=model, tools=5) + _assert_answered(family, response, marker) + assert _tool_values(_one_delivery(rig, marker, family.target(False))) == frozenset() + + +def test_a_drop_params_deployment_keeps_its_own_tools_when_it_drops_malformed_request_tools(rig: _Rig) -> None: + marker: Final = _marker() + companion, value = CHAT.tool_choice + response: Final = _send( + rig, CHAT, marker, False, model=rig.tools_drop_group, tools=5, parallel_tool_calls=False, **{companion: value} + ) + _assert_answered(CHAT, response, marker) + sent: Final = _one_delivery(rig, marker, "/v1/chat/completions") + assert _tool_names(CHAT, sent) == _tool_names(CHAT, {"tools": [CHAT.tool]}), sent + assert _tool_values(sent) == frozenset({"tools"}), sent + + +@pytest.mark.parametrize("drop", [False, None], ids=["false", "null"]) +def test_an_explicit_request_drop_params_overrides_the_deployment_group(rig: _Rig, drop: JsonValue) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, model=rig.drop_group, tools=5, drop_params=drop) + _assert_rejected(response, _NOT_A_LIST) + assert _delivered(rig, marker) == () + + +def test_a_mixed_group_answers_malformed_tools_by_the_drop_params_of_the_deployment_it_picks(rig: _Rig) -> None: + markers: Final = tuple(_marker() for _ in range(20)) + statuses: Final = tuple( + _send(rig, CHAT, marker, False, model=rig.mixed_group, tools=5).status_code for marker in markers + ) + sent: Final = tuple(request.body.decode() for request in rig.wire.drain()) + deliveries: Final = tuple( + tuple(_JSON_OBJECT.validate_json(body) for body in sent if marker in body) for marker in markers + ) + assert sorted(set(statuses)) == [200, 400], statuses + assert all((status, len(bodies)) in ((200, 1), (400, 0)) for status, bodies in zip(statuses, deliveries)), ( + statuses, + deliveries, + ) + assert all(all("tools" not in body for body in bodies) for bodies in deliveries), deliveries + + +def test_a_deployment_dropping_tools_drops_every_tool(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, model=rig.field_group, tools=[CHAT.tool, 5]) + _assert_answered(CHAT, response, marker) + assert "tools" not in _one_delivery(rig, marker, "/v1/chat/completions") + + +def test_a_request_additional_drop_params_list_keeps_tools_a_deployment_list_drops(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + model=rig.field_group, + tools=[CHAT.tool], + additional_drop_params=["web_search_options"], + ) + _assert_answered(CHAT, response, marker) + assert _tool_names(CHAT, _one_delivery(rig, marker, "/v1/chat/completions")) == ("get_weather",) + + +def test_a_request_additional_drop_params_list_replaces_the_deployment_list(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + model=rig.field_group, + tools=[CHAT.tool, 5], + additional_drop_params=["web_search_options"], + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize( + ("family", "agent"), + [(MESSAGES, _CLAUDE_CODE_AGENT), (RESPONSES, _CODEX_AGENT), (CHAT, _CLAUDE_CODE_AGENT)], + ids=["claude_code_messages", "codex_responses", "claude_code_chat"], +) +def test_an_agentic_cli_drops_only_the_bad_entry(rig: _Rig, family: _Family, agent: str) -> None: + marker: Final = _marker() + response: Final = _send(rig, family, marker, False, user_agent=agent, tools=[family.tool, 5]) + _assert_answered(family, response, marker) + assert _one_delivery(rig, marker, family.target(False))["tools"] == [family.tool] + + +def test_an_agentic_cli_with_explicit_drop_params_false_rejects(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send( + rig, MESSAGES, marker, False, user_agent=_CLAUDE_CODE_AGENT, tools=[MESSAGES.tool, 5], drop_params=False + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize( + "fields", [{"tools": None}, {"tools": []}, {}], ids=["null_tools", "empty_list_tools", "missing_tools"] +) +def test_empty_null_and_missing_tools_still_answer(rig: _Rig, fields: dict[str, JsonValue]) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, **fields) + _assert_answered(CHAT, response, marker) + assert not _one_delivery(rig, marker, "/v1/chat/completions").get("tools") + + +def test_null_web_search_options_answers_and_logs_success(rig: _Rig) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = _send(rig, CHAT, marker, False, call_id=call_id, web_search_options=None) + _assert_answered(CHAT, response, marker) + assert _one_delivery(rig, marker, "/v1/chat/completions").get("web_search_options") is None + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == ("acompletion", "success"), row + + +@_EACH_FAMILY +def test_request_drop_params_with_well_formed_tools_logs_one_success_row(rig: _Rig, family: _Family) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = _send(rig, family, marker, False, call_id=call_id, tools=[family.tool], drop_params=True) + _assert_answered(family, response, marker) + assert _tool_names(family, _one_delivery(rig, marker, family.target(False))) == ("get_weather",) + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (family.call_type(False), "success"), row + + +def test_repeated_identical_dropped_requests_each_log_once(rig: _Rig) -> None: + marker: Final = _marker() + call_ids: Final = tuple(_call_id() for _ in range(3)) + for call_id in call_ids: + response = _send(rig, GENERATE, marker, False, call_id=call_id, tools=[GENERATE.tool, 5], drop_params=True) + _assert_answered(GENERATE, response, marker) + delivered: Final = _delivered(rig, marker) + assert ( + tuple(_tool_names(GENERATE, _JSON_OBJECT.validate_json(sent.body)) for sent in delivered) + == (("get_weather",),) * 3 + ) + rows: Final = tuple(_logged_once(call_id) for call_id in call_ids) + assert tuple((row["call_type"], row["status"]) for row in rows) == (("agenerate_content", "success"),) * 3 + + +@pytest.mark.parametrize("family", (CHAT, MESSAGES, RESPONSES), ids=("chat", "messages", "responses")) +def test_a_cache_hit_twin_answers_from_the_cache_and_logs_once_each(rig: _Rig, family: _Family) -> None: + marker: Final = _marker() + first_id: Final = _call_id() + second_id: Final = _call_id() + first: Final = _send(rig, family, marker, False, call_id=first_id, tools=[family.tool]) + _assert_answered(family, first, marker) + await_entry(redis_store(), _answer(_tag(marker))) + second: Final = _send(rig, family, marker, False, call_id=second_id, tools=[family.tool]) + _assert_answered(family, second, marker) + assert len(_delivered(rig, marker)) == 1 + first_row: Final = _logged_once(first_id) + second_row: Final = _logged_once(second_id) + assert (first_row["call_type"], second_row["call_type"]) == (family.call_type(False),) * 2, (first_row, second_row) + assert second_row["cache_hit"] == "True", second_row + + +def test_a_dropped_request_shares_the_cache_entry_of_its_clean_twin(rig: _Rig) -> None: + marker: Final = _marker() + dropped: Final = _send(rig, CHAT, marker, False, tools=[CHAT.tool, 5], drop_params=True) + _assert_answered(CHAT, dropped, marker) + await_entry(redis_store(), _answer(_tag(marker))) + clean: Final = _send(rig, CHAT, marker, False, tools=[CHAT.tool], drop_params=True) + _assert_answered(CHAT, clean, marker) + (sent,) = _delivered(rig, marker) + assert _JSON_OBJECT.validate_json(sent.body)["tools"] == [CHAT.tool] + + +def _field_group(rig: _Rig, groups: _FieldGroups, family: _Family) -> str: + return {MESSAGES.name: groups.anthropic, GENERATE.name: groups.gemini}.get(family.name, rig.field_group) + + +_FIELD_GROUP_FAMILIES: Final = pytest.mark.parametrize( + "family", (MESSAGES, RESPONSES, GENERATE), ids=("messages", "responses", "generate_content") +) + + +@_FIELD_GROUP_FAMILIES +def test_each_family_rejects_malformed_tools_when_a_request_list_replaces_the_deployment_list( + rig: _Rig, field_groups: _FieldGroups, family: _Family +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + family, + marker, + False, + model=_field_group(rig, field_groups, family), + tools=[family.tool, 5], + additional_drop_params=["web_search_options"], + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize( + ("family", "forwarded"), + [(MESSAGES, ()), (RESPONSES, ()), (GENERATE, ("get_weather",))], + ids=("messages", "responses", "generate_content"), +) +def test_each_family_answers_malformed_tools_a_deployment_list_drops( + rig: _Rig, field_groups: _FieldGroups, family: _Family, forwarded: tuple[str, ...] +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, family, marker, False, model=_field_group(rig, field_groups, family), tools=[family.tool, 5] + ) + _assert_answered(family, response, marker) + assert _tool_names(family, _one_delivery(rig, marker, family.target(False))) == forwarded + + +@_FIELD_GROUP_FAMILIES +def test_each_family_forwards_well_formed_tools_a_request_list_keeps( + rig: _Rig, field_groups: _FieldGroups, family: _Family +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + family, + marker, + False, + model=_field_group(rig, field_groups, family), + tools=[family.tool], + additional_drop_params=["web_search_options"], + ) + _assert_answered(family, response, marker) + assert _tool_names(family, _one_delivery(rig, marker, family.target(False))) == ("get_weather",) + + +_NOT_A_NAME_LIST: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "empty_list": [], + "null": None, + "string": "tools", + "empty_string": "", + "five_kb_string": "x" * 5120, + "int": 5, + } +) + + +@pytest.mark.parametrize("value", _NOT_A_NAME_LIST.values(), ids=_NOT_A_NAME_LIST.keys()) +def test_a_request_additional_drop_params_that_is_no_list_of_names_still_replaces_the_deployment_list( + rig: _Rig, value: JsonValue +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, CHAT, marker, False, model=rig.field_group, tools=[CHAT.tool, 5], additional_drop_params=value + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize( + "value", + [[], None, "tools", "", "x" * 5120], + ids=["empty_list", "null", "string", "empty_string", "five_kb_string"], +) +def test_a_request_additional_drop_params_replacing_the_deployment_list_keeps_well_formed_tools( + rig: _Rig, value: JsonValue +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, CHAT, marker, False, model=rig.field_group, tools=[CHAT.tool], additional_drop_params=value + ) + _assert_answered(CHAT, response, marker) + assert _one_delivery(rig, marker, "/v1/chat/completions")["tools"] == [CHAT.tool] + + +@pytest.mark.parametrize( + ("first", "last", "status"), + [('["tools"]', '["web_search_options"]', 400), ('["web_search_options"]', '["tools"]', 200)], + ids=["other_field_last", "tools_last"], +) +def test_duplicate_additional_drop_params_keys_use_the_last_one(rig: _Rig, first: str, last: str, status: int) -> None: + marker: Final = _marker() + prefix: Final = json.dumps( + {"model": rig.openai, "messages": [{"role": "user", "content": marker}], "tools": [CHAT.tool, 5]} + )[:-1] + response: Final = rig.gateway.client.post( + "/v1/chat/completions", + content=f'{prefix}, "additional_drop_params": {first}, "additional_drop_params": {last}}}'.encode(), + headers={"Authorization": f"Bearer {rig.gateway.key}", "Content-Type": "application/json"}, + ) + assert response.status_code == status, response.text + if status == 400: + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + return + _assert_answered(CHAT, response, marker) + assert "tools" not in _one_delivery(rig, marker, "/v1/chat/completions") + + +def test_count_tokens_drops_the_bad_entry_a_deployment_list_names(rig: _Rig, field_groups: _FieldGroups) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + {**_count_body(rig, marker, tools=[MESSAGES.tool, 5]), "model": field_groups.anthropic}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT}, response.text + assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool] + + +def test_count_tokens_rejects_malformed_tools_when_a_request_list_replaces_the_deployment_list( + rig: _Rig, field_groups: _FieldGroups +) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + { + **_count_body(rig, marker, tools=[MESSAGES.tool, 5], additional_drop_params=["web_search_options"]), + "model": field_groups.anthropic, + }, + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +def test_responses_input_tokens_drops_the_bad_entry_a_deployment_list_names( + rig: _Rig, field_groups: _FieldGroups +) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": field_groups.anthropic, "input": marker, "tools": [RESPONSES.tool, 5]}, + ) + assert response.status_code == 200, response.text + tools: Final = _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] + assert isinstance(tools, list) and len(tools) == 1 and "get_weather" in json.dumps(tools), tools + + +def test_responses_input_tokens_rejects_malformed_tools_when_a_request_list_replaces_the_deployment_list( + rig: _Rig, field_groups: _FieldGroups +) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/responses/input_tokens", + { + "model": field_groups.anthropic, + "input": marker, + "tools": [RESPONSES.tool, 5], + "additional_drop_params": ["web_search_options"], + }, + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +def _text_items(marker: str) -> list[JsonValue]: + return [{"role": "user", "content": marker}] + + +@pytest.mark.parametrize("compact_input", [_text_items, lambda marker: marker], ids=["list_input", "string_input"]) +def test_a_compact_rejection_logs_one_failure_row_as_compact( + rig: _Rig, compact_input: Callable[[str], JsonValue] +) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = rig.gateway.request( + "POST", + "/v1/responses/compact", + {"model": rig.openai, "input": compact_input(marker), "tools": 5}, + headers={"x-litellm-call-id": call_id}, + ) + _assert_rejected(response, _NOT_A_LIST) + assert _delivered(rig, marker) == () + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == ("acompact_responses", "failure"), row + + +@dataclass(frozen=True, slots=True) +class _Limited: + path: Callable[[_Rig], str] + body: Callable[[_Rig, str], dict[str, JsonValue]] + call_type: str + + +_LIMITED: Final[Mapping[str, _Limited]] = MappingProxyType( + { + "chat": _Limited( + lambda rig: "/v1/chat/completions", lambda rig, k: _chat_body(rig.openai, k, False), "acompletion" + ), + "messages": _Limited( + lambda rig: "/v1/messages", lambda rig, k: _messages_body(rig.anthropic, k, False), "anthropic_messages" + ), + "responses_list_input": _Limited( + lambda rig: "/v1/responses", lambda rig, k: {"model": rig.openai, "input": _text_items(k)}, "aresponses" + ), + "compact_list_input": _Limited( + lambda rig: "/v1/responses/compact", + lambda rig, k: {"model": rig.openai, "input": _text_items(k)}, + "acompact_responses", + ), + "generate_content": _Limited( + lambda rig: _generate_path(rig.gemini, False), + lambda rig, k: _generate_body(rig.gemini, k, False), + "agenerate_content", + ), + "interactions_list_input": _Limited( + lambda rig: "/v1beta/interactions", + lambda rig, k: {"model": rig.gemini, "input": [{"type": "text", "text": k}]}, + "acreate_interaction", + ), + "video_generation": _Limited( + lambda rig: "/v1/videos", lambda rig, k: {"model": rig.openai, "prompt": k}, "avideo_generation" + ), + } +) + + +@pytest.mark.parametrize("limited", _LIMITED.values(), ids=_LIMITED.keys()) +def test_a_rate_limited_request_logs_one_failure_row_with_its_endpoint_call_type( + rig: _Rig, limited_key: str, limited: _Limited +) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = rig.gateway.request( + "POST", + limited.path(rig), + limited.body(rig, marker), + key=limited_key, + headers={"x-litellm-call-id": call_id}, + ) + assert response.status_code == 429, response.text + assert _delivered(rig, marker) == () + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (limited.call_type, "failure"), row + + +@pytest.mark.parametrize("family", (CHAT, MESSAGES), ids=("chat", "messages")) +def test_a_rate_limited_twin_of_a_cached_request_logs_its_endpoint_call_type( + rig: _Rig, limited_key: str, family: _Family +) -> None: + marker: Final = _marker() + _assert_answered(family, _send(rig, family, marker, False, tools=[family.tool]), marker) + await_entry(redis_store(), _answer(_tag(marker))) + call_id: Final = _call_id() + response: Final = _send(rig, family, marker, False, key=limited_key, call_id=call_id, tools=[family.tool]) + assert response.status_code == 429, response.text + assert len(_delivered(rig, marker)) == 1 + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (family.call_type(False), "failure"), row + + +def test_a_rate_limited_image_edit_logs_one_failure_row_as_an_image_edit(rig: _Rig, limited_key: str) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = rig.gateway.client.post( + "/v1/images/edits", + data={"model": rig.openai, "prompt": marker}, + files={"image": ("edit.png", b"\x89PNG\r\n\x1a\n", "image/png")}, + headers={"Authorization": f"Bearer {limited_key}", "x-litellm-call-id": call_id}, + ) + assert response.status_code == 429, response.text + assert _delivered(rig, marker) == () + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == ("aimage_edit", "failure"), row + + +_Sender = Callable[[_Rig, str, str, Mapping[str, JsonValue]], httpx.Response] + + +@dataclass(frozen=True, slots=True) +class _Checked: + provider: str + tool: dict[str, JsonValue] + named_as: _Family + send: _Sender + target: str + answered: Callable[[httpx.Response, str], None] + kept_by_a_tools_list: tuple[str, ...] + bad_entry: dict[str, JsonValue] + + +def _family_sender(family: _Family) -> _Sender: + return lambda rig, model, marker, fields: _send(rig, family, marker, False, model=model, **fields) + + +def _family_answered(family: _Family) -> Callable[[httpx.Response, str], None]: + return lambda response, marker: _assert_answered(family, response, marker) + + +def _count_tokens(rig: _Rig, model: str, marker: str, fields: Mapping[str, JsonValue]) -> httpx.Response: + return rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": [{"role": "user", "content": marker}], **fields}, + ) + + +def _input_tokens(rig: _Rig, model: str, marker: str, fields: Mapping[str, JsonValue]) -> httpx.Response: + return rig.gateway.request("POST", "/v1/responses/input_tokens", {"model": model, "input": marker, **fields}) + + +def _counted(expected: Mapping[str, JsonValue]) -> Callable[[httpx.Response, str], None]: + def check(response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content) == expected, (marker, response.text) + + return check + + +_CHAT_BAD_ENTRY: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "parameters": "bad"}, +} +_MESSAGES_BAD_ENTRY: Final[dict[str, JsonValue]] = {"name": "get_weather", "input_schema": [1]} +_RESPONSES_BAD_ENTRY: Final[dict[str, JsonValue]] = {"type": "function", "name": "get_weather", "parameters": 5} +_CHECKED: Final[Mapping[str, _Checked]] = MappingProxyType( + { + "chat": _Checked( + provider="openai", + tool=CHAT.tool, + named_as=CHAT, + send=_family_sender(CHAT), + target="/v1/chat/completions", + answered=_family_answered(CHAT), + kept_by_a_tools_list=(), + bad_entry=_CHAT_BAD_ENTRY, + ), + "messages": _Checked( + provider="anthropic", + tool=MESSAGES.tool, + named_as=MESSAGES, + send=_family_sender(MESSAGES), + target="/v1/messages", + answered=_family_answered(MESSAGES), + kept_by_a_tools_list=(), + bad_entry=_MESSAGES_BAD_ENTRY, + ), + "responses": _Checked( + provider="openai", + tool=RESPONSES.tool, + named_as=RESPONSES, + send=_family_sender(RESPONSES), + target="/v1/responses", + answered=_family_answered(RESPONSES), + kept_by_a_tools_list=(), + bad_entry=_RESPONSES_BAD_ENTRY, + ), + "generate_content": _Checked( + provider="gemini", + tool=GENERATE.tool, + named_as=GENERATE, + send=_family_sender(GENERATE), + target=":generateContent", + answered=_family_answered(GENERATE), + kept_by_a_tools_list=("get_weather",), + bad_entry={"parameters": 5}, + ), + "count_tokens": _Checked( + provider="anthropic", + tool=MESSAGES.tool, + named_as=MESSAGES, + send=_count_tokens, + target="/v1/messages/count_tokens", + answered=_counted({"input_tokens": _COUNT}), + kept_by_a_tools_list=("get_weather",), + bad_entry=_MESSAGES_BAD_ENTRY, + ), + "input_tokens": _Checked( + provider="anthropic", + tool=RESPONSES.tool, + named_as=MESSAGES, + send=_input_tokens, + target="/v1/messages/count_tokens", + answered=_counted({"object": "response.input_tokens", "input_tokens": _COUNT}), + kept_by_a_tools_list=("get_weather",), + bad_entry=_RESPONSES_BAD_ENTRY, + ), + } +) +_EACH_CHECKED: Final = pytest.mark.parametrize("checked", _CHECKED.values(), ids=_CHECKED.keys()) +_CHAT_AND_COUNT: Final = pytest.mark.parametrize( + "checked", (_CHECKED["chat"], _CHECKED["count_tokens"]), ids=("chat", "count_tokens") +) +_ENTRY_SHAPES: Final = pytest.mark.parametrize( + "bad_entry", + [lambda checked: 5, lambda checked: checked.bad_entry], + ids=["int_entry", "dict_entry"], +) +_BadEntry = Callable[[_Checked], JsonValue] + + +def _assert_forwarded( + rig: _Rig, checked: _Checked, response: httpx.Response, marker: str, names: tuple[str, ...] +) -> None: + checked.answered(response, marker) + assert _tool_names(checked.named_as, _one_delivery(rig, marker, checked.target)) == names + + +def _assert_held_back(rig: _Rig, response: httpx.Response, marker: str) -> None: + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +_PROVIDERS: Final[Mapping[str, Callable[[str], dict[str, JsonValue]]]] = MappingProxyType( + { + "openai": lambda wire: { + "model": f"openai/{_OPENAI_BACKEND}", + "api_base": f"{wire}/v1", + "api_key": _PROVIDER_KEY, + }, + "anthropic": lambda wire: {"model": f"anthropic/{_CLAUDE_BACKEND}", "api_base": wire, "api_key": _PROVIDER_KEY}, + "gemini": lambda wire: {"model": f"gemini/{_GEMINI_BACKEND}", "api_base": wire, "api_key": _PROVIDER_KEY}, + } +) +_DROP_VARIANTS: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + "plain": {}, + "lists_web_search": {"additional_drop_params": ["web_search_options"]}, + "lists_tools": {"additional_drop_params": ["tools"]}, + "keeps": {"drop_params": False}, + } +) + + +@dataclass(frozen=True, slots=True) +class _Owned: + rig: _Rig + suffix: str + + +def _owned_name(suffix: str, variant: str, provider: str) -> str: + return f"malformed-tools-{variant}-{provider}-{suffix}" + + +def _owned_id(suffix: str, variant: str, provider: str) -> str: + return f"{_owned_name(suffix, variant, provider)}-id" + + +@dataclass(frozen=True, slots=True) +class _OwnedSetup: + router_settings: Callable[[str], dict[str, JsonValue]] + general: Mapping[str, JsonValue] = MappingProxyType({}) + extra_models: Callable[[str, str], tuple[dict[str, JsonValue], ...]] = lambda wire, suffix: () + + +def _router_defaults(defaults: dict[str, JsonValue]) -> _OwnedSetup: + return _OwnedSetup(router_settings=lambda suffix: {"default_litellm_params": defaults}) + + +def _owned_config(path: Path, wire: str, suffix: str, setup: _OwnedSetup) -> Path: + path.write_text( + json.dumps( + { + "model_list": [ + *( + { + "model_name": _owned_name(suffix, variant, provider), + "litellm_params": {**deployment(wire), **extra}, + "model_info": {"id": _owned_id(suffix, variant, provider)}, + } + for (variant, extra), (provider, deployment) in product( + _DROP_VARIANTS.items(), _PROVIDERS.items() + ) + ), + *setup.extra_models(wire, suffix), + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + **setup.general, + }, + "router_settings": setup.router_settings(suffix), + } + ) + ) + return path + + +@contextmanager +def _owned_proxy(rig: _Rig, directory: Path, setup: _OwnedSetup) -> Iterator[_Owned]: + suffix: Final = uuid.uuid4().hex[:12] + config: Final = _owned_config(directory / "owned.yaml", rig.wire.url, suffix, setup) + with owned_proxy_process(rig.gateway, directory, {}, config=config, workers=2) as proxy: + yield _Owned(replace(rig, gateway=proxy.gateway), suffix) + + +@pytest.fixture(scope="module") +def tools_default(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-tools-default") + with _owned_proxy(rig, directory, _router_defaults({"additional_drop_params": ["tools"]})) as owned: + yield owned + + +@pytest.fixture(scope="module") +def drop_default(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-drop-default") + with _owned_proxy(rig, directory, _router_defaults({"drop_params": True, "additional_drop_params": None})) as owned: + yield owned + + +@pytest.fixture(scope="module") +def client_side(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-client-side") + setup: Final = replace( + _router_defaults({"additional_drop_params": "tools"}), + general=MappingProxyType({"allow_client_side_credentials": True}), + ) + with _owned_proxy(rig, directory, setup) as owned: + yield owned + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +@_EACH_CHECKED +def test_a_router_default_list_naming_tools_answers_malformed_tools( + tools_default: _Owned, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + model: Final = _owned_name(tools_default.suffix, "plain", checked.provider) + response: Final = checked.send(tools_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(tools_default.rig, checked, response, marker, checked.kept_by_a_tools_list) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_EACH_CHECKED +def test_a_request_list_replaces_a_router_default_list(tools_default: _Owned, checked: _Checked) -> None: + marker: Final = _marker() + model: Final = _owned_name(tools_default.suffix, "plain", checked.provider) + response: Final = checked.send( + tools_default.rig, + model, + marker, + {"tools": [checked.tool, 5], "additional_drop_params": ["web_search_options"]}, + ) + _assert_held_back(tools_default.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +@_EACH_CHECKED +def test_a_router_default_list_replaces_a_deployment_list( + tools_default: _Owned, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + model: Final = _owned_name(tools_default.suffix, "lists_web_search", checked.provider) + response: Final = checked.send(tools_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(tools_default.rig, checked, response, marker, checked.kept_by_a_tools_list) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +def test_a_router_default_list_replaces_the_list_of_a_deployment_named_by_id( + tools_default: _Owned, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["chat"] + model: Final = _owned_id(tools_default.suffix, "lists_web_search", checked.provider) + response: Final = checked.send(tools_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(tools_default.rig, checked, response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize("value", [None, [], 5], ids=["null", "empty_list", "int"]) +def test_a_request_list_naming_nothing_replaces_a_router_default_list(tools_default: _Owned, value: JsonValue) -> None: + marker: Final = _marker() + model: Final = _owned_name(tools_default.suffix, "plain", "openai") + response: Final = _send( + tools_default.rig, CHAT, marker, False, model=model, tools=[CHAT.tool, 5], additional_drop_params=value + ) + _assert_held_back(tools_default.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_a_router_default_list_naming_tools_drops_well_formed_tools(tools_default: _Owned) -> None: + marker: Final = _marker() + model: Final = _owned_name(tools_default.suffix, "plain", "openai") + response: Final = _send(tools_default.rig, CHAT, marker, False, model=model, tools=[CHAT.tool]) + _assert_forwarded(tools_default.rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +@_EACH_CHECKED +def test_a_router_default_drop_params_keeps_the_good_entries( + drop_default: _Owned, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + model: Final = _owned_name(drop_default.suffix, "plain", checked.provider) + response: Final = checked.send(drop_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(drop_default.rig, checked, response, marker, ("get_weather",)) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize("drop", [False, None], ids=["false", "null"]) +def test_a_request_drop_params_replaces_a_router_default_drop_params(drop_default: _Owned, drop: JsonValue) -> None: + marker: Final = _marker() + model: Final = _owned_name(drop_default.suffix, "plain", "openai") + response: Final = _send(drop_default.rig, CHAT, marker, False, model=model, tools=[CHAT.tool, 5], drop_params=drop) + _assert_held_back(drop_default.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +@_CHAT_AND_COUNT +def test_a_router_default_drop_params_replaces_a_deployment_drop_params( + drop_default: _Owned, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + model: Final = _owned_name(drop_default.suffix, "keeps", checked.provider) + response: Final = checked.send(drop_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(drop_default.rig, checked, response, marker, ("get_weather",)) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +@_CHAT_AND_COUNT +def test_a_null_router_default_list_leaves_the_deployment_list_in_force( + drop_default: _Owned, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + model: Final = _owned_name(drop_default.suffix, "lists_tools", checked.provider) + response: Final = checked.send( + drop_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)], "drop_params": False} + ) + _assert_forwarded(drop_default.rig, checked, response, marker, checked.kept_by_a_tools_list) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_a_router_default_list_that_is_a_string_still_rejects_a_bad_tool_entry(client_side: _Owned) -> None: + marker: Final = _marker() + model: Final = _owned_name(client_side.suffix, "plain", "openai") + response: Final = _send(client_side.rig, CHAT, marker, False, model=model, tools=[CHAT.tool, 5]) + _assert_held_back(client_side.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_a_router_default_list_that_is_a_string_keeps_well_formed_tools(client_side: _Owned) -> None: + marker: Final = _marker() + model: Final = _owned_name(client_side.suffix, "plain", "openai") + response: Final = _send(client_side.rig, CHAT, marker, False, model=model, tools=[CHAT.tool]) + _assert_answered(CHAT, response, marker) + assert _one_delivery(client_side.rig, marker, "/v1/chat/completions")["tools"] == [CHAT.tool] + + +_USER_CONFIG_DROPS: Final[Mapping[str, tuple[dict[str, JsonValue], dict[str, JsonValue], tuple[str, ...]]]] = ( + MappingProxyType( + { + "deployment_list": ({"additional_drop_params": ["tools"]}, {}, ()), + "router_default_list": ({}, {"additional_drop_params": ["tools"]}, ()), + "router_default_drop_params": ({}, {"drop_params": True}, ("get_weather",)), + } + ) +) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize(("deployment", "defaults", "kept"), _USER_CONFIG_DROPS.values(), ids=_USER_CONFIG_DROPS.keys()) +def test_a_user_config_router_decides_which_malformed_tools_drop( + client_side: _Owned, + deployment: dict[str, JsonValue], + defaults: dict[str, JsonValue], + kept: tuple[str, ...], +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["chat"] + name: Final = f"malformed-tools-user-config-{uuid.uuid4().hex}" + user_config: Final[dict[str, JsonValue]] = { + "model_list": [ + {"model_name": name, "litellm_params": {**_PROVIDERS["openai"](client_side.rig.wire.url), **deployment}} + ], + "default_litellm_params": defaults, + } + response: Final = checked.send( + client_side.rig, name, marker, {"tools": [checked.tool, checked.bad_entry], "user_config": user_config} + ) + _assert_forwarded(client_side.rig, checked, response, marker, kept) + + +@dataclass(frozen=True, slots=True) +class _ById: + lists_tools: Mapping[str, str] + drops: str + mixed_drops: str + mixed_keeps: str + shadowed: str + + +@pytest.fixture(scope="module") +def by_id(rig: _Rig) -> Iterator[_ById]: + openai_deployment: Final = _PROVIDERS["openai"](rig.wire.url) + lists_tools: Final[dict[str, JsonValue]] = {"additional_drop_params": ["tools"]} + names: Final = MappingProxyType({provider: f"integration-by-id-{uuid.uuid4().hex}" for provider in _PROVIDERS}) + drops_name: Final = f"integration-by-id-drops-{uuid.uuid4().hex}" + mixed: Final = f"integration-by-id-mixed-{uuid.uuid4().hex}" + shadow_name: Final = f"integration-by-id-shadow-{uuid.uuid4().hex}" + with rig.gateway.scenario() as scenario: + listed: Final = MappingProxyType( + { + provider: _deployment_ids(scenario, names[provider], ({**deployment(rig.wire.url), **lists_tools},))[0] + for provider, deployment in _PROVIDERS.items() + } + ) + (drops,) = _deployment_ids(scenario, drops_name, ({**openai_deployment, "drop_params": True},)) + mixed_drops, mixed_keeps = _deployment_ids( + scenario, mixed, ({**openai_deployment, **lists_tools}, openai_deployment) + ) + (shadowed,) = _deployment_ids(scenario, shadow_name, ({**openai_deployment, **lists_tools},)) + _deployment_ids(scenario, shadowed, (openai_deployment,)) + for group, deployments in ( + *((name, 1) for name in names.values()), + (drops_name, 1), + (mixed, 2), + (shadow_name, 1), + (shadowed, 1), + ): + wait_for_worker_readiness(rig.gateway, group, deployments, _PROXY_WORKERS) + yield _ById(listed, drops, mixed_drops, mixed_keeps, shadowed) + + +@_ENTRY_SHAPES +@_EACH_CHECKED +def test_a_deployment_id_answers_malformed_tools_its_deployment_list_drops( + rig: _Rig, by_id: _ById, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + response: Final = checked.send( + rig, by_id.lists_tools[checked.provider], marker, {"tools": [checked.tool, bad_entry(checked)]} + ) + _assert_forwarded(rig, checked, response, marker, checked.kept_by_a_tools_list) + + +@_EACH_CHECKED +def test_a_request_list_replaces_the_list_of_a_deployment_named_by_id( + rig: _Rig, by_id: _ById, checked: _Checked +) -> None: + marker: Final = _marker() + response: Final = checked.send( + rig, + by_id.lists_tools[checked.provider], + marker, + {"tools": [checked.tool, 5], "additional_drop_params": ["web_search_options"]}, + ) + _assert_held_back(rig, response, marker) + + +@_ENTRY_SHAPES +def test_the_id_of_the_dropping_deployment_in_a_mixed_group_drops_malformed_tools( + rig: _Rig, by_id: _ById, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["chat"] + response: Final = checked.send(rig, by_id.mixed_drops, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(rig, checked, response, marker, ()) + + +def test_the_keeping_id_of_a_mixed_group_rejects_malformed_tools(rig: _Rig, by_id: _ById) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, model=by_id.mixed_keeps, tools=[CHAT.tool, 5]) + _assert_held_back(rig, response, marker) + + +@_ENTRY_SHAPES +def test_the_id_of_a_drop_params_deployment_keeps_the_good_entries( + rig: _Rig, by_id: _ById, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["chat"] + response: Final = checked.send(rig, by_id.drops, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(rig, checked, response, marker, ("get_weather",)) + + +def test_a_model_name_wins_over_an_equal_deployment_id(rig: _Rig, by_id: _ById) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, model=by_id.shadowed, tools=[CHAT.tool, 5]) + _assert_held_back(rig, response, marker) + + +def test_a_deployment_id_with_a_tools_list_drops_well_formed_tools(rig: _Rig, by_id: _ById) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, model=by_id.lists_tools["openai"], tools=[CHAT.tool]) + _assert_forwarded(rig, _CHECKED["chat"], response, marker, ()) + + +def _failed_over(request: Request) -> bool: + return request.target.startswith((_DOWN, *(prefix for prefix, _ in _REFUSALS.values()))) + + +def _assert_failed_over( + rig: _Rig, checked: _Checked, response: httpx.Response, marker: str, names: tuple[str, ...] +) -> None: + checked.answered(response, marker) + delivered: Final = _delivered(rig, marker) + (served,) = tuple(request for request in delivered if not _failed_over(request)) + assert checked.target in served.target, served.target + assert _tool_names(checked.named_as, _JSON_OBJECT.validate_json(served.body)) == names + + +_FALLBACK_CHECKED: Final = pytest.mark.parametrize( + "checked", + tuple(_CHECKED[name] for name in ("chat", "messages", "responses")), + ids=("chat", "messages", "responses"), +) +_TOKEN_COUNTS: Final = pytest.mark.parametrize( + "checked", (_CHECKED["count_tokens"], _CHECKED["input_tokens"]), ids=("count_tokens", "input_tokens") +) + + +@dataclass(frozen=True, slots=True) +class _Reroute: + down: Mapping[str, str] + lists_tools: Mapping[str, str] + refusing: Mapping[str, str] + fallback_keys: Mapping[str, str] + team_key: str + team_models: Mapping[str, str] + + +@pytest.fixture(scope="module") +def reroute(rig: _Rig, field_groups: _FieldGroups) -> Iterator[_Reroute]: + lists_tools: Final = MappingProxyType({"openai": rig.field_group, "anthropic": field_groups.anthropic}) + with rig.gateway.scenario() as scenario: + down: Final = MappingProxyType( + {provider: scenario.model(**_PROVIDERS[provider](f"{rig.wire.url}/down")) for provider in lists_tools} + ) + refusing: Final = MappingProxyType( + { + chain: scenario.model(**_PROVIDERS["openai"](f"{rig.wire.url}{prefix.rstrip('/')}")) + for chain, (prefix, _) in _REFUSALS.items() + } + ) + router_settings: Final[dict[str, JsonValue]] = {"fallbacks": [{down["openai"]: [rig.field_group]}]} + public_team: Final = scenario.team() + team_models: Final = MappingProxyType( + { + provider: scenario.model( + **_PROVIDERS[provider](rig.wire.url), + additional_drop_params=["tools"], + model_info={"team_id": public_team}, + ) + for provider in lists_tools + } + ) + built: Final = _Reroute( + down=down, + lists_tools=lists_tools, + refusing=refusing, + fallback_keys=MappingProxyType( + { + "key": scenario.key(router_settings=router_settings), + "team": scenario.key(team_id=scenario.team(router_settings=router_settings)), + } + ), + team_key=scenario.key(team_id=public_team), + team_models=team_models, + ) + for group in (*down.values(), *refusing.values(), *team_models.values()): + wait_for_worker_readiness(rig.gateway, group, 1, _PROXY_WORKERS) + yield built + + +@_FALLBACK_CHECKED +def test_a_request_fallback_to_a_group_listing_tools_answers_a_bad_tool_entry( + rig: _Rig, reroute: _Reroute, checked: _Checked +) -> None: + marker: Final = _marker() + response: Final = checked.send( + rig, + reroute.down[checked.provider], + marker, + {"tools": [checked.tool, checked.bad_entry], "fallbacks": [reroute.lists_tools[checked.provider]]}, + ) + _assert_failed_over(rig, checked, response, marker, ()) + + +def test_a_generic_request_fallback_to_a_group_listing_tools_answers_a_bad_tool_entry( + rig: _Rig, reroute: _Reroute +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["chat"] + response: Final = checked.send( + rig, + reroute.down["openai"], + marker, + {"tools": [checked.tool, checked.bad_entry], "fallbacks": [{"*": [rig.field_group]}]}, + ) + _assert_failed_over(rig, checked, response, marker, ()) + + +@pytest.mark.parametrize( + "fallbacks", + [[], None, lambda rig: [{rig.drop_group: [rig.field_group]}]], + ids=["empty_list", "null", "chain_of_another_group"], +) +def test_a_request_fallback_that_cannot_apply_still_rejects_a_bad_tool_entry( + rig: _Rig, fallbacks: JsonValue | Callable[[_Rig], JsonValue] +) -> None: + marker: Final = _marker() + chains: Final = fallbacks(rig) if callable(fallbacks) else fallbacks + response: Final = _send(rig, CHAT, marker, False, tools=[CHAT.tool, 5], fallbacks=chains) + _assert_held_back(rig, response, marker) + + +@pytest.mark.parametrize("owner", ["key", "team"]) +def test_router_settings_fallbacks_of_a_key_or_team_answer_a_bad_tool_entry( + rig: _Rig, reroute: _Reroute, owner: str +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + model=reroute.down["openai"], + key=reroute.fallback_keys[owner], + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + ) + _assert_failed_over(rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.parametrize("chain", tuple(_REFUSALS)) +def test_an_error_specific_fallback_never_sees_a_bad_tool_entry_the_first_deployment_rejects( + rig: _Rig, reroute: _Reroute, chain: str +) -> None: + marker: Final = _marker() + refusing: Final = reroute.refusing[chain] + response: Final = _send( + rig, + CHAT, + marker, + False, + model=refusing, + tools=[CHAT.tool, 5], + **{chain: [{refusing: [rig.field_group]}]}, + ) + _assert_held_back(rig, response, marker) + + +def test_a_comma_batch_lets_each_group_decide_on_a_bad_tool_entry(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send( + rig, CHAT, marker, False, model=f"{rig.field_group},{rig.openai}", tools=[CHAT.tool, _CHAT_BAD_ENTRY] + ) + assert response.status_code == 200, response.text + assert _answer(_tag(marker)) in response.text, response.text + delivered: Final = sorted( + (_JSON_OBJECT.validate_json(request.body).get("tools") for request in _delivered(rig, marker)), key=bool + ) + assert delivered == [None] + + +def test_a_team_public_model_name_answers_a_bad_tool_entry_its_deployment_list_drops( + rig: _Rig, reroute: _Reroute +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + model=reroute.team_models["openai"], + key=reroute.team_key, + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + ) + _assert_forwarded(rig, _CHECKED["chat"], response, marker, ()) + + +def test_count_tokens_by_a_team_public_model_name_rejects_a_bad_tool_entry(rig: _Rig, reroute: _Reroute) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + { + "model": reroute.team_models["anthropic"], + "messages": [{"role": "user", "content": marker}], + "tools": [MESSAGES.tool, 5], + }, + key=reroute.team_key, + ) + _assert_held_back(rig, response, marker) + + +def test_a_request_list_naming_web_search_options_on_a_fallback_chain_still_rejects_a_bad_tool_entry( + rig: _Rig, +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + tools=[CHAT.tool, 5], + web_search_options=5, + additional_drop_params=["web_search_options"], + fallbacks=[{rig.openai: [rig.field_group]}], + ) + _assert_held_back(rig, response, marker) + + +def test_request_drop_params_on_a_fallback_chain_keeps_the_good_entries(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + drop_params=True, + fallbacks=[{rig.openai: [rig.field_group]}], + ) + _assert_forwarded(rig, _CHECKED["chat"], response, marker, ("get_weather",)) + + +@_ENTRY_SHAPES +def test_a_request_list_with_a_non_string_entry_drops_tools_it_names_on_messages( + rig: _Rig, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["messages"] + response: Final = checked.send( + rig, + rig.anthropic, + marker, + {"tools": [checked.tool, bad_entry(checked)], "additional_drop_params": ["tools", 5]}, + ) + _assert_forwarded(rig, checked, response, marker, ()) + + +@_ENTRY_SHAPES +def test_a_deployment_list_with_a_non_string_entry_drops_tools_it_names_on_messages( + rig: _Rig, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["messages"] + with rig.gateway.scenario() as scenario: + group: Final = scenario.model(**_PROVIDERS["anthropic"](rig.wire.url), additional_drop_params=["tools", 5]) + wait_for_worker_readiness(rig.gateway, group, 1, _PROXY_WORKERS) + response: Final = checked.send(rig, group, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(rig, checked, response, marker, ()) + + +@_TOKEN_COUNTS +def test_token_counting_with_a_request_list_with_a_non_string_entry_drops_the_bad_entry( + rig: _Rig, checked: _Checked +) -> None: + marker: Final = _marker() + response: Final = checked.send( + rig, rig.anthropic, marker, {"tools": [checked.tool, 5], "additional_drop_params": ["tools", 5]} + ) + _assert_forwarded(rig, checked, response, marker, ("get_weather",)) + + +def _alias(suffix: str, variant: str) -> str: + return f"malformed-tools-alias-{variant}-{suffix}" + + +def _routing_group(suffix: str, provider: str) -> str: + return f"malformed-tools-routed-{provider}-{suffix}" + + +def _wildcard(suffix: str) -> str: + return f"malformed-tools-wild-{suffix}" + + +def _direct(suffix: str) -> str: + return f"malformed-tools-direct-{suffix}" + + +def _down_models(wire: str, suffix: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + {"model_name": _owned_name(suffix, "down", provider), "litellm_params": _PROVIDERS[provider](f"{wire}/down")} + for provider in ("openai", "anthropic") + ) + + +def _rerouting_models(wire: str, suffix: str) -> tuple[dict[str, JsonValue], ...]: + lists_tools: Final = _owned_name(suffix, "lists_tools", "openai") + return ( + *_down_models(wire, suffix), + { + "model_name": _owned_name(suffix, "auto", "openai"), + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": lists_tools, + "complexity_router_config": { + "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), lists_tools) + }, + }, + }, + { + "model_name": _owned_name(suffix, "direct", "openai"), + "litellm_params": { + **_PROVIDERS["openai"](wire), + "model": f"openai/{_direct(suffix)}", + "additional_drop_params": ["tools"], + }, + }, + { + "model_name": f"{_wildcard(suffix)}/*", + "litellm_params": {**_PROVIDERS["openai"](wire), "model": "openai/*", "additional_drop_params": ["tools"]}, + }, + ) + + +def _rerouting_settings(suffix: str) -> dict[str, JsonValue]: + return { + "fallbacks": [ + {_owned_name(suffix, "down", provider): [_owned_name(suffix, "lists_tools", provider)]} + for provider in ("openai", "anthropic") + ], + "model_group_alias": { + _alias(suffix, variant): _owned_name(suffix, variant, "openai") for variant in ("lists_tools", "plain") + }, + "routing_groups": [ + { + "group_name": _routing_group(suffix, provider), + "models": [_owned_name(suffix, "lists_tools", provider)], + "routing_strategy": "simple-shuffle", + } + for provider in ("openai", "anthropic") + ], + "num_retries": 0, + "disable_cooldowns": True, + } + + +@pytest.fixture(scope="module") +def rerouting(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-rerouting") + setup: Final = _OwnedSetup(router_settings=_rerouting_settings, extra_models=_rerouting_models) + with _owned_proxy(rig, directory, setup) as owned: + yield owned + + +@pytest.fixture(scope="module") +def default_fallbacks(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-default-fallbacks") + setup: Final = _OwnedSetup( + router_settings=lambda suffix: { + "default_fallbacks": [_owned_name(suffix, "lists_tools", "openai")], + "num_retries": 0, + "disable_cooldowns": True, + }, + extra_models=_down_models, + ) + with _owned_proxy(rig, directory, setup) as owned: + yield owned + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize("checked", (_CHECKED["chat"], _CHECKED["messages"]), ids=("chat", "messages")) +def test_a_router_fallback_to_a_group_listing_tools_answers_a_bad_tool_entry( + rerouting: _Owned, checked: _Checked +) -> None: + marker: Final = _marker() + model: Final = _owned_name(rerouting.suffix, "down", checked.provider) + response: Final = checked.send(rerouting.rig, model, marker, {"tools": [checked.tool, checked.bad_entry]}) + _assert_failed_over(rerouting.rig, checked, response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize( + ("variant", "fields"), + [("plain", {}), ("down", {"fallbacks": []})], + ids=["group_outside_the_chain", "chain_cleared_by_the_request"], +) +def test_a_router_fallback_that_cannot_apply_still_rejects_a_bad_tool_entry( + rerouting: _Owned, variant: str, fields: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + model: Final = _owned_name(rerouting.suffix, variant, "openai") + response: Final = _send(rerouting.rig, CHAT, marker, False, model=model, tools=[CHAT.tool, 5], **fields) + _assert_held_back(rerouting.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_TOKEN_COUNTS +def test_token_counting_on_a_group_with_a_router_fallback_rejects_a_bad_tool_entry( + rerouting: _Owned, checked: _Checked +) -> None: + marker: Final = _marker() + model: Final = _owned_name(rerouting.suffix, "down", "anthropic") + response: Final = checked.send(rerouting.rig, model, marker, {"tools": [checked.tool, 5]}) + _assert_held_back(rerouting.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_an_alias_of_a_group_listing_tools_answers_a_bad_tool_entry(rerouting: _Owned) -> None: + marker: Final = _marker() + response: Final = _send( + rerouting.rig, + CHAT, + marker, + False, + model=_alias(rerouting.suffix, "lists_tools"), + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + ) + _assert_forwarded(rerouting.rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_an_alias_of_a_plain_group_rejects_a_bad_tool_entry(rerouting: _Owned) -> None: + marker: Final = _marker() + response: Final = _send( + rerouting.rig, CHAT, marker, False, model=_alias(rerouting.suffix, "plain"), tools=[CHAT.tool, 5] + ) + _assert_held_back(rerouting.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_a_routing_group_over_a_group_listing_tools_answers_a_bad_tool_entry(rerouting: _Owned) -> None: + marker: Final = _marker() + response: Final = _send( + rerouting.rig, + CHAT, + marker, + False, + model=_routing_group(rerouting.suffix, "openai"), + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + ) + _assert_forwarded(rerouting.rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_count_tokens_by_a_routing_group_rejects_a_bad_tool_entry(rerouting: _Owned) -> None: + marker: Final = _marker() + response: Final = _count_tokens( + rerouting.rig, _routing_group(rerouting.suffix, "anthropic"), marker, {"tools": [MESSAGES.tool, 5]} + ) + _assert_held_back(rerouting.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize( + ("model", "fields"), + [ + (lambda suffix: _owned_name(suffix, "auto", "openai"), {}), + (lambda suffix: f"{_wildcard(suffix)}/{_OPENAI_BACKEND}", {}), + (lambda suffix: f"openai/{_direct(suffix)}", {}), + (lambda suffix: f"openai/{_direct(suffix)}", {"specific_deployment": True}), + ], + ids=["auto_router", "wildcard_route", "model_string", "specific_deployment"], +) +def test_a_model_the_router_resolves_later_answers_a_bad_tool_entry_its_deployment_list_drops( + rerouting: _Owned, model: Callable[[str], str], fields: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + response: Final = _send( + rerouting.rig, + CHAT, + marker, + False, + model=model(rerouting.suffix), + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + **fields, + ) + _assert_forwarded(rerouting.rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_default_fallbacks_to_a_group_listing_tools_answer_a_bad_tool_entry(default_fallbacks: _Owned) -> None: + marker: Final = _marker() + response: Final = _send( + default_fallbacks.rig, + CHAT, + marker, + False, + model=_owned_name(default_fallbacks.suffix, "down", "openai"), + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + ) + _assert_failed_over(default_fallbacks.rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_a_user_config_model_named_like_a_proxy_group_drops_tools_its_own_list_names(client_side: _Owned) -> None: + marker: Final = _marker() + name: Final = _owned_name(client_side.suffix, "plain", "openai") + user_config: Final[dict[str, JsonValue]] = { + "model_list": [ + { + "model_name": name, + "litellm_params": { + **_PROVIDERS["openai"](client_side.rig.wire.url), + "additional_drop_params": ["tools"], + }, + } + ], + } + response: Final = _send( + client_side.rig, + CHAT, + marker, + False, + model=name, + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + user_config=user_config, + ) + _assert_forwarded(client_side.rig, _CHECKED["chat"], response, marker, ()) + + +_GUARDED_NAMES: Final = ("chat", "messages", "responses", "generate_content") +_EACH_GUARDED: Final = pytest.mark.parametrize( + "checked", tuple(_CHECKED[name] for name in _GUARDED_NAMES), ids=_GUARDED_NAMES +) +_GUARDRAIL_VENDOR_PATH: Final = "/beta/litellm_basic_guardrail_api" + + +@dataclass(frozen=True, slots=True) +class _Guarded: + rig: _Rig + vendor: Wire + models: Mapping[str, str] + + +def _guarded_config(path: Path, wire: str, vendor: str, models: Mapping[str, str]) -> Path: + path.write_text( + json.dumps( + { + "model_list": [ + {"model_name": name, "litellm_params": _PROVIDERS[provider](wire)} + for provider, name in models.items() + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + "guardrails": [ + { + "guardrail_name": "malformed-tools-vendor", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": True, + "api_base": vendor, + "api_key": "synthetic-guardrail-key", + }, + } + ], + } + ) + ) + return path + + +def _guardrail_verdict(request: Request) -> Reply: + assert request.target == _GUARDRAIL_VENDOR_PATH, request.target + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +@pytest.fixture(scope="module") +def guarded(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Guarded]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-guarded") + suffix: Final = uuid.uuid4().hex[:12] + models: Final = MappingProxyType( + {provider: f"malformed-tools-guarded-{provider}-{suffix}" for provider in _PROVIDERS} + ) + with wire_server(_guardrail_verdict) as vendor: + config: Final = _guarded_config(directory / "guarded.yaml", rig.wire.url, vendor.url, models) + with owned_proxy_process(rig.gateway, directory, {}, config=config, workers=2) as proxy: + yield _Guarded(replace(rig, gateway=proxy.gateway), vendor, models) + + +def _vendor_payloads(guarded: _Guarded, marker: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + _JSON_OBJECT.validate_json(request.body) + for request in guarded.vendor.drain() + if marker in request.body.decode() + ) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_EACH_GUARDED +def test_a_pre_call_guardrail_passes_well_formed_tools_through_to_the_provider( + guarded: _Guarded, checked: _Checked +) -> None: + marker: Final = _marker() + response: Final = checked.send(guarded.rig, guarded.models[checked.provider], marker, {"tools": [checked.tool]}) + _assert_forwarded(guarded.rig, checked, response, marker, ("get_weather",)) + payloads: Final = _vendor_payloads(guarded, marker) + assert len(payloads) == 1, payloads + assert all(payload.get("tools") is None or isinstance(payload["tools"], list) for payload in payloads), payloads + + +_PROVIDER_WIRES: Final[Mapping[str, tuple[_Family, Callable[[bool], str]]]] = MappingProxyType( + { + "openai": (CHAT, lambda stream: "/v1/chat/completions"), + "openai_responses": (RESPONSES, lambda stream: "/v1/responses"), + "anthropic": (MESSAGES, lambda stream: "/v1/messages"), + "gemini": (GENERATE, lambda stream: ":streamGenerateContent" if stream else ":generateContent"), + } +) +_NATIVE_PROVIDERS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( + { + "chat": ("openai",), + "messages": ("anthropic",), + "responses": ("openai", "openai_responses"), + "generate_content": ("gemini",), + } +) +_ROUTED_ELSEWHERE: Final[Mapping[tuple[str, str], str]] = MappingProxyType({("messages", "openai"): "openai_responses"}) +_BRIDGES: Final = tuple( + (family, provider) + for family, provider in product(_FAMILIES, _PROVIDER_WIRES) + if provider not in _NATIVE_PROVIDERS[family.name] +) + + +@pytest.fixture(scope="module") +def bridge_groups(rig: _Rig) -> Iterator[Mapping[str, str]]: + with rig.gateway.scenario() as scenario: + responses_group: Final = scenario.model( + model=f"openai/responses/{_OPENAI_BACKEND}", api_base=f"{rig.wire.url}/v1", api_key=_PROVIDER_KEY + ) + wait_for_worker_readiness(rig.gateway, responses_group, 1, _PROXY_WORKERS) + yield MappingProxyType( + { + "openai": rig.openai, + "openai_responses": responses_group, + "anthropic": rig.anthropic, + "gemini": rig.gemini, + } + ) + + +@_STREAMS +@pytest.mark.parametrize( + ("family", "provider"), _BRIDGES, ids=tuple(f"{family.name}_to_{provider}" for family, provider in _BRIDGES) +) +def test_well_formed_tools_cross_the_bridge_to_another_provider( + rig: _Rig, bridge_groups: Mapping[str, str], family: _Family, provider: str, stream: bool +) -> None: + marker: Final = _marker() + response: Final = _send(rig, family, marker, stream, model=bridge_groups[provider], tools=[family.tool]) + assert response.status_code == 200, response.text + assert _answer(_tag(marker)) in response.text, response.text + named_as, target = _PROVIDER_WIRES[_ROUTED_ELSEWHERE.get((family.name, provider), provider)] + assert _tool_names(named_as, _one_delivery(rig, marker, target(stream))) == ("get_weather",) diff --git a/tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py b/tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py index fedfc2b0848..8eefee1dfcf 100644 --- a/tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py +++ b/tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py @@ -9,12 +9,14 @@ response. A failure raised after the router picked a deployment (a post-call guardrail block, a provider error) carries no flag and keeps its deployment labels. """ +from typing import Final + import pytest from fastapi import HTTPException from prometheus_client import REGISTRY from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY -from litellm.exceptions import GuardrailRaisedException +from litellm.exceptions import GuardrailRaisedException, MalformedRequestFieldError from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -99,3 +101,15 @@ def test_failure_after_routing_still_marks_deployment_partial_outage(failure: Ex assert _model_id_values(logger.litellm_deployment_state) == {"dep-1"} assert _model_id_values(logger.litellm_deployment_failure_responses) == {"dep-1"} + + +def test_malformed_field_rejection_after_routing_leaves_deployment_healthy() -> None: + logger: Final = PrometheusLogger() + rejection: Final = MalformedRequestFieldError( + message="tools: Input should be a valid list", model="internal-model", llm_provider="openai" + ) + + logger.set_llm_deployment_failure_metrics(_attributed_failure_kwargs(rejection, rejected_before_routing=False)) + + assert logger.litellm_deployment_state._metrics == {} + assert _model_id_values(logger.litellm_deployment_failure_responses) == {""} diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index c8dbbd00d66..7291c4858f6 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,4 +1,6 @@ +from typing import Final + import pytest import litellm @@ -49,11 +51,12 @@ def test_web_search_cost_high(): ) -def test_get_web_search_options_preserves_explicit_none_error(): - with pytest.raises(TypeError): - StandardBuiltInToolCostTracking.get_web_search_options( - {"web_search_options": None, "tools": [{"type": "web_search_preview"}]} - ) +def test_get_web_search_options_reads_tools_when_options_are_null() -> None: + options: Final = StandardBuiltInToolCostTracking.get_web_search_options( + {"web_search_options": None, "tools": [{"type": "web_search_preview"}]} + ) + + assert options == {"type": "web_search_preview"} def test_get_web_search_options_accepts_non_list_tool_iterables(): @@ -1070,3 +1073,13 @@ def test_azure_ai_non_claude_web_search_keeps_generic_fallback() -> None: usage = Usage(prompt_tokens=1840, completion_tokens=412, total_tokens=2252) assert get_cost_for_web_search_request(custom_llm_provider="azure_ai", usage=usage, model_info=model_info) is None + + +@pytest.mark.parametrize( + "kwargs", + [{"tools": 5}, {"web_search_options": None}, {"web_search_options": 5}], + ids=["int_tools", "null_web_search_options", "int_web_search_options"], +) +def test_built_in_tool_lookups_ignore_malformed_request_values(kwargs: dict[str, object]) -> None: + assert StandardBuiltInToolCostTracking.get_web_search_options(kwargs) is None + assert StandardBuiltInToolCostTracking.get_file_search_tool_call(kwargs) is None diff --git a/tests/unit/litellm_core_utils/test_tool_fields.py b/tests/unit/litellm_core_utils/test_tool_fields.py new file mode 100644 index 00000000000..23d7de98fbc --- /dev/null +++ b/tests/unit/litellm_core_utils/test_tool_fields.py @@ -0,0 +1,187 @@ +import datetime +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm import Router +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.utils import async_pre_call_deployment_hook + +_GOOD_TOOL: Final = {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}} +_MALFORMED_TOOLS: Final = ( + 5, + "get_weather", + ["get_weather"], + {"type": "function"}, + [{"type": "function", "function": {"name": "get_weather", "parameters": "object"}}], +) + + +def _deployment(name: str, **params: object) -> Mapping[str, object]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-5.2", "api_key": "sk-test", "mock_response": f"served by {name}", **params}, + } + + +def _proxy_request(model: str, route: str = "acompletion", **fields: object) -> Mapping[str, object]: + return { + "model": model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "proxy_server_request": {"url": "/v1/chat/completions", "body": {}}, + "litellm_logging_obj": Logging( + model=model, + messages=[], + stream=False, + call_type=route, + start_time=datetime.datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ), + "num_retries": 0, + **fields, + } + + +@pytest.mark.parametrize("tools", _MALFORMED_TOOLS) +@pytest.mark.asyncio +async def test_malformed_tools_fall_back_to_a_deployment_that_drops_them(tools: object) -> None: + router: Final = Router( + model_list=[_deployment("primary"), _deployment("backup", additional_drop_params=["tools"])], + fallbacks=[{"primary": ["backup"]}], + ) + + response: Final = await router.acompletion(**_proxy_request("primary", tools=tools)) + + assert response.choices[0].message.content == "served by backup" + + +@pytest.mark.asyncio +async def test_fallback_on_unsupported_params_false_keeps_the_400() -> None: + router: Final = Router( + model_list=[_deployment("primary"), _deployment("backup", additional_drop_params=["tools"])], + fallbacks=[{"primary": ["backup"]}], + fallback_on_unsupported_params=False, + ) + + with pytest.raises(litellm.BadRequestError): + await router.acompletion(**_proxy_request("primary", tools=5)) + + +@pytest.mark.asyncio +async def test_a_rejected_request_cools_down_no_deployment() -> None: + cools_on_any_400: Final = {"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}, "cooldown_time": 120} + router: Final = Router( + model_list=[ + {**_deployment("primary"), "model_info": {"id": "primary", **cools_on_any_400}}, + {**_deployment("backup"), "model_info": {"id": "backup", **cools_on_any_400}}, + ], + fallbacks=[{"primary": ["backup"]}], + ) + + with pytest.raises(litellm.BadRequestError): + await router.acompletion(**_proxy_request("primary", tools=5)) + + primary: Final = await router.acompletion(**_proxy_request("primary")) + backup: Final = await router.acompletion(**_proxy_request("backup")) + assert (primary.choices[0].message.content, backup.choices[0].message.content) == ( + "served by primary", + "served by backup", + ) + + +@pytest.mark.parametrize("tools", _MALFORMED_TOOLS) +@pytest.mark.parametrize( + "deployment", + [_deployment("keeps"), _deployment("keeps", tools=[_GOOD_TOOL])], + ids=["plain", "with_configured_tools"], +) +@pytest.mark.asyncio +async def test_malformed_tools_answer_400_when_the_served_deployment_keeps_tools( + tools: object, deployment: Mapping[str, object] +) -> None: + router: Final = Router(model_list=[deployment]) + + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion(**_proxy_request("keeps", tools=tools)) + + assert (raised.value.status_code, raised.value.param, raised.value.type) == ( + 400, + "tools", + "invalid_request_error", + ) + + +@pytest.mark.asyncio +async def test_a_request_drop_params_false_overrides_a_deployment_that_drops() -> None: + router: Final = Router(model_list=[_deployment("drops", drop_params=True)]) + + with pytest.raises(litellm.BadRequestError): + await router.acompletion(**_proxy_request("drops", tools=5, drop_params=False)) + + +@pytest.mark.parametrize( + "route", ["acompletion", "aresponses", "anthropic_messages", "agenerate_content", "acreate_interaction"] +) +@pytest.mark.asyncio +async def test_drop_params_drops_every_malformed_tool_field(route: str) -> None: + request: Final = _proxy_request( + "drops", route, drop_params=True, tools=5, tool_choice="auto", web_search_options="high" + ) + + sent: Final = await async_pre_call_deployment_hook(dict(request), route) + + assert {"tools", "tool_choice", "web_search_options"}.isdisjoint(sent) + assert sent["messages"] == request["messages"] + + +@pytest.mark.asyncio +async def test_additional_drop_params_drops_only_the_malformed_tool_entries() -> None: + request: Final = _proxy_request("drops", additional_drop_params=["tools"], tools=[_GOOD_TOOL, "get_weather"]) + + sent: Final = await async_pre_call_deployment_hook(dict(request), "acompletion") + + assert sent["tools"] == [_GOOD_TOOL] + + +@pytest.mark.asyncio +async def test_an_sdk_call_keeps_the_tools_it_passed() -> None: + tools: Final = (_GOOD_TOOL,) + request: Final = { + key: value for key, value in _proxy_request("sdk", tools=tools).items() if key != "proxy_server_request" + } + + sent: Final = await async_pre_call_deployment_hook(dict(request), "acompletion") + + assert sent["tools"] is tools + + +@pytest.mark.parametrize("schema", ["", [], False, 0, {}, '{"type": "object"}']) +@pytest.mark.parametrize( + "tool", + [ + lambda schema: {"type": "function", "function": {"name": "get_weather", "parameters": schema}}, + lambda schema: {"type": "function", "name": "get_weather", "parameters": schema}, + lambda schema: {"name": "get_weather", "input_schema": schema}, + ], + ids=["chat", "responses", "messages"], +) +@pytest.mark.asyncio +async def test_empty_and_json_string_schemas_reach_the_deployment( + tool: Callable[[object], Mapping[str, object]], schema: object +) -> None: + tools: Final = [tool(schema)] + + sent: Final = await async_pre_call_deployment_hook(dict(_proxy_request("keeps", tools=tools)), "acompletion") + + assert sent["tools"] is tools + + +@pytest.mark.parametrize("tools", [(_GOOD_TOOL,), [{"type": "web_search", "function": None}]], ids=["tuple", "null"]) +@pytest.mark.asyncio +async def test_tools_core_accepts_reach_the_deployment(tools: object) -> None: + sent: Final = await async_pre_call_deployment_hook(dict(_proxy_request("keeps", tools=tools)), "acompletion") + + assert sent["tools"] is tools diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index e4129ebbc39..c7bb05e5d30 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -5,13 +5,16 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objec import json import logging import unittest +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +from litellm import Router from litellm._logging import verbose_proxy_logger from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.types.utils import TokenCountResponse class TestAnthropicEndpoints(unittest.TestCase): @@ -467,6 +470,151 @@ class TestErrorLogCarriesCallId: assert call_id in record.getMessage() +def _claude_router() -> Router: + return Router( + model_list=[ + { + "model_name": "claude-sonnet", + "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "sk-test"}, + } + ] + ) + + +_MALFORMED_TOOL_ENTRY: Final = [{"name": "get_weather", "input_schema": "object"}] +_SEVEN_TOKENS: Final = TokenCountResponse( + total_tokens=7, request_model="claude-sonnet", model_used="claude-sonnet", tokenizer_type="test" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "tools"), + [ + ("claude-sonnet", 5), + ("claude-sonnet", ""), + ("claude-sonnet", ["get_weather"]), + ("claude-sonnet", _MALFORMED_TOOL_ENTRY), + ("team-public-name", 5), + ], + ids=["int", "empty_string", "list_of_strings", "input_schema_string", "unnamed_model_int"], +) +async def test_count_tokens_answers_400_for_malformed_tools_without_counting(model: str, tools: object) -> None: + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + request: Final = MagicMock() + request.headers = {} + counter: Final = AsyncMock() + body: Final = {"model": model, "messages": [{"role": "user", "content": "hi"}], "tools": tools} + + with ( + patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts it is never reached + patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route + ): + response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + assert response.status_code == 400 + assert json.loads(response.body)["error"]["type"] == "invalid_request_error" + counter.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_count_tokens_counts_malformed_tool_entries_for_a_model_the_router_cannot_name() -> None: + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + request: Final = MagicMock() + request.headers = {} + counter: Final = AsyncMock(return_value=_SEVEN_TOKENS) + body: Final = {"model": "team-public-name", "messages": [{"role": "user", "content": "hi"}], "tools": _MALFORMED_TOOL_ENTRY} + + with ( + patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives + patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route + ): + response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + assert response == {"input_tokens": 7} + assert counter.await_args.kwargs["request"].tools == _MALFORMED_TOOL_ENTRY + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("dropping", "headers"), + [ + ({"drop_params": True}, {}), + ({"additional_drop_params": ["tools"]}, {}), + ({}, {"user-agent": "claude-cli/2.1.0 (external, cli)"}), + ], + ids=["drop_params", "additional_drop_params", "claude_code_user_agent"], +) +async def test_count_tokens_drops_malformed_tools_when_drop_params_is_set( + dropping: dict[str, object], headers: dict[str, str] +) -> None: + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + request: Final = MagicMock() + request.headers = headers + counter: Final = AsyncMock(return_value=_SEVEN_TOKENS) + messages: Final = [{"role": "user", "content": "hi"}] + body: Final = {"model": "claude-sonnet", "messages": messages, "tools": 5, **dropping} + + with ( + patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives + ): + response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + assert response == {"input_tokens": 7} + counted: Final = counter.await_args.kwargs["request"] + assert counted.tools is None + assert counted.messages == messages + + +@pytest.mark.asyncio +async def test_count_tokens_drops_malformed_tools_the_counted_deployment_drops_even_with_fallbacks() -> None: + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm import Router + from litellm.proxy._types import UserAPIKeyAuth + + router: Final = Router( + model_list=[ + { + "model_name": "claude-sonnet", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_key": "sk-test", + "additional_drop_params": ["tools"], + }, + }, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-haiku-5-5", "api_key": "sk-test"}}, + ], + fallbacks=[{"claude-sonnet": ["backup"]}], + ) + request: Final = MagicMock() + request.headers = {} + counter: Final = AsyncMock(return_value=_SEVEN_TOKENS) + body: Final = {"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}], "tools": 5} + + with ( + patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives + patch.object(proxy_server, "llm_router", new=router), # test-quality-ok: module global imported at call time; no injection seam on this route + ): + response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + assert response == {"input_tokens": 7} + assert counter.await_args.kwargs["request"].tools is None + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index b047c29243a..a7b955f5b84 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3454,10 +3454,11 @@ class TestResponsesInputTokens: body: dict[str, Any], path: str = "/v1/responses/input_tokens", counter: AsyncMock | None = None, + llm_router: litellm.Router | None = None, ) -> tuple[Response, AsyncMock]: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - from litellm.proxy.response_api_endpoints.endpoints import _proxy_token_counter + from litellm.proxy.response_api_endpoints.endpoints import _proxy_llm_router, _proxy_token_counter from litellm.types.utils import TokenCountResponse token_counter_mock = ( @@ -3474,6 +3475,7 @@ class TestResponsesInputTokens: ) app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(api_key="sk-test", request_route=path) app.dependency_overrides[_proxy_token_counter] = lambda: token_counter_mock + app.dependency_overrides[_proxy_llm_router] = lambda: llm_router try: client = TestClient(app) response = client.post(path, json=body, headers={"Authorization": "Bearer sk-9876"}) @@ -3481,6 +3483,7 @@ class TestResponsesInputTokens: finally: app.dependency_overrides.pop(user_api_key_auth, None) app.dependency_overrides.pop(_proxy_token_counter, None) + app.dependency_overrides.pop(_proxy_llm_router, None) def test_string_input_returns_openai_input_tokens_shape(self): response, counter = self._post_input_tokens({"model": "gpt-4o", "input": "Hello, how are you?"}) @@ -3594,6 +3597,48 @@ class TestResponsesInputTokens: assert response.status_code == 400, response.text error = response.json()["error"] assert error["type"] == "invalid_request_error" + assert error["param"] == "tools" + assert error["message"].startswith("tools: ") + counter.assert_not_awaited() + + def test_invalid_tools_are_dropped_under_drop_params(self) -> None: + response, counter = self._post_input_tokens( + {"model": "gpt-4o", "input": "hi", "tools": [5, {"type": "web_search"}], "drop_params": True} + ) + + assert response.status_code == 200, response.text + assert counter.await_args.kwargs["request"].tools == [{"type": "web_search"}] + + def test_malformed_tools_follow_the_counted_deployment_even_when_fallbacks_are_configured(self) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-5.2", + "litellm_params": {"model": "openai/gpt-5.2", "api_key": "sk-test", "additional_drop_params": ["tools"]}, + }, + {"model_name": "backup", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-test"}}, + ], + fallbacks=[{"gpt-5.2": ["backup"]}], + ) + + response, counter = self._post_input_tokens( + {"model": "gpt-5.2", "input": "hi", "tools": 5}, llm_router=router + ) + + assert response.status_code == 200, response.text + assert counter.await_args.kwargs["request"].tools is None + + def test_malformed_tools_for_a_model_outside_the_router_get_a_clean_400(self) -> None: + router: Final = litellm.Router( + model_list=[{"model_name": "gpt-5.2", "litellm_params": {"model": "openai/gpt-5.2", "api_key": "sk-test"}}] + ) + + response, counter = self._post_input_tokens( + {"model": "openai/gpt-5.4-mini", "input": "hi", "tools": 5}, llm_router=router + ) + + assert response.status_code == 400, response.text + assert response.json()["error"]["param"] == "tools" counter.assert_not_awaited() def test_provider_error_maps_status_code(self): diff --git a/tests/unit/proxy/test_tools_allowlist_enforcement.py b/tests/unit/proxy/test_tools_allowlist_enforcement.py index e855ef47dd5..e306f7be4f8 100644 --- a/tests/unit/proxy/test_tools_allowlist_enforcement.py +++ b/tests/unit/proxy/test_tools_allowlist_enforcement.py @@ -145,6 +145,21 @@ class TestExtractRequestToolNames: data = {"tools": [{"type": "function", "function": {"name": "x"}}]} assert extract_request_tool_names("/v1/embeddings", data) == [] + @pytest.mark.parametrize( + ("route", "data"), + [ + ("/v1/chat/completions", {"tools": 5}), + ("/v1/chat/completions", {"functions": 5}), + ("/v1/responses", {"tools": 5}), + ("/v1/messages", {"tools": 5}), + ("/generate_content", {"tools": 5}), + ("/generate_content", {"tools": [{"functionDeclarations": 5}]}), + ], + ids=["chat_tools", "chat_functions", "responses", "messages", "generate_content", "function_declarations"], + ) + def test_malformed_tool_lists_yield_no_names(self, route: str, data: dict[str, object]) -> None: + assert extract_request_tool_names(route, data) == [] + class TestCheckToolsAllowlist: """Test allowlist enforcement in auth (no DB in hot path).""" diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 84aa1798da4..e83efa21dcc 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -4,6 +4,7 @@ and ``_handle_logging_proxy_only_error``.""" from __future__ import annotations import asyncio +from collections.abc import Callable from datetime import datetime from types import MappingProxyType from typing import Final @@ -921,6 +922,48 @@ async def test_post_call_failure_hook_falls_back_to_body_shape_without_a_route(p assert request_data["call_type"] == "acompletion" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type, body", + [ + ("anthropic_messages", {"model": "m", "messages": [{"role": "user", "content": "hi"}]}), + ("aresponses", {"model": "m", "input": [{"role": "user", "content": "hi"}]}), + ("acompact_responses", {"model": "m", "input": [{"role": "user", "content": "hi"}]}), + ], +) +async def test_post_call_failure_hook_keeps_the_call_type_request_setup_logged( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + call_type: str, + body: dict[str, object], +) -> None: + """A request rejected after request setup already carries its endpoint's call type, so the body + shape must not relabel a /v1/messages rejection as a chat completion or a list-input /v1/responses + rejection as an embedding.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + proxy_logging.alert_types = [] + request_data: Final = { + **body, + "litellm_logging_obj": LiteLLMLoggingObj( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type=call_type, + start_time=datetime.now(), + litellm_call_id=f"test_keeps_{call_type}", + function_id=f"test_keeps_{call_type}", + ), + "metadata": {}, + } + await proxy_logging.post_call_failure_hook( + request_data=request_data, + original_exception=HTTPException(status_code=400, detail="tools must be a list"), + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + ) + assert request_data["call_type"] == call_type + + @pytest.mark.asyncio @pytest.mark.parametrize("route", ["/v1/files", "/files/file-abc", "/v1/containers"]) async def test_post_call_failure_hook_keeps_the_route_for_multi_operation_routes( diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index af8fca72252..3532870df1c 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -287,6 +287,19 @@ class TestFilterWebSearchDeployments: assert "deployment-2" in result_ids assert "deployment-3" not in result_ids + def test_web_search_in_a_tuple_filters_unsupported(self, sample_deployments: list[dict[str, object]]) -> None: + result: Final = filter_web_search_deployments(sample_deployments, {"tools": ({"type": "web_search"},)}) + assert [d["model_info"]["id"] for d in result] == ["deployment-1", "deployment-2"] + + @pytest.mark.parametrize("tool_type", [{}, [], None, 5], ids=["object", "list", "null", "int"]) + def test_a_non_string_tool_type_is_no_web_search_request( + self, sample_deployments: list[dict[str, object]], tool_type: object + ) -> None: + result: Final = filter_web_search_deployments( + sample_deployments, {"tools": [{"type": tool_type}, {"type": "function"}]} + ) + assert result == sample_deployments + def test_web_search_preview_filters_unsupported(self, sample_deployments): """web_search_preview type should also trigger filtering""" request_kwargs = {"tools": [{"type": "web_search_preview"}]} diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index c8ecf901602..50a9d98385d 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -9740,6 +9740,96 @@ def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice assert kwargs["tool_choice"] == "none" +@pytest.mark.parametrize( + "deployment_params, router_defaults, request_params, expected_tools", + [ + ({"drop_params": True}, {}, {}, [{"type": "web_search"}]), + ({}, {"drop_params": True}, {}, [{"type": "web_search"}]), + ({}, {}, {"additional_drop_params": ["tools"]}, [{"type": "web_search"}]), + ({"drop_params": True}, {}, {"drop_params": False}, 5), + ({}, {}, {}, 5), + ], + ids=[ + "deployment_drop_params", + "router_default_drop_params", + "request_drops_tools", + "request_overrides_deployment_drop", + "no_drop_policy", + ], +) +def test_update_kwargs_with_deployment_keeps_deployment_tools_when_malformed_request_tools_drop( + deployment_params: dict[str, object], + router_defaults: dict[str, object], + request_params: dict[str, object], + expected_tools: object, +) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "o3-deep-research", + "litellm_params": { + "model": "openai/o3-deep-research", + "api_key": "fake-key", + "tools": [{"type": "web_search"}], + **deployment_params, + }, + }, + ], + default_litellm_params={**router_defaults}, + ) + kwargs: Final[dict[str, object]] = {"metadata": {}, "tools": 5, **request_params} + router._update_kwargs_with_deployment( + deployment=router.get_deployment_by_model_group_name(model_group_name="o3-deep-research"), kwargs=kwargs + ) + + assert kwargs["tools"] == expected_tools + + +_NAMED_REQUEST_TOOL_CHOICE: Final = {"type": "function", "function": {"name": "lookup_order"}} + + +@pytest.mark.parametrize( + "deployment_params, expected_tool_fields", + [ + ( + {"drop_params": True, "tool_choice": "required"}, + {"tools": [{"type": "web_search"}], "tool_choice": "required"}, + ), + ({"drop_params": True}, {"tools": [{"type": "web_search"}]}), + ({}, {"tools": 5, "tool_choice": _NAMED_REQUEST_TOOL_CHOICE, "parallel_tool_calls": False}), + ], + ids=["deployment_tool_choice_applies", "no_deployment_tool_choice", "no_drop_policy"], +) +def test_update_kwargs_with_deployment_drops_request_tool_settings_with_malformed_request_tools( + deployment_params: dict[str, object], expected_tool_fields: dict[str, object] +) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "o3-deep-research", + "litellm_params": { + "model": "openai/o3-deep-research", + "api_key": "fake-key", + "tools": [{"type": "web_search"}], + **deployment_params, + }, + }, + ], + ) + kwargs: Final[dict[str, object]] = { + "metadata": {}, + "tools": 5, + "tool_choice": _NAMED_REQUEST_TOOL_CHOICE, + "parallel_tool_calls": False, + } + router._update_kwargs_with_deployment( + deployment=router.get_deployment_by_model_group_name(model_group_name="o3-deep-research"), kwargs=kwargs + ) + + tool_fields: Final = ("tools", "tool_choice", "parallel_tool_calls") + assert {name: kwargs[name] for name in tool_fields if name in kwargs} == expected_tool_fields + + def test_credential_name_injected_as_tag(): """ Test that litellm_credential_name from deployment litellm_params diff --git a/tests/unit/vector_stores/test_vector_store_registry.py b/tests/unit/vector_stores/test_vector_store_registry.py index f6c888d20ff..83b59d77aaf 100644 --- a/tests/unit/vector_stores/test_vector_store_registry.py +++ b/tests/unit/vector_stores/test_vector_store_registry.py @@ -349,6 +349,56 @@ async def test_malformed_vector_store_index_row_from_the_db_raises_a_validation_ await VectorStoreIndexRegistry.get_vector_store_indexes_from_db(_database_listing_index_rows([row])) +@pytest.mark.parametrize( + ("tools", "expected"), + [ + (5, []), + (["vector_store_ids", {"type": "file_search", "vector_store_ids": ["vs_kept"]}], ["vs_kept"]), + ( + [{"type": "file_search", "vector_store_ids": 5}, {"type": "file_search", "vector_store_ids": ["vs_kept"]}], + ["vs_kept"], + ), + ], + ids=["int", "string_entry", "int_vector_store_ids"], +) +def test_vector_store_ids_to_run_skips_malformed_tools(tools: object, expected: list[str]) -> None: + assert VectorStoreRegistry().get_vector_store_ids_to_run(non_default_params={}, tools=tools) == expected + + +@pytest.mark.parametrize( + "vector_store_ids", + ["vs_requested", ["vs_requested", 5]], + ids=["string", "list_with_a_non_string"], +) +def test_vector_store_ids_to_run_reads_every_store_id_a_malformed_vector_store_ids_field_names( + vector_store_ids: object, +) -> None: + tools: Final = [{"type": "file_search", "vector_store_ids": vector_store_ids}] + assert VectorStoreRegistry().get_vector_store_ids_to_run(non_default_params={}, tools=tools) == ["vs_requested"] + + +@pytest.mark.parametrize( + "vector_store_ids", + ["vs_allowed", ["vs_allowed", 12345], [12345], ["vs_allowed"]], + ids=["string", "list_with_a_non_string", "only_non_strings", "list_of_strings"], +) +def test_vector_store_tools_pop_exactly_the_store_ids_the_registry_reports_to_run(vector_store_ids: object) -> None: + now: Final = datetime.now(timezone.utc) + registry: Final = VectorStoreRegistry( + [ + LiteLLM_ManagedVectorStore(vector_store_id=vector_store_id, created_at=now, updated_at=now) + for vector_store_id in ("vs_allowed", 12345) + ] + ) + tools: Final = [{"type": "file_search", "vector_store_ids": vector_store_ids}] + checked: Final = registry.get_vector_store_ids_to_run(non_default_params={}, tools=tools) + ran: Final[list[str]] = [] + + registry.get_and_pop_recognised_vector_store_tools(tools=[dict(tool) for tool in tools], vector_store_ids=ran) + + assert ran == checked + + class TestRegistriesDoNotShareDefaultState: def test_adding_to_one_registry_leaves_the_next_one_empty(self): VectorStoreRegistry().add_vector_store_to_registry(