From 00b73b0d08dada4c386cf72f022e7ce62f773acd Mon Sep 17 00:00:00 2001 From: Paul Hernandez <60959+phernandez@users.noreply.github.com> Date: Sat, 11 Oct 2025 09:11:46 -0500 Subject: [PATCH] feat: Optimize directory operations for 10-100x performance improvement (#350) Signed-off-by: phernandez Co-authored-by: Claude --- .../api/routers/directory_router.py | 21 ++ .../repository/entity_repository.py | 60 +++++ src/basic_memory/repository/repository.py | 19 +- .../services/directory_service.py | 127 ++++++++++- src/basic_memory/services/project_service.py | 8 +- tests/api/test_directory_router.py | 133 +++++++++++ tests/repository/test_entity_repository.py | 212 ++++++++++++++++++ tests/services/test_directory_service.py | 48 ++++ 8 files changed, 621 insertions(+), 7 deletions(-) diff --git a/src/basic_memory/api/routers/directory_router.py b/src/basic_memory/api/routers/directory_router.py index b7d8a089..89ce48a6 100644 --- a/src/basic_memory/api/routers/directory_router.py +++ b/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, diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index 854a0930..486b23e6 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/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 diff --git a/src/basic_memory/repository/repository.py b/src/basic_memory/repository/repository.py index c878eca7..1cf69dd5 100644 --- a/src/basic_memory/repository/repository.py +++ b/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) diff --git a/src/basic_memory/services/directory_service.py b/src/basic_memory/services/directory_service.py index e5a1b669..cefdcf3d 100644 --- a/src/basic_memory/services/directory_service.py +++ b/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]: diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py index c5bf2502..e4646ce1 100644 --- a/src/basic_memory/services/project_service.py +++ b/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.""" diff --git a/tests/api/test_directory_router.py b/tests/api/test_directory_router.py index e9e050af..a688430d 100644 --- a/tests/api/test_directory_router.py +++ b/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"] == [] diff --git a/tests/repository/test_entity_repository.py b/tests/repository/test_entity_repository.py index 77f3698e..9b585cae 100644 --- a/tests/repository/test_entity_repository.py +++ b/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 diff --git a/tests/services/test_directory_service.py b/tests/services/test_directory_service.py index a40e4c3c..565a559d 100644 --- a/tests/services/test_directory_service.py +++ b/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