deepagents 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

deepagents/middleware/memory.py

@@ -0,0 +1,369 @@
+"""Middleware for loading agent memory/context from AGENTS.md files.
+
+This module implements support for the AGENTS.md specification (https://agents.md/),
+loading memory/context from configurable sources and injecting into the system prompt.
+
+## Overview
+
+AGENTS.md files provide project-specific context and instructions to help AI agents
+work effectively. Unlike skills (which are on-demand workflows), memory is always
+loaded and provides persistent context.
+
+## Usage
+
+```python
+from deepagents import MemoryMiddleware
+from deepagents.backends.filesystem import FilesystemBackend
+
+# Security: FilesystemBackend allows reading/writing from the entire filesystem.
+# Either ensure the agent is running within a sandbox OR add human-in-the-loop (HIL)
+# approval to file operations.
+backend = FilesystemBackend(root_dir="/")
+
+middleware = MemoryMiddleware(
+    backend=backend,
+    sources=[
+        "~/.deepagents/AGENTS.md",
+        "./.deepagents/AGENTS.md",
+    ],
+)
+
+agent = create_deep_agent(middleware=[middleware])
+```
+
+## Memory Sources
+
+Sources are simply paths to AGENTS.md files that are loaded in order and combined.
+Multiple sources are concatenated in order, with all content included.
+Later sources appear after earlier ones in the combined prompt.
+
+## File Format
+
+AGENTS.md files are standard Markdown with no required structure.
+Common sections include:
+- Project overview
+- Build/test commands
+- Code style guidelines
+- Architecture notes
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING, Annotated, NotRequired, TypedDict
+
+from langchain.messages import SystemMessage
+from langchain_core.runnables import RunnableConfig
+
+if TYPE_CHECKING:
+    from deepagents.backends.protocol import BACKEND_TYPES, BackendProtocol
+
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+    ModelRequest,
+    ModelResponse,
+    PrivateStateAttr,
+)
+from langchain.tools import ToolRuntime
+from langgraph.runtime import Runtime
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryState(AgentState):
+    """State schema for MemoryMiddleware.
+
+    Attributes:
+        memory_contents: Dict mapping source paths to their loaded content.
+            Marked as private so it's not included in the final agent state.
+    """
+
+    memory_contents: NotRequired[Annotated[dict[str, str], PrivateStateAttr]]
+
+
+class MemoryStateUpdate(TypedDict):
+    """State update for MemoryMiddleware."""
+
+    memory_contents: dict[str, str]
+
+
+# Default system prompt template for memory
+MEMORY_SYSTEM_PROMPT = """
+## Agent Memory
+
+You have access to persistent memory that provides context and instructions.
+
+{memory_locations}
+
+{memory_contents}
+
+**Memory Guidelines:**
+- Memory content above provides project-specific context and instructions
+- Follow any guidelines, conventions, or patterns described in memory
+- Memory is read-only during this session (loaded at startup)
+- If you need to update memory, use the appropriate file editing tools
+"""
+
+
+class MemoryMiddleware(AgentMiddleware):
+    """Middleware for loading agent memory from AGENTS.md files.
+
+    Loads memory content from configured sources and injects into the system prompt.
+    Supports multiple sources that are combined together.
+
+    Args:
+        backend: Backend instance or factory function for file operations.
+        sources: List of MemorySource configurations specifying paths and names.
+    """
+
+    state_schema = MemoryState
+
+    def __init__(
+        self,
+        *,
+        backend: BACKEND_TYPES,
+        sources: list[str],
+    ) -> None:
+        """Initialize the memory middleware.
+
+        Args:
+            backend: Backend instance or factory function that takes runtime
+                     and returns a backend. Use a factory for StateBackend.
+            sources: List of memory file paths to load (e.g., ["~/.deepagents/AGENTS.md",
+                     "./.deepagents/AGENTS.md"]). Display names are automatically derived
+                     from the paths. Sources are loaded in order.
+        """
+        self._backend = backend
+        self.sources = sources
+        self.system_prompt_template = MEMORY_SYSTEM_PROMPT
+
+    def _get_backend(self, state: MemoryState, runtime: Runtime, config: RunnableConfig) -> BackendProtocol:
+        """Resolve backend from instance or factory.
+
+        Args:
+            state: Current agent state.
+            runtime: Runtime context for factory functions.
+            config: Runnable config to pass to backend factory.
+
+        Returns:
+            Resolved backend instance.
+        """
+        if callable(self._backend):
+            # Construct an artificial tool runtime to resolve backend factory
+            tool_runtime = ToolRuntime(
+                state=state,
+                context=runtime.context,
+                stream_writer=runtime.stream_writer,
+                store=runtime.store,
+                config=config,
+                tool_call_id=None,
+            )
+            return self._backend(tool_runtime)
+        return self._backend
+
+    def _format_memory_locations(self) -> str:
+        """Format memory source locations for display."""
+        if not self.sources:
+            return "**Memory Sources:** None configured"
+
+        lines = ["**Memory Sources:**"]
+        for path in self.sources:
+            lines.append(f"- `{path}`")
+        return "\n".join(lines)
+
+    def _format_memory_contents(self, contents: dict[str, str]) -> str:
+        """Format loaded memory contents for injection into prompt.
+
+        Args:
+            contents: Dict mapping source paths to content.
+
+        Returns:
+            Formatted string with all memory contents concatenated.
+        """
+        if not contents:
+            return "(No memory loaded)"
+
+        sections = []
+        for path in self.sources:
+            if contents.get(path):
+                sections.append(contents[path])
+
+        if not sections:
+            return "(No memory loaded)"
+
+        return "\n\n".join(sections)
+
+    async def _load_memory_from_backend(
+        self,
+        backend: BackendProtocol,
+        path: str,
+    ) -> str | None:
+        """Load memory content from a backend path.
+
+        Args:
+            backend: Backend to load from.
+            path: Path to the AGENTS.md file.
+
+        Returns:
+            File content if found, None otherwise.
+        """
+        results = await backend.adownload_files([path])
+        # Should get exactly one response for one path
+        if len(results) != 1:
+            raise AssertionError(f"Expected 1 response for path {path}, got {len(results)}")
+        response = results[0]
+
+        if response.error is not None:
+            raise ValueError(f"Failed to download {path}: {response.error}")
+
+        if response.content is not None:
+            return response.content.decode("utf-8")
+
+        return None
+
+    def _load_memory_from_backend_sync(
+        self,
+        backend: BackendProtocol,
+        path: str,
+    ) -> str | None:
+        """Load memory content from a backend path synchronously.
+
+        Args:
+            backend: Backend to load from.
+            path: Path to the AGENTS.md file.
+
+        Returns:
+            File content if found, None otherwise.
+        """
+        results = backend.download_files([path])
+        # Should get exactly one response for one path
+        if len(results) != 1:
+            raise AssertionError(f"Expected 1 response for path {path}, got {len(results)}")
+        response = results[0]
+
+        if response.error is not None:
+            raise ValueError(f"Failed to download {path}: {response.error}")
+
+        if response.content is not None:
+            return response.content.decode("utf-8")
+
+        return None
+
+    def before_agent(self, state: MemoryState, runtime: Runtime, config: RunnableConfig) -> MemoryStateUpdate | None:
+        """Load memory content before agent execution (synchronous).
+
+        Loads memory from all configured sources and stores in state.
+        Only loads if not already present in state.
+
+        Args:
+            state: Current agent state.
+            runtime: Runtime context.
+            config: Runnable config.
+
+        Returns:
+            State update with memory_contents populated.
+        """
+        # Skip if already loaded
+        if "memory_contents" in state:
+            return None
+
+        backend = self._get_backend(state, runtime, config)
+        contents: dict[str, str] = {}
+
+        for path in self.sources:
+            content = self._load_memory_from_backend_sync(backend, path)
+            if content:
+                contents[path] = content
+                logger.debug(f"Loaded memory from: {path}")
+
+        return MemoryStateUpdate(memory_contents=contents)
+
+    async def abefore_agent(self, state: MemoryState, runtime: Runtime, config: RunnableConfig) -> MemoryStateUpdate | None:
+        """Load memory content before agent execution.
+
+        Loads memory from all configured sources and stores in state.
+        Only loads if not already present in state.
+
+        Args:
+            state: Current agent state.
+            runtime: Runtime context.
+            config: Runnable config.
+
+        Returns:
+            State update with memory_contents populated.
+        """
+        # Skip if already loaded
+        if "memory_contents" in state:
+            return None
+
+        backend = self._get_backend(state, runtime, config)
+        contents: dict[str, str] = {}
+
+        for path in self.sources:
+            content = await self._load_memory_from_backend(backend, path)
+            if content:
+                contents[path] = content
+                logger.debug(f"Loaded memory from: {path}")
+
+        return MemoryStateUpdate(memory_contents=contents)
+
+    def modify_request(self, request: ModelRequest) -> ModelRequest:
+        """Inject memory content into the system prompt.
+
+        Args:
+            request: Model request to modify.
+
+        Returns:
+            Modified request with memory injected into system prompt.
+        """
+        contents = request.state.get("memory_contents", {})
+        memory_locations = self._format_memory_locations()
+        memory_contents = self._format_memory_contents(contents)
+
+        memory_section = self.system_prompt_template.format(
+            memory_locations=memory_locations,
+            memory_contents=memory_contents,
+        )
+
+        if request.system_prompt:
+            system_prompt = memory_section + "\n\n" + request.system_prompt
+        else:
+            system_prompt = memory_section
+
+        return request.override(system_message=SystemMessage(system_prompt))
+
+    def wrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], ModelResponse],
+    ) -> ModelResponse:
+        """Wrap model call to inject memory into system prompt.
+
+        Args:
+            request: Model request being processed.
+            handler: Handler function to call with modified request.
+
+        Returns:
+            Model response from handler.
+        """
+        modified_request = self.modify_request(request)
+        return handler(modified_request)
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelResponse:
+        """Async wrap model call to inject memory into system prompt.
+
+        Args:
+            request: Model request being processed.
+            handler: Async handler function to call with modified request.
+
+        Returns:
+            Model response from handler.
+        """
+        modified_request = self.modify_request(request)
+        return await handler(modified_request)