From 192e38fa7ba2f529cfaad3bcfc28d2613aca28d5 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 7 Sep 2026 12:28:38 -0700 Subject: [PATCH] feat(skills): semantic search over the LiteLLM-hosted skill registry (#39401) * feat(skills): semantic search over the LiteLLM-hosted skill registry Adds GET /v1/skills?query= (custom_llm_provider=litellm_proxy) and a skill_search MCP virtual tool, ranking the caller's accessible skills by semantic similarity, mirroring the A2A agent registry search (LIT-6309). Also fixes a pre-existing bug where create_skill() dropped description and instructions for the litellm_proxy provider, which left every LiteLLM-hosted skill with no searchable text. * fix(mcp): coerce skill_search top_k instead of raising 500 on malformed input The MCP-REST skill_search dispatch validated raw tool arguments through a pydantic model directly, so a non-numeric top_k raised a ValidationError that the endpoint's catch-all turned into an HTTP 500. Mirrors the agent_search branch's tolerant coerce_top_k handling instead. * fix(skills): enforce key limits on search embeddings and bound the semantic index Semantic search embeddings now run the same pre_call_hook the /embeddings route runs, so key rate limits, budgets and guardrails apply before the embedding model is called. The shared SemanticTextIndex caps cached vectors and evicts the least recently searched entries, and each skill's embedded text is capped so one skill cannot inflate the embedding batch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(skills): surface proxy 429s from search embeddings instead of a 503 ProxyRateLimitError is also an OpenAIError, so the search engine was folding a key rate limit into skill_search_unavailable. Proxy HTTPExceptions now propagate so the caller gets the same 429 the /embeddings route returns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(skills): import assert_never from typing_extensions for Python 3.10 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(skills): embed the request as the pre-call hooks returned it, not the original text Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(skills): keep the litellm_proxy provider check for GET /v1/skills?query= inside llms/ Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(skills): move the GET /v1/skills?query= endpoint tests under tests/test_litellm/proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + litellm/__init__.py | 1 + .../llms/litellm_proxy/skills/constants.py | 5 + litellm/llms/litellm_proxy/skills/handler.py | 19 +- .../llms/litellm_proxy/skills/skill_search.py | 161 +++++++ .../litellm_proxy/skills/transformation.py | 11 +- .../mcp_server/rest_endpoints.py | 13 + .../proxy/_experimental/mcp_server/server.py | 9 + .../_experimental/mcp_server/tool_search.py | 68 ++- litellm/proxy/_lazy_openapi_snapshot.json | 57 ++- litellm/proxy/agent_endpoints/agent_search.py | 7 +- litellm/proxy/agent_endpoints/endpoints.py | 3 +- .../anthropic_endpoints/skills_endpoints.py | 97 +++- .../proxy/common_utils/semantic_text_index.py | 79 +++- litellm/skills/main.py | 6 +- litellm/types/llms/anthropic_skills.py | 9 + .../litellm_proxy/skills/test_skill_search.py | 436 ++++++++++++++++++ .../mcp_server/test_mcp_tool_search.py | 50 +- .../agent_endpoints/test_agent_search.py | 27 +- .../test_skills_endpoints.py | 175 +++++++ tests/test_litellm/skills/test_skills_main.py | 57 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 14 + 22 files changed, 1259 insertions(+), 46 deletions(-) create mode 100644 litellm/llms/litellm_proxy/skills/skill_search.py create mode 100644 tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py create mode 100644 tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py create mode 100644 tests/test_litellm/skills/test_skills_main.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 33245ec5b5f..cc606339a20 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -116,6 +116,7 @@ jobs: tests/test_litellm/rerank_api tests/test_litellm/rust_bridge tests/test_litellm/sandbox + tests/test_litellm/skills tests/test_litellm/test_router tests/test_litellm/vector_stores tests/test_litellm/videos diff --git a/litellm/__init__.py b/litellm/__init__.py index 42c0ea881fd..dc2f40af46e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -495,6 +495,7 @@ public_model_groups: Optional[List[str]] = None public_agent_groups: Optional[List[str]] = None agent_search_embedding_model: Optional[str] = None mcp_tool_search: Optional[Mapping[str, object]] = None +skill_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/llms/litellm_proxy/skills/constants.py b/litellm/llms/litellm_proxy/skills/constants.py index a6c88718f11..04a3a7dbc91 100644 --- a/litellm/llms/litellm_proxy/skills/constants.py +++ b/litellm/llms/litellm_proxy/skills/constants.py @@ -16,3 +16,8 @@ DEFAULT_MAX_ITERATIONS: Final[int] = 10 DEFAULT_SANDBOX_TIMEOUT: Final[int] = 120 """Default timeout in seconds for sandbox code execution.""" + +MAX_SKILLS_PER_SEARCH: Final[int] = 5000 +"""Upper bound on how many of the caller's accessible skills a single semantic +search embeds. Ranking runs in memory over this candidate set (no tsvector/DB-side +filtering yet), so this caps worst-case embedding cost per search request.""" diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 73f6ed23092..9b625cb0571 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -6,11 +6,15 @@ Used by the transformation layer and skills injection hook. """ import uuid +from collections.abc import Sequence from typing import Final from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache -from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX +from litellm.llms.litellm_proxy.skills.constants import ( + LITELLM_SKILL_ID_PREFIX, + MAX_SKILLS_PER_SEARCH, +) from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, @@ -131,6 +135,19 @@ class LiteLLMSkillsHandler: ) return [_prisma_skill_to_litellm(s) for s in skills] + @staticmethod + async def list_skills_for_search( + user_api_key_dict: UserAPIKeyAuth | None = None, + ) -> Sequence[LiteLLM_SkillsTable]: + """Every skill the caller can access, for ranking. Same owner-scope filter as + ``list_skills``, but unpaginated (up to ``MAX_SKILLS_PER_SEARCH``) since a query + must be scored against the whole accessible set, not one page of it.""" + return await LiteLLMSkillsHandler.list_skills( + limit=MAX_SKILLS_PER_SEARCH, + offset=0, + user_api_key_dict=user_api_key_dict, + ) + @staticmethod async def _load_skill(skill_id: str) -> object | None: """Cache-first read of the Prisma skill row. Owner-scope filtering diff --git a/litellm/llms/litellm_proxy/skills/skill_search.py b/litellm/llms/litellm_proxy/skills/skill_search.py new file mode 100644 index 00000000000..f975c6c4cab --- /dev/null +++ b/litellm/llms/litellm_proxy/skills/skill_search.py @@ -0,0 +1,161 @@ +"""Semantic ranking over the LiteLLM-hosted skill registry, shared by GET /v1/skills?query= and the skill_search MCP tool.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import BaseModel, ConfigDict + +from litellm.llms.litellm_proxy.skills.constants import MAX_SKILLS_PER_SEARCH +from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler +from litellm.proxy.common_utils.semantic_text_index import ( + Embedder, + EmbeddingFailed, + SemanticTextIndex, + router_embedder, +) +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + +DEFAULT_SKILL_SEARCH_TOP_K: Final = 5 +MAX_SKILL_SEARCH_TOP_K: Final = 100 +"""Matches the ``le=100`` bound GET /v1/skills?query= enforces via FastAPI's Query +validation, so the MCP tool can't return a larger payload than the REST endpoint allows.""" +MAX_SKILL_SEARCH_TEXT_CHARS: Final = 4000 +"""Per-skill cap on the title + description + instructions text that gets embedded, so one +search embeds at most ``MAX_SKILLS_PER_SEARCH * MAX_SKILL_SEARCH_TEXT_CHARS`` characters no +matter how long the stored instructions are.""" + + +@dataclass(frozen=True, slots=True) +class SkillSearchHit: + skill: LiteLLM_SkillsTable + score: float + + +@dataclass(frozen=True, slots=True) +class SkillSearchHits: + hits: tuple[SkillSearchHit, ...] + + +@dataclass(frozen=True, slots=True) +class SkillSearchNotConfigured: + reason: str + + +@dataclass(frozen=True, slots=True) +class SkillSearchEmbeddingFailed: + reason: str + + +@dataclass(frozen=True, slots=True) +class SkillSearchUnsupportedProvider: + reason: str + + +SkillSearchOutcome: TypeAlias = SkillSearchHits | SkillSearchNotConfigured | SkillSearchEmbeddingFailed +HostedSkillSearchOutcome: TypeAlias = SkillSearchOutcome | SkillSearchUnsupportedProvider + + +class SkillSearchResult(BaseModel): + model_config = ConfigDict(frozen=True) + + skill_id: str + display_title: str | None + description: str | None + score: float + + +def skill_search_text(skill: LiteLLM_SkillsTable) -> str: + joined: Final = "\n".join(part for part in (skill.display_title, skill.description, skill.instructions) if part) + return joined[:MAX_SKILL_SEARCH_TEXT_CHARS] + + +def skill_search_result(hit: SkillSearchHit) -> SkillSearchResult: + return SkillSearchResult( + skill_id=hit.skill.skill_id, + display_title=hit.skill.display_title, + description=hit.skill.description, + score=hit.score, + ) + + +class SkillSearchIndex: + """Caches one vector per distinct skill text per embedding model, so repeat searches only embed the query.""" + + def __init__(self, max_entries: int = MAX_SKILLS_PER_SEARCH) -> None: + self._index: Final = SemanticTextIndex(max_entries=max_entries) + + async def search( + self, + query: str, + skills: Sequence[LiteLLM_SkillsTable], + top_k: int, + embed: Embedder, + embedding_model: str, + ) -> SkillSearchHits | SkillSearchEmbeddingFailed: + texts: Final = tuple(skill_search_text(skill) for skill in skills) + scores: Final = await self._index.scores(query, texts, embed, embedding_model) + if isinstance(scores, EmbeddingFailed): + return SkillSearchEmbeddingFailed(reason=scores.reason) + ranked: Final = sorted( + (SkillSearchHit(skill=skill, score=score) for skill, score in zip(skills, scores, strict=True)), + key=lambda hit: hit.score, + reverse=True, + ) + return SkillSearchHits(hits=tuple(ranked[:top_k])) + + +global_skill_search_index: Final = SkillSearchIndex() + + +async def search_skills( + query: str, + skills: Sequence[LiteLLM_SkillsTable], + top_k: int, + router: Router | None, + embedding_model: str | None, + index: SkillSearchIndex, + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> SkillSearchOutcome: + if embedding_model is None: + return SkillSearchNotConfigured( + reason="skill search needs litellm_settings.skill_search_embedding_model set to an embedding model from model_list" + ) + if router is None: + return SkillSearchNotConfigured(reason="skill search needs a model_list so the embedding model can be called") + embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj) + return await index.search(query, skills, top_k, embed, embedding_model) + + +async def search_hosted_skills( + custom_llm_provider: str | None, + query: str, + top_k: int, + router: Router | None, + embedding_model: str | None, + index: SkillSearchIndex, + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> HostedSkillSearchOutcome: + """GET /v1/skills?query= for the skills LiteLLM hosts itself: only ``litellm_proxy`` has a registry to rank.""" + if custom_llm_provider != LlmProviders.LITELLM_PROXY.value: + return SkillSearchUnsupportedProvider(reason="query is only supported for custom_llm_provider=litellm_proxy") + skills: Final = await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict=user_api_key_dict) + return await search_skills( + query=query, + skills=skills, + top_k=top_k, + router=router, + embedding_model=embedding_model, + index=index, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) diff --git a/litellm/llms/litellm_proxy/skills/transformation.py b/litellm/llms/litellm_proxy/skills/transformation.py index c972dc349c9..9fc2d2cbb45 100644 --- a/litellm/llms/litellm_proxy/skills/transformation.py +++ b/litellm/llms/litellm_proxy/skills/transformation.py @@ -154,7 +154,7 @@ class LiteLLMSkillsTransformationHandler: user_api_key_dict=user_api_key_dict, ) - return self._db_skill_to_response(db_skill) + return self.db_skill_to_response(db_skill) def list_skills_handler( self, @@ -222,7 +222,9 @@ class LiteLLMSkillsTransformationHandler: user_api_key_dict=user_api_key_dict, ) - skills: Final = [self._db_skill_to_response(s) for s in db_skills] + skills: Final = [ # mutable-ok: ListSkillsResponse.data needs list[Skill]; never mutated after + self.db_skill_to_response(s) for s in db_skills + ] return ListSkillsResponse( data=skills, has_more=len(skills) >= limit, @@ -288,7 +290,7 @@ class LiteLLMSkillsTransformationHandler: skill_id=skill_id, user_api_key_dict=user_api_key_dict, ) - return self._db_skill_to_response(db_skill) + return self.db_skill_to_response(db_skill) def delete_skill_handler( self, @@ -354,7 +356,7 @@ class LiteLLMSkillsTransformationHandler: type=result.get("type", "skill_deleted"), ) - def _db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill: + def db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill: """ Convert a database skill record to Anthropic-compatible Skill response. @@ -375,4 +377,5 @@ class LiteLLMSkillsTransformationHandler: latest_version=db_skill.latest_version, source=db_skill.source or "custom", type="skill", + description=db_skill.description, ) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 5fbfad54a39..102129ffdd0 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -104,6 +104,9 @@ if MCP_AVAILABLE: from mcp.types import Tool as MCPTool from litellm.experimental_mcp_client.client import MCPClient + from litellm.llms.litellm_proxy.skills.skill_search import ( + DEFAULT_SKILL_SEARCH_TOP_K, + ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, global_mcp_server_manager, @@ -188,10 +191,12 @@ if MCP_AVAILABLE: AGENT_SEARCH_TOOL_NAME, DEFAULT_AGENT_SEARCH_TOP_K, MCP_TOOL_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, coerce_top_k, handle_agent_search, handle_mcp_tool_call, handle_mcp_tool_search, + handle_skill_search, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.proxy_server import general_settings, proxy_config, proxy_logging_obj @@ -210,6 +215,14 @@ if MCP_AVAILABLE: ), user_api_key_dict=user_api_key_dict, ) + if tool_name == SKILL_SEARCH_TOOL_NAME: + return await handle_skill_search( + query=str(tool_arguments.get("query", "")), + top_k=coerce_top_k( + tool_arguments.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K + ), + user_api_key_dict=user_api_key_dict, + ) rest_client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) ( virtual_mcp_auth_header, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 60a9af89cc3..0b424c31c4b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -911,15 +911,18 @@ if MCP_AVAILABLE: Returns a CallToolResult when ``name`` is a virtual tool, else ``None`` so the caller falls through to normal tool routing. """ + from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K from litellm.proxy._experimental.mcp_server.tool_search import ( AGENT_SEARCH_TOOL_NAME, DEFAULT_AGENT_SEARCH_TOP_K, MCP_TOOL_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, VIRTUAL_TOOL_NAMES, coerce_top_k, handle_agent_search, handle_mcp_tool_call, handle_mcp_tool_search, + handle_skill_search, ) if name not in VIRTUAL_TOOL_NAMES: @@ -961,6 +964,12 @@ if MCP_AVAILABLE: 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, ) + if name == SKILL_SEARCH_TOOL_NAME: + return await handle_skill_search( + query=str(args.get("query", "")), + top_k=coerce_top_k(args.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_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 af02c11ad86..f19340d30cb 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -11,6 +11,7 @@ from pydantic import ValidationError from typing_extensions import ReadOnly, Required, assert_never import litellm +from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K from litellm.proxy.common_utils.semantic_text_index import ( Embedder, @@ -30,7 +31,10 @@ MCP_TOOL_SEARCH_SETTINGS_KEY: Final[str] = "mcp_tool_search" 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)) +SKILL_SEARCH_TOOL_NAME: Final[str] = "skill_search" +VIRTUAL_TOOL_NAMES: Final = frozenset( + (MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, SKILL_SEARCH_TOOL_NAME) +) def coerce_top_k(value: Any, default: int = 5) -> int: @@ -199,8 +203,28 @@ _AGENT_SEARCH_DEFINITION: Final[VirtualToolDefinition] = { } +_SKILL_SEARCH_DEFINITION: Final[VirtualToolDefinition] = { + "name": SKILL_SEARCH_TOOL_NAME, + "description": "Find registered skills by describing what you need in natural language. Returns the best " + "matching skills you can access, ranked by semantic similarity, each with its skill_id, display_title, " + "description, and score.", + "inputSchema": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "What you need the skill to do, in natural language."}, + "top_k": { + "type": "integer", + "description": "Maximum number of skills to return.", + "default": DEFAULT_SKILL_SEARCH_TOP_K, + }, + }, + "required": _json_array("query"), + }, +} + + def get_virtual_tool_definitions() -> tuple[VirtualToolDefinition, ...]: - return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION) + return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION, _SKILL_SEARCH_DEFINITION) def _text_tool_result(text: str, is_error: bool) -> CallToolResult: @@ -223,7 +247,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI ) 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 + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj await check_feature_access_for_user(user_api_key_dict, "agents") outcome: Final = await search_agents( @@ -234,6 +258,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI embedding_model=litellm.agent_search_embedding_model, index=global_agent_search_index, user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) match outcome: case AgentSearchHits(hits): @@ -245,6 +270,39 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI assert_never(outcome) +async def handle_skill_search(query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth) -> CallToolResult: + from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler + from litellm.llms.litellm_proxy.skills.skill_search import ( + MAX_SKILL_SEARCH_TOP_K, + SkillSearchEmbeddingFailed, + SkillSearchHits, + SkillSearchNotConfigured, + global_skill_search_index, + search_skills, + skill_search_result, + ) + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj + + outcome: Final = await search_skills( + query=query, + skills=await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict), + top_k=min(max(top_k, 1), MAX_SKILL_SEARCH_TOP_K), + router=llm_router, + embedding_model=litellm.skill_search_embedding_model, + index=global_skill_search_index, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + match outcome: + case SkillSearchHits(hits): + results: Final = tuple(skill_search_result(hit).model_dump() for hit in hits) + return _text_tool_result(json.dumps(results), is_error=False) + case SkillSearchNotConfigured(reason) | SkillSearchEmbeddingFailed(reason): + return _text_tool_result(reason, is_error=True) + case _: + assert_never(outcome) + + async def handle_mcp_tool_search( query: str, top_k: int, @@ -257,7 +315,7 @@ async def handle_mcp_tool_search( raw_headers: dict[str, str] | None = None, ) -> CallToolResult: from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj settings: Final = mcp_tool_search_settings() if isinstance(settings, ValidationError): @@ -271,7 +329,7 @@ async def handle_mcp_tool_search( ) ranker: Final = ( SemanticToolRanker( - embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict), + embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict, proxy_logging_obj), embedding_model=settings.embedding_model, index=global_mcp_tool_search_index, ) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 91a97ad6544..4093f2c5248 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -4687,6 +4687,17 @@ "title": "Created At", "type": "string" }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, "display_title": { "anyOf": [ { @@ -4713,6 +4724,17 @@ ], "title": "Latest Version" }, + "search_score": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Search Score" + }, "source": { "title": "Source", "type": "string" @@ -4781,7 +4803,7 @@ "paths": { "/v1/skills": { "get": { - "description": "List skills on Anthropic.\n\nRequires `?beta=true` query parameter.\n\nModel-based routing (for multi-account support):\n- Pass model via header: `x-litellm-model: claude-account-1`\n- Pass model via query: `?model=claude-account-1`\n- Pass model via body: `{\"model\": \"claude-account-1\"}`\n\nExample usage:\n```bash\n# Basic usage\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\"\n\n# With model-based routing\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\" -H \"x-litellm-model: claude-account-1\"\n```\n\nReturns: ListSkillsResponse with list of skills", + "description": "List skills on Anthropic.\n\nRequires `?beta=true` query parameter.\n\nModel-based routing (for multi-account support):\n- Pass model via header: `x-litellm-model: claude-account-1`\n- Pass model via query: `?model=claude-account-1`\n- Pass model via body: `{\"model\": \"claude-account-1\"}`\n\nExample usage:\n```bash\n# Basic usage\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\"\n\n# With model-based routing\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\" -H \"x-litellm-model: claude-account-1\"\n```\n\nPass `?custom_llm_provider=litellm_proxy&query=` to rank the LiteLLM-hosted skills you can\naccess by semantic similarity instead of paging through the whole registry:\n```bash\ncurl \"http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5\" -H \"Authorization: Bearer your-key\"\n```\n\nReturns: ListSkillsResponse with list of skills", "operationId": "list_skills_v1_skills_get", "parameters": [ { @@ -4849,6 +4871,39 @@ "default": "anthropic", "title": "Custom Llm Provider" } + }, + { + "description": "Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model.", + "in": "query", + "name": "query", + "required": false, + "schema": { + "anyOf": [ + { + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model.", + "title": "Query" + } + }, + { + "description": "With query: the maximum number of ranked skills to return.", + "in": "query", + "name": "top_k", + "required": false, + "schema": { + "default": 5, + "description": "With query: the maximum number of ranked skills 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 index 76e3fe6c5ad..65a89bb2c7a 100644 --- a/litellm/proxy/agent_endpoints/agent_search.py +++ b/litellm/proxy/agent_endpoints/agent_search.py @@ -18,6 +18,7 @@ from litellm.types.agents import AgentResponse if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging from litellm.router import Router DEFAULT_AGENT_SEARCH_TOP_K: Final = 5 @@ -132,6 +133,7 @@ async def search_agents( embedding_model: str | None, index: AgentSearchIndex, user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, ) -> AgentSearchOutcome: if embedding_model is None: return AgentSearchNotConfigured( @@ -139,6 +141,5 @@ async def search_agents( ) 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 - ) + embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj) + return await index.search(query, agents, top_k, embed, embedding_model) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index cc17672553b..aa8979a73c6 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -249,7 +249,7 @@ def _agent_search_error(status_code: int, error: str, message: str) -> HTTPExcep 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 + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj outcome: Final = await search_agents( query=query, @@ -259,6 +259,7 @@ async def _rank_agents_by_query( embedding_model=litellm.agent_search_embedding_model, index=global_agent_search_index, user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) match outcome: case AgentSearchHits(hits): diff --git a/litellm/proxy/anthropic_endpoints/skills_endpoints.py b/litellm/proxy/anthropic_endpoints/skills_endpoints.py index 9390bf4c537..4426c0b547a 100644 --- a/litellm/proxy/anthropic_endpoints/skills_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/skills_endpoints.py @@ -2,11 +2,23 @@ Anthropic Skills API endpoints - /v1/skills """ -from typing import Final +from types import MappingProxyType +from typing import Annotated, Final import orjson -from fastapi import APIRouter, Depends, Request, Response +from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response +from typing_extensions import ReadOnly, TypedDict, assert_never +import litellm +from litellm.llms.litellm_proxy.skills.skill_search import ( + DEFAULT_SKILL_SEARCH_TOP_K, + SkillSearchEmbeddingFailed, + SkillSearchHits, + SkillSearchNotConfigured, + SkillSearchUnsupportedProvider, + global_skill_search_index, + search_hosted_skills, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -23,6 +35,51 @@ from litellm.types.llms.anthropic_skills import ( router: Final = APIRouter() +class _SkillSearchErrorDetail(TypedDict): + error: ReadOnly[str] + message: ReadOnly[str] + + +def _skill_search_error(status_code: int, error: str, message: str) -> HTTPException: + detail: Final[_SkillSearchErrorDetail] = {"error": error, "message": message} + return HTTPException(status_code=status_code, detail=detail) + + +async def _search_skills( + custom_llm_provider: str | None, query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth +) -> ListSkillsResponse: + from litellm.llms.litellm_proxy.skills.transformation import ( + LiteLLMSkillsTransformationHandler, + ) + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj + + outcome: Final = await search_hosted_skills( + custom_llm_provider=custom_llm_provider, + query=query, + top_k=top_k, + router=llm_router, + embedding_model=litellm.skill_search_embedding_model, + index=global_skill_search_index, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + to_response: Final = LiteLLMSkillsTransformationHandler().db_skill_to_response + match outcome: + case SkillSearchHits(hits): + skills: Final = [ # mutable-ok: ListSkillsResponse.data requires list[Skill]; never mutated after + to_response(hit.skill).model_copy(update=MappingProxyType({"search_score": hit.score})) for hit in hits + ] + return ListSkillsResponse(data=skills, has_more=False, next_page=None) + case SkillSearchUnsupportedProvider(reason): + raise _skill_search_error(400, "skill_search_unsupported_provider", reason) + case SkillSearchNotConfigured(reason): + raise _skill_search_error(400, "skill_search_not_configured", reason) + case SkillSearchEmbeddingFailed(reason): + raise _skill_search_error(503, "skill_search_unavailable", reason) + case _: + assert_never(outcome) + + @router.post( "/v1/skills", tags=["[beta] Anthropic Skills API"], @@ -134,32 +191,58 @@ async def list_skills( after_id: str | None = None, before_id: str | None = None, custom_llm_provider: str | None = "anthropic", + query: Annotated[ + str | None, + Query( + min_length=1, + description="Describe what you need in natural language to rank the skills you can access by " + "semantic similarity over their title and description. Each result carries a search_score. " + "Only supported for custom_llm_provider=litellm_proxy. Requires " + "litellm_settings.skill_search_embedding_model.", + ), + ] = None, + top_k: Annotated[ + int, + Query(ge=1, le=100, description="With query: the maximum number of ranked skills to return."), + ] = DEFAULT_SKILL_SEARCH_TOP_K, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ List skills on Anthropic. - + Requires `?beta=true` query parameter. - + Model-based routing (for multi-account support): - Pass model via header: `x-litellm-model: claude-account-1` - Pass model via query: `?model=claude-account-1` - Pass model via body: `{"model": "claude-account-1"}` - + Example usage: ```bash # Basic usage curl "http://localhost:4000/v1/skills?beta=true&limit=10" \ -H "Authorization: Bearer your-key" - + # With model-based routing curl "http://localhost:4000/v1/skills?beta=true&limit=10" \ -H "Authorization: Bearer your-key" \ -H "x-litellm-model: claude-account-1" ``` - + + Pass `?custom_llm_provider=litellm_proxy&query=` to rank the LiteLLM-hosted skills you can + access by semantic similarity instead of paging through the whole registry: + ```bash + curl "http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5" \ + -H "Authorization: Bearer your-key" + ``` + Returns: ListSkillsResponse with list of skills """ + if query is not None: + return await _search_skills( + custom_llm_provider=custom_llm_provider, query=query, top_k=top_k, user_api_key_dict=user_api_key_dict + ) + from litellm.proxy.proxy_server import ( general_settings, llm_router, diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py index 0820459af49..b8d3595163e 100644 --- a/litellm/proxy/common_utils/semantic_text_index.py +++ b/litellm/proxy/common_utils/semantic_text_index.py @@ -5,10 +5,11 @@ from __future__ import annotations import math from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass -from itertools import chain +from itertools import chain, islice from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, TypeAlias +from fastapi import HTTPException from openai import OpenAIError from pydantic import BaseModel, ConfigDict @@ -16,10 +17,14 @@ from litellm.exceptions import BudgetExceededError if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging from litellm.router import Router Vector: TypeAlias = tuple[float, ...] +DEFAULT_MAX_CACHED_VECTORS: Final = 5000 +"""Ceiling on how many (embedding model, text) vectors one index keeps; the least recently searched are evicted first.""" + class Embedder(Protocol): def __call__(self, texts: Sequence[str]) -> Awaitable[Sequence[Vector]]: ... @@ -42,6 +47,16 @@ class _EmbeddingData(BaseModel): data: tuple[_EmbeddingItem, ...] +class _EmbeddingRequest(BaseModel): + """The /embeddings-shaped request as the pre-call hooks (rate limits, budgets, guardrails) hand it back.""" + + model_config = ConfigDict(frozen=True, extra="ignore") + + model: str + input: tuple[str, ...] + metadata: dict[str, object] # mutable-ok: the router mutates the metadata dict it is handed + + 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)) @@ -57,23 +72,40 @@ def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, obj } -def router_embedder(router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth) -> Embedder: +def router_embedder( + router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging +) -> Embedder: + """Embeds through the router after the same key rate-limit, budget and guardrail pre-call hooks /embeddings runs.""" + async def embed(texts: Sequence[str]) -> Sequence[Vector]: - batch: Final = list(texts) # mutable-ok: Router.aembedding accepts only str | list input + request: Final = { # mutable-ok: pre_call_hook mutates the request dict in place + "model": embedding_model, + "input": list(texts), # mutable-ok: Router.aembedding accepts only str | list input + "metadata": embedding_spend_metadata(user_api_key_dict), + } + processed: Final = _EmbeddingRequest.model_validate( + await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=request, call_type="aembedding" + ) + ) response: Final = await router.aembedding( - model=embedding_model, input=batch, metadata=embedding_spend_metadata(user_api_key_dict) + model=processed.model, + input=list(processed.input), # mutable-ok: Router.aembedding accepts only str | list input + metadata=processed.metadata, ) return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data) return embed -_NO_VECTORS: Final[Mapping[str, Vector]] = MappingProxyType({}) +_CacheKey: TypeAlias = tuple[str, str] async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed: try: vectors: Final = tuple(await embed(texts)) + except HTTPException: + raise except (OpenAIError, ValueError, BudgetExceededError) as exc: return EmbeddingFailed(reason=f"embedding the search query failed: {exc}") if len(vectors) != len(texts): @@ -111,20 +143,34 @@ async def _embed_query_and_texts( class SemanticTextIndex: - """Caches one vector per distinct text per embedding model, so repeat searches only embed the query.""" + """Caches one vector per distinct text per embedding model, so repeat searches only embed the query. - def __init__(self) -> None: - self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({}) + Holds at most ``max_entries`` vectors across all models: once full, the texts no recent search touched go first.""" - def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]: - kept: Final = MappingProxyType( + def __init__(self, max_entries: int = DEFAULT_MAX_CACHED_VECTORS) -> None: + self._max_entries: Final = max_entries + self._vectors: Mapping[_CacheKey, Vector] = MappingProxyType({}) + + def _cached(self, embedding_model: str) -> Mapping[str, Vector]: + return MappingProxyType( + {text: vector for (model, text), vector in self._vectors.items() if model == embedding_model} + ) + + def _merged(self, embedding_model: str, embedded: _Embedded, texts: Sequence[str]) -> Mapping[_CacheKey, Vector]: + dimension: Final = len(embedded.query_vector) + touched: Final = MappingProxyType({(embedding_model, text): embedded.vectors[text] for text in texts}) + untouched: Final = MappingProxyType( { - text: vector - for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items() - if len(vector) == len(embedded.query_vector) + key: vector + for key, vector in chain( + self._vectors.items(), + (((embedding_model, text), vector) for text, vector in embedded.vectors.items()), + ) + if key not in touched and (key[0] != embedding_model or len(vector) == dimension) } ) - return MappingProxyType({**kept, **embedded.vectors}) + ordered: Final = MappingProxyType({**untouched, **touched}) + return MappingProxyType(dict(islice(ordered.items(), max(len(ordered) - self._max_entries, 0), None))) async def scores( self, query: str, texts: Sequence[str], embed: Embedder, embedding_model: str @@ -132,11 +178,10 @@ class SemanticTextIndex: """Cosine similarity of `query` to each entry of `texts`, in the same order.""" if not texts: return () - cached: Final = self._vectors.get(embedding_model, _NO_VECTORS) - embedded: Final = await _embed_query_and_texts(embed, query, texts, cached) + embedded: Final = await _embed_query_and_texts(embed, query, texts, self._cached(embedding_model)) if isinstance(embedded, EmbeddingFailed): return embedded if not _same_dimension(embedded.query_vector, embedded.vectors, texts): return EmbeddingFailed(reason=f"embedding model {embedding_model} returned vectors of mixed dimensions") - self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)}) + self._vectors = self._merged(embedding_model, embedded, texts) return tuple(cosine_similarity(embedded.query_vector, embedded.vectors[text]) for text in texts) diff --git a/litellm/skills/main.py b/litellm/skills/main.py index 002419dbad4..71fd78f11a3 100644 --- a/litellm/skills/main.py +++ b/litellm/skills/main.py @@ -182,10 +182,14 @@ def create_skill( if extra_body: create_request.update(extra_body) - # Route to LiteLLM DB if custom_llm_provider="litellm_proxy" + # Route to LiteLLM DB if custom_llm_provider="litellm_proxy". description/instructions + # arrive as top-level kwargs from the REST form endpoint, or nested in extra_body from + # the SDK convention used by other providers' create_request above. if custom_llm_provider == LlmProviders.LITELLM_PROXY.value: return _get_litellm_skills_handler().create_skill_handler( display_title=display_title, + description=kwargs.get("description") or (extra_body.get("description") if extra_body else None), + instructions=kwargs.get("instructions") or (extra_body.get("instructions") if extra_body else None), files=files, metadata=_get_skill_request_metadata(kwargs, extra_body), user_id=kwargs.get("user_id"), diff --git a/litellm/types/llms/anthropic_skills.py b/litellm/types/llms/anthropic_skills.py index 51eefe7154f..4b27f9b17ef 100644 --- a/litellm/types/llms/anthropic_skills.py +++ b/litellm/types/llms/anthropic_skills.py @@ -57,6 +57,15 @@ class Skill(BaseModel): updated_at: str """ISO 8601 timestamp of when the skill was last updated""" + description: str | None = None + """Description of the skill. Populated for the LiteLLM-hosted registry + (custom_llm_provider="litellm_proxy"); Anthropic's list endpoint does not + return a description, so this is None there.""" + + search_score: float | None = None + """Semantic similarity to the ``query`` passed to ``GET /v1/skills``. None + unless a query was given.""" + class ListSkillsResponse(BaseModel): """Response from listing skills""" diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py new file mode 100644 index 00000000000..a0f22a59f0c --- /dev/null +++ b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py @@ -0,0 +1,436 @@ +import asyncio +import json +from collections.abc import Sequence +from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from openai import APIConnectionError + +import litellm +from litellm.llms.litellm_proxy.skills.skill_search import ( + MAX_SKILL_SEARCH_TEXT_CHARS, + SkillSearchEmbeddingFailed, + SkillSearchHits, + SkillSearchIndex, + SkillSearchNotConfigured, + search_skills, + skill_search_text, +) +from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.common_utils.semantic_text_index import Vector, cosine_similarity + +CALLER: Final = UserAPIKeyAuth(api_key="hashed-caller-key", team_id="team-1", user_id="user-1") + +TRANSLATOR: Final = LiteLLM_SkillsTable( + skill_id="translate-file", + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", +) +SQL_ANALYST: Final = LiteLLM_SkillsTable( + skill_id="warehouse-sql-analyst", + display_title="Warehouse SQL Analyst", + description="Runs SQL against the inventory database", +) +TRIP_PLANNER: Final = LiteLLM_SkillsTable( + skill_id="trip-planner", + display_title="Trip Planner", + description="Books flights and hotels", +) +SKILLS: Final = (TRANSLATOR, SQL_ANALYST, TRIP_PLANNER) + +VECTORS: Final = MappingProxyType( + { + "language translation": (1.0, 0.0, 0.0), + skill_search_text(TRANSLATOR): (0.9, 0.1, 0.0), + skill_search_text(SQL_ANALYST): (0.0, 1.0, 0.0), + skill_search_text(TRIP_PLANNER): (0.3, 0.0, 1.0), + } +) + + +def _pass_through_key_limits() -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + return limits + + +def _embedding_router() -> 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)], + ) + ) + return router + + +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 TestSkillSearchText: + def test_joins_title_description_and_instructions(self) -> None: + assert skill_search_text(TRANSLATOR) == ( + "Document Translator\n" + "Converts files from one language into another\n" + "Take an uploaded document and produce it in the target language" + ) + + def test_missing_fields_fall_back_to_whatever_is_present(self) -> None: + assert skill_search_text(LiteLLM_SkillsTable(skill_id="bare", display_title="bare")) == "bare" + + def test_all_fields_absent_is_an_empty_string(self) -> None: + assert skill_search_text(LiteLLM_SkillsTable(skill_id="empty")) == "" + + def test_oversized_instructions_are_cut_so_one_skill_cannot_blow_up_the_embedding_batch(self) -> None: + bloated = LiteLLM_SkillsTable( + skill_id="bloated", display_title="Bloated", instructions="x" * (MAX_SKILL_SEARCH_TEXT_CHARS * 3) + ) + text = skill_search_text(bloated) + assert len(text) == MAX_SKILL_SEARCH_TEXT_CHARS + assert text.startswith("Bloated\n") + + +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 TestSkillSearchIndex: + @pytest.mark.asyncio + async def test_ranks_by_similarity_and_truncates_to_top_k(self) -> None: + outcome = await SkillSearchIndex().search( + "language translation", SKILLS, top_k=2, embed=FakeEmbedder(), embedding_model="m" + ) + assert isinstance(outcome, SkillSearchHits) + assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file", "trip-planner"] + assert outcome.hits[0].score > outcome.hits[1].score + + @pytest.mark.asyncio + async def test_second_search_only_embeds_the_query(self) -> None: + index = SkillSearchIndex() + embedder = FakeEmbedder() + await index.search("language translation", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("language translation", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert len(embedder.calls[0]) == 1 + len(SKILLS) + assert embedder.calls[1] == ("language translation",) + + @pytest.mark.asyncio + async def test_switching_embedding_models_does_not_reuse_cached_vectors(self) -> None: + index = SkillSearchIndex() + await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="small") + wide = FixedDimensionEmbedder(2) + outcome = await index.search("language translation", SKILLS, top_k=5, embed=wide, embedding_model="wide") + assert isinstance(outcome, SkillSearchHits) + assert len(wide.calls[0]) == 1 + len(SKILLS) + + @pytest.mark.asyncio + async def test_cached_vectors_of_another_dimension_are_re_embedded(self) -> None: + index = SkillSearchIndex() + await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="m") + fallback = FixedDimensionEmbedder(2) + outcome = await index.search("language translation", SKILLS, top_k=5, embed=fallback, embedding_model="m") + assert isinstance(outcome, SkillSearchHits) + assert fallback.calls == [ + ("language translation",), + ("language translation", *(skill_search_text(skill) for skill in SKILLS)), + ] + + @pytest.mark.asyncio + async def test_re_embedding_a_subset_drops_the_other_skills_old_vectors(self) -> None: + index = SkillSearchIndex() + await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="m") + wide = FixedDimensionEmbedder(2) + await index.search("language translation", SKILLS[:1], top_k=5, embed=wide, embedding_model="m") + await index.search("language translation", SKILLS, top_k=5, embed=wide, embedding_model="m") + assert wide.calls[-1] == ("language translation", *(skill_search_text(skill) for skill in SKILLS[1:])) + + @pytest.mark.asyncio + async def test_concurrent_searches_keep_each_others_vectors(self) -> None: + index = SkillSearchIndex() + embedder = FixedDimensionEmbedder(3) + await asyncio.gather( + index.search("q", SKILLS[:1], top_k=5, embed=embedder, embedding_model="m"), + index.search("q", SKILLS[1:], top_k=5, embed=embedder, embedding_model="m"), + ) + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q",) + + @pytest.mark.asyncio + async def test_least_recently_searched_skills_are_evicted_once_the_index_is_full(self) -> None: + index = SkillSearchIndex(max_entries=len(SKILLS)) + embedder = FixedDimensionEmbedder(3) + newcomer = LiteLLM_SkillsTable(skill_id="newcomer", display_title="Newcomer") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS[:1], top_k=5, embed=embedder, embedding_model="m") + await index.search("q", (newcomer,), top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q", skill_search_text(SKILLS[1])) + + @pytest.mark.asyncio + async def test_deleted_skills_stop_occupying_the_index_after_enough_new_ones(self) -> None: + index = SkillSearchIndex(max_entries=2) + embedder = FixedDimensionEmbedder(3) + for generation in range(50): + skill = LiteLLM_SkillsTable(skill_id=f"gen-{generation}", display_title=f"Generation {generation}") + await index.search("q", (skill,), top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q", skill_search_text(SKILLS[0])) + + @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 SkillSearchIndex().search("q", SKILLS, top_k=5, embed=mixed, embedding_model="m") + assert isinstance(outcome, SkillSearchEmbeddingFailed) + assert "mixed dimensions" in outcome.reason + + @pytest.mark.asyncio + async def test_no_accessible_skills_returns_no_hits_without_embedding(self) -> None: + embedder = FakeEmbedder() + outcome = await SkillSearchIndex().search("anything", (), top_k=5, embed=embedder, embedding_model="m") + assert outcome == SkillSearchHits(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 SkillSearchIndex().search("q", SKILLS, top_k=5, embed=failing, embedding_model="m") + assert isinstance(outcome, SkillSearchEmbeddingFailed) + 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 SkillSearchIndex().search("q", SKILLS, top_k=5, embed=short, embedding_model="m") + assert isinstance(outcome, SkillSearchEmbeddingFailed) + + +class TestSearchSkills: + @pytest.mark.asyncio + async def test_no_embedding_model_is_not_configured(self) -> None: + outcome = await search_skills( + "q", + SKILLS, + 5, + router=MagicMock(), + embedding_model=None, + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + assert isinstance(outcome, SkillSearchNotConfigured) + assert "skill_search_embedding_model" in outcome.reason + + @pytest.mark.asyncio + async def test_no_router_is_not_configured(self) -> None: + outcome = await search_skills( + "q", + SKILLS, + 5, + router=None, + embedding_model="m", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + assert isinstance(outcome, SkillSearchNotConfigured) + + @pytest.mark.asyncio + async def test_router_embeddings_are_read_from_the_response(self) -> None: + router = _embedding_router() + outcome = await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + assert isinstance(outcome, SkillSearchHits) + assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file"] + 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 = _embedding_router() + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + 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" + + @pytest.mark.asyncio + async def test_key_limits_are_checked_against_the_real_embedding_call_before_it_runs(self) -> None: + router = _embedding_router() + key_limits = _pass_through_key_limits() + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + checked = key_limits.pre_call_hook.await_args.kwargs + assert checked["user_api_key_dict"] is CALLER + assert checked["call_type"] == "aembedding" + assert checked["data"]["model"] == "text-embedding-3-small" + assert checked["data"]["input"] == router.aembedding.await_args.kwargs["input"] + assert checked["data"]["metadata"]["user_api_key"] == "hashed-caller-key" + + @pytest.mark.asyncio + async def test_the_embedding_model_sees_the_request_as_the_guardrails_rewrote_it(self) -> None: + router = MagicMock() + router.aembedding = AsyncMock( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": i, "embedding": [1.0, 0.0, 0.0]} for i in range(len(input))], + ) + ) + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type: { + **data, + "input": ["[MASKED]" for _ in data["input"]], + "metadata": {**data["metadata"], "guardrail": "masked"}, + } + ) + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + sent = tuple(call.kwargs for call in router.aembedding.await_args_list) + assert sent + assert all(set(call["input"]) == {"[MASKED]"} for call in sent) + assert all(call["metadata"]["guardrail"] == "masked" for call in sent) + + @pytest.mark.asyncio + async def test_a_key_over_its_limit_never_reaches_the_embedding_model(self) -> None: + router = _embedding_router() + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock(side_effect=ProxyRateLimitError(detail="rpm exceeded")) + with pytest.raises(ProxyRateLimitError) as raised: + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + assert raised.value.status_code == 429 + router.aembedding.assert_not_awaited() + + +@pytest.fixture +def accessible_skills(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + list_for_search = AsyncMock(return_value=list(SKILLS)) + monkeypatch.setattr( + "litellm.llms.litellm_proxy.skills.handler.LiteLLMSkillsHandler.list_skills_for_search", list_for_search + ) + return list_for_search + + +@pytest.fixture +def key_limits(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + limits = _pass_through_key_limits() + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", limits) + return limits + + +@pytest.fixture +def embedding_router(monkeypatch: pytest.MonkeyPatch, key_limits: MagicMock) -> MagicMock: + router = _embedding_router() + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "skill_search_embedding_model", "text-embedding-3-small") + return router + + +class TestHandleSkillSearchMCP: + @pytest.mark.asyncio + async def test_top_k_is_clamped_to_the_same_ceiling_as_rest( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + from litellm.llms.litellm_proxy.skills.skill_search import MAX_SKILL_SEARCH_TOP_K + from litellm.proxy._experimental.mcp_server.tool_search import handle_skill_search + + many_skills: Final = tuple( + LiteLLM_SkillsTable( + skill_id=f"skill-{i}", display_title=TRIP_PLANNER.display_title, description=TRIP_PLANNER.description + ) + for i in range(MAX_SKILL_SEARCH_TOP_K + 50) + ) + accessible_skills.return_value = list(many_skills) + + result = await handle_skill_search( + query="language translation", top_k=10_000, user_api_key_dict=UserAPIKeyAuth(user_id="u") + ) + assert result.isError is False + assert len(json.loads(result.content[0].text)) == MAX_SKILL_SEARCH_TOP_K + + @pytest.mark.asyncio + async def test_top_k_below_one_is_raised_to_one( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + from litellm.proxy._experimental.mcp_server.tool_search import handle_skill_search + + result = await handle_skill_search( + query="language translation", top_k=0, user_api_key_dict=UserAPIKeyAuth(user_id="u") + ) + assert result.isError is False + assert len(json.loads(result.content[0].text)) == 1 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 239f89ebd90..798b0001af1 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 @@ -24,6 +24,7 @@ from litellm.proxy._experimental.mcp_server.tool_search import ( AGENT_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, MCP_TOOL_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, SemanticToolRanker, ToolSearchResult, coerce_top_k, @@ -272,8 +273,8 @@ class TestSearchTools: class TestGetVirtualToolDefinitions: - def test_returns_three_tools(self) -> None: - assert len(get_virtual_tool_definitions()) == 3 + def test_returns_four_tools(self) -> None: + assert len(get_virtual_tool_definitions()) == 4 def test_agent_search_schema_requires_query(self) -> None: tools = get_virtual_tool_definitions() @@ -330,6 +331,7 @@ class TestGetVirtualToolDefinitions: MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, } @@ -364,7 +366,12 @@ 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, AGENT_SEARCH_TOOL_NAME} + assert set(tool_names) == { + MCP_TOOL_SEARCH_TOOL_NAME, + MCP_TOOL_CALL_TOOL_NAME, + AGENT_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, + } @pytest.mark.asyncio async def test_returns_full_catalog_when_flag_disabled(self) -> None: @@ -737,6 +744,31 @@ class TestCallToolRestApiVirtualTools: assert mock_search.await_args.kwargs["top_k"] == 1 assert mock_search.await_args.kwargs["agents"] == (translator,) + @pytest.mark.asyncio + async def test_skill_search_call_tolerates_malformed_top_k(self) -> None: + """Regression: a caller-supplied non-numeric top_k must be coerced to the default, + the same as agent_search, instead of raising a pydantic ValidationError that the + endpoint's catch-all turns into an HTTP 500.""" + from mcp.types import CallToolResult, TextContent + + from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K + + user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) + request = self._make_request( + {"name": SKILL_SEARCH_TOOL_NAME, "arguments": {"query": "translate a document", "top_k": "not-a-number"}} + ) + fake_result = CallToolResult(content=[TextContent(type="text", text="[]")], isError=False) + with patch( # test-quality-ok: the embedding router only resolves via proxy_server globals, no injection seam + "litellm.proxy._experimental.mcp_server.tool_search.handle_skill_search", + new_callable=AsyncMock, + return_value=fake_result, + ) as mock_search: + result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict) + + assert result.isError is False + assert mock_search.await_args.kwargs["top_k"] == DEFAULT_SKILL_SEARCH_TOP_K + assert mock_search.await_args.kwargs["query"] == "translate a document" + @pytest.mark.asyncio async def test_agent_search_call_reports_missing_embedding_model_as_tool_error(self) -> None: from litellm.proxy.agent_endpoints.agent_search import AgentSearchNotConfigured @@ -782,10 +814,15 @@ class TestCallToolRestApiVirtualTools: router = MagicMock() router.aembedding = AsyncMock(side_effect=fake_aembedding) + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) with ( patch( # test-quality-ok: the proxy's router is a module global; the handler reaches it the way production does "litellm.proxy.proxy_server.llm_router", router ), + patch( # test-quality-ok: the proxy's key-limit hooks are a module global; the embedding call runs them like /embeddings does + "litellm.proxy.proxy_server.proxy_logging_obj", key_limits + ), patch( # test-quality-ok: the authorized catalog is the seam every virtual tool shares; the ranking under test stays real "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", new_callable=AsyncMock, @@ -795,6 +832,8 @@ class TestCallToolRestApiVirtualTools: result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) assert mock_list.await_args.kwargs["user_api_key_auth"] is user_api_key_dict + assert key_limits.pre_call_hook.await_args.kwargs["call_type"] == "aembedding" + assert key_limits.pre_call_hook.await_args.kwargs["data"]["model"] == "emb" assert result.isError is False assert [t["name"] for t in json.loads(result.content[0].text)] == [FX_TOOL.name] @@ -810,7 +849,9 @@ class TestCallToolRestApiVirtualTools: assert "mcp_tool_search.embedding_model" in result.content[0].text @pytest.mark.asyncio - async def test_mcp_tool_search_reports_invalid_settings_as_tool_error(self, monkeypatch: pytest.MonkeyPatch) -> None: + async def test_mcp_tool_search_reports_invalid_settings_as_tool_error( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: monkeypatch.setattr(litellm, "mcp_tool_search", {"top_k": 0}) user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) @@ -1196,6 +1237,7 @@ class TestHandleListToolsVirtual: MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, + SKILL_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 index daca244c0a1..c02ed1f37e5 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py @@ -211,11 +211,24 @@ class TestAgentSearchIndex: assert isinstance(outcome, AgentSearchEmbeddingFailed) +def _pass_through_key_limits() -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + return limits + + 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 + "q", + AGENTS, + 5, + router=MagicMock(), + embedding_model=None, + index=AgentSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchNotConfigured) assert "agent_search_embedding_model" in outcome.reason @@ -223,7 +236,14 @@ class TestSearchAgents: @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 + "q", + AGENTS, + 5, + router=None, + embedding_model="m", + index=AgentSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchNotConfigured) @@ -244,6 +264,7 @@ class TestSearchAgents: embedding_model="text-embedding-3-small", index=AgentSearchIndex(), user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchHits) assert [hit.agent.agent_id for hit in outcome.hits] == ["translator"] @@ -266,6 +287,7 @@ class TestSearchAgents: embedding_model="text-embedding-3-small", index=AgentSearchIndex(), user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) metadata = router.aembedding.await_args.kwargs["metadata"] assert metadata["user_api_key"] == "hashed-caller-key" @@ -302,6 +324,7 @@ def embedding_router(monkeypatch: pytest.MonkeyPatch) -> MagicMock: ) ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", _pass_through_key_limits()) monkeypatch.setattr(litellm, "agent_search_embedding_model", "text-embedding-3-small") return router diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py new file mode 100644 index 00000000000..ae6b1471c93 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py @@ -0,0 +1,175 @@ +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.llms.litellm_proxy.skills.skill_search import skill_search_text +from litellm.proxy._types import LiteLLM_SkillsTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.skills_endpoints import router +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + +TRANSLATOR: Final = LiteLLM_SkillsTable( + skill_id="translate-file", + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", +) +SQL_ANALYST: Final = LiteLLM_SkillsTable( + skill_id="warehouse-sql-analyst", + display_title="Warehouse SQL Analyst", + description="Runs SQL against the inventory database", +) +TRIP_PLANNER: Final = LiteLLM_SkillsTable( + skill_id="trip-planner", + display_title="Trip Planner", + description="Books flights and hotels", +) +SKILLS: Final = (TRANSLATOR, SQL_ANALYST, TRIP_PLANNER) + +VECTORS: Final = MappingProxyType( + { + "language translation": (1.0, 0.0, 0.0), + skill_search_text(TRANSLATOR): (0.9, 0.1, 0.0), + skill_search_text(SQL_ANALYST): (0.0, 1.0, 0.0), + skill_search_text(TRIP_PLANNER): (0.3, 0.0, 1.0), + } +) + + +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 accessible_skills(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + list_for_search = AsyncMock(return_value=list(SKILLS)) + monkeypatch.setattr( + "litellm.llms.litellm_proxy.skills.handler.LiteLLMSkillsHandler.list_skills_for_search", list_for_search + ) + return list_for_search + + +@pytest.fixture +def key_limits(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", limits) + return limits + + +@pytest.fixture +def embedding_router(monkeypatch: pytest.MonkeyPatch, key_limits: MagicMock) -> MagicMock: + embedding_router = MagicMock() + embedding_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", embedding_router) + monkeypatch.setattr(litellm, "skill_search_embedding_model", "text-embedding-3-small") + return embedding_router + + +class TestGetSkillsQuery: + def test_query_ranks_and_scores_and_truncates( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "language translation", "top_k": 2}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 200 + body = response.json()["data"] + assert [skill["id"] for skill in body] == ["translate-file", "trip-planner"] + 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_restricted_key_only_ranks_the_skills_it_can_access( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + accessible_skills.return_value = [SQL_ANALYST] + response = _client(LitellmUserRoles.INTERNAL_USER).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "language translation"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 200 + assert [skill["id"] for skill in response.json()["data"]] == ["warehouse-sql-analyst"] + + def test_no_accessible_skills_is_a_no_match_empty_result( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + accessible_skills.return_value = [] + response = _client(LitellmUserRoles.INTERNAL_USER).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 200 + assert response.json()["data"] == [] + embedding_router.aembedding.assert_not_awaited() + + def test_query_is_unsupported_for_the_anthropic_passthrough_provider( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", params={"query": "anything"}, headers={"Authorization": "Bearer k"} + ) + assert response.status_code == 400 + assert response.json()["detail"]["error"] == "skill_search_unsupported_provider" + accessible_skills.assert_not_awaited() + + def test_missing_embedding_model_is_a_400( + self, accessible_skills: AsyncMock, embedding_router: MagicMock, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "skill_search_embedding_model", None) + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 400 + assert response.json()["detail"]["error"] == "skill_search_not_configured" + + def test_embedding_provider_failure_is_a_503( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + embedding_router.aembedding = AsyncMock(side_effect=APIConnectionError(request=MagicMock())) + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 503 + assert response.json()["detail"]["error"] == "skill_search_unavailable" + + def test_a_key_over_its_rate_limit_gets_a_429_without_embedding( + self, accessible_skills: AsyncMock, embedding_router: MagicMock, key_limits: MagicMock + ) -> None: + key_limits.pre_call_hook = AsyncMock(side_effect=ProxyRateLimitError(detail="rpm exceeded")) + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "language translation"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 429 + embedding_router.aembedding.assert_not_awaited() + + def test_top_k_is_validated(self, accessible_skills: AsyncMock, embedding_router: MagicMock) -> None: + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything", "top_k": 0}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 422 diff --git a/tests/test_litellm/skills/test_skills_main.py b/tests/test_litellm/skills/test_skills_main.py new file mode 100644 index 00000000000..e1c66c8d9ea --- /dev/null +++ b/tests/test_litellm/skills/test_skills_main.py @@ -0,0 +1,57 @@ +from unittest.mock import MagicMock + +import litellm.skills.main as skills_main +from litellm.types.utils import LlmProviders + + +def test_create_skill_forwards_description_and_instructions_from_top_level_kwargs( + monkeypatch, +) -> None: + """The REST /v1/skills form endpoint passes description/instructions as top-level + kwargs (not extra_body). Regression for a bug where the litellm_proxy dispatch + branch of create_skill() dropped both, so every LiteLLM-hosted skill was created + with description=None and instructions=None regardless of what the caller sent.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Converts files from one language into another" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == ( + "Take an uploaded document and produce it in the target language" + ) + + +def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: + """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under + extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Warehouse SQL Analyst", + extra_body={"description": "Runs SQL against the inventory database", "instructions": "Summarize results"}, + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Runs SQL against the inventory database" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == "Summarize results" + + +def test_create_skill_without_description_or_instructions_passes_none(monkeypatch) -> None: + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill(display_title="Bare Skill", custom_llm_provider=LlmProviders.LITELLM_PROXY.value) + + assert handler.create_skill_handler.call_args.kwargs["description"] is None + assert handler.create_skill_handler.call_args.kwargs["instructions"] is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 842a3da4122..12e7eb801e6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -19872,6 +19872,12 @@ export interface paths { * curl "http://localhost:4000/v1/skills?beta=true&limit=10" -H "Authorization: Bearer your-key" -H "x-litellm-model: claude-account-1" * ``` * + * Pass `?custom_llm_provider=litellm_proxy&query=` to rank the LiteLLM-hosted skills you can + * access by semantic similarity instead of paging through the whole registry: + * ```bash + * curl "http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5" -H "Authorization: Bearer your-key" + * ``` + * * Returns: ListSkillsResponse with list of skills */ get: operations["list_skills_v1_skills_get"]; @@ -36050,12 +36056,16 @@ export interface components { Skill: { /** Created At */ created_at: string; + /** Description */ + description?: string | null; /** Display Title */ display_title?: string | null; /** Id */ id: string; /** Latest Version */ latest_version?: string | null; + /** Search Score */ + search_score?: number | null; /** Source */ source: string; /** @@ -64556,6 +64566,10 @@ export interface operations { after_id?: string | null; before_id?: string | null; custom_llm_provider?: string | null; + /** @description Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model. */ + query?: string | null; + /** @description With query: the maximum number of ranked skills to return. */ + top_k?: number; }; header?: never; path?: never;