Newer
Older
navi-1 / navi / core / registry.py
@Eugene Sukhodolskiy Eugene Sukhodolskiy 3 hours ago 10 KB profiles: tool_developer folds into developer
"""Registries for tools, profiles, and LLM backends."""

import structlog

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,
    NotifyTool,
    PeerTool,
    ReflectTool,
    PlanTool,
    ScheduleRecallTool,
    SpawnAgentTool,
    SshExecTool,
    ScratchpadTool,
    SwitchProfileTool,
    SynapseInstructionsTool,
    TasksTool,
    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

log = structlog.get_logger()


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 clear_external_tools(self) -> list[str]:
        """Drop every external tool and return the names removed.

        External registrations are rebuilt from scratch on reload: an MCP
        server removed from the config (or a tool it stopped exposing) has no
        other way to leave the registry.
        """
        removed = sorted(self._external_names)
        for name in removed:
            self.unregister_external(name)
        return removed

    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)
    # profile_registry + mcp_manager let it resolve a name against the profile's own
    # toolset the same way the executor will — and say "registered but not enabled
    # for this profile" instead of a bare "no manual found".
    manual_tool = ToolManualTool(
        registry=tools, profile_registry=profiles, mcp_manager=mcp_manager
    )
    memory_tool = MemoryTool(memory_store) if memory_store else None
    mcp_status_tool = McpStatusTool(mcp_manager=mcp_manager)
    create_mcp_server_tool = CreateMcpServerTool()
    test_mcp_tool_tool = TestMcpToolTool(mcp_manager=mcp_manager)
    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,
        tool_registry=tools,
        mcp_manager=mcp_manager,
    )
    list_profiles_tool = ListProfilesTool(profile_registry=profiles)
    synapse_instructions_tool = SynapseInstructionsTool(
        pool_provider=getattr(session_store, "_get_pool", None),
    )
    notify_tool = NotifyTool()

    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),
                PlanTool(),
                PeerTool(),
                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,
                synapse_instructions_tool, notify_tool,
                TasksTool()]
    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