summaryrefslogtreecommitdiff
path: root/src/ollama_web/tools/registry.py
diff options
context:
space:
mode:
authortanaka takashi <isazaka@gmail.com>2026-06-26 23:54:17 +0900
committertanaka takashi <isazaka@gmail.com>2026-06-26 23:54:17 +0900
commit3a9e8e18df4c596360da54000f65e1bebc6039c2 (patch)
tree0c6dc5fd75cc8f751ca6d7dca01c7e7b5fb3c9f5 /src/ollama_web/tools/registry.py
parenta296c2c3a294648c04286a5c027deddbbdc8d5ae (diff)
MCP機能追加
Diffstat (limited to 'src/ollama_web/tools/registry.py')
-rw-r--r--src/ollama_web/tools/registry.py43
1 files changed, 42 insertions, 1 deletions
diff --git a/src/ollama_web/tools/registry.py b/src/ollama_web/tools/registry.py
index 2b6c7ae..17c755d 100644
--- a/src/ollama_web/tools/registry.py
+++ b/src/ollama_web/tools/registry.py
@@ -5,6 +5,7 @@ from __future__ import annotations
import json
from collections.abc import Callable, Sequence
from dataclasses import dataclass, field
+from typing import Any
from .fetch import search_and_fetch
from .pdf import pdf_to_text
@@ -17,12 +18,35 @@ class ToolRegistry:
"""A simple registry mapping tool names to callable implementations."""
_tools: dict[str, Callable[..., object]] = field(default_factory=dict)
+ _ollama_tool_defs: list[dict[str, Any]] = field(default_factory=list)
+ _mcp_tool_names: set[str] = field(default_factory=set)
def register(self, func: Callable[..., object]) -> Callable[..., object]:
"""Register a callable under its ``__name__``."""
self._tools[func.__name__] = func
return func
+ def register_mcp_tool(
+ self,
+ name: str,
+ tool_def: dict[str, Any],
+ executor: Callable[..., object],
+ ) -> Callable[..., object]:
+ """Register a remote MCP tool under a normalized name.
+
+ Args:
+ name: The normalized ollama-visible tool name.
+ tool_def: The ollama-compatible tool definition dict.
+ executor: A callable that accepts ``(full_name, arguments)`` and
+ returns a JSON-encoded result string.
+ """
+ self._tools[name] = executor
+ self._mcp_tool_names.add(name)
+ # Avoid duplicate definitions if a registry is rebuilt.
+ self._ollama_tool_defs = [d for d in self._ollama_tool_defs if d.get("name") != name]
+ self._ollama_tool_defs.append(tool_def)
+ return executor
+
@property
def names(self) -> list[str]:
return list(self._tools.keys())
@@ -31,6 +55,19 @@ class ToolRegistry:
def callables(self) -> list[Callable[..., object]]:
return list(self._tools.values())
+ @property
+ def ollama_tools(self) -> list[dict[str, Any] | Callable[..., object]]:
+ """Return a list suitable for ollama's ``tools`` argument.
+
+ Built-in tools are passed as their callables so ollama can infer schemas
+ from docstrings. MCP tools are passed as pre-built dict definitions.
+ """
+ # Only return built-in callables and dict definitions; skip the MCP
+ # executor callables stored in the tools dict.
+ return [
+ func for name, func in self._tools.items() if name not in self._mcp_tool_names
+ ] + self._ollama_tool_defs
+
def get(self, name: str) -> Callable[..., object] | None:
return self._tools.get(name)
@@ -57,7 +94,11 @@ class ToolRegistry:
return json.dumps({"error": f"Invalid JSON arguments for {name}: {arguments}"})
try:
- result = func(**arguments)
+ if name in self._mcp_tool_names:
+ # MCP executor expects (full_name, arguments).
+ result = func(name, arguments)
+ else:
+ result = func(**arguments)
except Exception as exc: # noqa: BLE001
return json.dumps({"error": f"{type(exc).__name__}: {exc}"})