mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-17 18:12:43 +00:00
a297d27d48
# What does this PR do? Adds a new brave tool provider ## Test Plan ``` curl -X POST 'http://localhost:5000/alpha/toolgroups/register' \ -H 'Content-Type: application/json' \ -d '{ "name": "search", "tool_group": { "type": "user_defined", "tools": [ { "name": "brave_search", "description": "A web search tool", "parameters": [ { "name": "query", "parameter_type": "string", "description": "The query to search" } ], "metadata": {}, "tool_prompt_format": "json" } ] } }' curl -X POST 'http://localhost:5000/alpha/tool-runtime/invoke' \ -H 'Content-Type: application/json' \ -d '{ "tool_id": "brave_search", "args": { "query": "who is meta ceo" } }' | jq .content % Total % Received % Xferd Average Speed Time Time Time Current Dload Upload Total Spent Left Speed 100 1973 100 1884 100 89 11288 533 --:--:-- --:--:-- --:--:-- 11885 "{'title': 'Mark Zuckerberg, Founder, Chairman and Chief Executive ...', 'url': 'https://about.meta.com/media-gallery/executives/mark-zuckerberg/', 'description': 'Not Logged In · Please log in to see this page', 'type': 'search_result'}\n{'title': 'Meta - Leadership & Governance', 'url': 'https://investor.fb.com/leadership-and-governance/', 'description': '<strong>Mark Zuckerberg</strong> is the founder, chairman and CEO of Meta, which he originally founded as Facebook in 2004. Mark is responsible for setting the overall direction and product strategy for the company. He leads the design of Meta's services and development of its core technology and infrastructure.', 'type': 'search_result'}\n[{'type': 'video_result', 'url': '2372542949/', 'title': 'Mark Zuckerberg, the CEO of Meta, has officially joined the ...', 'description': \"Express Tribune, Karachi, Pakistan. 2,334,400 likes · 36,360 talking about this · 205 were here. The Express Tribune is Pakistan's #1 brand for breaking news in politics, sports, business, lifestyle\"}, {'type': 'video_result', 'url': 'https://www.youtube.com/watch?v=Y3oeQqtRvqk', 'title': \"Meta CEO: Mark Zuckerberg becomes World's Second Richest Person!\", 'description': 'Try VectorVest Risk-Free ➥➥➥ https://www.vectorvest.com/YTUse this link for a FREE Stock Analysis Report ➥➥➥ vectorvest.com/YTFSAVectorVest Merch Store ➥➥➥'}, {'type': 'video_result', 'url': '5348412224/', 'title': '#WATCH | Meta founder and CEO Mark Zuckerberg recently ...', 'description': 'See posts, photos and more on Facebook'}]" curl -X POST 'http://localhost:5000/alpha/tool-runtime/invoke' \ -H 'Content-Type: application/json' -H 'X-LlamaStack-ProviderData: {"api_key": "<KEY>"}' \ -d '{ "tool_id": "brave_search", "args": { "query": "who is meta ceo" } }' ```
127 lines
3.5 KiB
Python
127 lines
3.5 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
from typing import Annotated, Any, Dict, List, Literal, Optional, Union
|
|
|
|
from llama_models.llama3.api.datatypes import ToolPromptFormat
|
|
from llama_models.schema_utils import json_schema_type, register_schema, webmethod
|
|
from pydantic import BaseModel, Field
|
|
from typing_extensions import Protocol, runtime_checkable
|
|
|
|
from llama_stack.apis.common.content_types import InterleavedContent, URL
|
|
from llama_stack.apis.resource import Resource, ResourceType
|
|
from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
|
|
|
|
|
|
@json_schema_type
|
|
class ToolParameter(BaseModel):
|
|
name: str
|
|
parameter_type: str
|
|
description: str
|
|
|
|
|
|
@json_schema_type
|
|
class Tool(Resource):
|
|
type: Literal[ResourceType.tool.value] = ResourceType.tool.value
|
|
tool_group: str
|
|
description: str
|
|
parameters: List[ToolParameter]
|
|
provider_id: Optional[str] = None
|
|
metadata: Optional[Dict[str, Any]] = None
|
|
tool_prompt_format: Optional[ToolPromptFormat] = Field(
|
|
default=ToolPromptFormat.json
|
|
)
|
|
|
|
|
|
@json_schema_type
|
|
class ToolDef(BaseModel):
|
|
name: str
|
|
description: str
|
|
parameters: List[ToolParameter]
|
|
metadata: Dict[str, Any]
|
|
tool_prompt_format: Optional[ToolPromptFormat] = Field(
|
|
default=ToolPromptFormat.json
|
|
)
|
|
|
|
|
|
@json_schema_type
|
|
class MCPToolGroup(BaseModel):
|
|
"""
|
|
A tool group that is defined by in a model context protocol server.
|
|
Refer to https://modelcontextprotocol.io/docs/concepts/tools for more information.
|
|
"""
|
|
|
|
type: Literal["model_context_protocol"] = "model_context_protocol"
|
|
name: str
|
|
endpoint: URL
|
|
|
|
|
|
@json_schema_type
|
|
class UserDefinedToolGroup(BaseModel):
|
|
type: Literal["user_defined"] = "user_defined"
|
|
name: str
|
|
tools: List[ToolDef]
|
|
|
|
|
|
ToolGroup = register_schema(
|
|
Annotated[Union[MCPToolGroup, UserDefinedToolGroup], Field(discriminator="type")],
|
|
name="ToolGroup",
|
|
)
|
|
|
|
|
|
@json_schema_type
|
|
class ToolInvocationResult(BaseModel):
|
|
content: InterleavedContent
|
|
error_message: Optional[str] = None
|
|
error_code: Optional[int] = None
|
|
|
|
|
|
class ToolStore(Protocol):
|
|
def get_tool(self, tool_id: str) -> Tool: ...
|
|
|
|
|
|
@runtime_checkable
|
|
@trace_protocol
|
|
class Tools(Protocol):
|
|
@webmethod(route="/toolgroups/register", method="POST")
|
|
async def register_tool_group(
|
|
self,
|
|
tool_group: ToolGroup,
|
|
provider_id: Optional[str] = None,
|
|
) -> None:
|
|
"""Register a tool group"""
|
|
...
|
|
|
|
@webmethod(route="/tools/get", method="GET")
|
|
async def get_tool(
|
|
self,
|
|
tool_id: str,
|
|
) -> Tool: ...
|
|
|
|
@webmethod(route="/tools/list", method="GET")
|
|
async def list_tools(self) -> List[Tool]:
|
|
"""List tools with optional provider"""
|
|
|
|
@webmethod(route="/tools/unregister", method="POST")
|
|
async def unregister_tool(self, tool_id: str) -> None:
|
|
"""Unregister a tool"""
|
|
...
|
|
|
|
|
|
@runtime_checkable
|
|
@trace_protocol
|
|
class ToolRuntime(Protocol):
|
|
tool_store: ToolStore
|
|
|
|
@webmethod(route="/tool-runtime/discover", method="POST")
|
|
async def discover_tools(self, tool_group: ToolGroup) -> List[Tool]: ...
|
|
|
|
@webmethod(route="/tool-runtime/invoke", method="POST")
|
|
async def invoke_tool(
|
|
self, tool_id: str, args: Dict[str, Any]
|
|
) -> ToolInvocationResult:
|
|
"""Run a tool with the given arguments"""
|
|
...
|