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;