simplify backend protocol

src/deepagents/backends/__init__.py

@@ -1,18 +1,16 @@
 """Memory backends for pluggable file storage."""
 
-from deepagents.backends.composite import CompositeBackend, CompositeStateBackendProvider
+from deepagents.backends.composite import CompositeBackend
 from deepagents.backends.filesystem import FilesystemBackend
-from deepagents.backends.state import StateBackend, StateBackendProvider
-from deepagents.backends.store import StoreBackend, StoreBackendProvider
-from deepagents.backends.protocol import BackendProtocol
+from deepagents.backends.protocol import Backend, BackendProvider
+from deepagents.backends.state import StateBackend
+from deepagents.backends.store import StoreBackend
 
 __all__ = [
-    "BackendProtocol",
+    "Backend",
+    "BackendProvider",
     "CompositeBackend",
-    "CompositeStateBackendProvider",
     "FilesystemBackend",
     "StateBackend",
-    "StateBackendProvider",
     "StoreBackend",
-    "StoreBackendProvider",
 ]

src/deepagents/backends/composite.py

@@ -1,36 +1,45 @@
 """CompositeBackend: Route operations to different backends based on path prefix."""
 
-from typing import Any, Literal, Optional, TYPE_CHECKING
+from deepagents.backends.protocol import Backend, EditResult, WriteResult
 
-from langchain.tools import ToolRuntime
 
-from deepagents.backends.protocol import BackendProtocol, BackendProvider, StateBackendProvider, StateBackendProtocol
-from deepagents.backends.state import StateBackend, StateBackendProvider
-from langgraph.types import Command
+class CompositeBackend(Backend):
+    """Composite backend that routes operations to different backends based on path prefix.
 
+    This backend allows combining multiple backends (checkpoint storage and external storage)
+    under different path prefixes. For example:
+    - /memories/* → StoreBackend (persistent across conversations)
+    - /* → StateBackend (checkpoint storage per conversation)
+
+    Storage Model: Mixed (routes to checkpoint or external backends)
+    """
 
-class _CompositeBackend:
-
     def __init__(
         self,
-        default: BackendProtocol | StateBackend,
-        routes: dict[str, BackendProtocol],
+        default: Backend,
+        routes: dict[str, Backend],
     ) -> None:
+        """Initialize CompositeBackend.
+
+        Args:
+            default: Default backend for paths that don't match any route
+            routes: Dictionary mapping path prefixes to backends
+        """
         # Default backend
         self.default = default
 
-        # Virtual routes
+        # Routed backends
         self.routes = routes
-        
+
         # Sort routes by length (longest first) for correct prefix matching
         self.sorted_routes = sorted(routes.items(), key=lambda x: len(x[0]), reverse=True)
 
-    def _get_backend_and_key(self, key: str) -> tuple[BackendProtocol, str]:
+    def _get_backend_and_key(self, key: str) -> tuple[Backend, str]:
         """Determine which backend handles this key and strip prefix.
-        
+
         Args:
             key: Original file path
-        
+
         Returns:
             Tuple of (backend, stripped_key) where stripped_key has the route
             prefix removed (but keeps leading slash).
@@ -40,39 +49,39 @@ def _get_backend_and_key(self, key: str) -> tuple[BackendProtocol, str]:
             if key.startswith(prefix):
                 # Strip prefix but keep leading slash
                 # e.g., "/memories/notes.txt" → "/notes.txt"
-                stripped_key = key[len(prefix) - 1:] if key[len(prefix) - 1:] else "/"
+                stripped_key = key[len(prefix) - 1 :] if key[len(prefix) - 1 :] else "/"
                 return backend, stripped_key
-        
+
         return self.default, key
-    
+
     def ls(self, path: str) -> list[str]:
         """List files from backends, with appropriate prefixes.
-        
+
         Args:
             path: Absolute path to directory.
-        
+
         Returns:
             List of file paths with route prefixes added.
         """
         # Check if path matches a specific route
         for route_prefix, backend in self.sorted_routes:
             if path.startswith(route_prefix.rstrip("/")):
                 # Query only the matching routed backend
-                search_path = path[len(route_prefix) - 1:]
+                search_path = path[len(route_prefix) - 1 :]
                 keys = backend.ls(search_path if search_path else "/")
                 return [f"{route_prefix[:-1]}{key}" for key in keys]
-        
+
         # Path doesn't match a route: query only default backend
         return self.default.ls(path)
-    
+
     def read(
-        self, 
+        self,
         file_path: str,
         offset: int = 0,
         limit: int = 2000,
     ) -> str:
         """Read file content, routing to appropriate backend.
-        
+
         Args:
             file_path: Absolute file path
             offset: Line offset to start reading from (0-indexed)
