aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/middleware/filesystem.py

@@ -0,0 +1,690 @@
+"""Filesystem middleware for agent file operations.
+
+Provides file storage and management capabilities for agents, enabling them
+to read, write, edit, and search files within isolated thread-specific storage.
+
+TODO: When migrating to deepagents:
+  1. Import BackendProtocol from deepagents.backends.protocol
+  2. Keep InMemoryBackend or use deepagents.backends.StateBackend
+  3. Tools will import from deepagents middleware tools
+
+Authors:
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from functools import lru_cache
+from pathlib import Path
+from typing import Any
+
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, ConfigDict, Field
+
+from aip_agents.common.callable_introspection import callable_accepts_keyword
+from aip_agents.middleware.backends.in_memory import InMemoryBackend
+from aip_agents.middleware.backends.local_disk import LocalDiskBackend
+from aip_agents.middleware.backends.protocol import BackendProtocol
+from aip_agents.middleware.backends.utils import sanitize_tool_call_id
+from aip_agents.middleware.base import AgentMiddleware, ModelRequest
+from aip_agents.middleware.tools.edit_file import EditFileTool
+from aip_agents.middleware.tools.execute import ExecuteTool
+from aip_agents.middleware.tools.grep_file import GrepTool
+from aip_agents.middleware.tools.ls import LsTool
+from aip_agents.middleware.tools.read_file import ReadFileTool
+from aip_agents.middleware.tools.write_file import WriteFileTool
+from aip_agents.sandbox.defaults import DEFAULT_SANDBOX_TEMPLATE
+from aip_agents.utils.logger import get_logger
+
+# Execution system prompt addition (shown when execute tool is available)
+EXECUTION_SYSTEM_PROMPT = """## Execute Tool `execute`
+
+You have access to an `execute` tool for running shell commands.
+Use this tool to run commands, scripts, tests, builds, and other shell operations.
+
+- execute: run a shell command (returns output and exit code)
+
+⚠️ Always handle errors gracefully. Read stderr if execution fails."""
+
+
+def _backend_supports_execution(backend: BackendProtocol) -> bool:
+    """Check if backend supports command execution.
+
+    Handles both direct backends and CompositeBackend which may wrap
+    a sandbox backend as its default.
+
+    Args:
+        backend: The backend to check.
+
+    Returns:
+        True if the backend supports execution, False otherwise.
+    """
+    try:
+        composite_module = __import__(
+            "aip_agents.middleware.backends.composite",
+            fromlist=["CompositeBackend"],
+        )
+        composite_backend_cls = getattr(composite_module, "CompositeBackend")
+
+        if isinstance(backend, composite_backend_cls):
+            default = getattr(backend, "default", None)
+            if default is None:
+                return False
+            return _backend_supports_execution(default)
+    except ImportError:
+        # CompositeBackend is optional in this repository state.
+        # Keep compatibility with existing tests that use a class-name-based mock.
+        if getattr(backend, "__class__", None) and backend.__class__.__name__ == "CompositeBackend":
+            default = getattr(backend, "default", None)
+            if default is None:
+                return False
+            return _backend_supports_execution(default)
+
+    # For other backends, check if execute method exists
+    return hasattr(backend, "execute") and callable(getattr(backend, "execute"))
+
+
+def _is_execution_enabled(backend: BackendProtocol) -> bool:
+    """Check if execution is enabled for this backend.
+
+    Even if a backend supports execution (has execute method),
+    it may be disabled via allow_execute flag.
+
+    Args:
+        backend: The backend to check.
+
+    Returns:
+        True if execution is enabled, False otherwise.
+    """
+    # Check if backend supports execution first
+    if not _backend_supports_execution(backend):
+        return False
+
+    try:
+        composite_module = __import__(
+            "aip_agents.middleware.backends.composite",
+            fromlist=["CompositeBackend"],
+        )
+        composite_backend_cls = getattr(composite_module, "CompositeBackend")
+
+        if isinstance(backend, composite_backend_cls):
+            default = getattr(backend, "default", None)
+            if default is None:
+                return False
+            return _is_execution_enabled(default)
+    except ImportError:
+        if getattr(backend, "__class__", None) and backend.__class__.__name__ == "CompositeBackend":
+            default = getattr(backend, "default", None)
+            if default is None:
+                return False
+            return _is_execution_enabled(default)
+
+    # Check if allow_execute is set to True (defaults to False for security)
+    return getattr(backend, "allow_execute", False)
+
+
+@lru_cache(maxsize=128)
+def _execute_accepts_timeout(backend_class: type) -> bool:
+    """Check whether a backend class's `execute` accepts a `timeout` kwarg.
+
+    Args:
+        backend_class: The backend class to inspect.
+
+    Returns:
+        True if the execute method accepts a timeout parameter, False otherwise.
+    """
+    execute_fn = getattr(backend_class, "execute", None)
+    return callable_accepts_keyword(execute_fn, "timeout")
+
+
+__all__ = [
+    "FilesystemConfig",
+    "InMemoryConfig",
+    "LocalDiskConfig",
+    "SandboxConfig",
+    "FilesystemMiddleware",
+    "FILESYSTEM_SYSTEM_PROMPT",
+    "EXECUTION_SYSTEM_PROMPT",
+    "_backend_supports_execution",
+    "_is_execution_enabled",
+    "_execute_accepts_timeout",
+]
+
+logger = get_logger(__name__)
+
+
+class FilesystemConfig(BaseModel):
+    """Abstract base class for filesystem backend configurations.
+
+    Subclasses implement create_backend() to instantiate specific backends.
+    This allows users to configure filesystem storage with custom settings.
+
+    Example:
+        >>> config = InMemoryConfig()
+        >>> backend = config.create_backend()
+        >>> middleware = FilesystemMiddleware(backend=backend)
+
+        >>> config = LocalDiskConfig(base_directory="/tmp/agent_files")
+        >>> backend = config.create_backend()
+        >>> middleware = FilesystemMiddleware(backend=backend)
+    """
+
+    @abstractmethod
+    def create_backend(self) -> BackendProtocol:
+        """Create and return a backend instance based on this configuration.
+
+        Returns:
+            BackendProtocol: Configured backend instance ready for use.
+        """
+        ...
+
+
+class InMemoryConfig(FilesystemConfig):
+    """Configuration for in-memory filesystem backend.
+
+    Stores all files in memory. Files are lost when the process exits.
+    Best for temporary storage and testing scenarios.
+
+    Example:
+        >>> config = InMemoryConfig()
+        >>> agent = LangGraphReactAgent(
+        ...     name="test_agent",
+        ...     instruction="Test",
+        ...     filesystem=config
+        ... )
+    """
+
+    def create_backend(self) -> BackendProtocol:
+        """Create an in-memory backend.
+
+        Returns:
+            BackendProtocol: InMemoryBackend instance.
+        """
+        return InMemoryBackend()
+
+    model_config = ConfigDict(frozen=True)
+
+
+class LocalDiskConfig(FilesystemConfig):
+    """Configuration for local disk filesystem backend.
+
+    Stores files on the local filesystem at the specified base path.
+    Files persist between process restarts.
+    Best for production use cases requiring file persistence.
+
+
+    Args:
+        base_directory: Root directory for file storage.
+
+    Example:
+        >>> import tempfile
+        >>> with tempfile.TemporaryDirectory() as tmpdir:
+        ...     config = LocalDiskConfig(base_directory=tmpdir)
+        ...     agent = LangGraphReactAgent(
+        ...         name="file_agent",
+        ...         instruction="Process files",
+        ...         filesystem=config
+        ...     )
+    """
+
+    base_directory: str | Path = Field(
+        ...,
+        description="Root directory for file storage.",
+    )
+    allow_execute: bool = Field(
+        default=False,
+        description="Whether execute() is enabled for this local disk backend.",
+    )
+    env: dict[str, str] | None = Field(
+        default=None,
+        description="Optional environment variables overlay for execute() calls.",
+    )
+
+    def create_backend(self) -> BackendProtocol:
+        """Create a local disk backend.
+
+        Returns:
+            BackendProtocol: LocalDiskBackend instance configured with base_path.
+        """
+        return LocalDiskBackend(
+            base_directory=Path(self.base_directory),
+            allow_execute=self.allow_execute,
+            env=self.env,
+        )
+
+    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
+
+
+class SandboxConfig(FilesystemConfig):
+    """Configuration for sandbox filesystem backend.
+
+    Stores files and command execution inside E2B sandbox environments.
+    """
+
+    template: str = Field(
+        default=DEFAULT_SANDBOX_TEMPLATE,
+        description="E2B template alias. Defaults to DEFAULT_SANDBOX_TEMPLATE.",
+    )
+    base_dir: str | None = Field(
+        default=None,
+        description=(
+            "Base workspace directory inside the sandbox. When None, resolved at "
+            "runtime to $HOME/workspace (falling back to /home/user/workspace) so "
+            "it is writable by the sandbox user."
+        ),
+    )
+    allow_execute: bool = Field(
+        default=True,
+        description="Whether execute() is enabled for this sandbox backend.",
+    )
+    env: dict[str, str] | None = Field(
+        default=None,
+        description="Optional environment variables forwarded to sandbox execute() calls.",
+    )
+    timeout_seconds: int = Field(
+        default=300,
+        ge=1,
+        description="Default runtime timeout in seconds for sandbox commands.",
+    )
+    auto_build_template: bool = Field(
+        default=True,
+        description="Whether to auto-build missing template alias before fallback.",
+    )
+
+    def create_backend(self) -> BackendProtocol:
+        """Create a sandbox backend.
+
+        Returns:
+            BackendProtocol: SandboxBackend configured from this model.
+
+        Raises:
+            ImportError: If sandbox dependencies (e.g., e2b) are not installed.
+        """
+        try:
+            from aip_agents.middleware.backends.sandbox import SandboxBackend
+        except ModuleNotFoundError as e:
+            missing_root = (e.name or "").split(".")[0]
+            if missing_root in {"e2b", "e2b_code_interpreter"}:
+                raise ImportError(
+                    "SandboxBackend requires sandbox dependencies. Install with: pip install 'aip-agents[local]'"
+                ) from e
+            raise
+        return SandboxBackend(
+            template=self.template,
+            base_dir=self.base_dir,
+            allow_execute=self.allow_execute,
+            env=self.env,
+            timeout_seconds=self.timeout_seconds,
+            auto_build_template=self.auto_build_template,
+        )
+
+    model_config = ConfigDict(frozen=True)
+
+
+NUM_CHARS_PER_TOKEN = 4
+TOOL_TOKEN_LIMIT_BEFORE_EVICT = 20000
+PREVIEW_HEAD_SIZE = 500
+PREVIEW_TAIL_SIZE = 500
+AUTO_SAVE_DIR = "/tool_outputs"
+
+TOOLS_EXCLUDED_FROM_EVICTION = (
+    "ls",
+    "grep",
+    "read_file",
+    "edit_file",
+    "write_file",
+)
+
+
+FILESYSTEM_SYSTEM_PROMPT = """## Filesystem Tools `read_file`, `grep`, `ls`, `write_file`, `edit_file`
+
+You have access to a filesystem which you can interact with using these tools.
+All file paths must start with a /.
+
+- read_file: read a file from the filesystem
+- grep: search for text within files
+- ls: list files in a directory (requires absolute path)
+- write_file: write to a file in the filesystem
+- edit_file: perform exact string replacements in existing files
+
+## Pagination Strategy
+
+**CRITICAL: Choose ONE mode per file and never mix them in the same session.**
+
+**Line-based** — for normal code/text files:
+- `read_file(path, line_offset=0, limit=100)` → first 100 lines
+- `read_file(path, line_offset=100, limit=100)` → next 100 lines
+- Do NOT set `char_offset` or `max_chars`
+
+**Character-based** — for large/long-line files (minified JS, JSON, CSV, etc.):
+- `read_file(path, char_offset=0, max_chars=20000)` → first 20K chars
+- `read_file(path, char_offset=20000, max_chars=20000)` → next 20K chars
+- Do NOT set `line_offset` or `limit`
+- Default chunk: 20000 characters; hard limit: 80000 characters per call
+
+**How to decide:**
+1. If you know the file has long lines (JSON, minified, CSV) → use char-based immediately
+2. If a line-based read returns a truncation warning → copy the exact `char_offset` from that warning and continue with char-based
+3. Once you pick a mode for a file, continue with that mode until done
+
+## Auto-Save
+
+Large tool outputs may be automatically saved to files. You will receive:
+- The file path where content was saved
+- A preview (head + tail) of the content
+- Use `read_file` to access specific sections
+"""
+
+
+class FilesystemMiddleware(AgentMiddleware):
+    """Middleware providing filesystem capabilities for agents.
+
+    Manages file storage and retrieval for agents using a pluggable backend.
+    Provides hooks for lifecycle management and system prompt enhancement
+    with file context information.
+
+    Thread Safety:
+        Uses the underlying backend's thread safety mechanisms.
+        Each thread_id has isolated storage.
+
+    Attributes:
+        tools: List of file operation tools (empty in PR 001, populated in 002-004).
+        system_prompt_additions: System prompt text explaining file access.
+        backend: Storage backend for file operations.
+    """
+
+    def __init__(
+        self,
+        backend: BackendProtocol | None = None,
+        tool_token_limit_before_evict: int | None = TOOL_TOKEN_LIMIT_BEFORE_EVICT,
+        max_execute_timeout: int = 300,
+    ) -> None:
+        """Initialize filesystem tools and tool-output eviction policy.
+
+        Args:
+            backend: Filesystem backend implementation.
+                If ``None``, ``InMemoryBackend`` is used.
+            tool_token_limit_before_evict: Token threshold used to
+                decide when tool output should be evicted to
+                ``/tool_outputs/<tool_call_id>.txt``. The middleware uses
+                ``NUM_CHARS_PER_TOKEN`` as an approximation.
+                Set to ``None`` to disable eviction.
+            max_execute_timeout: Maximum timeout in seconds for execute commands.
+                Defaults to 300. Only used if backend supports execution.
+
+        Notes:
+            Eviction is skipped for tools listed in
+            ``TOOLS_EXCLUDED_FROM_EVICTION``.
+        """
+        self.backend = backend if backend is not None else InMemoryBackend()
+        self.tools: list[BaseTool] = [
+            ReadFileTool(backend=self.backend),
+            GrepTool(backend=self.backend),
+            LsTool(backend=self.backend),
+            WriteFileTool(backend=self.backend),
+            EditFileTool(backend=self.backend),
+        ]
+
+        # Add execute tool if backend supports it AND execution is enabled
+        if _is_execution_enabled(self.backend):
+            self.tools.append(
+                ExecuteTool(
+                    backend=self.backend,
+                    max_execute_timeout=max_execute_timeout,
+                )
+            )
+            self.system_prompt_additions = FILESYSTEM_SYSTEM_PROMPT + "\n\n" + EXECUTION_SYSTEM_PROMPT
+        else:
+            self.system_prompt_additions = FILESYSTEM_SYSTEM_PROMPT
+
+        self.tool_token_limit_before_evict = tool_token_limit_before_evict
+
+    def replace_backend(self, backend: BackendProtocol, max_execute_timeout: int = 300) -> None:
+        """Replace the backend and rebuild all tools to point to it.
+
+        Called by the coordinator during sub-agent registration to inject a
+        shared backend. After this call, all tools will operate on the new
+        backend.
+
+        Args:
+            backend: The new backend instance to use. Must implement BackendProtocol.
+            max_execute_timeout: Maximum timeout in seconds for execute commands.
+                Defaults to 300. Only used if backend supports execution.
+        """
+        self.backend = backend
+        self.tools = [
+            ReadFileTool(backend=self.backend),
+            GrepTool(backend=self.backend),
+            LsTool(backend=self.backend),
+            WriteFileTool(backend=self.backend),
+            EditFileTool(backend=self.backend),
+        ]
+
+        # Add execute tool if backend supports it AND execution is enabled
+        if _is_execution_enabled(self.backend):
+            self.tools.append(
+                ExecuteTool(
+                    backend=self.backend,
+                    max_execute_timeout=max_execute_timeout,
+                )
+            )
+            self.system_prompt_additions = FILESYSTEM_SYSTEM_PROMPT + "\n\n" + EXECUTION_SYSTEM_PROMPT
+        else:
+            self.system_prompt_additions = FILESYSTEM_SYSTEM_PROMPT
+
+    def _is_tool_excluded_from_eviction(self, message: ToolMessage) -> bool:
+        """Check if the tool is excluded from eviction."""
+        tool_name = None
+        tool_calls = getattr(message, "tool_calls", None)
+        if isinstance(tool_calls, dict):
+            tool_name = tool_calls.get("name")
+
+        if not tool_name:
+            tool_name = getattr(message, "name", None)
+
+        return isinstance(tool_name, str) and tool_name in TOOLS_EXCLUDED_FROM_EVICTION
+
+    def _exceeds_eviction_limit(self, content: str) -> bool:
+        """Check if the content exceeds the eviction limit."""
+        if self.tool_token_limit_before_evict is None:
+            return False
+        return len(content) > NUM_CHARS_PER_TOKEN * self.tool_token_limit_before_evict
+
+    def _save_output(self, path: str, content: str, thread_id: str) -> None:
+        """Save the output to the backend.
+
+        Auto-saved files are saved to the ``"default"`` thread storage because
+        file tools (ReadFileTool, GrepTool, etc.) do not currently pass
+        ``thread_id`` to the backend.  See issue #1158.
+        """
+        if isinstance(self.backend, InMemoryBackend):
+            self.backend.write(path, content, thread_id="default")
+            return
+        self.backend.write(path, content)
+
+    def _format_auto_saved_preview(self, file_path: str, content: str) -> str:
+        """Format the auto-saved preview."""
+        head = content[:PREVIEW_HEAD_SIZE]
+        tail = content[-PREVIEW_TAIL_SIZE:] if len(content) > PREVIEW_TAIL_SIZE else content
+        return (
+            f"Output saved to: {file_path}\n"
+            f"Preview (first {PREVIEW_HEAD_SIZE} chars): {head}...\n"
+            f"Preview (last {PREVIEW_TAIL_SIZE} chars): ...{tail}\n"
+            "Use read_file(path, char_offset=0, max_chars=20000) to read this file in chunks."
+        )
+
+    def _replace_tool_message_content(self, message: ToolMessage, content: str) -> ToolMessage:
+        """Replace the content of a tool message and update tool_calls output.
+
+        Args:
+            message: The original ToolMessage to update.
+            content: The new content to set in the message.
+
+        Returns:
+            A new ToolMessage instance with updated content and tool_calls.output.
+        """
+        # Update content and tool_calls.output for consistency
+        update_dict: dict[str, Any] = {"content": content}
+        if hasattr(message, "tool_calls") and isinstance(message.tool_calls, dict):
+            updated_tool_calls = dict(message.tool_calls)
+            updated_tool_calls["output"] = content
+            update_dict["tool_calls"] = updated_tool_calls
+        return message.model_copy(update=update_dict)
+
+    def before_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Hook executed before model invocation.
+
+        Loads files from state into the backend storage for the current thread.
+        This restores the filesystem state from checkpointed data.
+
+        Args:
+            state: Current agent state containing messages and other context.
+                Expected keys:
+                - "files": Dictionary mapping file paths to file data (optional).
+                - "configurable": Dict containing "thread_id" for isolation.
+
+        Returns:
+            Empty dict (no state updates needed from this hook).
+        """
+        thread_id = state.get("configurable", {}).get("thread_id", "default")
+
+        if "files" in state and state["files"] is not None:
+            self.backend.set_files(state["files"], thread_id=thread_id)
+
+        return {}
+
+    def _evict_tool_messages(
+        self,
+        messages: list[Any],
+        thread_id: str,
+    ) -> tuple[list[Any], bool]:
+        """Evict large tool messages and return updated list.
+
+        Args:
+            messages: List of messages to process.
+            thread_id: Thread ID for backend storage.
+
+        Returns:
+            Tuple of (updated_messages, messages_changed).
+        """
+        updated_messages = []
+        messages_changed = False
+
+        for message in messages:
+            if not isinstance(message, ToolMessage):
+                updated_messages.append(message)
+                continue
+
+            content = message.content if isinstance(message.content, str) else str(message.content)
+            if self._is_tool_excluded_from_eviction(message):
+                updated_messages.append(message)
+                continue
+
+            tool_name = getattr(message, "name", "unknown")
+            threshold = (
+                NUM_CHARS_PER_TOKEN * self.tool_token_limit_before_evict
+                if self.tool_token_limit_before_evict is not None
+                else None
+            )
+            logger.debug(
+                "Checking tool output for eviction: tool=%s size=%d threshold=%s",
+                tool_name,
+                len(content),
+                threshold,
+            )
+
+            if not self._exceeds_eviction_limit(content):
+                updated_messages.append(message)
+                continue
+
+            raw_call_id = getattr(message, "tool_call_id", "") or "unknown"
+            safe_call_id = sanitize_tool_call_id(str(raw_call_id))
+            file_path = f"{AUTO_SAVE_DIR}/{safe_call_id}.txt"
+
+            self._save_output(file_path, content, thread_id)
+            preview_text = self._format_auto_saved_preview(file_path=file_path, content=content)
+            updated_messages.append(self._replace_tool_message_content(message=message, content=preview_text))
+            messages_changed = True
+
+            logger.info(
+                "Evicted tool output: tool=%s size=%d chars path=%s",
+                tool_name,
+                len(content),
+                file_path,
+            )
+
+        return updated_messages, messages_changed
+
+    def after_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Hook executed after model invocation.
+
+        Persists files to backend for checkpointing. Note: message eviction
+        is handled in modify_model_request hook before LLM invocation.
+
+        Args:
+            state: Current agent state after model invocation.
+
+        Returns:
+            Dict with "files" key if files exist, or empty dict if no updates needed.
+        """
+        thread_id = state.get("configurable", {}).get("thread_id", "default")
+
+        updates: dict[str, Any] = {}
+
+        # Persist files to checkpoint only (message eviction happens in modify_model_request)
+        files = self.backend.get_files(thread_id=thread_id)
+        if files:
+            updates["files"] = files
+
+        return updates
+
+    def modify_model_request(self, request: ModelRequest, state: dict[str, Any]) -> ModelRequest:
+        """Modify the model request before invocation."""
+        messages = request.get("messages", [])
+        thread_id = state.get("configurable", {}).get("thread_id", "default")
+
+        updated_messages, messages_changed = self._evict_tool_messages(messages, thread_id)
+
+        if messages_changed:
+            request["messages"] = updated_messages
+
+        return request
+
+    async def abefore_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Async version of before_model.
+
+        Args:
+            state: Current agent state containing messages and other context.
+
+        Returns:
+            Dict of state updates to merge into the agent state. Return empty dict
+            if no updates are needed.
+        """
+        return self.before_model(state)
+
+    async def aafter_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Async version of after_model.
+
+        Args:
+            state: Current agent state after model invocation.
+
+        Returns:
+            Dict of state updates to merge into the agent state. Return empty dict
+            if no updates are needed.
+        """
+        thread_id = state.get("configurable", {}).get("thread_id", "default")
+
+        updates: dict[str, Any] = {}
+
+        aget_files = getattr(self.backend, "aget_files", None)
+        if callable(aget_files):
+            files = await aget_files(thread_id=thread_id)
+        else:
+            files = self.backend.get_files(thread_id=thread_id)
+
+        if files:
+            updates["files"] = files
+
+        return updates