Newer
Older
navi-1 / navi / tools / _internal / base.py
"""
Base class for all tools.

Each tool is self-describing: name, description, and parameters (JSON Schema).
The schema() method builds the LLM-facing function spec automatically.

current_session_id — ContextVar set by Agent before every tool call.
Tools that need per-session state (e.g. SSH connection pool) read it here.
"""

import asyncio
from abc import ABC, abstractmethod
from contextvars import ContextVar
from dataclasses import dataclass, field

from navi.llm.base import ToolSchema

# Set by Agent before every tool call. Tools that need per-session state read this.
current_session_id: ContextVar[str | None] = ContextVar("current_session_id", default=None)

# Set by run_stream() before executing a tool. run_ephemeral() reads this to forward
# sub-agent tool events up to the parent WS stream.
current_event_sink: ContextVar[asyncio.Queue | None] = ContextVar("current_event_sink", default=None)

# Set by _run_agent() before run_stream(). Cooperative stop: when set, the agent
# breaks out of LLM loops cleanly (aclose() is called → Ollama stream closes gracefully,
# model stays in VRAM). Never use task.cancel() for stopping generation.
current_stop_event: ContextVar[asyncio.Event | None] = ContextVar("current_stop_event", default=None)

# Set by run_stream() / run_ephemeral() to expose the current profile's model name
# to tools that need to make their own LLM calls (e.g. AIHelper-powered tools).
current_model: ContextVar[list[str] | str | None] = ContextVar("current_model", default=None)

# Set by run_stream() alongside current_user_id. Admins bypass sandbox restrictions.
current_user_id: ContextVar[str | None] = ContextVar("current_user_id", default=None)

# Set by run_stream() alongside current_user_id. Admins bypass sandbox restrictions.
current_user_role: ContextVar[str] = ContextVar("current_user_role", default="user")

# Set by run_stream() / messages endpoint to expose the full user profile to
# ContextBuilder so the LLM receives [User context] with name, email, locale, etc.
current_user_info: ContextVar[dict | None] = ContextVar("current_user_info", default=None)


@dataclass
class ToolContext:
    """Explicit execution context passed to tools instead of hidden ContextVars."""

    session_id: str | None = None
    event_sink: asyncio.Queue | None = None
    stop_event: asyncio.Event | None = None
    model: list[str] | str | None = None
    user_id: str | None = None
    user_role: str = "user"
    user_info: dict | None = None


@dataclass
class ToolResult:
    success: bool
    output: str  # always a string — LLM consumes this
    error: str | None = None
    metadata: dict = field(default_factory=dict)

    def to_message_content(self) -> str:
        if self.success:
            return self.output
        # Always include output — it contains tracebacks and details the model needs.
        if self.error and self.output:
            return f"Error: {self.error}\n{self.output}"
        if self.error:
            return f"Error: {self.error}"
        return self.output


class Tool(ABC):
    """Abstract base for all tools."""

    # Override in subclasses:
    name: str
    description: str
    parameters: dict  # JSON Schema object

    @abstractmethod
    async def execute(self, params: dict, ctx: ToolContext | None = None) -> ToolResult:
        """Execute the tool with given parameters."""

    def schema(self) -> ToolSchema:
        return ToolSchema(
            type="function",
            function={
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        )