fix yml parsing/writing in document service

This commit is contained in:
phernandez
2025-01-03 23:59:09 -06:00
parent f4c5327889
commit eda2a30099
8 changed files with 495 additions and 151 deletions
+10 -10
View File
@@ -11,7 +11,6 @@ from sqlalchemy.ext.asyncio import (
from basic_memory import db
from basic_memory.config import ProjectConfig, config
from basic_memory.db import DatabaseType
from basic_memory.markdown.knowledge_writer import KnowledgeWriter
from basic_memory.repository.document_repository import DocumentRepository
from basic_memory.repository.entity_repository import EntityRepository
@@ -107,6 +106,12 @@ DocumentRepositoryDep = Annotated[DocumentRepository, Depends(get_document_repos
## services
async def get_file_service() -> FileService:
return FileService()
FileServiceDep = Annotated[FileService, Depends(get_file_service)]
async def get_entity_service(entity_repository: EntityRepositoryDep) -> EntityService:
"""Create EntityService with repository."""
@@ -135,10 +140,10 @@ RelationServiceDep = Annotated[RelationService, Depends(get_relation_service)]
async def get_document_service(
document_repository: DocumentRepositoryDep, project_config: ProjectConfigDep
document_repository: DocumentRepositoryDep, project_config: ProjectConfigDep, file_service: FileServiceDep,
) -> DocumentService:
"""Create RelationService with repository."""
return DocumentService(document_repository, project_config.documents_dir)
return DocumentService(document_repository, project_config.documents_dir, file_service)
DocumentServiceDep = Annotated[DocumentService, Depends(get_document_service)]
@@ -153,18 +158,13 @@ async def get_activity_service(
return ActivityService(
entity_service=entity_service,
document_service=document_service,
relation_service=relation_service
relation_service=relation_service,
)
ActivityServiceDep = Annotated[ActivityService, Depends(get_activity_service)]
async def get_file_service() -> FileService:
return FileService()
FileServiceDep = Annotated[FileService, Depends(get_file_service)]
async def get_knowledge_writer() -> KnowledgeWriter:
@@ -193,4 +193,4 @@ async def get_knowledge_service(
)
KnowledgeServiceDep = Annotated[KnowledgeService, Depends(get_knowledge_service)]
KnowledgeServiceDep = Annotated[KnowledgeService, Depends(get_knowledge_service)]
+2 -2
View File
@@ -6,7 +6,7 @@ from typing import List, TypeVar, Generic, Optional, Dict, Any, Tuple
from loguru import logger
from basic_memory.utils.file_utils import parse_frontmatter, ParseError, FileError
from basic_memory.utils.file_utils import parse_content_with_frontmatter, ParseError, FileError
T = TypeVar("T") # The parsed document type
@@ -62,7 +62,7 @@ class MarkdownParser(ABC, Generic[T]):
"""
try:
# Split into frontmatter and content
frontmatter, markdown = await parse_frontmatter(content)
frontmatter, markdown = await parse_content_with_frontmatter(content)
# Parse frontmatter
parsed_frontmatter = await self.parse_frontmatter(frontmatter)
+86 -62
View File
@@ -1,34 +1,30 @@
"""Service for managing documents in the system."""
import hashlib
from datetime import datetime, UTC
from pathlib import Path
from typing import Optional, Dict, Any, Tuple, Sequence
import yaml
from loguru import logger
from basic_memory.models import Document
from basic_memory.repository.document_repository import DocumentRepository
from basic_memory.services.service import BaseService
from basic_memory.utils.file_utils import compute_checksum, ensure_directory, add_frontmatter
from basic_memory.services.file_service import FileService
from basic_memory.services.exceptions import FileOperationError
class DocumentError(Exception):
"""Base exception for document operations."""
pass
class DocumentNotFoundError(DocumentError):
"""Raised when a document doesn't exist."""
pass
class DocumentWriteError(DocumentError):
"""Raised when document file operations fail."""
pass
@@ -40,38 +36,35 @@ class DocumentService(BaseService[DocumentRepository]):
them in sync. The filesystem is the source of truth.
"""
def __init__(self, document_repository: DocumentRepository, documents_path: Path):
def __init__(
self,
document_repository: DocumentRepository,
documents_path: Path,
file_service: FileService,
):
super().__init__(document_repository)
self.documents_base_path = documents_path
self.file_service = file_service
def get_document_path(self, path_id: str) -> Path:
"""Get full filesystem path for a document."""
doc_path = Path(path_id)
if doc_path.is_absolute():
raise DocumentError(f"Document path {path_id} must be relative")
document_path = Path(self.documents_base_path / path_id)
document_path = self.documents_base_path / path_id
logger.debug(f"Document path: '{path_id}' file_path: {document_path}")
return document_path
async def add_frontmatter(
self, content: str, path_id: str, metadata: Optional[Dict[str, Any]] = None
) -> str:
"""Add frontmatter to document content."""
# Generate frontmatter with timestamps
now = datetime.now(UTC).isoformat()
frontmatter = {"id": path_id, "created": now, "modified": now}
if metadata:
frontmatter.update(metadata)
return await add_frontmatter(content, frontmatter)
async def list_documents(self) -> Sequence[Document]:
"""List all documents."""
return await self.repository.find_all()
async def create_document(
self, path_id: str, content: str, metadata: Optional[Dict[str, Any]] = None
self,
path_id: str,
content: str,
metadata: Optional[Dict[str, Any]] = None,
) -> Document:
"""
Create a new document.
@@ -89,24 +82,28 @@ class DocumentService(BaseService[DocumentRepository]):
"""
logger.debug(f"Creating document path_id: {path_id}")
# Ensure parent directories exist
file_path = self.get_document_path(path_id)
await ensure_directory(file_path.parent)
try:
# 1. Add frontmatter with path_id
content_with_frontmatter = await self.add_frontmatter(content, path_id, metadata)
# Prepare frontmatter
now = datetime.now(UTC).isoformat()
frontmatter = {
"id": path_id,
"created": now,
"modified": now
}
if metadata:
frontmatter.update(metadata)
# 2. Compute checksum
checksum = await compute_checksum(content_with_frontmatter)
# 3. Write complete file
file_path.write_text(content_with_frontmatter)
# Let FileService handle the write with frontmatter
checksum = await self.file_service.write_with_frontmatter(
path=file_path,
content=content,
frontmatter=frontmatter
)
# 4. Create DB record with checksum
# Create DB record
document = await self.repository.create({
"path_id": path_id,
# for create the file_path will be the same as path_id
"file_path": path_id,
"checksum": checksum,
"doc_metadata": metadata
@@ -114,9 +111,11 @@ class DocumentService(BaseService[DocumentRepository]):
return document
except FileOperationError as e:
raise DocumentWriteError(f"Failed to create document: {e}")
except Exception as e:
# Clean up on any failure
file_path.unlink(missing_ok=True)
await self.file_service.delete_file(file_path)
raise DocumentWriteError(f"Failed to create document: {e}")
async def read_document_by_path_id(self, path_id: str) -> Tuple[Document, str]:
@@ -140,17 +139,20 @@ class DocumentService(BaseService[DocumentRepository]):
if not document:
raise DocumentNotFoundError(f"Document not found: {path_id}")
# Read content since file is source of truth
# use the actual file_path of the doc on the filesystem
file_path = self.get_document_path(document.file_path)
try:
content = file_path.read_text()
# Read content using FileService
file_path = self.get_document_path(document.file_path)
content, _ = await self.file_service.read_file(file_path)
return document, content
except Exception as e:
raise DocumentError(f"Failed to read document {path_id}: {e} at path {file_path}")
except FileOperationError as e:
raise DocumentError(f"Failed to read document {path_id}: {e}")
async def update_document_by_path_id(
self, path_id: str, content: str, metadata: Optional[Dict[str, Any]] = None
self,
path_id: str,
content: str,
metadata: Optional[Dict[str, Any]] = None,
) -> Document:
"""
Update a document using its PathId.
@@ -174,24 +176,45 @@ class DocumentService(BaseService[DocumentRepository]):
if not document:
raise DocumentNotFoundError(f"Document not found: {path_id}")
# Add frontmatter with metadata
content_with_frontmatter = await self.add_frontmatter(content, path_id, metadata)
# Write new content first
try:
# Read existing content to preserve frontmatter
file_path = self.get_document_path(document.file_path)
file_path.write_text(content_with_frontmatter)
except Exception as e:
raise DocumentWriteError(f"Failed to write document: {e}")
old_content, _ = await self.file_service.read_file(file_path)
try:
existing_frontmatter = await self.file_service.parse_frontmatter(old_content)
except:
# If we can't parse existing frontmatter, start fresh
existing_frontmatter = {
"id": path_id,
"created": datetime.now(UTC).isoformat(),
}
# Update DB record
checksum = await compute_checksum(content_with_frontmatter)
update_data = {"checksum": checksum}
if metadata is not None:
update_data["doc_metadata"] = metadata # pyright: ignore [reportArgumentType]
# Update frontmatter
now = datetime.now(UTC).isoformat()
frontmatter = {
**existing_frontmatter,
"modified": now,
}
if metadata:
frontmatter.update(metadata)
updated_document = await self.repository.update(document.id, update_data)
return updated_document
# Update file using FileService
checksum = await self.file_service.write_with_frontmatter(
path=file_path,
content=content,
frontmatter=frontmatter
)
# Update DB record
update_data = {"checksum": checksum}
if metadata is not None:
update_data["doc_metadata"] = metadata
updated_document = await self.repository.update(document.id, update_data)
return updated_document
except FileOperationError as e:
raise DocumentWriteError(f"Failed to update document: {e}")
async def delete_document_by_path_id(self, path_id: str) -> None:
"""
@@ -211,12 +234,13 @@ class DocumentService(BaseService[DocumentRepository]):
if not document:
raise DocumentNotFoundError(f"Document not found: {path_id}")
# Delete file first since it's source of truth
try:
# Delete file using FileService
file_path = self.get_document_path(document.path_id)
file_path.unlink(missing_ok=True)
except Exception as e:
raise DocumentWriteError(f"Failed to delete document {path_id}: {e}")
await self.file_service.delete_file(file_path)
# Delete database record
await self.repository.delete(document.id)
# Delete database record
await self.repository.delete(document.id)
except FileOperationError as e:
raise DocumentWriteError(f"Failed to delete document {path_id}: {e}")
+120 -5
View File
@@ -106,19 +106,88 @@ class FileService:
logger.error(f"Failed to delete file {path}: {e}")
raise FileOperationError(f"Failed to delete file: {e}")
async def has_frontmatter(self, content: str) -> bool:
"""
Check if content has frontmatter markers.
Args:
content: Content to check
Returns:
True if content appears to have frontmatter
"""
try:
return file_utils.has_frontmatter(content)
except Exception as e:
logger.error(f"Failed to check frontmatter: {e}")
return False
async def parse_frontmatter(self, content: str) -> Dict[str, Any]:
"""
Parse frontmatter from content.
Args:
content: Content containing frontmatter
Returns:
Parsed frontmatter as dict
Raises:
FileOperationError: If parsing fails
"""
try:
return file_utils.parse_frontmatter(content)
except Exception as e:
logger.error(f"Failed to parse frontmatter: {e}")
raise FileOperationError(f"Failed to parse frontmatter: {e}")
async def remove_frontmatter(self, content: str) -> str:
"""
Remove frontmatter from content.
Args:
content: Content with frontmatter
Returns:
Content with frontmatter removed
Raises:
FileOperationError: If removal fails
"""
try:
return file_utils.remove_frontmatter(content)
except Exception as e:
logger.error(f"Failed to remove frontmatter: {e}")
raise FileOperationError(f"Failed to remove frontmatter: {e}")
async def remove_frontmatter_lenient(self, content: str) -> str:
"""
Remove frontmatter without validation.
Args:
content: Content that may contain frontmatter
Returns:
Content with potential frontmatter removed
"""
try:
return file_utils.remove_frontmatter_lenient(content)
except Exception as e:
logger.error(f"Failed to remove frontmatter leniently: {e}")
raise FileOperationError(f"Failed to remove frontmatter: {e}")
async def add_frontmatter(
self,
*,
content: str,
frontmatter: Dict[str, Any],
metadata: Optional[Dict[str, Any]] = None,
content: str,
) -> str:
"""
Add YAML frontmatter to content.
Args:
frontmatter: frontmatter info
content: Content to add frontmatter to
frontmatter: Frontmatter to add
metadata: Optional additional metadata
Returns:
@@ -132,7 +201,53 @@ class FileService:
frontmatter.update(metadata)
return await file_utils.add_frontmatter(content, frontmatter)
except Exception as e:
logger.error(f"Failed to add frontmatter: {e}")
raise FileOperationError(f"Failed to add frontmatter: {e}")
async def write_with_frontmatter(
self,
path: Path,
content: str,
frontmatter: Dict[str, Any],
) -> str:
"""
Write content to file with frontmatter, properly handling existing frontmatter.
If content already has frontmatter, it will be updated with new values.
If not, frontmatter will be added.
Args:
path: Path where to write
content: Content to write
frontmatter: Frontmatter to add/update
Returns:
Checksum of written content
Raises:
FileOperationError: If operation fails
"""
try:
final_content: str
if await self.has_frontmatter(content):
try:
# Try to parse and merge existing frontmatter
existing_frontmatter = await self.parse_frontmatter(content)
content_only = await self.remove_frontmatter(content)
merged_frontmatter = {**existing_frontmatter, **frontmatter}
final_content = await self.add_frontmatter(content_only, merged_frontmatter)
except FileOperationError:
# If parsing fails, just strip any frontmatter-like content and start fresh
content_only = await self.remove_frontmatter_lenient(content)
final_content = await self.add_frontmatter(content_only, frontmatter)
else:
# No existing frontmatter, just add new
final_content = await self.add_frontmatter(content, frontmatter)
# Write and return checksum
return await self.write_file(path, final_content)
except Exception as e:
logger.error(f"Failed to write file with frontmatter {path}: {e}")
raise FileOperationError(f"Failed to write file with frontmatter: {e}")
+131 -30
View File
@@ -80,6 +80,121 @@ async def write_file_atomic(path: Path, content: str) -> None:
raise FileWriteError(f"Failed to write file {path}: {e}")
def has_frontmatter(content: str) -> bool:
"""
Check if content contains YAML frontmatter.
Args:
content: Content to check
Returns:
True if content has frontmatter delimiter (---), False otherwise
"""
content = content.strip()
return content.startswith("---") and "---" in content[3:]
def parse_frontmatter(content: str) -> Dict[str, Any]:
"""
Parse YAML frontmatter from content.
Args:
content: Content with YAML frontmatter
Returns:
Dictionary of frontmatter values
Raises:
ParseError: If frontmatter is invalid or parsing fails
"""
try:
if not has_frontmatter(content):
raise ParseError("Content has no frontmatter")
# Split on first two occurrences of ---
parts = content.split("---", 2)
if len(parts) < 3:
raise ParseError("Invalid frontmatter format")
# Parse YAML
try:
frontmatter = yaml.safe_load(parts[1])
# Handle empty frontmatter (None from yaml.safe_load)
if frontmatter is None:
return {}
if not isinstance(frontmatter, dict):
raise ParseError("Frontmatter must be a YAML dictionary")
return frontmatter
except yaml.YAMLError as e:
raise ParseError(f"Invalid YAML in frontmatter: {e}")
except Exception as e:
if not isinstance(e, ParseError):
logger.error(f"Failed to parse frontmatter: {e}")
raise ParseError(f"Failed to parse frontmatter: {e}")
raise
def remove_frontmatter(content: str) -> str:
"""
Remove YAML frontmatter from content.
Args:
content: Content with frontmatter
Returns:
Content with frontmatter removed
Raises:
ParseError: If frontmatter format is invalid
"""
try:
if not has_frontmatter(content):
return content.strip()
# Split on first two occurrences of ---
parts = content.split("---", 2)
if len(parts) < 3:
raise ParseError("Invalid frontmatter format")
return parts[2].strip()
except Exception as e:
if not isinstance(e, ParseError):
logger.error(f"Failed to remove frontmatter: {e}")
raise ParseError(f"Failed to remove frontmatter: {e}")
raise
def remove_frontmatter_lenient(content: str) -> str:
"""
Remove frontmatter markers and anything between them without validation.
This is a more permissive version of remove_frontmatter that doesn't
try to validate the YAML content. It simply removes everything between
the first two '---' markers if they exist.
Args:
content: Content that may contain frontmatter
Returns:
Content with any frontmatter markers and content removed
"""
content = content.strip()
if not content.startswith("---"):
return content
# Find the second marker
rest = content[3:].strip()
if "---" not in rest:
return content
# Split on the second marker and take everything after
parts = rest.split("---", 1)
return parts[1].strip()
async def add_frontmatter(content: str, frontmatter: Dict[str, Any]) -> str:
"""
Add YAML frontmatter to content.
@@ -96,49 +211,35 @@ async def add_frontmatter(content: str, frontmatter: Dict[str, Any]) -> str:
"""
try:
yaml_fm = yaml.dump(frontmatter, sort_keys=False)
return f"---\n{yaml_fm}---\n\n{content}"
return f"---\n{yaml_fm}---\n\n{content.strip()}"
except yaml.YAMLError as e:
logger.error(f"Failed to add frontmatter: {e}")
raise ParseError(f"Failed to add frontmatter: {e}")
async def parse_frontmatter(content: str) -> Tuple[Dict[str, Any], str]:
async def parse_content_with_frontmatter(content: str) -> Tuple[Dict[str, Any], str]:
"""
Parse YAML frontmatter from content.
Parse both frontmatter and content.
Args:
content: Text content with optional frontmatter
Returns:
Tuple of (frontmatter dict, remaining content)
Tuple of (frontmatter dict, content without frontmatter)
Raises:
ParseError: If frontmatter parsing fails
ParseError: If parsing fails
"""
try:
# Ensure we have frontmatter
if not content.strip().startswith("---"):
if not has_frontmatter(content):
return {}, content.strip()
# Split on first two occurrences of ---
parts = content.split("---", 2)
if len(parts) < 3:
raise ParseError("Invalid frontmatter format")
# Parse YAML (skipping empty first part)
try:
frontmatter = yaml.safe_load(parts[1])
if not isinstance(frontmatter, dict):
raise ParseError("Frontmatter must be a YAML dictionary")
# Return parsed frontmatter and rest of content
return frontmatter, parts[2].strip()
except yaml.YAMLError as e:
raise ParseError(f"Invalid YAML in frontmatter: {e}")
frontmatter = parse_frontmatter(content)
remaining = remove_frontmatter(content)
return frontmatter, remaining
except Exception as e:
if not isinstance(e, ParseError):
logger.error(f"Failed to parse frontmatter: {e}")
raise ParseError(f"Failed to parse frontmatter: {e}") from e
logger.error(f"Failed to parse content with frontmatter: {e}")
raise ParseError(f"Failed to parse content with frontmatter: {e}")
raise
+2 -1
View File
@@ -86,9 +86,10 @@ async def document_repository(
async def document_service(
document_repository: DocumentRepository,
test_config: ProjectConfig,
file_service: FileService,
) -> DocumentService:
"""Create a DocumentService instance."""
return DocumentService(document_repository, test_config.documents_dir)
return DocumentService(document_repository, test_config.documents_dir, file_service)
@pytest_asyncio.fixture(scope="function")
-19
View File
@@ -24,27 +24,8 @@ def test_dir():
yield Path(tmpdir)
@pytest.fixture
def document_service(session_maker, test_dir):
"""Create document service."""
docs_path = test_dir / 'documents'
repository = DocumentRepository(session_maker)
return DocumentService(repository, docs_path)
@pytest.fixture
def entity_service(session_maker):
"""Create entity service."""
repository = EntityRepository(session_maker)
return EntityService(repository)
@pytest.fixture
def relation_service(session_maker):
"""Create relation service."""
repository = RelationRepository(session_maker)
return RelationService(repository)
@pytest.fixture
def activity_service(document_service, entity_service, relation_service):
+144 -22
View File
@@ -10,6 +10,9 @@ from basic_memory.utils.file_utils import (
write_file_atomic,
add_frontmatter,
parse_frontmatter,
has_frontmatter,
remove_frontmatter,
parse_content_with_frontmatter,
FileError,
FileWriteError,
ParseError,
@@ -83,47 +86,166 @@ async def test_add_frontmatter():
assert "- a\n- b" in result or "['a', 'b']" in result
# Should preserve content
assert result.endswith(content)
assert result.endswith(f"test content")
@pytest.mark.asyncio
async def test_parse_frontmatter():
def test_has_frontmatter():
"""Test frontmatter detection."""
# Valid frontmatter
assert has_frontmatter("""---
title: Test
---
content""")
# Just content
assert not has_frontmatter("Just content")
# Empty content
assert not has_frontmatter("")
# Just delimiter
assert not has_frontmatter("---")
# Delimiter not at start
assert not has_frontmatter("""
Some text
---
title: Test
---""")
# Invalid format
assert not has_frontmatter("--title: test--")
def test_parse_frontmatter():
"""Test parsing frontmatter."""
# Valid frontmatter
content = """---
title: Test
tags:
- a
- b
---
content"""
result = parse_frontmatter(content)
assert result == {"title": "Test", "tags": ["a", "b"]}
# Empty frontmatter
content = """---
---
content"""
result = parse_frontmatter(content)
assert result == None or result == {}
# Invalid YAML
with pytest.raises(ParseError):
parse_frontmatter("""---
[invalid yaml]
---
content""")
# No frontmatter
with pytest.raises(ParseError):
parse_frontmatter("Just content")
# Incomplete frontmatter
with pytest.raises(ParseError):
parse_frontmatter("""---
title: Test
content""")
def test_remove_frontmatter():
"""Test removing frontmatter."""
# With frontmatter
content = """---
title: Test
---
test content"""
assert remove_frontmatter(content) == "test content"
# No frontmatter
content = "test content"
assert remove_frontmatter(content) == "test content"
# Only frontmatter
content = """---
title: Test
---
"""
assert remove_frontmatter(content) == ""
# frontmatter missing some fields
assert remove_frontmatter("""---
title: Test
content""") == "---\ntitle: Test\ncontent"
@pytest.mark.asyncio
async def test_parse_content_with_frontmatter():
"""Test combined frontmatter and content parsing."""
# Full document
content = """---
title: Test
tags:
- a
- b
---
test content"""
metadata, remaining = await parse_frontmatter(content)
frontmatter, body = await parse_content_with_frontmatter(content)
assert frontmatter == {"title": "Test", "tags": ["a", "b"]}
assert body == "test content"
assert metadata == {"title": "Test", "tags": ["a", "b"]}
assert remaining.strip() == "test content"
@pytest.mark.asyncio
async def test_parse_frontmatter_no_frontmatter():
"""Test parsing content without frontmatter."""
# No frontmatter
content = "test content"
metadata, remaining = await parse_frontmatter(content)
frontmatter, body = await parse_content_with_frontmatter(content)
assert frontmatter == {}
assert body == "test content"
assert metadata == {}
assert remaining == content
# Empty document
frontmatter, body = await parse_content_with_frontmatter("")
assert frontmatter == {}
assert body == ""
# Only frontmatter
content = """---
title: Test
---
"""
frontmatter, body = await parse_content_with_frontmatter(content)
assert frontmatter == {"title": "Test"}
assert body == ""
@pytest.mark.asyncio
async def test_parse_frontmatter_error():
"""Test frontmatter parse error handling."""
# Really invalid YAML frontmatter
async def test_frontmatter_whitespace_handling():
"""Test frontmatter handling with various whitespace."""
# Extra newlines before frontmatter
content = """
---
title: Test
---
content"""
assert has_frontmatter(content.strip())
frontmatter = parse_frontmatter(content.strip())
assert frontmatter == {"title": "Test"}
# Extra newlines after frontmatter
content = """---
[[ this is not valid yaml ]]
title:: [}
title: Test
---
test content"""
with pytest.raises(ParseError):
await parse_frontmatter(content)
content"""
result = await add_frontmatter("content", {"title": "Test"})
assert result.count("\n\n") == 1 # Should normalize to single blank line
# Spaces around content
content = """---
title: Test
---
content """
assert remove_frontmatter(content).strip() == "content"