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

Merge 4026aa6575 into 3300fc3a96

Browse source
This commit is contained in:
Mateo Wang 2026-08-28 03:58:19 +00:00 • committed by GitHub
commit 187301d88b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 1042 additions and 76 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,

156
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, Sequence
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,116 @@ 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[Sequence[str]]
class VirtualToolDefinition(TypedDict):
name: ReadOnly[str]
description: ReadOnly[str]
inputSchema: ReadOnly[_ToolInputSchema]
def _json_array(*items: str) -> Sequence[str]:
return list(items) # mutable-ok: jsonschema's metaschema only accepts a JSON array for required
_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": _json_array("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": _json_array("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": _json_array("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,
user_api_key_dict=user_api_key_dict,
)
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(

46
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": {

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

@ -0,0 +1,250 @@
"""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.proxy._types import UserAPIKeyAuth
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 embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]:
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
return { # mutable-ok: the router mutates the metadata dict it is handed
**LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict),
"user_api_key": user_api_key_dict.api_key,
}
def router_embedder(router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth) -> 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, metadata=embedding_spend_metadata(user_api_key_dict)
)
return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data)
return embed
_NO_VECTORS: Final[Mapping[str, Vector]] = MappingProxyType({})
async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | AgentSearchEmbeddingFailed:
try:
vectors: Final = tuple(await embed(texts))
except (OpenAIError, ValueError, BudgetExceededError) as exc:
return AgentSearchEmbeddingFailed(reason=f"embedding the search query failed: {exc}")
if len(vectors) != len(texts):
return AgentSearchEmbeddingFailed(
reason=f"embedding model returned {len(vectors)} vectors for {len(texts)} inputs"
)
return vectors
@dataclass(frozen=True, slots=True)
class _Embedded:
query_vector: Vector
vectors: Mapping[str, Vector]
def _same_dimension(query_vector: Vector, vectors: Mapping[str, Vector], texts: Sequence[str]) -> bool:
return all(len(vectors[text]) == len(query_vector) for text in texts)
async def _embed_query_and_agents(
embed: Embedder, query: str, texts: Sequence[str], cached: Mapping[str, Vector]
) -> _Embedded | AgentSearchEmbeddingFailed:
missing: Final = tuple(dict.fromkeys(text for text in texts if text not in cached))
embedded: Final = await _embed_all(embed, (query, *missing))
if isinstance(embedded, AgentSearchEmbeddingFailed):
return embedded
vectors: Final = MappingProxyType(dict(chain(cached.items(), zip(missing, embedded[1:], strict=True))))
if _same_dimension(embedded[0], vectors, texts):
return _Embedded(query_vector=embedded[0], vectors=vectors)
unique: Final = tuple(dict.fromkeys(texts))
reembedded: Final = await _embed_all(embed, (query, *unique))
if isinstance(reembedded, AgentSearchEmbeddingFailed):
return reembedded
return _Embedded(
query_vector=reembedded[0], vectors=MappingProxyType(dict(zip(unique, reembedded[1:], strict=True)))
)
class AgentSearchIndex:
"""Caches one vector per distinct agent text per embedding model, so repeat searches only embed the query."""
def __init__(self) -> None:
self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({})
def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]:
kept: Final = {
text: vector
for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items()
if len(vector) == len(embedded.query_vector)
}
return MappingProxyType({**kept, **embedded.vectors})
async def search(
self, query: str, agents: Sequence[AgentResponse], top_k: int, embed: Embedder, embedding_model: str
) -> AgentSearchHits | AgentSearchEmbeddingFailed:
if not agents:
return AgentSearchHits(hits=())
texts: Final = tuple(agent_search_text(agent) for agent in agents)
cached: Final = self._vectors.get(embedding_model, _NO_VECTORS)
embedded: Final = await _embed_query_and_agents(embed, query, texts, cached)
if isinstance(embedded, AgentSearchEmbeddingFailed):
return embedded
if not _same_dimension(embedded.query_vector, embedded.vectors, texts):
return AgentSearchEmbeddingFailed(
reason=f"embedding model {embedding_model} returned vectors of mixed dimensions"
)
self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)})
ranked: Final = sorted(
(
AgentSearchHit(agent=agent, score=cosine_similarity(embedded.query_vector, embedded.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,
user_api_key_dict: UserAPIKeyAuth,
) -> 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, user_api_key_dict), 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)

95
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,41 @@ 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, user_api_key_dict: UserAPIKeyAuth
) -> 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,
user_api_key_dict=user_api_key_dict,
)
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 +268,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 +296,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 +377,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, user_api_key_dict)
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

138
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()]
@ -140,6 +148,17 @@ class TestGetVirtualToolDefinitions:
assert "arguments" in props
assert "tool_name" in call_tool["inputSchema"]["required"]
def test_input_schemas_validate_arguments_like_the_mcp_server_does(self) -> None:
from jsonschema import ValidationError, validate
from mcp.types import Tool
for definition in get_virtual_tool_definitions():
tool = Tool.model_validate(definition)
required_arguments = {name: "x" for name in tool.inputSchema["required"]}
validate(instance=required_arguments, schema=tool.inputSchema)
with pytest.raises(ValidationError):
validate(instance={}, schema=tool.inputSchema)
def test_all_tools_have_description(self) -> None:
for tool in get_virtual_tool_definitions():
assert tool.get("description"), f"{tool['name']} missing description"
@ -153,6 +172,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 +207,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 +537,82 @@ 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 mock_search.await_args.kwargs["user_api_key_dict"] is user_api_key_dict
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 +688,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 +981,7 @@ class TestHandleListToolsVirtual:
assert {t.name for t in tools} == {
MCP_TOOL_SEARCH_TOOL_NAME,
MCP_TOOL_CALL_TOOL_NAME,
AGENT_SEARCH_TOOL_NAME,
}

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

@ -0,0 +1,374 @@
import asyncio
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
CALLER: Final = UserAPIKeyAuth(api_key="hashed-caller-key", team_id="team-1", user_id="user-1")
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 FixedDimensionEmbedder:
def __init__(self, dimensions: int) -> None:
self.dimensions: Final = dimensions
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))
await asyncio.sleep(0)
return tuple((1.0,) * self.dimensions for _ 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(), embedding_model="m"
)
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, embedding_model="m")
await index.search("language translation", AGENTS, top_k=5, embed=embedder, embedding_model="m")
assert len(embedder.calls[0]) == 1 + len(AGENTS)
assert embedder.calls[1] == ("language translation",)
@pytest.mark.asyncio
async def test_switching_embedding_models_does_not_reuse_cached_vectors(self) -> None:
index = AgentSearchIndex()
await index.search("language translation", AGENTS, top_k=5, embed=FakeEmbedder(), embedding_model="small")
wide = FixedDimensionEmbedder(2)
outcome = await index.search("language translation", AGENTS, top_k=5, embed=wide, embedding_model="wide")
assert isinstance(outcome, AgentSearchHits)
assert len(wide.calls[0]) == 1 + len(AGENTS)
@pytest.mark.asyncio
async def test_cached_vectors_of_another_dimension_are_re_embedded(self) -> None:
index = AgentSearchIndex()
await index.search("language translation", AGENTS, top_k=5, embed=FakeEmbedder(), embedding_model="m")
fallback = FixedDimensionEmbedder(2)
outcome = await index.search("language translation", AGENTS, top_k=5, embed=fallback, embedding_model="m")
assert isinstance(outcome, AgentSearchHits)
assert fallback.calls == [
("language translation",),
("language translation", *(agent_search_text(agent) for agent in AGENTS)),
]
@pytest.mark.asyncio
async def test_re_embedding_a_subset_drops_the_other_agents_old_vectors(self) -> None:
index = AgentSearchIndex()
await index.search("language translation", AGENTS, top_k=5, embed=FakeEmbedder(), embedding_model="m")
wide = FixedDimensionEmbedder(2)
await index.search("language translation", AGENTS[:1], top_k=5, embed=wide, embedding_model="m")
await index.search("language translation", AGENTS, top_k=5, embed=wide, embedding_model="m")
assert wide.calls[-1] == ("language translation", *(agent_search_text(agent) for agent in AGENTS[1:]))
@pytest.mark.asyncio
async def test_concurrent_searches_keep_each_others_vectors(self) -> None:
index = AgentSearchIndex()
embedder = FixedDimensionEmbedder(3)
await asyncio.gather(
index.search("q", AGENTS[:1], top_k=5, embed=embedder, embedding_model="m"),
index.search("q", AGENTS[1:], top_k=5, embed=embedder, embedding_model="m"),
)
await index.search("q", AGENTS, top_k=5, embed=embedder, embedding_model="m")
assert embedder.calls[-1] == ("q",)
@pytest.mark.asyncio
async def test_mixed_dimensions_in_one_batch_become_embedding_failed(self) -> None:
async def mixed(texts: Sequence[str]) -> Sequence[Vector]:
return ((1.0, 0.0), *((1.0, 0.0, 0.0) for _ in texts[1:]))
outcome = await AgentSearchIndex().search("q", AGENTS, top_k=5, embed=mixed, embedding_model="m")
assert isinstance(outcome, AgentSearchEmbeddingFailed)
assert "mixed dimensions" in outcome.reason
@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, embedding_model="m")
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, embedding_model="m")
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, embedding_model="m")
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(), user_api_key_dict=CALLER
)
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(), user_api_key_dict=CALLER
)
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, metadata: 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(),
user_api_key_dict=CALLER,
)
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"
@pytest.mark.asyncio
async def test_embedding_spend_is_attributed_to_the_calling_key(self) -> None:
router = MagicMock()
router.aembedding = AsyncMock(
side_effect=lambda model, input, metadata: litellm.EmbeddingResponse(
model=model,
data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)],
)
)
await search_agents(
"language translation",
AGENTS,
1,
router=router,
embedding_model="text-embedding-3-small",
index=AgentSearchIndex(),
user_api_key_dict=CALLER,
)
metadata = router.aembedding.await_args.kwargs["metadata"]
assert metadata["user_api_key"] == "hashed-caller-key"
assert metadata["user_api_key_team_id"] == "team-1"
assert metadata["user_api_key_user_id"] == "user-1"
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, metadata: 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"]
assert embedding_router.aembedding.await_args.kwargs["metadata"]["user_api_key_user_id"] == "u"
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

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

@ -16586,6 +16586,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"];
@ -22498,6 +22502,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 */
@ -58582,6 +58588,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;