feat: Optimize directory operations for 10-100x performance improvement (#350)

Signed-off-by: phernandez <paul@basicmachines.co>
Co-authored-by: Claude <noreply@anthropic.com>

src/basic_memory/api/routers/directory_router.py

@@ -31,6 +31,27 @@ async def get_directory_tree(
     return tree
 
 
+@router.get("/structure", response_model=DirectoryNode)
+async def get_directory_structure(
+    directory_service: DirectoryServiceDep,
+    project_id: ProjectIdDep,
+):
+    """Get folder structure for navigation (no files).
+
+    Optimized endpoint for folder tree navigation. Returns only directory nodes
+    without file metadata. For full tree with files, use /directory/tree.
+
+    Args:
+        directory_service: Service for directory operations
+        project_id: ID of the current project
+
+    Returns:
+        DirectoryNode tree containing only folders (type="directory")
+    """
+    structure = await directory_service.get_directory_structure()
+    return structure
+
+
 @router.get("/list", response_model=List[DirectoryNode])
 async def list_directory(
     directory_service: DirectoryServiceDep,

src/basic_memory/repository/entity_repository.py

@@ -176,6 +176,66 @@ class EntityRepository(Repository[Entity]):
                     entity = await self._handle_permalink_conflict(entity, session)
                     return entity
 
+    async def get_distinct_directories(self) -> List[str]:
+        """Extract unique directory paths from file_path column.
+
+        Optimized method for getting directory structure without loading full entities
+        or relationships. Returns a sorted list of unique directory paths.
+
+        Returns:
+            List of unique directory paths (e.g., ["notes", "notes/meetings", "specs"])
+        """
+        # Query only file_path column, no entity objects or relationships
+        query = select(Entity.file_path).distinct()
+        query = self._add_project_filter(query)
+
+        # Execute with use_query_options=False to skip eager loading
+        result = await self.execute_query(query, use_query_options=False)
+        file_paths = [row for row in result.scalars().all()]
+
+        # Parse file paths to extract unique directories
+        directories = set()
+        for file_path in file_paths:
+            parts = [p for p in file_path.split("/") if p]
+            # Add all parent directories (exclude filename which is the last part)
+            for i in range(len(parts) - 1):
+                dir_path = "/".join(parts[: i + 1])
+                directories.add(dir_path)
+
+        return sorted(directories)
+
+    async def find_by_directory_prefix(self, directory_prefix: str) -> Sequence[Entity]:
+        """Find entities whose file_path starts with the given directory prefix.
+
+        Optimized method for listing directory contents without loading all entities.
+        Uses SQL LIKE pattern matching to filter entities by directory path.
+
+        Args:
+            directory_prefix: Directory path prefix (e.g., "docs", "docs/guides")
+                             Empty string returns all entities (root directory)
+
+        Returns:
+            Sequence of entities in the specified directory and subdirectories
+        """
+        # Build SQL LIKE pattern
+        if directory_prefix == "" or directory_prefix == "/":
+            # Root directory - return all entities
+            return await self.find_all()
+
+        # Remove leading/trailing slashes for consistency
+        directory_prefix = directory_prefix.strip("/")
+
+        # Query entities with file_path starting with prefix
+        # Pattern matches "prefix/" to ensure we get files IN the directory,
+        # not just files whose names start with the prefix
+        pattern = f"{directory_prefix}/%"
+
+        query = self.select().where(Entity.file_path.like(pattern))
+
+        # Skip eager loading - we only need basic entity fields for directory trees
+        result = await self.execute_query(query, use_query_options=False)
+        return list(result.scalars().all())
+
     async def _handle_permalink_conflict(self, entity: Entity, session: AsyncSession) -> Entity:
         """Handle permalink conflicts by generating a unique permalink."""
         base_permalink = entity.permalink

src/basic_memory/repository/repository.py

@@ -152,12 +152,25 @@ class Repository[T: Base]:
         # Add project filter if applicable
         return self._add_project_filter(query)
 
-    async def find_all(self, skip: int = 0, limit: Optional[int] = None) -> Sequence[T]:
-        """Fetch records from the database with pagination."""
+    async def find_all(
+        self, skip: int = 0, limit: Optional[int] = None, use_load_options: bool = True
+    ) -> Sequence[T]:
+        """Fetch records from the database with pagination.
+
+        Args:
+            skip: Number of records to skip
+            limit: Maximum number of records to return
+            use_load_options: Whether to apply eager loading options (default: True)
+        """
         logger.debug(f"Finding all {self.Model.__name__} (skip={skip}, limit={limit})")
 
         async with db.scoped_session(self.session_maker) as session:
-            query = select(self.Model).offset(skip).options(*self.get_load_options())
+            query = select(self.Model).offset(skip)
+
+            # Only apply load options if requested
+            if use_load_options:
+                query = query.options(*self.get_load_options())
+
             # Add project filter if applicable
             query = self._add_project_filter(query)
 

src/basic_memory/services/directory_service.py

@@ -3,8 +3,9 @@
 import fnmatch
 import logging
 import os
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Sequence
 
+from basic_memory.models import Entity
 from basic_memory.repository import EntityRepository
 from basic_memory.schemas.directory import DirectoryNode
 
@@ -89,6 +90,49 @@ class DirectoryService:
         # Return the root node with its children
         return root_node
 
+    async def get_directory_structure(self) -> DirectoryNode:
+        """Build a hierarchical directory structure without file details.
+
+        Optimized method for folder navigation that only returns directory nodes,
+        no file metadata. Much faster than get_directory_tree() for large knowledge bases.
+
+        Returns:
+            DirectoryNode tree containing only folders (type="directory")
+        """
+        # Get unique directories without loading entities
+        directories = await self.entity_repository.get_distinct_directories()
+
+        # Create a root directory node
+        root_node = DirectoryNode(name="Root", directory_path="/", type="directory")
+
+        # Map to store directory nodes by path for easy lookup
+        dir_map: Dict[str, DirectoryNode] = {"/": root_node}
+
+        # Build tree with just folders
+        for dir_path in directories:
+            parts = [p for p in dir_path.split("/") if p]
+            current_path = "/"
+
+            for i, part in enumerate(parts):
+                parent_path = current_path
+                # Build the directory path
+                current_path = (
+                    f"{current_path}{part}" if current_path == "/" else f"{current_path}/{part}"
+                )
+
+                # Create directory node if it doesn't exist
+                if current_path not in dir_map:
+                    dir_node = DirectoryNode(
+                        name=part, directory_path=current_path, type="directory"
+                    )
+                    dir_map[current_path] = dir_node
+
+                    # Add to parent's children
+                    if parent_path in dir_map:
+                        dir_map[parent_path].children.append(dir_node)
+
+        return root_node
+
     async def list_directory(
         self,
         dir_name: str = "/",
@@ -118,8 +162,13 @@ class DirectoryService:
         if dir_name != "/" and dir_name.endswith("/"):
             dir_name = dir_name.rstrip("/")
 
-        # Get the full directory tree
-        root_tree = await self.get_directory_tree()
+        # Optimize: Query only entities in the target directory
+        # instead of loading the entire tree
+        dir_prefix = dir_name.lstrip("/")
+        entity_rows = await self.entity_repository.find_by_directory_prefix(dir_prefix)
+
+        # Build a partial tree from only the relevant entities
+        root_tree = self._build_directory_tree_from_entities(entity_rows, dir_name)
 
         # Find the target directory node
         target_node = self._find_directory_node(root_tree, dir_name)
@@ -132,6 +181,78 @@ class DirectoryService:
 
         return result
 
+    def _build_directory_tree_from_entities(
+        self, entity_rows: Sequence[Entity], root_path: str
+    ) -> DirectoryNode:
+        """Build a directory tree from a subset of entities.
+
+        Args:
+            entity_rows: Sequence of entity objects to build tree from
+            root_path: Root directory path for the tree
+
+        Returns:
+            DirectoryNode representing the tree root
+        """
+        # Create a root directory node
+        root_node = DirectoryNode(name="Root", directory_path=root_path, type="directory")
+
+        # Map to store directory nodes by path for easy lookup
+        dir_map: Dict[str, DirectoryNode] = {root_path: root_node}
+
+        # First pass: create all directory nodes
+        for file in entity_rows:
+            # Process directory path components
+            parts = [p for p in file.file_path.split("/") if p]
+
+            # Create directory structure
+            current_path = "/"
+            for i, part in enumerate(parts[:-1]):  # Skip the filename
+                parent_path = current_path
+                # Build the directory path
+                current_path = (
+                    f"{current_path}{part}" if current_path == "/" else f"{current_path}/{part}"
+                )
+
+                # Create directory node if it doesn't exist
+                if current_path not in dir_map:
+                    dir_node = DirectoryNode(
+                        name=part, directory_path=current_path, type="directory"
+                    )
+                    dir_map[current_path] = dir_node
+
+                    # Add to parent's children
+                    if parent_path in dir_map:
+                        dir_map[parent_path].children.append(dir_node)
+
+        # Second pass: add file nodes to their parent directories
+        for file in entity_rows:
+            file_name = os.path.basename(file.file_path)
+            parent_dir = os.path.dirname(file.file_path)
+            directory_path = "/" if parent_dir == "" else f"/{parent_dir}"
+
+            # Create file node
+            file_node = DirectoryNode(
+                name=file_name,
+                file_path=file.file_path,
+                directory_path=f"/{file.file_path}",
+                type="file",
+                title=file.title,
+                permalink=file.permalink,
+                entity_id=file.id,
+                entity_type=file.entity_type,
+                content_type=file.content_type,
+                updated_at=file.updated_at,
+            )
+
+            # Add to parent directory's children
+            if directory_path in dir_map:
+                dir_map[directory_path].children.append(file_node)
+            elif root_path in dir_map:
+                # Fallback to root if parent not found
+                dir_map[root_path].children.append(file_node)
+
+        return root_node
+
     def _find_directory_node(
         self, root: DirectoryNode, target_path: str
     ) -> Optional[DirectoryNode]:

src/basic_memory/services/project_service.py

@@ -80,7 +80,13 @@ class ProjectService:
         return os.environ.get("BASIC_MEMORY_PROJECT", self.config_manager.default_project)
 
     async def list_projects(self) -> Sequence[Project]:
-        return await self.repository.find_all()
+        """List all projects without loading entity relationships.
+
+        Returns only basic project fields (name, path, etc.) without
+        eager loading the entities relationship which could load thousands
+        of entities for large knowledge bases.
+        """
+        return await self.repository.find_all(use_load_options=False)
 
     async def get_project(self, name: str) -> Optional[Project]:
         """Get the file path for a project by name or permalink."""

tests/api/test_directory_router.py

@@ -277,3 +277,136 @@ async def test_list_directory_endpoint_mocked(client, project_url):
         assert file_item["file_path"] == "file1.md"
         assert file_item["title"] == "File 1"
         assert file_item["permalink"] == "file-1"
+
+
+@pytest.mark.asyncio
+async def test_get_directory_structure_endpoint(test_graph, client, project_url):
+    """Test the get_directory_structure endpoint returns folders only."""
+    # Call the endpoint
+    response = await client.get(f"{project_url}/directory/structure")
+
+    # Verify response
+    assert response.status_code == 200
+    data = response.json()
+
+    # Check that the response is a valid directory tree
+    assert "name" in data
+    assert "directory_path" in data
+    assert "children" in data
+    assert "type" in data
+    assert data["type"] == "directory"
+
+    # Root should be present
+    assert data["name"] == "Root"
+    assert data["directory_path"] == "/"
+
+    # Should have the test directory
+    assert len(data["children"]) == 1
+    test_dir = data["children"][0]
+    assert test_dir["name"] == "test"
+    assert test_dir["type"] == "directory"
+    assert test_dir["directory_path"] == "/test"
+
+    # Should NOT have any files (test_graph has files but no subdirectories)
+    assert len(test_dir["children"]) == 0
+
+    # Verify no file metadata is present in directory nodes
+    assert test_dir.get("entity_id") is None
+    assert test_dir.get("content_type") is None
+    assert test_dir.get("title") is None
+    assert test_dir.get("permalink") is None
+
+
+@pytest.mark.asyncio
+async def test_get_directory_structure_empty(client, project_url):
+    """Test the get_directory_structure endpoint with empty database."""
+    # Call the endpoint
+    response = await client.get(f"{project_url}/directory/structure")
+
+    # Verify response
+    assert response.status_code == 200
+    data = response.json()
+
+    # Should return root with no children
+    assert data["name"] == "Root"
+    assert data["directory_path"] == "/"
+    assert data["type"] == "directory"
+    assert len(data["children"]) == 0
+
+
+@pytest.mark.asyncio
+async def test_get_directory_structure_mocked(client, project_url):
+    """Test the get_directory_structure endpoint with mocked service."""
+    # Create a mock directory structure (folders only, no files)
+    mock_structure = DirectoryNode(
+        name="Root",
+        directory_path="/",
+        type="directory",
+        children=[
+            DirectoryNode(
+                name="docs",
+                directory_path="/docs",
+                type="directory",
+                children=[
+                    DirectoryNode(
+                        name="guides",
+                        directory_path="/docs/guides",
+                        type="directory",
+                        children=[],
+                    ),
+                    DirectoryNode(
+                        name="api",
+                        directory_path="/docs/api",
+                        type="directory",
+                        children=[],
+                    ),
+                ],
+            ),
+            DirectoryNode(name="specs", directory_path="/specs", type="directory", children=[]),
+        ],
+    )
+
+    # Patch the directory service
+    with patch(
+        "basic_memory.services.directory_service.DirectoryService.get_directory_structure",
+        return_value=mock_structure,
+    ):
+        # Call the endpoint
+        response = await client.get(f"{project_url}/directory/structure")
+
+        # Verify response
+        assert response.status_code == 200
+        data = response.json()
+
+        # Check structure matches our mock (folders only)
+        assert data["name"] == "Root"
+        assert data["directory_path"] == "/"
+        assert data["type"] == "directory"
+        assert len(data["children"]) == 2
+
+        # Check docs directory
+        docs = data["children"][0]
+        assert docs["name"] == "docs"
+        assert docs["directory_path"] == "/docs"
+        assert docs["type"] == "directory"
+        assert len(docs["children"]) == 2
+
+        # Check subdirectories
+        guides = docs["children"][0]
+        assert guides["name"] == "guides"
+        assert guides["directory_path"] == "/docs/guides"
+        assert guides["type"] == "directory"
+        assert guides["children"] == []
+
+        api = docs["children"][1]
+        assert api["name"] == "api"
+        assert api["directory_path"] == "/docs/api"
+        assert api["type"] == "directory"
+        assert api["children"] == []
+
+        # Check specs directory
+        specs = data["children"][1]
+        assert specs["name"] == "specs"
+        assert specs["directory_path"] == "/specs"
+        assert specs["type"] == "directory"
+        assert specs["children"] == []

tests/repository/test_entity_repository.py

@@ -432,3 +432,215 @@ async def test_get_by_file_path(entity_repository: EntityRepository, session_mak
     # Test non-existent file_path
     found = await entity_repository.get_by_file_path("not/a/real/file.md")
     assert found is None
+
+
+@pytest.mark.asyncio
+async def test_get_distinct_directories(entity_repository: EntityRepository, session_maker):
+    """Test getting distinct directory paths from entity file paths."""
+    # Create test entities with various directory structures
+    async with db.scoped_session(session_maker) as session:
+        entities = [
+            Entity(
+                project_id=entity_repository.project_id,
+                title="File 1",
+                entity_type="test",
+                permalink="docs/guides/file1",
+                file_path="docs/guides/file1.md",
+                content_type="text/markdown",
+                created_at=datetime.now(timezone.utc),
+                updated_at=datetime.now(timezone.utc),
+            ),
+            Entity(
+                project_id=entity_repository.project_id,
+                title="File 2",
+                entity_type="test",
+                permalink="docs/guides/file2",
+                file_path="docs/guides/file2.md",
+                content_type="text/markdown",
+                created_at=datetime.now(timezone.utc),
+                updated_at=datetime.now(timezone.utc),
+            ),
+            Entity(
+                project_id=entity_repository.project_id,
+                title="File 3",
+                entity_type="test",
+                permalink="docs/api/file3",
+                file_path="docs/api/file3.md",
+                content_type="text/markdown",
+                created_at=datetime.now(timezone.utc),
+                updated_at=datetime.now(timezone.utc),
+            ),
+            Entity(
+                project_id=entity_repository.project_id,
+                title="File 4",
+                entity_type="test",
+                permalink="specs/file4",
+                file_path="specs/file4.md",
+                content_type="text/markdown",
+                created_at=datetime.now(timezone.utc),
+                updated_at=datetime.now(timezone.utc),
+            ),
+            Entity(
+                project_id=entity_repository.project_id,
+                title="File 5",
+                entity_type="test",
+                permalink="notes/2024/q1/file5",
+                file_path="notes/2024/q1/file5.md",
+                content_type="text/markdown",
+                created_at=datetime.now(timezone.utc),
+                updated_at=datetime.now(timezone.utc),
+            ),
+        ]
+        session.add_all(entities)
+        await session.flush()
+
+    # Get distinct directories
+    directories = await entity_repository.get_distinct_directories()
+
+    # Verify directories are extracted correctly
+    assert isinstance(directories, list)
+    assert len(directories) > 0
+
+    # Should include all parent directories but not filenames
+    expected_dirs = {
+        "docs",
+        "docs/guides",
+        "docs/api",
+        "notes",
+        "notes/2024",
+        "notes/2024/q1",
+        "specs",
+    }
+    assert set(directories) == expected_dirs
+
+    # Verify results are sorted
+    assert directories == sorted(directories)
+
+    # Verify no file paths are included
+    for dir_path in directories:
+        assert not dir_path.endswith(".md")
+
+
+@pytest.mark.asyncio
+async def test_get_distinct_directories_empty_db(entity_repository: EntityRepository):
+    """Test getting distinct directories when database is empty."""
+    directories = await entity_repository.get_distinct_directories()
+    assert directories == []
+
+
+@pytest.mark.asyncio
+async def test_find_by_directory_prefix(entity_repository: EntityRepository, session_maker):
+    """Test finding entities by directory prefix."""
+    # Create test entities in various directories
+    async with db.scoped_session(session_maker) as session:
+        entities = [
+            Entity(
+                project_id=entity_repository.project_id,
+                title="File 1",
+                entity_type="test",
+                permalink="docs/file1",
+                file_path="docs/file1.md",
+                content_type="text/markdown",
+                created_at=datetime.now(timezone.utc),
+                updated_at=datetime.now(timezone.utc),
+            ),
+            Entity(
+                project_id=entity_repository.project_id,
+                title="File 2",
+                entity_type="test",
+                permalink="docs/guides/file2",
+                file_path="docs/guides/file2.md",
+                content_type="text/markdown",
+                created_at=datetime.now(timezone.utc),
+                updated_at=datetime.now(timezone.utc),
+            ),
+            Entity(
+                project_id=entity_repository.project_id,
+                title="File 3",
+                entity_type="test",
+                permalink="docs/api/file3",
+                file_path="docs/api/file3.md",
+                content_type="text/markdown",
+                created_at=datetime.now(timezone.utc),
+                updated_at=datetime.now(timezone.utc),
+            ),
+            Entity(
+                project_id=entity_repository.project_id,
+                title="File 4",
+                entity_type="test",
+                permalink="specs/file4",
+                file_path="specs/file4.md",
+                content_type="text/markdown",
+                created_at=datetime.now(timezone.utc),
+                updated_at=datetime.now(timezone.utc),
+            ),
+        ]
+        session.add_all(entities)
+        await session.flush()
+
+    # Test finding all entities in "docs" directory and subdirectories
+    docs_entities = await entity_repository.find_by_directory_prefix("docs")
+    assert len(docs_entities) == 3
+    file_paths = {e.file_path for e in docs_entities}
+    assert file_paths == {"docs/file1.md", "docs/guides/file2.md", "docs/api/file3.md"}
+
+    # Test finding entities in "docs/guides" subdirectory
+    guides_entities = await entity_repository.find_by_directory_prefix("docs/guides")
+    assert len(guides_entities) == 1
+    assert guides_entities[0].file_path == "docs/guides/file2.md"
+
+    # Test finding entities in "specs" directory
+    specs_entities = await entity_repository.find_by_directory_prefix("specs")
+    assert len(specs_entities) == 1
+    assert specs_entities[0].file_path == "specs/file4.md"
+
+    # Test with root directory (empty string)
+    all_entities = await entity_repository.find_by_directory_prefix("")
+    assert len(all_entities) == 4
+
+    # Test with root directory (slash)
+    all_entities = await entity_repository.find_by_directory_prefix("/")
+    assert len(all_entities) == 4
+
+    # Test with non-existent directory
+    nonexistent = await entity_repository.find_by_directory_prefix("nonexistent")
+    assert len(nonexistent) == 0
+
+
+@pytest.mark.asyncio
+async def test_find_by_directory_prefix_basic_fields_only(
+    entity_repository: EntityRepository, session_maker
+):
+    """Test that find_by_directory_prefix returns basic entity fields.
+
+    Note: This method uses use_query_options=False for performance,
+    so it doesn't eager load relationships. Directory trees only need
+    basic entity fields.
+    """
+    # Create test entity
+    async with db.scoped_session(session_maker) as session:
+        entity = Entity(
+            project_id=entity_repository.project_id,
+            title="Test Entity",
+            entity_type="test",
+            permalink="docs/test",
+            file_path="docs/test.md",
+            content_type="text/markdown",
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+        session.add(entity)
+        await session.flush()
+
+    # Query entity by directory prefix
+    entities = await entity_repository.find_by_directory_prefix("docs")
+    assert len(entities) == 1
+
+    # Verify basic fields are present (all we need for directory trees)
+    entity = entities[0]
+    assert entity.title == "Test Entity"
+    assert entity.file_path == "docs/test.md"
+    assert entity.permalink == "docs/test"
+    assert entity.entity_type == "test"
+    assert entity.content_type == "text/markdown"
+    assert entity.updated_at is not None

tests/services/test_directory_service.py

@@ -208,3 +208,51 @@ async def test_list_directory_default_parameters(directory_service: DirectorySer
     assert len(result) == 1
     assert result[0].name == "test"
     assert result[0].type == "directory"
+
+
+@pytest.mark.asyncio
+async def test_directory_structure_empty(directory_service: DirectoryService):
+    """Test getting empty directory structure."""
+    # When no entities exist, result should just be the root
+    result = await directory_service.get_directory_structure()
+    assert result is not None
+    assert len(result.children) == 0
+
+    assert result.name == "Root"
+    assert result.directory_path == "/"
+    assert result.type == "directory"
+    assert result.has_children is False
+
+
+@pytest.mark.asyncio
+async def test_directory_structure(directory_service: DirectoryService, test_graph):
+    """Test getting directory structure with folders only (no files)."""
+    # test_graph files:
+    # /
+    # ├── test
+    # │   ├── Connected Entity 1.md
+    # │   ├── Connected Entity 2.md
+    # │   ├── Deep Entity.md
+    # │   ├── Deeper Entity.md
+    # │   └── Root.md
+
+    result = await directory_service.get_directory_structure()
+    assert result is not None
+    assert len(result.children) == 1
+
+    # Should only have the "test" directory, not the files
+    node_0 = result.children[0]
+    assert node_0.name == "test"
+    assert node_0.type == "directory"
+    assert node_0.directory_path == "/test"
+    assert node_0.has_children is False  # No subdirectories, only files
+
+    # Verify no file metadata is present
+    assert node_0.content_type is None
+    assert node_0.entity_id is None
+    assert node_0.entity_type is None
+    assert node_0.title is None
+    assert node_0.permalink is None
+
+    # No file nodes should be present
+    assert len(node_0.children) == 0