Newer
Older
navi-1 / navi / core / registry.py
"""Registries for tools, profiles, and LLM backends."""

from navi.config import settings
from navi.exceptions import ProfileNotFound, ToolNotFound
from navi.llm.base import LLMBackend
from navi.llm.ollama import OllamaBackend
from navi.profiles import ALL_PROFILES
from navi.profiles.base import AgentProfile
from navi.tools import (
    CodeExecTool,
    FilesystemTool,
    HttpRequestTool,
    ImageViewTool,
    MemoryForgetTool,
    MemorySearchTool,
    SpawnAgentTool,
    SshExecTool,
    TerminalTool,
    TodoTool,
    Tool,
    WebSearchTool,
    WebViewTool,
)
from navi.tools.list_tools import ListToolsTool
from navi.tools.reload_tools import ReloadToolsTool
from navi.tools.tool_manual import ToolManualTool
from navi.tools.write_tool import WriteToolTool
from navi.tools.loader import LoadResult, load_tools_from_dir


class ToolRegistry:
    def __init__(self) -> None:
        self._tools: dict[str, Tool] = {}
        self._builtin_names: set[str] = set()

    def register(self, tool: Tool, builtin: bool = False) -> None:
        self._tools[tool.name] = tool
        if builtin:
            self._builtin_names.add(tool.name)

    def get(self, name: str) -> Tool:
        if name not in self._tools:
            raise ToolNotFound(name)
        return self._tools[name]

    def resolve(self, names: list[str]) -> list[Tool]:
        return [self.get(n) for n in names]

    def all(self) -> list[Tool]:
        return list(self._tools.values())

    def reload_user_tools(self, tools_dir: str) -> LoadResult:
        """Remove all user tools and reload from disk. Safe: errors are isolated."""
        # Drop previously loaded user tools
        for name in list(self._tools):
            if name not in self._builtin_names:
                del self._tools[name]

        result = load_tools_from_dir(tools_dir)
        for tool in result.loaded:
            self._tools[tool.name] = tool
        return result


class ProfileRegistry:
    def __init__(self) -> None:
        self._profiles: dict[str, AgentProfile] = {}

    def register(self, profile: AgentProfile) -> None:
        self._profiles[profile.id] = profile

    def get(self, profile_id: str) -> AgentProfile:
        if profile_id not in self._profiles:
            raise ProfileNotFound(profile_id)
        return self._profiles[profile_id]

    def all(self) -> list[AgentProfile]:
        return list(self._profiles.values())


class BackendRegistry:
    def __init__(self) -> None:
        self._backends: dict[str, LLMBackend] = {}

    def register(self, key: str, backend: LLMBackend) -> None:
        self._backends[key] = backend

    def get(self, key: str) -> LLMBackend:
        backend = self._backends.get(key)
        if backend is None:
            raise KeyError(f"LLM backend '{key}' not registered")
        return backend

    def all_keys(self) -> list[str]:
        return list(self._backends.keys())


def build_default_registries(
    memory_store=None,
    session_store=None,
) -> tuple[ToolRegistry, ProfileRegistry, BackendRegistry]:
    """Build and populate registries with all built-in components."""

    tools = ToolRegistry()
    reload_tool = ReloadToolsTool(registry=tools)
    write_tool = WriteToolTool(registry=tools)
    list_tool = ListToolsTool(registry=tools)
    manual_tool = ToolManualTool(registry=tools)
    memory_search = MemorySearchTool(memory_store) if memory_store else None
    memory_forget = MemoryForgetTool(memory_store) if memory_store else None
    builtins = [WebSearchTool(), FilesystemTool(), HttpRequestTool(), WebViewTool(),
                CodeExecTool(), TerminalTool(), SshExecTool(), ImageViewTool(),
                TodoTool(), reload_tool, write_tool, list_tool, manual_tool]
    if memory_search:
        builtins.extend([memory_search, memory_forget])
    for builtin in builtins:
        tools.register(builtin, builtin=True)

    # User-defined tools loaded from tools_dir
    result = load_tools_from_dir(settings.tools_dir)
    for user_tool in result.loaded:
        tools.register(user_tool)

    profiles = ProfileRegistry()
    for p in ALL_PROFILES:
        profiles.register(p)

    # SpawnAgentTool registered after profiles so it can resolve profile names.
    # session_store may be None at build time (injected separately in deps.py).
    spawn_tool = SpawnAgentTool(
        profile_registry=profiles,
        tool_registry=tools,
        backend_registry=None,   # patched below after backends are built
        session_store=session_store,
        memory_store=memory_store,
    )
    tools.register(spawn_tool, builtin=True)

    backends = BackendRegistry()
    backends.register(
        "ollama",
        OllamaBackend(
            model=settings.ollama_default_model,
            host=settings.ollama_host,
        ),
    )
    # Patch backend registry into spawn_tool now that it's available
    spawn_tool._backend_registry = backends

    return tools, profiles, backends