diff --git a/litellm/__init__.py b/litellm/__init__.py index ec2960c196e..179cac47663 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3a8fd6de5e5..3efb6429326 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index c6b2ac489bb..989b08b929a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 3b0dd2071ae..f79765f6d01 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -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( diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 4ccfde18d36..582563842cf 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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=` 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": { diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py new file mode 100644 index 00000000000..46ab36d7b72 --- /dev/null +++ b/litellm/proxy/agent_endpoints/agent_search.py @@ -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 + ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 81586b4eef3..d0ac94d3710 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -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) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index d348bc01153..b6c41a17503 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -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=` 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: diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 7e499dde642..c85507b77c1 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 0e442102e53..16221f44efe 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -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, } diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py new file mode 100644 index 00000000000..3fb09076e5f --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py @@ -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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e80f9fdbefe..3d85a9eec9e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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=` 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;