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.fallback import FallbackOllamaBackend, ServerEntry, load_servers_from_file
from navi.profiles import ALL_PROFILES
from navi.profiles.base import AgentProfile
from navi.tools import (
    CodeExecTool,
    CreateMcpServerTool,
    FilesystemTool,
    ImageViewTool,
    ListProfilesTool,
    ManageRecallTool,
    MemoryTool,
    ReflectTool,
    ScheduleRecallTool,
    SpawnAgentTool,
    SshExecTool,
    ScratchpadTool,
    SwitchProfileTool,
    TerminalTool,
    TestMcpToolTool,
    TodoTool,
    Tool,
)
from navi.tools.list_tools import ListToolsTool
from navi.tools.reload_tools import ReloadToolsTool
from navi.tools.tool_manual import ToolManualTool
from navi.tools.share_file import ShareFileTool
from navi.tools.content_publish import ContentPublishTool
from navi.tools.mcp_status import McpStatusTool
from navi.tools._internal.loader import LoadResult, load_tools_from_dir
from navi.tools._internal.logging_middleware import LoggingMiddleware
from navi.context_providers._loader import ContextProviderRegistry


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

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

    def register_external(self, tool: Tool) -> None:
        """Register a tool from an external source (e.g. MCP server).

        External tools survive ``reload_user_tools()`` just like builtins.
        """
        self._tools[tool.name] = tool
        self._external_names.add(tool.name)

    def unregister_external(self, name: str) -> None:
        """Remove a previously registered external tool."""
        self._external_names.discard(name)
        self._tools.pop(name, None)

    def add_middleware(self, middleware) -> None:
        """Add a ToolMiddleware instance."""
        self._middlewares.append(middleware)

    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 (not builtin, not external)
        for name in list(self._tools):
            if name not in self._builtin_names and name not in self._external_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())

    def update(self, profile: AgentProfile) -> None:
        """Replace an existing profile in-memory."""
        if profile.id not in self._profiles:
            raise ProfileNotFound(profile.id)
        self._profiles[profile.id] = profile


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 _discover_backends() -> list[tuple[str, LLMBackend]]:
    """Auto-discover LLM backends from navi/llm/ modules."""
    discovered: list[tuple[str, LLMBackend]] = []
    from navi.llm.fallback import FallbackOllamaBackend
    from navi.llm.openai_backend import OpenAIBackend
    ollama_http_timeout = max(
        settings.ollama_request_timeout,
        settings.llm_complete_timeout,
        settings.llm_stream_first_chunk_timeout,
    )

    # Ollama backend (primary) — always use FallbackOllamaBackend so model
    # priority lists work regardless of multi-server vs single-server config.
    if settings.ollama_backends_file:
        servers = load_servers_from_file(settings.ollama_backends_file)
        if not servers and settings.ollama_host:
            log.warning("fallback.backends_file_empty", path=settings.ollama_backends_file)
    else:
        servers = []

    if not servers and settings.ollama_host:
        servers = [ServerEntry(
            host=settings.ollama_host,
            api_key=settings.ollama_api_key,
        )]

    if servers:
        discovered.append(("ollama", FallbackOllamaBackend(servers)))

    # OpenAI backend (if configured)
    if settings.openai_api_key:
        discovered.append(("openai", OpenAIBackend(
            model=settings.openai_model,
            api_key=settings.openai_api_key,
            base_url=settings.openai_base_url,
        )))

    return discovered


def build_default_registries(
    memory_store=None,
    session_store=None,
    scheduler=None,
    kv_store=None,
    mcp_manager=None,
    terminal_manager=None,
) -> tuple[ToolRegistry, ProfileRegistry, BackendRegistry, ContextProviderRegistry]:
    """Build and populate registries with all built-in components."""
    from navi.core.ai_helper import AIHelper

    # Phase 1: Create registries that have no cross-dependencies on tools.
    backends = BackendRegistry()
    backend_instances = _discover_backends()
    if not backend_instances:
        raise RuntimeError("No LLM backends discovered. Check OLLAMA_HOST or OPENAI_API_KEY.")

    for key, backend in backend_instances:
        backends.register(key, backend)

    # Use primary backend for AIHelper
    primary_backend = backend_instances[0][1]
    ai_helper = AIHelper(
        backend=primary_backend,
        default_model=settings.ollama_default_model,
    )

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

    cp_registry = ContextProviderRegistry()
    from pathlib import Path as _Path
    cp_registry.load_from_dir(str(_Path(__file__).parent.parent / "context_providers"), builtin=True)
    cp_registry.load_from_dir(settings.context_providers_dir, builtin=False)

    # Phase 2: Create tools with ALL dependencies already available.
    tools = ToolRegistry()

    reload_tool = ReloadToolsTool(registry=tools, cp_registry=cp_registry, mcp_manager=mcp_manager)
    list_tool = ListToolsTool(registry=tools, profile_registry=profiles, mcp_manager=mcp_manager)
    manual_tool = ToolManualTool(registry=tools)
    memory_tool = MemoryTool(memory_store) if memory_store else None
    mcp_status_tool = McpStatusTool()
    create_mcp_server_tool = CreateMcpServerTool()
    test_mcp_tool_tool = TestMcpToolTool()
    schedule_recall_tool = ScheduleRecallTool(scheduler)
    manage_recall_tool = ManageRecallTool(scheduler)
    spawn_tool = SpawnAgentTool(
        profile_registry=profiles,
        tool_registry=tools,
        backend_registry=backends,
        session_store=session_store,
        memory_store=memory_store,
        mcp_manager=mcp_manager,
    )
    switch_tool = SwitchProfileTool(
        session_store=session_store,
        profile_registry=profiles,
    )
    list_profiles_tool = ListProfilesTool(profile_registry=profiles)

    builtins = [FilesystemTool(ai_helper=ai_helper),
                CodeExecTool(), TerminalTool(terminal_manager=terminal_manager), SshExecTool(), ImageViewTool(),
                ShareFileTool(), ContentPublishTool(),
                TodoTool(kv_store=kv_store), ScratchpadTool(kv_store=kv_store),
                ReflectTool(ai_helper=ai_helper),
                reload_tool, list_tool, manual_tool,
                mcp_status_tool, create_mcp_server_tool, test_mcp_tool_tool,
                schedule_recall_tool, manage_recall_tool,
                spawn_tool, switch_tool, list_profiles_tool]
    if memory_tool:
        builtins.append(memory_tool)
    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)

    # Register built-in middleware
    tools.add_middleware(LoggingMiddleware())

    return tools, profiles, backends, cp_registry