@@ -81,16 +90,16 @@ def read(
         """
         backend, stripped_key = self._get_backend_and_key(file_path)
         return backend.read(stripped_key, offset=offset, limit=limit)
-    
+
     def grep(
         self,
         pattern: str,
-        path: Optional[str] = None,
-        glob: Optional[str] = None,
+        path: str | None = None,
+        glob: str | None = None,
         output_mode: str = "files_with_matches",
     ) -> str:
         """Search for a pattern in files, routing to appropriate backend(s).
-        
+
         Args:
             pattern: String pattern to search for
             path: Path to search in (default "/")
@@ -100,11 +109,11 @@ def grep(
         """
         for route_prefix, backend in self.sorted_routes:
             if path is not None and path.startswith(route_prefix.rstrip("/")):
-                search_path = path[len(route_prefix) - 1:]
+                search_path = path[len(route_prefix) - 1 :]
                 result = backend.grep(pattern, search_path if search_path else "/", glob, output_mode)
                 if result.startswith("No matches found"):
                     return result
-                
+
                 lines = result.split("\n")
                 prefixed_lines = []
                 for line in lines:
@@ -118,11 +127,11 @@ def grep(
                 return "\n".join(prefixed_lines)
 
         all_results = []
-        
+
         default_result = self.default.grep(pattern, path, glob, output_mode)
         if not default_result.startswith("No matches found"):
             all_results.append(default_result)
-        
+
         for route_prefix, backend in self.routes.items():
             result = backend.grep(pattern, None, glob, output_mode)
             if not result.startswith("No matches found"):
@@ -137,12 +146,12 @@ def grep(
                     else:
                         prefixed_lines.append(line)
                 all_results.append("\n".join(prefixed_lines))
-        
+
         if not all_results:
             return f"No matches found for pattern: '{pattern}'"
-        
+
         return "\n".join(all_results)
-    
+
     def glob(self, pattern: str, path: str = "/") -> list[str]:
         """Find files matching a glob pattern across all backends.
 
@@ -157,7 +166,7 @@ def glob(self, pattern: str, path: str = "/") -> list[str]:
         for route_prefix, backend in self.sorted_routes:
             if path.startswith(route_prefix.rstrip("/")):
                 # Path matches a specific route - search only that backend
-                search_path = path[len(route_prefix) - 1:]
+                search_path = path[len(route_prefix) - 1 :]
                 matches = backend.glob(pattern, search_path if search_path else "/")
                 results.extend(f"{route_prefix[:-1]}{match}" for match in matches)
                 return sorted(results)
@@ -173,99 +182,40 @@ def glob(self, pattern: str, path: str = "/") -> list[str]:
 
         return sorted(results)
 
-
-class CompositeStateBacked(_CompositeBackend):
-
-    def __init__(
-            self,
-            default: StateBackend,
-            routes: dict[str, BackendProtocol],
-    ) -> None:
-        self.default = default
-        self.routes = routes
-
-        # Sort routes by length (longest first) for correct prefix matching
-        self.sorted_routes = sorted(routes.items(), key=lambda x: len(x[0]), reverse=True)
-
     def write(
-            self,
-            file_path: str,
-            content: str,
-    ) -> Command | str:
+        self,
+        file_path: str,
+        content: str,
+    ) -> WriteResult:
         """Create a new file, routing to appropriate backend.
 
         Args:
             file_path: Absolute file path
-            content: File content as a stringReturns:
-            Success message or Command object, or error if file already exists.
-        """
-        backend, stripped_key = self._get_backend_and_key(file_path)
-        return backend.write(stripped_key, content)
-
-    def edit(
-            self,
-            file_path: str,
-            old_string: str,
-            new_string: str,
-            replace_all: bool = False,
-    ) -> Command | str:
-        """Edit a file, routing to appropriate backend.
-
-        Args:
-            file_path: Absolute file path
-            old_string: String to find and replace
-            new_string: Replacement string
-            replace_all: If True, replace all occurrencesReturns:
-            Success message or Command object, or error message on failure.
-        """
-        backend, stripped_key = self._get_backend_and_key(file_path)
-        return backend.edit(stripped_key, old_string, new_string, replace_all=replace_all)
-
+            content: File content as a string
 
-class CompositeBackend(_CompositeBackend):
-    def write(
-            self,
-            file_path: str,
-            content: str,
-    ) -> str:
-        """Create a new file, routing to appropriate backend.
-
-        Args:
-            file_path: Absolute file path
-            content: File content as a stringReturns:
-            Success message or Command object, or error if file already exists.
+        Returns:
+            WriteResult from the appropriate backend.
         """
         backend, stripped_key = self._get_backend_and_key(file_path)
         return backend.write(stripped_key, content)
 
     def edit(
-            self,
-            file_path: str,
-            old_string: str,
-            new_string: str,
-            replace_all: bool = False,
-    ) -> str:
+        self,
+        file_path: str,
+        old_string: str,
+        new_string: str,
+        replace_all: bool = False,
+    ) -> EditResult:
         """Edit a file, routing to appropriate backend.
 
         Args:
             file_path: Absolute file path
             old_string: String to find and replace
             new_string: Replacement string
-            replace_all: If True, replace all occurrencesReturns:
-            Success message or Command object, or error message on failure.
+            replace_all: If True, replace all occurrences
+
+        Returns:
+            EditResult from the appropriate backend.
         """
         backend, stripped_key = self._get_backend_and_key(file_path)
         return backend.edit(stripped_key, old_string, new_string, replace_all=replace_all)
-
-class CompositeStateBackendProvider(StateBackendProvider):
-
-    def __init__(self, routes: dict[str, BackendProtocol | BackendProvider]):
-        self.routes = routes
-
-    def get_backend(self, runtime: ToolRuntime) -> StateBackendProtocol:
-        return CompositeStateBacked(
-            default=StateBackendProvider.get_backend(runtime),
-            routes={
-                k: v if isinstance(v, BackendProtocol) else v.get_backend(runtime) for k,v in self.routes.items()
-            }
-        )