From 3a9e8e18df4c596360da54000f65e1bebc6039c2 Mon Sep 17 00:00:00 2001 From: tanaka takashi Date: Fri, 26 Jun 2026 23:54:17 +0900 Subject: MCP機能追加 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/ollama_web/mcp.py | 488 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 488 insertions(+) create mode 100644 src/ollama_web/mcp.py (limited to 'src/ollama_web/mcp.py') diff --git a/src/ollama_web/mcp.py b/src/ollama_web/mcp.py new file mode 100644 index 0000000..6cd3cc9 --- /dev/null +++ b/src/ollama_web/mcp.py @@ -0,0 +1,488 @@ +"""MCP client integration for ollama-web. + +This module loads MCP server configuration from ``data/mcpServers.json`` and +bridges remote MCP tools into the local ``ToolRegistry`` used by the ollama +tool-calling loop. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from .config import settings + +logger = logging.getLogger("ollama_web.mcp") + +DEFAULT_TIMEOUT = 30.0 + +_MCP_TOOL_PREFIX = "mcp__" + + +def mcp_servers_file(data_dir: str | Path | None = None) -> Path: + """Return the path to ``mcpServers.json`` under the configured data dir.""" + return Path(data_dir or settings.data_dir).resolve() / "mcpServers.json" + + +def load_mcp_config(data_dir: str | Path | None = None) -> dict[str, Any]: + """Load the MCP server configuration from disk. + + Returns: + A dict with a ``mcpServers`` key, or an empty default config when the + file does not exist or is malformed. + """ + path = mcp_servers_file(data_dir) + if not path.exists(): + return {"mcpServers": {}} + try: + data = json.loads(path.read_text(encoding="utf-8")) + except Exception as exc: # noqa: BLE001 + logger.warning("failed to read %s: %s", path, exc) + return {"mcpServers": {}} + + if not isinstance(data, dict): + return {"mcpServers": {}} + servers = data.get("mcpServers", {}) + if not isinstance(servers, dict): + return {"mcpServers": {}} + return {"mcpServers": servers} + + +def save_mcp_config( + data_dir: str | Path | None = None, + config: dict[str, Any] | None = None, +) -> None: + """Persist the MCP server configuration to disk. + + Args: + data_dir: Data directory root. Defaults to ``settings.data_dir``. + config: Configuration dict; must contain a ``mcpServers`` key. + """ + path = mcp_servers_file(data_dir) + path.parent.mkdir(parents=True, exist_ok=True) + config = config or {"mcpServers": {}} + if "mcpServers" not in config: + config = {"mcpServers": {}} + path.write_text( + json.dumps(config, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + +def _is_stdio(server: dict[str, Any]) -> bool: + """Return whether the server definition describes a stdio transport.""" + return "command" in server + + +def _is_http(server: dict[str, Any]) -> bool: + """Return whether the server definition describes an HTTP transport.""" + return "url" in server + + +def _normalized_tool_name(server_name: str, tool_name: str) -> str: + """Build a unique ollama-visible tool name for an MCP tool.""" + return f"{_MCP_TOOL_PREFIX}{server_name}__{tool_name}" + + +def _parse_tool_name(full_name: str) -> tuple[str, str] | None: + """Recover ``(server_name, tool_name)`` from a normalized MCP tool name.""" + if not full_name.startswith(_MCP_TOOL_PREFIX): + return None + rest = full_name[len(_MCP_TOOL_PREFIX) :] + parts = rest.split("__", 1) + if len(parts) != 2: + return None + return parts[0], parts[1] + + +def _to_ollama_tool(server_name: str, tool: Any) -> dict[str, Any]: + """Convert an MCP ``Tool`` object into an ollama-compatible tool dict. + + The resulting dict uses the OpenAI-compatible function-calling shape that + ollama accepts as an element of the ``tools`` argument. + """ + name = getattr(tool, "name", "") + description = getattr(tool, "description", "") or name + schema = getattr(tool, "inputSchema", None) or {} + + return { + "type": "function", + "function": { + "name": _normalized_tool_name(server_name, str(name)), + "description": str(description), + "parameters": _clean_schema(schema), + }, + } + + +def _clean_schema(schema: Any) -> dict[str, Any]: + """Return a JSON schema suitable for ollama, removing unsupported keys.""" + if not isinstance(schema, dict): + return {"type": "object", "properties": {}, "required": []} + + cleaned: dict[str, Any] = { + "type": schema.get("type", "object"), + "properties": {}, + "required": [], + } + properties = schema.get("properties") + if isinstance(properties, dict): + cleaned["properties"] = properties + required = schema.get("required") + if isinstance(required, list): + cleaned["required"] = [str(r) for r in required] + + # Preserve additional metadata ollama understands, if present. + for key in ("additionalProperties", "anyOf", "oneOf", "enum"): + if key in schema: + cleaned[key] = schema[key] + return cleaned + + +async def _list_tools_for_server( + server_name: str, + server: dict[str, Any], +) -> list[dict[str, Any]]: + """Connect to a single MCP server and return its tools in ollama form. + + Exceptions are caught and logged so that one misbehaving server does not + break the whole tool registry. + """ + tools: list[dict[str, Any]] = [] + if _is_stdio(server): + tools = await _list_tools_stdio(server_name, server) + elif _is_http(server): + tools = await _list_tools_http(server_name, server) + else: + logger.warning("MCP server %s has no recognized transport", server_name) + return tools + + +async def _list_tools_stdio( + server_name: str, + server: dict[str, Any], +) -> list[dict[str, Any]]: + """List tools from a stdio MCP server.""" + from mcp import ClientSession + from mcp.client.stdio import StdioServerParameters, stdio_client + + command = str(server["command"]) + args = [str(a) for a in server.get("args", [])] + env = server.get("env") + cwd = server.get("cwd") + + params = StdioServerParameters( + command=command, + args=args, + env=env, + cwd=cwd, + ) + + async with stdio_client(params) as (read_stream, write_stream): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + result = await session.list_tools() + return [_to_ollama_tool(server_name, t) for t in getattr(result, "tools", [])] + + +async def _list_tools_http( + server_name: str, + server: dict[str, Any], +) -> list[dict[str, Any]]: + """List tools from an HTTP (streamable) MCP server.""" + import httpx + from mcp import ClientSession + from mcp.client.streamable_http import streamable_http_client + + url = str(server["url"]) + headers: dict[str, str] | None = None + raw_headers = server.get("headers") + if isinstance(raw_headers, dict): + headers = {str(k): str(v) for k, v in raw_headers.items()} + + timeout = server.get("timeout", DEFAULT_TIMEOUT) + try: + timeout = float(timeout) + except (TypeError, ValueError): + timeout = DEFAULT_TIMEOUT + + sse_read_timeout = server.get("sse_read_timeout", DEFAULT_TIMEOUT) + try: + sse_read_timeout = float(sse_read_timeout) + except (TypeError, ValueError): + sse_read_timeout = DEFAULT_TIMEOUT + + client = httpx.AsyncClient( + headers=headers, + timeout=httpx.Timeout(timeout, read=sse_read_timeout), + ) + async with client: + async with streamable_http_client(url, http_client=client) as ( + read_stream, + write_stream, + _get_session_id, + ): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + result = await session.list_tools() + return [_to_ollama_tool(server_name, t) for t in getattr(result, "tools", [])] + + +async def collect_mcp_tools( + config: dict[str, Any] | None = None, +) -> list[dict[str, Any]]: + """Collect ollama-compatible tool definitions from all configured MCP servers. + + Args: + config: MCP configuration dict. If None, loaded from disk. + + Returns: + A list of ollama tool dicts (``{"type": "function", "function": {...}}``). + """ + if config is None: + config = load_mcp_config() + + servers = config.get("mcpServers", {}) + if not isinstance(servers, dict): + return [] + + tasks = [ + _list_tools_for_server(name, server) + for name, server in servers.items() + if isinstance(server, dict) + ] + if not tasks: + return [] + + results = await asyncio.gather(*tasks, return_exceptions=True) + all_tools: list[dict[str, Any]] = [] + for name, result in zip(servers.keys(), results, strict=True): + if isinstance(result, Exception): + logger.warning("failed to list tools from MCP server %s: %s", name, result) + continue + all_tools.extend(result) + return all_tools + + +def collect_mcp_tools_sync(config: dict[str, Any] | None = None) -> list[dict[str, Any]]: + """Synchronous wrapper around ``collect_mcp_tools``. + + Useful for callers running in a thread without an active event loop. + """ + return asyncio.run(collect_mcp_tools(config)) + + +async def _call_tool_on_server( + server_name: str, + server: dict[str, Any], + tool_name: str, + arguments: dict[str, Any], +) -> str: + """Call a single tool on the specified MCP server and return its result. + + The result content is flattened into a JSON string. + """ + if _is_stdio(server): + return await _call_tool_stdio(server_name, server, tool_name, arguments) + if _is_http(server): + return await _call_tool_http(server_name, server, tool_name, arguments) + raise ValueError(f"server {server_name} has no recognized transport") + + +async def _call_tool_stdio( + server_name: str, + server: dict[str, Any], + tool_name: str, + arguments: dict[str, Any], +) -> str: + from mcp import ClientSession + from mcp.client.stdio import StdioServerParameters, stdio_client + + command = str(server["command"]) + args = [str(a) for a in server.get("args", [])] + env = server.get("env") + cwd = server.get("cwd") + + params = StdioServerParameters( + command=command, + args=args, + env=env, + cwd=cwd, + ) + + async with stdio_client(params) as (read_stream, write_stream): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + result = await session.call_tool(tool_name, arguments) + return _format_tool_result(result) + + +async def _call_tool_http( + server_name: str, + server: dict[str, Any], + tool_name: str, + arguments: dict[str, Any], +) -> str: + import httpx + from mcp import ClientSession + from mcp.client.streamable_http import streamable_http_client + + url = str(server["url"]) + headers: dict[str, str] | None = None + raw_headers = server.get("headers") + if isinstance(raw_headers, dict): + headers = {str(k): str(v) for k, v in raw_headers.items()} + + timeout = server.get("timeout", DEFAULT_TIMEOUT) + try: + timeout = float(timeout) + except (TypeError, ValueError): + timeout = DEFAULT_TIMEOUT + + sse_read_timeout = server.get("sse_read_timeout", DEFAULT_TIMEOUT) + try: + sse_read_timeout = float(sse_read_timeout) + except (TypeError, ValueError): + sse_read_timeout = DEFAULT_TIMEOUT + + client = httpx.AsyncClient( + headers=headers, + timeout=httpx.Timeout(timeout, read=sse_read_timeout), + ) + async with client: + async with streamable_http_client(url, http_client=client) as ( + read_stream, + write_stream, + _get_session_id, + ): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + result = await session.call_tool(tool_name, arguments) + return _format_tool_result(result) + + +def _format_tool_result(result: Any) -> str: + """Convert a ``CallToolResult`` into a JSON string for ollama.""" + content: list[Any] = [] + is_error = False + if isinstance(result, dict): + content = result.get("content", []) + is_error = bool(result.get("isError")) + else: + content = getattr(result, "content", []) + is_error = getattr(result, "isError", False) + + parts: list[str] = [] + for item in content: + if isinstance(item, dict): + item_type = item.get("type", "") + if item_type == "text": + parts.append(str(item.get("text", ""))) + elif item_type == "image": + parts.append(f"[image/{item.get('mimeType', 'unknown')}]") + elif item_type == "resource": + resource = item.get("resource", {}) + parts.append(_format_resource(resource)) + else: + text = getattr(item, "text", None) + if text is not None: + parts.append(str(text)) + else: + data = getattr(item, "data", None) + mime = getattr(item, "mimeType", None) + if data is not None: + parts.append(f"[image/{mime or 'unknown'}]") + else: + resource = getattr(item, "resource", None) + if resource is not None: + parts.append(_format_resource(resource)) + + body = "\n\n".join(parts) + if is_error: + return json.dumps({"error": body or "MCP tool returned an error"}, ensure_ascii=False) + return json.dumps({"result": body}, ensure_ascii=False) + + +def _format_resource(resource: Any) -> str: + """Convert an embedded resource reference into a short text snippet.""" + if isinstance(resource, dict): + uri = resource.get("uri", "") + mime = resource.get("mimeType", "") + text = resource.get("text", "") + if text: + return f"[resource {uri} ({mime})]\n{text}" + return f"[resource {uri} ({mime})]" + return f"[resource {resource}]" + + +def call_mcp_tool_sync( + config: dict[str, Any] | None, + full_name: str, + arguments: dict[str, Any] | str = "", +) -> str: + """Synchronous wrapper for calling an MCP tool by its normalized name. + + This is registered as the executor in ``ToolRegistry`` so that the + synchronous ollama tool loop can invoke remote MCP tools. + + Args: + config: MCP configuration dict. If None, loaded from disk. + full_name: Normalized MCP tool name (``mcp__server__tool``). + arguments: Tool arguments as a dict or JSON string. + + Returns: + A JSON-encoded result string, or an error object on failure. + """ + parsed = _parse_tool_name(full_name) + if parsed is None: + return json.dumps({"error": f"Invalid MCP tool name: {full_name}"}, ensure_ascii=False) + + server_name, tool_name = parsed + + if config is None: + config = load_mcp_config() + server = config.get("mcpServers", {}).get(server_name) + if not isinstance(server, dict): + return json.dumps( + {"error": f"MCP server not found for tool: {full_name}"}, ensure_ascii=False + ) + + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) if arguments else {} + except json.JSONDecodeError: + return json.dumps( + {"error": f"Invalid JSON arguments for {full_name}: {arguments}"}, + ensure_ascii=False, + ) + + try: + return asyncio.run(_call_tool_on_server(server_name, server, tool_name, arguments)) + except Exception as exc: # noqa: BLE001 + logger.warning("MCP tool %s execution failed: %s", full_name, exc) + return json.dumps({"error": f"{type(exc).__name__}: {exc}"}, ensure_ascii=False) + + +def make_mcp_executor( + config: dict[str, Any] | None = None, +) -> Callable[..., object]: + """Return a callable that executes an MCP tool by normalized name. + + The returned callable accepts ``full_name`` (injected by the registry) and + optionally ``arguments`` (ollama may pass extra keyword arguments). + """ + if config is None: + config = load_mcp_config() + + def executor( + full_name: str, + arguments: dict[str, Any] | str = "", + **_kwargs: Any, + ) -> str: + return call_mcp_tool_sync(config, full_name, arguments) + + return executor \ No newline at end of file -- cgit v1.2.3