mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
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
This commit is contained in:
Mateo Wang
GitHub
parent
57e69ce713
commit
d76ec3a47e
34 changed files with 4435 additions and 92 deletions
|
|
@ -1416,6 +1416,7 @@ from .exceptions import (
|
|||
InvalidRequestError,
|
||||
BadRequestError,
|
||||
ImageFetchError,
|
||||
MalformedRequestFieldError as MalformedRequestFieldError,
|
||||
VectorStoreSearchError,
|
||||
NotFoundError,
|
||||
PermissionDeniedError,
|
||||
|
|
|
|||
|
|
@ -313,6 +313,10 @@ class VectorStoreSearchError(BadRequestError):
|
|||
)
|
||||
|
||||
|
||||
class MalformedRequestFieldError(BadRequestError):
|
||||
pass
|
||||
|
||||
|
||||
class UnprocessableEntityError(openai.UnprocessableEntityError):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
164
litellm/litellm_core_utils/tool_fields.py
Normal file
164
litellm/litellm_core_utils/tool_fields.py
Normal file
|
|
@ -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),))
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"):
|
||||
|
|
|
|||
|
|
@ -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": <number>}
|
||||
"""
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
619
tests/integration/providers/test_malformed_tools_chaos.py
Normal file
619
tests/integration/providers/test_malformed_tools_chaos.py
Normal file
|
|
@ -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
|
||||
2731
tests/integration/providers/test_malformed_tools_wire.py
Normal file
2731
tests/integration/providers/test_malformed_tools_wire.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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) == {""}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
187
tests/unit/litellm_core_utils/test_tool_fields.py
Normal file
187
tests/unit/litellm_core_utils/test_tool_fields.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)."""
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"}]}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue