mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
fix(mcp): coerce skill_search top_k instead of raising 500 on malformed input
The MCP-REST skill_search dispatch validated raw tool arguments through a pydantic model directly, so a non-numeric top_k raised a ValidationError that the endpoint's catch-all turned into an HTTP 500. Mirrors the agent_search branch's tolerant coerce_top_k handling instead.
This commit is contained in:
parent
cce83d649a
commit
6f3447bb67
2 changed files with 32 additions and 11 deletions
|
|
@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Any, Final, Literal
|
|||
import anyio
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT
|
||||
|
|
@ -163,12 +162,6 @@ if MCP_AVAILABLE:
|
|||
request_path=request.scope.get("_original_path") or request.url.path,
|
||||
)
|
||||
|
||||
class _SkillSearchToolArguments(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
query: str = ""
|
||||
top_k: int = DEFAULT_SKILL_SEARCH_TOP_K
|
||||
|
||||
async def _handle_virtual_mcp_tool(
|
||||
request: Request,
|
||||
data: dict[str, Any],
|
||||
|
|
@ -211,10 +204,11 @@ if MCP_AVAILABLE:
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
if tool_name == SKILL_SEARCH_TOOL_NAME:
|
||||
skill_search_args: Final = _SkillSearchToolArguments.model_validate(tool_arguments)
|
||||
return await handle_skill_search(
|
||||
query=skill_search_args.query,
|
||||
top_k=skill_search_args.top_k,
|
||||
query=str(tool_arguments.get("query", "")),
|
||||
top_k=coerce_top_k(
|
||||
tool_arguments.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K
|
||||
),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
rest_client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
|
||||
|
|
|
|||
|
|
@ -744,6 +744,31 @@ class TestCallToolRestApiVirtualTools:
|
|||
assert mock_search.await_args.kwargs["top_k"] == 1
|
||||
assert mock_search.await_args.kwargs["agents"] == (translator,)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skill_search_call_tolerates_malformed_top_k(self) -> None:
|
||||
"""Regression: a caller-supplied non-numeric top_k must be coerced to the default,
|
||||
the same as agent_search, instead of raising a pydantic ValidationError that the
|
||||
endpoint's catch-all turns into an HTTP 500."""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
|
||||
request = self._make_request(
|
||||
{"name": SKILL_SEARCH_TOOL_NAME, "arguments": {"query": "translate a document", "top_k": "not-a-number"}}
|
||||
)
|
||||
fake_result = CallToolResult(content=[TextContent(type="text", text="[]")], isError=False)
|
||||
with patch( # test-quality-ok: the embedding router only resolves via proxy_server globals, no injection seam
|
||||
"litellm.proxy._experimental.mcp_server.tool_search.handle_skill_search",
|
||||
new_callable=AsyncMock,
|
||||
return_value=fake_result,
|
||||
) as mock_search:
|
||||
result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
assert result.isError is False
|
||||
assert mock_search.await_args.kwargs["top_k"] == DEFAULT_SKILL_SEARCH_TOP_K
|
||||
assert mock_search.await_args.kwargs["query"] == "translate a document"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_search_call_reports_missing_embedding_model_as_tool_error(self) -> None:
|
||||
from litellm.proxy.agent_endpoints.agent_search import AgentSearchNotConfigured
|
||||
|
|
@ -817,7 +842,9 @@ class TestCallToolRestApiVirtualTools:
|
|||
assert "mcp_tool_search.embedding_model" in result.content[0].text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tool_search_reports_invalid_settings_as_tool_error(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_mcp_tool_search_reports_invalid_settings_as_tool_error(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "mcp_tool_search", {"top_k": 0})
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
|
||||
result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue