phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2026-01-01 02:24:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add inline mcp provider

Browse source
This commit is contained in:
Dinesh Yeduguru 2025-01-08 23:11:43 -08:00
parent ffc6bd4805
commit 2c265d803c
16 changed files with 398 additions and 49 deletions
Download patch file Download diff file Expand all files Collapse all files

6
llama_stack/providers/inline/tool_runtime/code_interpreter/code_interpreter.py
View file

@ -9,8 +9,8 @@ import logging
import tempfile
from typing import Any, Dict, List, Optional
from llama_stack.apis.common.content_types import URL
from llama_stack.apis.tools import (
MCPConfig,
Tool,
ToolDef,
ToolInvocationResult,
@ -43,7 +43,9 @@ class CodeInterpreterToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime):
return
async def list_runtime_tools(
self, tool_group_id: Optional[str] = None, mcp_endpoint: Optional[URL] = None
self,
tool_group_id: Optional[str] = None,
mcp_config: Optional[MCPConfig] = None,
) -> List[ToolDef]:
return [
ToolDef(

6
llama_stack/providers/inline/tool_runtime/memory/memory.py
View file

@ -10,11 +10,11 @@ import secrets
import string
from typing import Any, Dict, List, Optional
from llama_stack.apis.common.content_types import URL
from llama_stack.apis.inference import Inference, InterleavedContent
from llama_stack.apis.memory import Memory, QueryDocumentsResponse
from llama_stack.apis.memory_banks import MemoryBanks
from llama_stack.apis.tools import (
MCPConfig,
ToolDef,
ToolInvocationResult,
ToolParameter,
@ -52,7 +52,9 @@ class MemoryToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime):
pass
async def list_runtime_tools(
self, tool_group_id: Optional[str] = None, mcp_endpoint: Optional[URL] = None
self,
tool_group_id: Optional[str] = None,
mcp_config: Optional[MCPConfig] = None,
) -> List[ToolDef]:
return [
ToolDef(

20
llama_stack/providers/inline/tool_runtime/model_context_protocol/__init__.py Normal file
View file

@ -0,0 +1,20 @@
# 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 pydantic import BaseModel
from .config import ModelContextProtocolConfig
from .model_context_protocol import ModelContextProtocolToolRuntimeImpl
class ModelContextProtocolToolProviderDataValidator(BaseModel):
api_key: str
async def get_provider_impl(config: ModelContextProtocolConfig, _deps):
impl = ModelContextProtocolToolRuntimeImpl(config)
await impl.initialize()
return impl

11
llama_stack/providers/inline/tool_runtime/model_context_protocol/config.py Normal file
View file

@ -0,0 +1,11 @@
# 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 pydantic import BaseModel
class ModelContextProtocolConfig(BaseModel):
pass

98
llama_stack/providers/inline/tool_runtime/model_context_protocol/model_context_protocol.py Normal file
View file

@ -0,0 +1,98 @@
# 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.
import json
from typing import Any, Dict, List, Optional
from mcp import ClientSession
from mcp.client.stdio import stdio_client, StdioServerParameters
from pydantic import TypeAdapter
from llama_stack.apis.tools import (
MCPConfig,
ToolDef,
ToolInvocationResult,
ToolParameter,
ToolRuntime,
)
from llama_stack.providers.datatypes import ToolsProtocolPrivate
from .config import ModelContextProtocolConfig
class ModelContextProtocolToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime):
def __init__(self, config: ModelContextProtocolConfig):
self.config = config
async def initialize(self):
pass
async def list_runtime_tools(
self,
tool_group_id: Optional[str] = None,
mcp_config: Optional[MCPConfig] = None,
) -> List[ToolDef]:
if mcp_config is None:
raise ValueError("mcp_config is required")
tools = []
async with stdio_client(
StdioServerParameters(
command=mcp_config.command,
args=mcp_config.args,
env=mcp_config.env,
)
) as streams:
async with ClientSession(*streams) as session:
await session.initialize()
tools_result = await session.list_tools()
for tool in tools_result.tools:
parameters = []
for param_name, param_schema in tool.inputSchema.get(
"properties", {}
).items():
parameters.append(
ToolParameter(
name=param_name,
parameter_type=param_schema.get("type", "string"),
description=param_schema.get("description", ""),
)
)
tools.append(
ToolDef(
name=tool.name,
description=tool.description,
parameters=parameters,
metadata={
"mcp_config": mcp_config.model_dump_json(),
},
)
)
return tools
async def invoke_tool(
self, tool_name: str, args: Dict[str, Any]
) -> ToolInvocationResult:
tool = await self.tool_store.get_tool(tool_name)
if tool.metadata is None or tool.metadata.get("mcp_config") is None:
raise ValueError(f"Tool {tool_name} does not have metadata")
mcp_config_dict = json.loads(tool.metadata.get("mcp_config"))
mcp_config = TypeAdapter(MCPConfig).validate_python(mcp_config_dict)
async with stdio_client(
StdioServerParameters(
command=mcp_config.command,
args=mcp_config.args,
env=mcp_config.env,
)
) as streams:
async with ClientSession(*streams) as session:
await session.initialize()
result = await session.call_tool(tool.identifier, arguments=args)
return ToolInvocationResult(
content="\n".join([result.model_dump_json() for result in result.content]),
error_code=1 if result.isError else 0,
)

7
llama_stack/providers/registry/tool_runtime.py
View file

@ -32,6 +32,13 @@ def available_providers() -> List[ProviderSpec]:
module="llama_stack.providers.inline.tool_runtime.code_interpreter",
config_class="llama_stack.providers.inline.tool_runtime.code_interpreter.config.CodeInterpreterToolConfig",
),
InlineProviderSpec(
api=Api.tool_runtime,
provider_type="inline::model-context-protocol",
pip_packages=["mcp"],
module="llama_stack.providers.inline.tool_runtime.model_context_protocol",
config_class="llama_stack.providers.inline.tool_runtime.model_context_protocol.config.ModelContextProtocolConfig",
),
remote_provider_spec(
api=Api.tool_runtime,
adapter=AdapterSpec(

6
llama_stack/providers/remote/tool_runtime/brave_search/brave_search.py
View file

@ -9,8 +9,8 @@ from typing import Any, Dict, List, Optional
import requests
from llama_models.llama3.api.datatypes import BuiltinTool
from llama_stack.apis.common.content_types import URL
from llama_stack.apis.tools import (
MCPConfig,
Tool,
ToolDef,
ToolInvocationResult,
@ -50,7 +50,9 @@ class BraveSearchToolRuntimeImpl(
return provider_data.api_key
async def list_runtime_tools(
self, tool_group_id: Optional[str] = None, mcp_endpoint: Optional[URL] = None
self,
tool_group_id: Optional[str] = None,
mcp_config: Optional[MCPConfig] = None,
) -> List[ToolDef]:
return [
ToolDef(

26
llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py
View file

@ -4,14 +4,15 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
import json
from typing import Any, Dict, List, Optional
from urllib.parse import urlparse
from mcp import ClientSession
from mcp.client.sse import sse_client
from pydantic import TypeAdapter
from llama_stack.apis.common.content_types import URL
from llama_stack.apis.tools import (
MCPConfig,
ToolDef,
ToolInvocationResult,
ToolParameter,
@ -30,13 +31,15 @@ class ModelContextProtocolToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime):
pass
async def list_runtime_tools(
self, tool_group_id: Optional[str] = None, mcp_endpoint: Optional[URL] = None
self,
tool_group_id: Optional[str] = None,
mcp_config: Optional[MCPConfig] = None,
) -> List[ToolDef]:
if mcp_endpoint is None:
raise ValueError("mcp_endpoint is required")
if mcp_config is None:
raise ValueError("mcp_config is required")
tools = []
async with sse_client(mcp_endpoint.uri) as streams:
async with sse_client(mcp_config.mcp_endpoint.uri) as streams:
async with ClientSession(*streams) as session:
await session.initialize()
tools_result = await session.list_tools()
@ -58,7 +61,7 @@ class ModelContextProtocolToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime):
description=tool.description,
parameters=parameters,
metadata={
"endpoint": mcp_endpoint.uri,
"mcp_config": mcp_config,
},
)
)
@ -68,13 +71,12 @@ class ModelContextProtocolToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime):
self, tool_name: str, args: Dict[str, Any]
) -> ToolInvocationResult:
tool = await self.tool_store.get_tool(tool_name)
if tool.metadata is None or tool.metadata.get("endpoint") is None:
if tool.metadata is None or tool.metadata.get("mcp_config") is None:
raise ValueError(f"Tool {tool_name} does not have metadata")
endpoint = tool.metadata.get("endpoint")
if urlparse(endpoint).scheme not in ("http", "https"):
raise ValueError(f"Endpoint {endpoint} is not a valid HTTP(S) URL")
mcp_config_dict = json.loads(tool.metadata.get("mcp_config"))
mcp_config = TypeAdapter(MCPConfig).validate_python(mcp_config_dict)
async with sse_client(endpoint) as streams:
async with sse_client(mcp_config.mcp_endpoint.uri) as streams:
async with ClientSession(*streams) as session:
await session.initialize()
result = await session.call_tool(tool.identifier, args)

6
llama_stack/providers/remote/tool_runtime/tavily_search/tavily_search.py
View file

@ -9,8 +9,8 @@ from typing import Any, Dict, List, Optional
import requests
from llama_stack.apis.common.content_types import URL
from llama_stack.apis.tools import (
MCPConfig,
Tool,
ToolDef,
ToolInvocationResult,
@ -50,7 +50,9 @@ class TavilySearchToolRuntimeImpl(
return provider_data.api_key
async def list_runtime_tools(
self, tool_group_id: Optional[str] = None, mcp_endpoint: Optional[URL] = None
self,
tool_group_id: Optional[str] = None,
mcp_config: Optional[MCPConfig] = None,
) -> List[ToolDef]:
return [
ToolDef(

6
llama_stack/providers/remote/tool_runtime/wolfram_alpha/wolfram_alpha.py
View file

@ -9,8 +9,8 @@ from typing import Any, Dict, List, Optional
import requests
from llama_stack.apis.common.content_types import URL
from llama_stack.apis.tools import (
MCPConfig,
Tool,
ToolDef,
ToolInvocationResult,
@ -51,7 +51,9 @@ class WolframAlphaToolRuntimeImpl(
return provider_data.api_key
async def list_runtime_tools(
self, tool_group_id: Optional[str] = None, mcp_endpoint: Optional[URL] = None
self,
tool_group_id: Optional[str] = None,
mcp_config: Optional[MCPConfig] = None,
) -> List[ToolDef]:
return [
ToolDef(