632ea8c032
Co-authored-by: Antoine <33425718+anth2o@users.noreply.github.com> Co-authored-by: Bastien <bastien.baret@gmail.com> Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Kim-Adeline Miguel <51720070+kimadeline@users.noreply.github.com> Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Maxime Dolores <maxime.dolores@ext.mistral.ai> Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Nelson PROIA <144663685+Nelson-PROIA@users.noreply.github.com> Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai> Co-authored-by: Quentin <quentin.torroba@mistral.ai> Co-authored-by: Robin Gullo <robin.gullo@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from vibe.core.tools.mcp import MCPRegistry
|
|
from vibe.core.tools.mcp.tools import (
|
|
RemoteTool,
|
|
create_mcp_http_proxy_tool_class,
|
|
create_mcp_stdio_proxy_tool_class,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from vibe.core.config import MCPServer
|
|
from vibe.core.tools.base import BaseTool
|
|
|
|
|
|
_BROKEN_SERVER_NAME = "broken-server"
|
|
|
|
|
|
class FakeMCPRegistry(MCPRegistry):
|
|
async def get_tools_async(
|
|
self, servers: list[MCPServer]
|
|
) -> dict[str, type[BaseTool]]:
|
|
return self.get_tools(servers)
|
|
|
|
def set_tools(
|
|
self, servers: list[MCPServer], tools: dict[str, type[BaseTool]]
|
|
) -> None:
|
|
for srv in servers:
|
|
key = self._server_key(srv)
|
|
self._cache[key] = dict(tools)
|
|
|
|
def get_tools(self, servers: list[MCPServer]) -> dict[str, type[BaseTool]]:
|
|
result: dict[str, type[BaseTool]] = {}
|
|
for srv in servers:
|
|
key = self._server_key(srv)
|
|
if key not in self._cache:
|
|
remote = RemoteTool(
|
|
name="fake_tool", description=f"A fake tool for {srv.name}"
|
|
)
|
|
match srv.transport:
|
|
case "stdio":
|
|
proxy = create_mcp_stdio_proxy_tool_class(
|
|
command=["fake-cmd"], remote=remote, alias=srv.name
|
|
)
|
|
case "http" | "streamable-http":
|
|
proxy = create_mcp_http_proxy_tool_class(
|
|
url="http://fake-mcp-server", remote=remote, alias=srv.name
|
|
)
|
|
case _:
|
|
raise ValueError(
|
|
f"FakeMCPRegistry: unsupported transport {srv.transport!r}"
|
|
)
|
|
self._cache[key] = {proxy.get_name(): proxy}
|
|
result.update(self._cache[key])
|
|
return result
|
|
|
|
|
|
class FakeMCPRegistryWithBrokenServer(FakeMCPRegistry):
|
|
def get_tools(self, servers: list[MCPServer]) -> dict[str, type[BaseTool]]:
|
|
working = [s for s in servers if s.name != _BROKEN_SERVER_NAME]
|
|
return super().get_tools(working)
|