From 4b1330a84bef5c5c5246e4b7fdfe7e5bd5bac6f8 Mon Sep 17 00:00:00 2001 From: Lance Galletti Date: Mon, 23 Jun 2025 10:56:08 -0400 Subject: [PATCH] feat: adding tool annotation used by MCP --- docs/_static/llama-stack-spec.html | 100 ++++++++++++++++++ docs/_static/llama-stack-spec.yaml | 40 +++++++ llama_stack/apis/agents/openai_responses.py | 1 + llama_stack/apis/tools/tools.py | 11 ++ .../distribution/routing_tables/toolgroups.py | 1 + llama_stack/models/llama/datatypes.py | 3 + .../agents/meta_reference/agent_instance.py | 1 + .../agents/meta_reference/openai_responses.py | 1 + llama_stack/providers/utils/tools/mcp.py | 2 + 9 files changed, 160 insertions(+) diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html index 0a5caa3d1..cab579d1d 100644 --- a/docs/_static/llama-stack-spec.html +++ b/docs/_static/llama-stack-spec.html @@ -5547,6 +5547,31 @@ "additionalProperties": { "$ref": "#/components/schemas/ToolParamDefinition" } + }, + "annotations": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + } } }, "additionalProperties": false, @@ -6447,6 +6472,31 @@ } ] } + }, + "annotations": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + } } }, "additionalProperties": false, @@ -8309,6 +8359,31 @@ }, "description": { "type": "string" + }, + "annotations": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + } } }, "additionalProperties": false, @@ -10810,6 +10885,31 @@ } ] } + }, + "annotations": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + } } }, "additionalProperties": false, diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml index c115e1df2..608cf00a3 100644 --- a/docs/_static/llama-stack-spec.yaml +++ b/docs/_static/llama-stack-spec.yaml @@ -3912,6 +3912,16 @@ components: type: object additionalProperties: $ref: '#/components/schemas/ToolParamDefinition' + annotations: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object additionalProperties: false required: - tool_name @@ -4610,6 +4620,16 @@ components: - type: string - type: array - type: object + annotations: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object additionalProperties: false required: - name @@ -5855,6 +5875,16 @@ components: type: string description: type: string + annotations: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object additionalProperties: false required: - input_schema @@ -7615,6 +7645,16 @@ components: - type: string - type: array - type: object + annotations: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object additionalProperties: false required: - identifier diff --git a/llama_stack/apis/agents/openai_responses.py b/llama_stack/apis/agents/openai_responses.py index addb72f14..39e3c963d 100644 --- a/llama_stack/apis/agents/openai_responses.py +++ b/llama_stack/apis/agents/openai_responses.py @@ -116,6 +116,7 @@ class MCPListToolsTool(BaseModel): input_schema: dict[str, Any] name: str description: str | None = None + annotations: dict[str, Any] | None = None @json_schema_type diff --git a/llama_stack/apis/tools/tools.py b/llama_stack/apis/tools/tools.py index 0c8d47edf..f1e619b41 100644 --- a/llama_stack/apis/tools/tools.py +++ b/llama_stack/apis/tools/tools.py @@ -27,6 +27,15 @@ class ToolParameter(BaseModel): default: Any | None = None +@json_schema_type +class ToolAnnotations(BaseModel): + title: str | None = None + readOnlyHint: bool | None = None + destructiveHint: bool | None = None + idempotentHint: bool | None = None + openWorldHint: bool | None = None + + @json_schema_type class Tool(Resource): type: Literal[ResourceType.tool] = ResourceType.tool @@ -34,6 +43,7 @@ class Tool(Resource): description: str parameters: list[ToolParameter] metadata: dict[str, Any] | None = None + annotations: ToolAnnotations | None = None @json_schema_type @@ -42,6 +52,7 @@ class ToolDef(BaseModel): description: str | None = None parameters: list[ToolParameter] | None = None metadata: dict[str, Any] | None = None + annotations: ToolAnnotations | None = None @json_schema_type diff --git a/llama_stack/distribution/routing_tables/toolgroups.py b/llama_stack/distribution/routing_tables/toolgroups.py index b86f057bd..453e5c2e7 100644 --- a/llama_stack/distribution/routing_tables/toolgroups.py +++ b/llama_stack/distribution/routing_tables/toolgroups.py @@ -74,6 +74,7 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups): parameters=t.parameters or [], metadata=t.metadata, provider_id=toolgroup.provider_id, + annotations=t.annotations, ) ) diff --git a/llama_stack/models/llama/datatypes.py b/llama_stack/models/llama/datatypes.py index f9f094c3d..883585c77 100644 --- a/llama_stack/models/llama/datatypes.py +++ b/llama_stack/models/llama/datatypes.py @@ -11,6 +11,8 @@ from typing import Annotated, Any, Literal from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator +from apis.tools import ToolAnnotations + # The goal is that these set of types are relevant for all Llama models. # That isn't the current state yet -- e.g., BuiltinTool is somewhat specific to # the llama3 series of models. @@ -99,6 +101,7 @@ class ToolDefinition(BaseModel): tool_name: BuiltinTool | str description: str | None = None parameters: dict[str, ToolParamDefinition] | None = None + annotations: ToolAnnotations | None = None @field_validator("tool_name", mode="before") @classmethod diff --git a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py index 4d2b9f8bf..2eb10d7aa 100644 --- a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py +++ b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py @@ -843,6 +843,7 @@ class ChatAgent(ShieldRunnerMixin): ) for param in tool_def.parameters }, + annotations=tool_def.annotations, ) tool_name_to_args[tool_def.identifier] = toolgroup_to_args.get(toolgroup_name, {}) diff --git a/llama_stack/providers/inline/agents/meta_reference/openai_responses.py b/llama_stack/providers/inline/agents/meta_reference/openai_responses.py index 4465a32fe..ab99d9e67 100644 --- a/llama_stack/providers/inline/agents/meta_reference/openai_responses.py +++ b/llama_stack/providers/inline/agents/meta_reference/openai_responses.py @@ -665,6 +665,7 @@ class OpenAIResponsesImpl: }, "required": [p.name for p in t.parameters if p.required], }, + annotations=t.annotations, ) ) else: diff --git a/llama_stack/providers/utils/tools/mcp.py b/llama_stack/providers/utils/tools/mcp.py index f024693a0..c0402c37a 100644 --- a/llama_stack/providers/utils/tools/mcp.py +++ b/llama_stack/providers/utils/tools/mcp.py @@ -65,6 +65,7 @@ async def list_mcp_tools(endpoint: str, headers: dict[str, str]) -> ListToolDefs description=param_schema.get("description", ""), ) ) + annotations = getattr(tool, "annotations", None) tools.append( ToolDef( name=tool.name, @@ -73,6 +74,7 @@ async def list_mcp_tools(endpoint: str, headers: dict[str, str]) -> ListToolDefs metadata={ "endpoint": endpoint, }, + annotations=annotations, ) ) return ListToolDefsResponse(data=tools)