scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

feat(a2a): semantic search over the agent registry via GET /v1/agents?query and an agent_search MCP tool

Browse source
This commit is contained in:
mateo-berri 2026-08-27 18:31:15 -07:00
parent 7083c47998
commit ca21cf5773
12 changed files with 910 additions and 77 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/__init__.py
View file

@ -487,6 +487,7 @@ public_mcp_servers: Optional[List[str]] = None
public_mcp_hub_strict_whitelist: bool = True
public_model_groups: Optional[List[str]] = None
public_agent_groups: Optional[List[str]] = None
agent_search_embedding_model: Optional[str] = None
# Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]])
# New format: { "displayName": { "url": "...", "index": 0 } }
# Old format: { "displayName": "url" } (for backward compatibility)

18
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -168,8 +168,11 @@ if MCP_AVAILABLE:
MCPRequestHandler,
)
from litellm.proxy._experimental.mcp_server.tool_search import (
AGENT_SEARCH_TOOL_NAME,
DEFAULT_AGENT_SEARCH_TOP_K,
MCP_TOOL_SEARCH_TOOL_NAME,
coerce_top_k,
handle_agent_search,
handle_mcp_tool_call,
handle_mcp_tool_search,
)
@ -182,6 +185,14 @@ if MCP_AVAILABLE:
detail={"error": "forbidden", "message": f"{tool_name} requires mcp_tool_search_enabled on the key"},
)
tool_arguments: Final = data.get("arguments") or {}
if tool_name == AGENT_SEARCH_TOOL_NAME:
return await handle_agent_search(
query=str(tool_arguments.get("query", "")),
top_k=coerce_top_k(
tool_arguments.get("top_k", DEFAULT_AGENT_SEARCH_TOP_K), default=DEFAULT_AGENT_SEARCH_TOP_K
),
user_api_key_dict=user_api_key_dict,
)
rest_client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
(
virtual_mcp_auth_header,
@ -939,12 +950,9 @@ if MCP_AVAILABLE:
tool_name: Final[str | None] = data.get("name")
tool_arguments: Final[dict[str, object]] = data.get("arguments") or {}
from litellm.proxy._experimental.mcp_server.tool_search import (
MCP_TOOL_CALL_TOOL_NAME,
MCP_TOOL_SEARCH_TOOL_NAME,
)
from litellm.proxy._experimental.mcp_server.tool_search import VIRTUAL_TOOL_NAMES
if tool_name in (MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME):
if tool_name in VIRTUAL_TOOL_NAMES:
return await _handle_virtual_mcp_tool(request, data, tool_name, user_api_key_dict)
# Validate required parameters early

13
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -912,14 +912,17 @@ if MCP_AVAILABLE:
the caller falls through to normal tool routing.
"""
from litellm.proxy._experimental.mcp_server.tool_search import (
MCP_TOOL_CALL_TOOL_NAME,
AGENT_SEARCH_TOOL_NAME,
DEFAULT_AGENT_SEARCH_TOP_K,
MCP_TOOL_SEARCH_TOOL_NAME,
VIRTUAL_TOOL_NAMES,
coerce_top_k,
handle_agent_search,
handle_mcp_tool_call,
handle_mcp_tool_search,
)
if name not in (MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME):
if name not in VIRTUAL_TOOL_NAMES:
return None
if not getattr(
@ -952,6 +955,12 @@ if MCP_AVAILABLE:
)
assert user_api_key_auth is not None # guaranteed by the flag check above
if name == AGENT_SEARCH_TOOL_NAME:
return await handle_agent_search(
query=str(args.get("query", "")),
top_k=coerce_top_k(args.get("top_k", DEFAULT_AGENT_SEARCH_TOP_K), default=DEFAULT_AGENT_SEARCH_TOP_K),
user_api_key_dict=user_api_key_auth,
)
virtual_logging_obj: Final = await _build_virtual_call_logging_obj(
name=name,
arguments=args,

151
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

@ -1,8 +1,14 @@
from __future__ import annotations
import json
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, TypedDict, assert_never
from typing_extensions import ReadOnly, Required
import litellm
from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K
if TYPE_CHECKING:
from mcp.types import CallToolResult
@ -12,6 +18,8 @@ if TYPE_CHECKING:
MCP_TOOL_SEARCH_TOOL_NAME: Final[str] = "mcp_tool_search"
MCP_TOOL_CALL_TOOL_NAME: Final[str] = "mcp_tool_call"
AGENT_SEARCH_TOOL_NAME: Final[str] = "agent_search"
VIRTUAL_TOOL_NAMES: Final = frozenset((MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME))
def coerce_top_k(value: Any, default: int = 5) -> int:
@ -34,46 +42,111 @@ def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = 5) -> lis
return [tool for _, tool in sorted(scored, key=lambda x: x[0], reverse=True)[:top_k]]
def get_virtual_tool_definitions() -> list[dict[str, Any]]:
return [
{
"name": MCP_TOOL_SEARCH_TOOL_NAME,
"description": "Search for MCP tools by keyword. Returns top matching tools with names, descriptions, and input schemas.",
"inputSchema": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Keywords to search for in tool names and descriptions.",
},
"top_k": {
"type": "integer",
"description": "Maximum number of results to return.",
"default": 5,
},
},
"required": ["query"],
class _ToolParamSchema(TypedDict, total=False):
type: Required[ReadOnly[str]]
description: Required[ReadOnly[str]]
default: ReadOnly[int]
class _ToolInputSchema(TypedDict):
type: ReadOnly[str]
properties: ReadOnly[Mapping[str, _ToolParamSchema]]
required: ReadOnly[tuple[str, ...]]
class VirtualToolDefinition(TypedDict):
name: ReadOnly[str]
description: ReadOnly[str]
inputSchema: ReadOnly[_ToolInputSchema]
_MCP_TOOL_SEARCH_DEFINITION: Final[VirtualToolDefinition] = {
"name": MCP_TOOL_SEARCH_TOOL_NAME,
"description": "Search for MCP tools by keyword. Returns top matching tools with names, descriptions, and input schemas.",
"inputSchema": {
"type": "object",
"properties": {
"query": {"type": "string", "description": "Keywords to search for in tool names and descriptions."},
"top_k": {"type": "integer", "description": "Maximum number of results to return.", "default": 5},
},
"required": ("query",),
},
}
_MCP_TOOL_CALL_DEFINITION: Final[VirtualToolDefinition] = {
"name": MCP_TOOL_CALL_TOOL_NAME,
"description": "Call an MCP tool by name with the given arguments.",
"inputSchema": {
"type": "object",
"properties": {
"tool_name": {"type": "string", "description": "The exact name of the MCP tool to call."},
"arguments": {"type": "object", "description": "Arguments to pass to the tool."},
},
"required": ("tool_name",),
},
}
_AGENT_SEARCH_DEFINITION: Final[VirtualToolDefinition] = {
"name": AGENT_SEARCH_TOOL_NAME,
"description": "Find A2A agents by describing the task in natural language. Returns the best matching agents you can access, ranked by semantic similarity, each with its agent_id, name, description, skills, and score.",
"inputSchema": {
"type": "object",
"properties": {
"query": {"type": "string", "description": "The task the agent should be able to do, in natural language."},
"top_k": {
"type": "integer",
"description": "Maximum number of agents to return.",
"default": DEFAULT_AGENT_SEARCH_TOP_K,
},
},
{
"name": MCP_TOOL_CALL_TOOL_NAME,
"description": "Call an MCP tool by name with the given arguments.",
"inputSchema": {
"type": "object",
"properties": {
"tool_name": {
"type": "string",
"description": "The exact name of the MCP tool to call.",
},
"arguments": {
"type": "object",
"description": "Arguments to pass to the tool.",
},
},
"required": ["tool_name"],
},
},
]
"required": ("query",),
},
}
def get_virtual_tool_definitions() -> tuple[VirtualToolDefinition, ...]:
return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION)
def _text_tool_result(text: str, is_error: bool) -> CallToolResult:
from mcp.types import CallToolResult, TextContent
return CallToolResult(
content=[TextContent(type="text", text=text)], # mutable-ok: CallToolResult accepts only list content
isError=is_error,
)
async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth) -> CallToolResult:
from litellm.proxy.agent_endpoints.agent_search import (
AgentSearchEmbeddingFailed,
AgentSearchHits,
AgentSearchNotConfigured,
agent_search_result,
global_agent_search_index,
search_agents,
)
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
from litellm.proxy.proxy_server import llm_router
await check_feature_access_for_user(user_api_key_dict, "agents")
outcome: Final = await search_agents(
query=query,
agents=await accessible_agents(user_api_key_dict),
top_k=max(top_k, 1),
router=llm_router,
embedding_model=litellm.agent_search_embedding_model,
index=global_agent_search_index,
)
match outcome:
case AgentSearchHits(hits):
results: Final = tuple(agent_search_result(hit).model_dump() for hit in hits)
return _text_tool_result(json.dumps(results), is_error=False)
case AgentSearchNotConfigured(reason) | AgentSearchEmbeddingFailed(reason):
return _text_tool_result(reason, is_error=True)
case _:
assert_never(outcome)
async def handle_mcp_tool_search(

90
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -2377,6 +2377,17 @@
],
"title": "Rpm Limit"
},
"search_score": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Search Score"
},
"session_rpm_limit": {
"anyOf": [
{
@ -3404,7 +3415,7 @@
},
"/v1/agents": {
"get": {
"description": "Example usage:\n```\ncurl -X GET \"http://localhost:4000/v1/agents\" -H \"Content-Type: application/json\" -H \"Authorization: Bearer your-key\" ```\n\nPass `?health_check=true` to filter out agents whose URL is unreachable:\n```\ncurl -X GET \"http://localhost:4000/v1/agents?health_check=true\" -H \"Content-Type: application/json\" -H \"Authorization: Bearer your-key\" ```\n\nReturns: List[AgentResponse]",
"description": "Example usage:\n```\ncurl -X GET \"http://localhost:4000/v1/agents\" -H \"Content-Type: application/json\" -H \"Authorization: Bearer your-key\" ```\n\nPass `?health_check=true` to filter out agents whose URL is unreachable:\n```\ncurl -X GET \"http://localhost:4000/v1/agents?health_check=true\" -H \"Content-Type: application/json\" -H \"Authorization: Bearer your-key\" ```\n\nPass `?query=<task>` to get the best matching agents ranked by semantic similarity:\n```\ncurl -X GET \"http://localhost:4000/v1/agents?query=translate+a+PDF+document&top_k=5\" -H \"Content-Type: application/json\" -H \"Authorization: Bearer your-key\" ```\n\nReturns: List[AgentResponse]",
"operationId": "get_agents_v1_agents_get",
"parameters": [
{
@ -3418,6 +3429,39 @@
"title": "Health Check",
"type": "boolean"
}
},
{
"description": "Describe the task in natural language to rank the agents you can reach by semantic similarity over their name, description, and skills. Each result carries a search_score. Requires litellm_settings.agent_search_embedding_model.",
"in": "query",
"name": "query",
"required": false,
"schema": {
"anyOf": [
{
"minLength": 1,
"type": "string"
},
{
"type": "null"
}
],
"description": "Describe the task in natural language to rank the agents you can reach by semantic similarity over their name, description, and skills. Each result carries a search_score. Requires litellm_settings.agent_search_embedding_model.",
"title": "Query"
}
},
{
"description": "With query: the maximum number of ranked agents to return.",
"in": "query",
"name": "top_k",
"required": false,
"schema": {
"default": 5,
"description": "With query: the maximum number of ranked agents to return.",
"maximum": 100,
"minimum": 1,
"title": "Top K",
"type": "integer"
}
}
],
"responses": {
@ -15038,6 +15082,17 @@
}
],
"title": "Upstream Resource"
},
"upstream_token_header": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Upstream Token Header"
}
},
"title": "MCPCredentials",
@ -17518,6 +17573,17 @@
}
],
"title": "Upstream Resource"
},
"upstream_token_header": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Upstream Token Header"
}
},
"title": "MCPCredentials",
@ -20352,6 +20418,17 @@
}
],
"title": "Upstream Resource"
},
"upstream_token_header": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Upstream Token Header"
}
},
"title": "MCPCredentials",
@ -23699,6 +23776,17 @@
}
],
"title": "Upstream Resource"
},
"upstream_token_header": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Upstream Token Header"
}
},
"title": "MCPCredentials",

184
litellm/proxy/agent_endpoints/agent_search.py Normal file
View file

@ -0,0 +1,184 @@
"""Semantic ranking over the in-memory A2A agent registry, shared by GET /v1/agents?query= and the agent_search MCP tool."""
from __future__ import annotations
import math
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
from openai import OpenAIError
from pydantic import BaseModel, ConfigDict, ValidationError
from litellm.exceptions import BudgetExceededError
from litellm.types.agents import AgentResponse
if TYPE_CHECKING:
from litellm.router import Router
DEFAULT_AGENT_SEARCH_TOP_K: Final = 5
Vector: TypeAlias = tuple[float, ...]
class Embedder(Protocol):
def __call__(self, texts: Sequence[str]) -> Awaitable[Sequence[Vector]]: ...
@dataclass(frozen=True, slots=True)
class AgentSearchHit:
agent: AgentResponse
score: float
@dataclass(frozen=True, slots=True)
class AgentSearchHits:
hits: tuple[AgentSearchHit, ...]
@dataclass(frozen=True, slots=True)
class AgentSearchNotConfigured:
reason: str
@dataclass(frozen=True, slots=True)
class AgentSearchEmbeddingFailed:
reason: str
AgentSearchOutcome: TypeAlias = AgentSearchHits | AgentSearchNotConfigured | AgentSearchEmbeddingFailed
class _SearchableSkill(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
name: str = ""
description: str = ""
tags: tuple[str, ...] = ()
class _SearchableCard(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
description: str = ""
skills: tuple[_SearchableSkill, ...] = ()
class _EmbeddingItem(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
embedding: tuple[float, ...]
class _EmbeddingData(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
data: tuple[_EmbeddingItem, ...]
class AgentSearchResult(BaseModel):
model_config = ConfigDict(frozen=True)
agent_id: str
agent_name: str
description: str
skills: tuple[_SearchableSkill, ...]
score: float
def _searchable_card(agent: AgentResponse) -> _SearchableCard:
try:
return _SearchableCard.model_validate(agent.agent_card_params)
except ValidationError:
return _SearchableCard()
def _skill_text(skill: _SearchableSkill) -> str:
return " ".join(part for part in (skill.name, skill.description, " ".join(skill.tags)) if part)
def agent_search_text(agent: AgentResponse) -> str:
card: Final = _searchable_card(agent)
skill_lines: Final = tuple(_skill_text(skill) for skill in card.skills)
return "\n".join(part for part in (agent.agent_name, card.description, *skill_lines) if part)
def agent_search_result(hit: AgentSearchHit) -> AgentSearchResult:
card: Final = _searchable_card(hit.agent)
return AgentSearchResult(
agent_id=hit.agent.agent_id,
agent_name=hit.agent.agent_name,
description=card.description,
skills=card.skills,
score=hit.score,
)
def cosine_similarity(left: Vector, right: Vector) -> float:
dot: Final = sum(a * b for a, b in zip(left, right, strict=True))
norms: Final = math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right))
return dot / norms if norms else 0.0
def router_embedder(router: Router, embedding_model: str) -> Embedder:
async def embed(texts: Sequence[str]) -> Sequence[Vector]:
batch: Final = list(texts) # mutable-ok: Router.aembedding accepts only str | list input
response: Final = await router.aembedding(model=embedding_model, input=batch)
return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data)
return embed
class AgentSearchIndex:
"""Caches one vector per distinct agent text, so repeat searches only embed the query."""
def __init__(self) -> None:
self._vectors: Mapping[str, Vector] = MappingProxyType({})
async def search(
self, query: str, agents: Sequence[AgentResponse], top_k: int, embed: Embedder
) -> AgentSearchHits | AgentSearchEmbeddingFailed:
if not agents:
return AgentSearchHits(hits=())
texts: Final = tuple(agent_search_text(agent) for agent in agents)
missing: Final = tuple(dict.fromkeys(text for text in texts if text not in self._vectors))
try:
vectors: Final = await embed((query, *missing))
except (OpenAIError, ValueError, BudgetExceededError) as exc:
return AgentSearchEmbeddingFailed(reason=f"embedding the search query failed: {exc}")
if len(vectors) != len(missing) + 1:
return AgentSearchEmbeddingFailed(
reason=f"embedding model returned {len(vectors)} vectors for {len(missing) + 1} inputs"
)
self._vectors = MappingProxyType(dict(chain(self._vectors.items(), zip(missing, vectors[1:], strict=True))))
ranked: Final = sorted(
(
AgentSearchHit(agent=agent, score=cosine_similarity(vectors[0], self._vectors[text]))
for agent, text in zip(agents, texts, strict=True)
),
key=lambda hit: hit.score,
reverse=True,
)
return AgentSearchHits(hits=tuple(ranked[:top_k]))
global_agent_search_index: Final = AgentSearchIndex()
async def search_agents(
query: str,
agents: Sequence[AgentResponse],
top_k: int,
router: Router | None,
embedding_model: str | None,
index: AgentSearchIndex,
) -> AgentSearchOutcome:
if embedding_model is None:
return AgentSearchNotConfigured(
reason="agent search needs litellm_settings.agent_search_embedding_model set to an embedding model from model_list"
)
if router is None:
return AgentSearchNotConfigured(reason="agent search needs a model_list so the embedding model can be called")
return await index.search(query, agents, top_k, router_embedder(router, embedding_model))

16
litellm/proxy/agent_endpoints/auth/agent_permission_handler.py
View file

@ -13,9 +13,11 @@ from litellm.proxy._types import (
UI_TEAM_ID,
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTable,
LitellmUserRoles,
UserAPIKeyAuth,
)
from litellm.repositories.table_repositories import AgentsRepository
from litellm.types.agents import AgentResponse
@dataclass(frozen=True, slots=True)
@ -439,3 +441,17 @@ class AgentRequestHandler:
except Exception as e:
verbose_logger.warning("Failed to get agent access groups for team: %s", e)
return []
async def accessible_agents(user_api_key_auth: UserAPIKeyAuth) -> tuple[AgentResponse, ...]:
"""Every registry agent for proxy admins, else the agents the key's and team's grants reach."""
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
all_agents: Final = global_agent_registry.get_agent_list()
if user_api_key_auth.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN.value):
return all_agents
match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_auth):
case UnrestrictedAgentAccess():
return all_agents
case RestrictedAgentAccess(allowed_agent_ids):
return tuple(agent for agent in all_agents if agent.agent_id in allowed_agent_ids)

92
litellm/proxy/agent_endpoints/endpoints.py
View file

@ -12,10 +12,11 @@ import asyncio
import os
import uuid
from collections.abc import Mapping, Sequence
from typing import Final, TypedDict
from types import MappingProxyType
from typing import Annotated, Final, TypedDict, assert_never
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from typing_extensions import Required
from typing_extensions import ReadOnly, Required
import litellm
from litellm._logging import verbose_proxy_logger
@ -32,6 +33,15 @@ from litellm.proxy.a2a.agent_card import (
merge_agent_card,
normalize_protocol_version,
)
from litellm.proxy.agent_endpoints.agent_search import (
DEFAULT_AGENT_SEARCH_TOP_K,
AgentSearchEmbeddingFailed,
AgentSearchHits,
AgentSearchNotConfigured,
global_agent_search_index,
search_agents,
)
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
@ -211,6 +221,38 @@ async def _check_agent_url_health(
}
class _AgentSearchErrorDetail(TypedDict):
error: ReadOnly[str]
message: ReadOnly[str]
def _agent_search_error(status_code: int, error: str, message: str) -> HTTPException:
detail: Final[_AgentSearchErrorDetail] = {"error": error, "message": message}
return HTTPException(status_code=status_code, detail=detail)
async def _rank_agents_by_query(query: str, agents: Sequence[AgentResponse], top_k: int) -> tuple[AgentResponse, ...]:
from litellm.proxy.proxy_server import llm_router
outcome: Final = await search_agents(
query=query,
agents=agents,
top_k=top_k,
router=llm_router,
embedding_model=litellm.agent_search_embedding_model,
index=global_agent_search_index,
)
match outcome:
case AgentSearchHits(hits):
return tuple(hit.agent.model_copy(update=MappingProxyType({"search_score": hit.score})) for hit in hits)
case AgentSearchNotConfigured(reason):
raise _agent_search_error(400, "agent_search_not_configured", reason)
case AgentSearchEmbeddingFailed(reason):
raise _agent_search_error(503, "agent_search_unavailable", reason)
case _:
assert_never(outcome)
@router.get(
"/v1/agents",
tags=["[beta] A2A Agents"],
@ -223,6 +265,17 @@ async def get_agents(
False,
description="When true, performs a GET request to each agent's URL. Agents with reachable URLs (HTTP status < 500) and agents without a URL are returned; unreachable agents are filtered out.",
),
query: Annotated[
str | None,
Query(
min_length=1,
description="Describe the task in natural language to rank the agents you can reach by semantic similarity over their name, description, and skills. Each result carries a search_score. Requires litellm_settings.agent_search_embedding_model.",
),
] = None,
top_k: Annotated[
int,
Query(ge=1, le=100, description="With query: the maximum number of ranked agents to return."),
] = DEFAULT_AGENT_SEARCH_TOP_K,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # Used for auth
):
"""
@ -240,37 +293,22 @@ async def get_agents(
-H "Authorization: Bearer your-key" \
```
Pass `?query=<task>` to get the best matching agents ranked by semantic similarity:
```
curl -X GET "http://localhost:4000/v1/agents?query=translate+a+PDF+document&top_k=5" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer your-key" \
```
Returns: List[AgentResponse]
"""
await check_feature_access_for_user(user_api_key_dict, "agents")
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
AgentRequestHandler,
RestrictedAgentAccess,
UnrestrictedAgentAccess,
)
try:
returned_agents: Sequence[AgentResponse] = ()
# Admin users get all agents
if (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
):
returned_agents = global_agent_registry.get_agent_list()
else:
# Get allowed agents from object_permission (key/team level)
agent_access: Final = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict)
all_agents: Final = global_agent_registry.get_agent_list()
match agent_access:
case UnrestrictedAgentAccess():
returned_agents = all_agents
case RestrictedAgentAccess(allowed_agent_ids):
returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids]
returned_agents: Sequence[AgentResponse] = await accessible_agents(user_api_key_dict)
# Fetch current spend from DB for all returned agents
from litellm.proxy.proxy_server import prisma_client
@ -336,7 +374,9 @@ async def get_agents(
healthy_ids: Final = {result["agent_id"] for result in health_results if result["healthy"]}
returned_agents = [agent for agent in agents_with_url if agent.agent_id in healthy_ids] + agents_without_url
return returned_agents
if query is None:
return returned_agents
return await _rank_agents_by_query(query, returned_agents, top_k)
except HTTPException:
raise
except Exception as e:

1
litellm/types/agents.py
View file

@ -225,6 +225,7 @@ class AgentResponse(BaseModel):
static_headers: dict[str, str] | None = None
extra_headers: list[str] | None = None
keys: list[AgentKeySummary] | None = None
search_score: float | None = None
created_at: datetime | None = None
updated_at: datetime | None = None
created_by: str | None = None

128
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
View file

@ -18,6 +18,7 @@ import pytest
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing
from litellm.proxy._experimental.mcp_server.tool_search import (
AGENT_SEARCH_TOOL_NAME,
MCP_TOOL_CALL_TOOL_NAME,
MCP_TOOL_SEARCH_TOOL_NAME,
coerce_top_k,
@ -114,8 +115,15 @@ class TestSearchTools:
class TestGetVirtualToolDefinitions:
def test_returns_two_tools(self) -> None:
assert len(get_virtual_tool_definitions()) == 2
def test_returns_three_tools(self) -> None:
assert len(get_virtual_tool_definitions()) == 3
def test_agent_search_schema_requires_query(self) -> None:
tools = get_virtual_tool_definitions()
agent_tool = next(t for t in tools if t["name"] == AGENT_SEARCH_TOOL_NAME)
props = agent_tool["inputSchema"]["properties"]
assert set(props) == {"query", "top_k"}
assert agent_tool["inputSchema"]["required"] == ("query",)
def test_has_mcp_tool_search(self) -> None:
names = [t["name"] for t in get_virtual_tool_definitions()]
@ -130,7 +138,7 @@ class TestGetVirtualToolDefinitions:
search_tool = next(t for t in tools if t["name"] == MCP_TOOL_SEARCH_TOOL_NAME)
props = search_tool["inputSchema"]["properties"]
assert "query" in props
assert search_tool["inputSchema"]["required"] == ["query"]
assert search_tool["inputSchema"]["required"] == ("query",)
def test_mcp_tool_call_schema_has_tool_name_and_arguments(self) -> None:
tools = get_virtual_tool_definitions()
@ -153,6 +161,7 @@ class TestGetVirtualToolDefinitions:
assert {t.name for t in built} == {
MCP_TOOL_SEARCH_TOOL_NAME,
MCP_TOOL_CALL_TOOL_NAME,
AGENT_SEARCH_TOOL_NAME,
}
@ -187,7 +196,7 @@ class TestListToolRestApiWithToolSearch:
assert result["error"] is None
tool_names = [t["name"] for t in result["tools"]]
assert set(tool_names) == {MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME}
assert set(tool_names) == {MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME}
@pytest.mark.asyncio
async def test_returns_full_catalog_when_flag_disabled(self) -> None:
@ -517,6 +526,81 @@ class TestCallToolRestApiVirtualTools:
mock_list.assert_awaited_once()
assert mock_list.await_args.kwargs["client_ip"] == "203.0.113.7"
@pytest.mark.asyncio
async def test_agent_search_call_ranks_accessible_agents(self) -> None:
from litellm.proxy.agent_endpoints.agent_search import AgentSearchHit, AgentSearchHits
from litellm.types.agents import AgentResponse
user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
request = self._make_request(
{"name": AGENT_SEARCH_TOOL_NAME, "arguments": {"query": "translate a document", "top_k": "1"}}
)
translator = AgentResponse(
agent_id="translator",
agent_name="document-translator",
agent_card_params={"description": "Translates files", "skills": [{"id": "t", "name": "Translate"}]},
)
with (
patch( # test-quality-ok: the tool resolves agent access through proxy_server globals, no injection seam
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.accessible_agents",
new_callable=AsyncMock,
return_value=(translator,),
),
patch( # test-quality-ok: the embedding router only resolves via proxy_server globals, no injection seam
"litellm.proxy.agent_endpoints.agent_search.search_agents",
new_callable=AsyncMock,
return_value=AgentSearchHits(hits=(AgentSearchHit(agent=translator, score=0.91),)),
) 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 json.loads(result.content[0].text) == [
{
"agent_id": "translator",
"agent_name": "document-translator",
"description": "Translates files",
"skills": [{"name": "Translate", "description": "", "tags": []}],
"score": 0.91,
}
]
assert mock_search.await_args.kwargs["query"] == "translate a document"
assert mock_search.await_args.kwargs["top_k"] == 1
assert mock_search.await_args.kwargs["agents"] == (translator,)
@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
user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
request = self._make_request({"name": AGENT_SEARCH_TOOL_NAME, "arguments": {"query": "anything"}})
with (
patch( # test-quality-ok: the tool resolves agent access through proxy_server globals, no injection seam
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.accessible_agents",
new_callable=AsyncMock,
return_value=(),
),
patch( # test-quality-ok: the embedding router only resolves via proxy_server globals, no injection seam
"litellm.proxy.agent_endpoints.agent_search.search_agents",
new_callable=AsyncMock,
return_value=AgentSearchNotConfigured(reason="set agent_search_embedding_model"),
),
):
result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict)
assert result.isError is True
assert result.content[0].text == "set agent_search_embedding_model"
@pytest.mark.asyncio
async def test_agent_search_requires_flag_enabled(self) -> None:
from fastapi import HTTPException
user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=False))
request = self._make_request({"name": AGENT_SEARCH_TOOL_NAME, "arguments": {"query": "anything"}})
with pytest.raises(HTTPException) as exc_info:
await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_mcp_tool_search_requires_flag_enabled(self) -> None:
from fastapi import HTTPException
@ -592,6 +676,41 @@ class TestDispatchVirtualMcpTool:
assert mock_search.await_args.kwargs["query"] == "q"
assert mock_search.await_args.kwargs["top_k"] == 3
@pytest.mark.asyncio
async def test_routes_agent_search_to_its_handler(self) -> None:
from litellm.proxy._experimental.mcp_server import server as srv
uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
with patch( # test-quality-ok: dispatch routing is the subject; the handler is faked like its siblings here
"litellm.proxy._experimental.mcp_server.tool_search.handle_agent_search",
new_callable=AsyncMock,
return_value="AGENT_RESULT",
) as mock_agent_search:
result = await srv._dispatch_virtual_mcp_tool(
name=AGENT_SEARCH_TOOL_NAME,
arguments={"query": "translate a document", "top_k": "2"},
user_api_key_auth=uak,
client_ip=None,
)
assert result == "AGENT_RESULT"
assert mock_agent_search.await_args.kwargs == {
"query": "translate a document",
"top_k": 2,
"user_api_key_dict": uak,
}
@pytest.mark.asyncio
async def test_agent_search_rejected_when_flag_disabled(self) -> None:
from litellm.proxy._experimental.mcp_server.server import _dispatch_virtual_mcp_tool
uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=False))
result = await _dispatch_virtual_mcp_tool(
name=AGENT_SEARCH_TOOL_NAME, arguments={"query": "x"}, user_api_key_auth=uak, client_ip=None
)
assert result is not None
assert result.isError is True
@pytest.mark.asyncio
async def test_routes_call_with_client_ip(self) -> None:
from litellm.proxy._experimental.mcp_server import server as srv
@ -850,6 +969,7 @@ class TestHandleListToolsVirtual:
assert {t.name for t in tools} == {
MCP_TOOL_SEARCH_TOOL_NAME,
MCP_TOOL_CALL_TOOL_NAME,
AGENT_SEARCH_TOOL_NAME,
}

281
tests/test_litellm/proxy/agent_endpoints/test_agent_search.py Normal file
View file

@ -0,0 +1,281 @@
from collections.abc import Sequence
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from openai import APIConnectionError
import litellm
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.agent_endpoints.agent_search import (
AgentSearchEmbeddingFailed,
AgentSearchHits,
AgentSearchIndex,
AgentSearchNotConfigured,
Vector,
agent_search_text,
cosine_similarity,
search_agents,
)
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import RestrictedAgentAccess
from litellm.proxy.agent_endpoints.endpoints import router, user_api_key_auth
from litellm.types.agents import AgentResponse
TRANSLATOR: Final = AgentResponse(
agent_id="translator",
agent_name="document-translator",
agent_card_params={
"name": "Document Translator",
"description": "Converts files from one language into another",
"skills": [
{
"id": "t",
"name": "Translate a file",
"description": "Produce the document in the target language",
"tags": ["localization", "documents"],
}
],
},
)
SQL_ANALYST: Final = AgentResponse(
agent_id="sql",
agent_name="warehouse-sql-analyst",
agent_card_params={
"name": "Warehouse SQL Analyst",
"description": "Runs SQL against the inventory database",
"skills": [],
},
)
TRIP_PLANNER: Final = AgentResponse(
agent_id="trip",
agent_name="trip-planner",
agent_card_params={"name": "Trip Planner", "description": "Books flights and hotels"},
)
AGENTS: Final = (TRANSLATOR, SQL_ANALYST, TRIP_PLANNER)
VECTORS: Final = MappingProxyType(
{
"language translation": (1.0, 0.0, 0.0),
agent_search_text(TRANSLATOR): (0.9, 0.1, 0.0),
agent_search_text(SQL_ANALYST): (0.0, 1.0, 0.0),
agent_search_text(TRIP_PLANNER): (0.3, 0.0, 1.0),
}
)
class FakeEmbedder:
def __init__(self) -> None:
self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs
async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
self.calls.append(tuple(texts))
return tuple(VECTORS[text] for text in texts)
class TestAgentSearchText:
def test_joins_name_description_and_skills_with_tags(self) -> None:
assert agent_search_text(TRANSLATOR) == (
"document-translator\n"
"Converts files from one language into another\n"
"Translate a file Produce the document in the target language localization documents"
)
def test_missing_card_fields_fall_back_to_the_name(self) -> None:
assert agent_search_text(AgentResponse(agent_id="x", agent_name="bare", agent_card_params={})) == "bare"
def test_malformed_skills_do_not_break_the_text(self) -> None:
agent = AgentResponse(
agent_id="x", agent_name="odd", agent_card_params={"skills": "not-a-list", "description": "d"}
)
assert agent_search_text(agent) == "odd"
class TestCosineSimilarity:
def test_identical_direction_scores_one(self) -> None:
assert cosine_similarity((2.0, 0.0), (1.0, 0.0)) == pytest.approx(1.0)
def test_orthogonal_scores_zero(self) -> None:
assert cosine_similarity((1.0, 0.0), (0.0, 1.0)) == pytest.approx(0.0)
def test_zero_vector_scores_zero_instead_of_dividing(self) -> None:
assert cosine_similarity((0.0, 0.0), (1.0, 0.0)) == 0.0
class TestAgentSearchIndex:
@pytest.mark.asyncio
async def test_ranks_by_similarity_and_truncates_to_top_k(self) -> None:
outcome = await AgentSearchIndex().search("language translation", AGENTS, top_k=2, embed=FakeEmbedder())
assert isinstance(outcome, AgentSearchHits)
assert [hit.agent.agent_id for hit in outcome.hits] == ["translator", "trip"]
assert outcome.hits[0].score > outcome.hits[1].score
@pytest.mark.asyncio
async def test_second_search_only_embeds_the_query(self) -> None:
index = AgentSearchIndex()
embedder = FakeEmbedder()
await index.search("language translation", AGENTS, top_k=5, embed=embedder)
await index.search("language translation", AGENTS, top_k=5, embed=embedder)
assert len(embedder.calls[0]) == 1 + len(AGENTS)
assert embedder.calls[1] == ("language translation",)
@pytest.mark.asyncio
async def test_empty_registry_returns_no_hits_without_embedding(self) -> None:
embedder = FakeEmbedder()
outcome = await AgentSearchIndex().search("anything", (), top_k=5, embed=embedder)
assert outcome == AgentSearchHits(hits=())
assert embedder.calls == []
@pytest.mark.asyncio
async def test_provider_error_becomes_embedding_failed(self) -> None:
async def failing(texts: Sequence[str]) -> Sequence[Vector]:
raise APIConnectionError(request=MagicMock())
outcome = await AgentSearchIndex().search("q", AGENTS, top_k=5, embed=failing)
assert isinstance(outcome, AgentSearchEmbeddingFailed)
assert "embedding the search query failed" in outcome.reason
@pytest.mark.asyncio
async def test_wrong_vector_count_becomes_embedding_failed(self) -> None:
async def short(texts: Sequence[str]) -> Sequence[Vector]:
return ((1.0, 0.0, 0.0),)
outcome = await AgentSearchIndex().search("q", AGENTS, top_k=5, embed=short)
assert isinstance(outcome, AgentSearchEmbeddingFailed)
class TestSearchAgents:
@pytest.mark.asyncio
async def test_no_embedding_model_is_not_configured(self) -> None:
outcome = await search_agents(
"q", AGENTS, 5, router=MagicMock(), embedding_model=None, index=AgentSearchIndex()
)
assert isinstance(outcome, AgentSearchNotConfigured)
assert "agent_search_embedding_model" in outcome.reason
@pytest.mark.asyncio
async def test_no_router_is_not_configured(self) -> None:
outcome = await search_agents("q", AGENTS, 5, router=None, embedding_model="m", index=AgentSearchIndex())
assert isinstance(outcome, AgentSearchNotConfigured)
@pytest.mark.asyncio
async def test_router_embeddings_are_read_from_the_response(self) -> None:
router = MagicMock()
router.aembedding = AsyncMock(
side_effect=lambda model, input: litellm.EmbeddingResponse(
model=model,
data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)],
)
)
outcome = await search_agents(
"language translation",
AGENTS,
1,
router=router,
embedding_model="text-embedding-3-small",
index=AgentSearchIndex(),
)
assert isinstance(outcome, AgentSearchHits)
assert [hit.agent.agent_id for hit in outcome.hits] == ["translator"]
assert router.aembedding.await_args.kwargs["model"] == "text-embedding-3-small"
def _client(role: LitellmUserRoles) -> TestClient:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="u", user_role=role)
return TestClient(app)
@pytest.fixture
def registry(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
from litellm.proxy.agent_endpoints import agent_registry as registry_module
mock_registry = MagicMock()
mock_registry.get_agent_list = MagicMock(return_value=AGENTS)
mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id}))
monkeypatch.setattr(registry_module, "global_agent_registry", mock_registry)
monkeypatch.setattr("litellm.proxy.agent_endpoints.endpoints.global_agent_search_index", AgentSearchIndex())
return mock_registry
@pytest.fixture
def embedding_router(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
router = MagicMock()
router.aembedding = AsyncMock(
side_effect=lambda model, input: litellm.EmbeddingResponse(
model=model,
data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)],
)
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
monkeypatch.setattr(litellm, "agent_search_embedding_model", "text-embedding-3-small")
return router
@pytest.fixture
def no_db(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
class TestGetAgentsQuery:
def test_query_ranks_and_scores_and_truncates(
self, registry: MagicMock, embedding_router: MagicMock, no_db: None
) -> None:
response = _client(LitellmUserRoles.PROXY_ADMIN).get(
"/v1/agents", params={"query": "language translation", "top_k": 2}, headers={"Authorization": "Bearer k"}
)
assert response.status_code == 200
body = response.json()
assert [agent["agent_id"] for agent in body] == ["translator", "trip"]
assert body[0]["search_score"] > body[1]["search_score"]
def test_without_query_the_list_is_unchanged_and_unscored(
self, registry: MagicMock, embedding_router: MagicMock, no_db: None
) -> None:
response = _client(LitellmUserRoles.PROXY_ADMIN).get("/v1/agents", headers={"Authorization": "Bearer k"})
assert response.status_code == 200
assert [agent["agent_id"] for agent in response.json()] == ["translator", "sql", "trip"]
assert all(agent["search_score"] is None for agent in response.json())
embedding_router.aembedding.assert_not_awaited()
def test_restricted_key_only_ranks_its_own_agents(
self, registry: MagicMock, embedding_router: MagicMock, no_db: None, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access",
AsyncMock(return_value=RestrictedAgentAccess(frozenset({"sql"}))),
)
response = _client(LitellmUserRoles.INTERNAL_USER).get(
"/v1/agents", params={"query": "language translation"}, headers={"Authorization": "Bearer k"}
)
assert response.status_code == 200
assert [agent["agent_id"] for agent in response.json()] == ["sql"]
def test_missing_embedding_model_is_a_400(
self, registry: MagicMock, embedding_router: MagicMock, no_db: None, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "agent_search_embedding_model", None)
response = _client(LitellmUserRoles.PROXY_ADMIN).get(
"/v1/agents", params={"query": "anything"}, headers={"Authorization": "Bearer k"}
)
assert response.status_code == 400
assert response.json()["detail"]["error"] == "agent_search_not_configured"
def test_embedding_provider_failure_is_a_503(
self, registry: MagicMock, embedding_router: MagicMock, no_db: None
) -> None:
embedding_router.aembedding = AsyncMock(side_effect=APIConnectionError(request=MagicMock()))
response = _client(LitellmUserRoles.PROXY_ADMIN).get(
"/v1/agents", params={"query": "anything"}, headers={"Authorization": "Bearer k"}
)
assert response.status_code == 503
assert response.json()["detail"]["error"] == "agent_search_unavailable"
def test_top_k_is_validated(self, registry: MagicMock, embedding_router: MagicMock, no_db: None) -> None:
response = _client(LitellmUserRoles.PROXY_ADMIN).get(
"/v1/agents", params={"query": "anything", "top_k": 0}, headers={"Authorization": "Bearer k"}
)
assert response.status_code == 422

12
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -16575,6 +16575,10 @@ export interface paths {
* ```
* curl -X GET "http://localhost:4000/v1/agents?health_check=true" -H "Content-Type: application/json" -H "Authorization: Bearer your-key" ```
*
* Pass `?query=<task>` to get the best matching agents ranked by semantic similarity:
* ```
* curl -X GET "http://localhost:4000/v1/agents?query=translate+a+PDF+document&top_k=5" -H "Content-Type: application/json" -H "Authorization: Bearer your-key" ```
*
* Returns: List[AgentResponse]
*/
get: operations["get_agents_v1_agents_get"];
@ -22487,6 +22491,8 @@ export interface components {
} | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Search Score */
search_score?: number | null;
/** Session Rpm Limit */
session_rpm_limit?: number | null;
/** Session Tpm Limit */
@ -30208,6 +30214,8 @@ export interface components {
token_exchange_profile?: string | null;
/** Upstream Resource */
upstream_resource?: string | null;
/** Upstream Token Header */
upstream_token_header?: string | null;
};
/**
* MCPEnvVar
@ -58515,6 +58523,10 @@ export interface operations {
query?: {
/** @description When true, performs a GET request to each agent's URL. Agents with reachable URLs (HTTP status < 500) and agents without a URL are returned; unreachable agents are filtered out. */
health_check?: boolean;
/** @description Describe the task in natural language to rank the agents you can reach by semantic similarity over their name, description, and skills. Each result carries a search_score. Requires litellm_settings.agent_search_embedding_model. */
query?: string | null;
/** @description With query: the maximum number of ranked agents to return. */
top_k?: number;
};
header?: never;
path?: never;