mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
add knowledge_service, file_service
This commit is contained in:
phernandez
@@ -0,0 +1,120 @@
|
||||
"""Service for file operations with checksum tracking."""
|
||||
|
||||
import hashlib
|
||||
from datetime import datetime, UTC
|
||||
from pathlib import Path
|
||||
from typing import Optional, Dict, Any, Tuple
|
||||
|
||||
import yaml
|
||||
from loguru import logger
|
||||
|
||||
from basic_memory.services.exceptions import FileOperationError
|
||||
from basic_memory.utils import file_utils
|
||||
|
||||
|
||||
class FileService:
|
||||
"""
|
||||
Service for handling file operations.
|
||||
|
||||
Features:
|
||||
- Consistent file writing with checksums
|
||||
- Frontmatter management
|
||||
- Atomic operations
|
||||
- Error handling
|
||||
"""
|
||||
|
||||
async def write_file(self, path: Path, content: str) -> str:
|
||||
"""
|
||||
Write content to file and return checksum.
|
||||
|
||||
Args:
|
||||
path: Path where to write
|
||||
content: Content to write
|
||||
|
||||
Returns:
|
||||
Checksum of written content
|
||||
|
||||
Raises:
|
||||
FileOperationError: If write fails
|
||||
"""
|
||||
try:
|
||||
# Ensure parent directory exists
|
||||
await file_utils.ensure_directory(path.parent)
|
||||
|
||||
# Write content atomically
|
||||
await file_utils.write_file_atomic(path, content)
|
||||
|
||||
# Compute and return checksum
|
||||
return await file_utils.compute_checksum(content)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to write file {path}: {e}")
|
||||
raise FileOperationError(f"Failed to write file: {e}")
|
||||
|
||||
async def read_file(self, path: Path) -> Tuple[str, str]:
|
||||
"""
|
||||
Read file and compute checksum.
|
||||
|
||||
Args:
|
||||
path: Path to read
|
||||
|
||||
Returns:
|
||||
Tuple of (content, checksum)
|
||||
|
||||
Raises:
|
||||
FileOperationError: If read fails
|
||||
"""
|
||||
try:
|
||||
content = path.read_text()
|
||||
checksum = await file_utils.compute_checksum(content)
|
||||
return content, checksum
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to read file {path}: {e}")
|
||||
raise FileOperationError(f"Failed to read file: {e}")
|
||||
|
||||
async def delete_file(self, path: Path) -> None:
|
||||
"""
|
||||
Delete file if it exists.
|
||||
|
||||
Args:
|
||||
path: Path to delete
|
||||
|
||||
Raises:
|
||||
FileOperationError: If deletion fails
|
||||
"""
|
||||
try:
|
||||
path.unlink(missing_ok=True)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete file {path}: {e}")
|
||||
raise FileOperationError(f"Failed to delete file: {e}")
|
||||
|
||||
async def add_frontmatter(
|
||||
self, content: str, id: int, metadata: Optional[Dict[str, Any]] = None
|
||||
) -> str:
|
||||
"""
|
||||
Add YAML frontmatter to content.
|
||||
|
||||
Args:
|
||||
content: Content to add frontmatter to
|
||||
id: ID to include in frontmatter
|
||||
metadata: Optional additional metadata
|
||||
|
||||
Returns:
|
||||
Content with frontmatter added
|
||||
|
||||
Raises:
|
||||
FileOperationError: If frontmatter creation fails
|
||||
"""
|
||||
try:
|
||||
# Generate frontmatter with timestamps
|
||||
now = datetime.now(UTC).isoformat()
|
||||
frontmatter = {"id": id, "created": now, "modified": now}
|
||||
if metadata:
|
||||
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}")
|
||||
@@ -0,0 +1,187 @@
|
||||
"""Service for managing knowledge graph entities and their file persistence."""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Sequence, List
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from basic_memory.markdown.knowledge_parser import KnowledgeParser
|
||||
from basic_memory.models import Entity as EntityModel
|
||||
from basic_memory.schemas import Entity as EntitySchema, Relation as RelationSchema
|
||||
from basic_memory.services.entity_service import EntityService
|
||||
from basic_memory.services.exceptions import EntityNotFoundError, FileOperationError
|
||||
from basic_memory.services.observation_service import ObservationService
|
||||
from basic_memory.services.relation_service import RelationService
|
||||
|
||||
|
||||
class KnowledgeService:
|
||||
"""
|
||||
Service for managing knowledge graph entities and their persistence.
|
||||
|
||||
Orchestrates operations between:
|
||||
- EntityService for core entity operations
|
||||
- ObservationService for atomic facts
|
||||
- RelationService for entity connections
|
||||
- FileService for persistence
|
||||
- KnowledgeParser for file formatting
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
entity_service: EntityService,
|
||||
observation_service: ObservationService,
|
||||
relation_service: RelationService,
|
||||
file_service, # FileService
|
||||
knowledge_parser: KnowledgeParser,
|
||||
):
|
||||
self.entity_service = entity_service
|
||||
self.observation_service = observation_service
|
||||
self.relation_service = relation_service
|
||||
self.file_service = file_service
|
||||
self.knowledge_parser = knowledge_parser
|
||||
|
||||
def get_entity_path(self, entity: EntityModel) -> Path:
|
||||
"""Generate filesystem path for entity."""
|
||||
# Store in entities/[type]/[name].md
|
||||
return Path("entities") / entity.entity_type / f"{entity.name}.md"
|
||||
|
||||
async def write_entity_file(self, entity: EntityModel) -> str:
|
||||
"""Write entity to filesystem and return checksum."""
|
||||
try:
|
||||
# Format content
|
||||
path = self.get_entity_path(entity)
|
||||
content = await self.knowledge_parser.format_entity(entity)
|
||||
|
||||
# Write and get checksum
|
||||
return await self.file_service.write_file(path, content)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to write entity file: {e}")
|
||||
raise FileOperationError(f"Failed to write entity file: {e}")
|
||||
|
||||
async def create_entity(self, entity: EntitySchema) -> EntityModel:
|
||||
"""Create a new entity and write to filesystem."""
|
||||
logger.debug(f"Creating entity: {entity}")
|
||||
try:
|
||||
# 1. Create entity in DB
|
||||
db_entity = await self.entity_service.create_entity(entity)
|
||||
|
||||
# 2. Write file and get checksum
|
||||
checksum = await self.write_entity_file(db_entity)
|
||||
|
||||
# 3. Update DB with checksum
|
||||
updated = await self.entity_service.update_entity(db_entity.id, {"checksum": checksum})
|
||||
|
||||
return updated
|
||||
|
||||
except Exception as e:
|
||||
# Clean up on any failure
|
||||
if "db_entity" in locals():
|
||||
await self.entity_service.delete_entity(db_entity.id) # pyright: ignore [reportPossiblyUnboundVariable]
|
||||
if "path" in locals():
|
||||
await self.file_service.delete_file(path) # pyright: ignore [reportUndefinedVariable]
|
||||
logger.error(f"Failed to create entity: {e}")
|
||||
raise
|
||||
|
||||
async def create_entities(self, entities: List[EntitySchema]) -> Sequence[EntityModel]:
|
||||
"""Create multiple entities."""
|
||||
logger.debug(f"Creating {len(entities)} entities")
|
||||
created = []
|
||||
|
||||
for entity in entities:
|
||||
try:
|
||||
created_entity = await self.create_entity(entity)
|
||||
created.append(created_entity)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create entity {entity.name}: {e}")
|
||||
continue
|
||||
|
||||
return created
|
||||
|
||||
async def create_relations(self, relations: List[RelationSchema]) -> Sequence[RelationSchema]:
|
||||
"""Create relations and update affected entity files."""
|
||||
logger.debug(f"Creating {len(relations)} relations")
|
||||
created = []
|
||||
|
||||
for relation in relations:
|
||||
try:
|
||||
# Create relation in DB
|
||||
db_relation = await self.relation_service.create_relation(relation)
|
||||
|
||||
# Get updated entities to write
|
||||
from_entity = await self.entity_service.get_entity(relation.from_id)
|
||||
to_entity = await self.entity_service.get_entity(relation.to_id)
|
||||
|
||||
# Update files with their new relations
|
||||
for entity in [from_entity, to_entity]:
|
||||
checksum = await self.write_entity_file(entity)
|
||||
await self.entity_service.update_entity(entity.id, {"checksum": checksum})
|
||||
|
||||
created.append(db_relation)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create relation: {e}")
|
||||
continue
|
||||
|
||||
return created
|
||||
|
||||
async def add_observations(
|
||||
self, entity_id: int, observations: List[str], context: str | None = None
|
||||
) -> EntityModel:
|
||||
"""Add observations to entity and update its file."""
|
||||
logger.debug(f"Adding observations to entity {entity_id}")
|
||||
|
||||
try:
|
||||
# Add observations to DB
|
||||
await self.observation_service.add_observations(entity_id, observations, context)
|
||||
|
||||
# Get updated entity
|
||||
entity = await self.entity_service.get_entity(entity_id)
|
||||
if not entity:
|
||||
raise EntityNotFoundError(f"Entity not found: {entity_id}")
|
||||
|
||||
# Write updated file
|
||||
checksum = await self.write_entity_file(entity)
|
||||
|
||||
# Update checksum in DB
|
||||
return await self.entity_service.update_entity(entity_id, {"checksum": checksum})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to add observations: {e}")
|
||||
raise
|
||||
|
||||
async def delete_entity(self, entity_id: int) -> bool:
|
||||
"""Delete entity and its file."""
|
||||
logger.debug(f"Deleting entity: {entity_id}")
|
||||
|
||||
try:
|
||||
# Get entity first for file deletion
|
||||
entity = await self.entity_service.get_entity(entity_id)
|
||||
if not entity:
|
||||
return True # Already deleted
|
||||
|
||||
# Delete file first (it's source of truth)
|
||||
path = self.get_entity_path(entity)
|
||||
await self.file_service.delete_file(path)
|
||||
|
||||
# Delete from DB (this will cascade to observations/relations)
|
||||
return await self.entity_service.delete_entity(entity_id)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete entity: {e}")
|
||||
raise
|
||||
|
||||
async def delete_entities(self, entity_ids: List[int]) -> bool:
|
||||
"""Delete multiple entities and their files."""
|
||||
logger.debug(f"Deleting entities: {entity_ids}")
|
||||
success = True
|
||||
|
||||
for entity_id in entity_ids:
|
||||
try:
|
||||
await self.delete_entity(entity_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete entity {entity_id}: {e}")
|
||||
success = False
|
||||
continue
|
||||
|
||||
return success
|
||||
@@ -0,0 +1,178 @@
|
||||
"""Tests for file operations service."""
|
||||
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from basic_memory.services.exceptions import FileOperationError
|
||||
from basic_memory.services.file_service import FileService
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def file_service():
|
||||
"""Create FileService instance."""
|
||||
return FileService()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_read_file(tmp_path: Path, file_service: FileService):
|
||||
"""Test basic write/read operations with checksums."""
|
||||
test_path = tmp_path / "test.md"
|
||||
test_content = "test content\nwith multiple lines"
|
||||
|
||||
# Write file and get checksum
|
||||
checksum = await file_service.write_file(test_path, test_content)
|
||||
assert test_path.exists()
|
||||
|
||||
# Read back and verify content/checksum
|
||||
content, read_checksum = await file_service.read_file(test_path)
|
||||
assert content == test_content
|
||||
assert read_checksum == checksum
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_creates_directories(tmp_path: Path, file_service: FileService):
|
||||
"""Test directory creation on write."""
|
||||
test_path = tmp_path / "subdir" / "nested" / "test.md"
|
||||
test_content = "test content"
|
||||
|
||||
# Write should create directories
|
||||
await file_service.write_file(test_path, test_content)
|
||||
assert test_path.exists()
|
||||
assert test_path.parent.is_dir()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_atomic(tmp_path: Path, file_service: FileService):
|
||||
"""Test atomic write with no partial files."""
|
||||
test_path = tmp_path / "test.md"
|
||||
temp_path = test_path.with_suffix(".tmp")
|
||||
|
||||
# Mock write_file_atomic to raise an error
|
||||
with patch('basic_memory.utils.file_utils.write_file_atomic') as mock_write:
|
||||
mock_write.side_effect = Exception("Write failed")
|
||||
|
||||
# Attempt write that will fail
|
||||
with pytest.raises(FileOperationError):
|
||||
await file_service.write_file(test_path, "test content")
|
||||
|
||||
# No partial files should exist
|
||||
assert not test_path.exists()
|
||||
assert not temp_path.exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_file(tmp_path: Path, file_service: FileService):
|
||||
"""Test file deletion."""
|
||||
test_path = tmp_path / "test.md"
|
||||
test_content = "test content"
|
||||
|
||||
# Create then delete
|
||||
await file_service.write_file(test_path, test_content)
|
||||
assert test_path.exists()
|
||||
|
||||
await file_service.delete_file(test_path)
|
||||
assert not test_path.exists()
|
||||
|
||||
# Delete non-existent file should not error
|
||||
await file_service.delete_file(test_path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_frontmatter(file_service: FileService):
|
||||
"""Test frontmatter addition."""
|
||||
test_content = "# Test\nSome content"
|
||||
test_metadata = {"type": "test", "tags": ["one", "two"]}
|
||||
|
||||
# Add frontmatter
|
||||
content_with_fm = await file_service.add_frontmatter(
|
||||
test_content, id=123, metadata=test_metadata
|
||||
)
|
||||
|
||||
# Verify structure
|
||||
assert content_with_fm.startswith("---\n")
|
||||
assert "id: 123" in content_with_fm
|
||||
assert "type: test" in content_with_fm
|
||||
assert "created:" in content_with_fm
|
||||
assert "modified:" in content_with_fm
|
||||
assert "tags:" in content_with_fm
|
||||
assert test_content in content_with_fm
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checksum_consistency(tmp_path: Path, file_service: FileService):
|
||||
"""Test checksum remains consistent."""
|
||||
test_path = tmp_path / "test.md"
|
||||
test_content = "test content\n" * 10
|
||||
|
||||
# Get checksum from write
|
||||
checksum1 = await file_service.write_file(test_path, test_content)
|
||||
|
||||
# Get checksum from read
|
||||
_, checksum2 = await file_service.read_file(test_path)
|
||||
|
||||
# Write again and get new checksum
|
||||
checksum3 = await file_service.write_file(test_path, test_content)
|
||||
|
||||
# All should match
|
||||
assert checksum1 == checksum2 == checksum3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_handling_missing_file(tmp_path: Path, file_service: FileService):
|
||||
"""Test error handling for missing files."""
|
||||
test_path = tmp_path / "missing.md"
|
||||
|
||||
with pytest.raises(FileOperationError):
|
||||
await file_service.read_file(test_path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_handling_invalid_path(tmp_path: Path, file_service: FileService):
|
||||
"""Test error handling for invalid paths."""
|
||||
# Try to write to a directory instead of file
|
||||
test_path = tmp_path / "test.md"
|
||||
test_path.mkdir() # Create a directory instead of a file
|
||||
|
||||
with pytest.raises(FileOperationError):
|
||||
await file_service.write_file(test_path, "test")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_frontmatter_invalid_metadata(file_service: FileService):
|
||||
"""Test error handling for invalid frontmatter metadata."""
|
||||
# Create an object that can't be serialized to YAML
|
||||
class NonSerializable:
|
||||
def __getstate__(self):
|
||||
raise ValueError("Can't serialize me!")
|
||||
|
||||
bad_metadata = {"bad": NonSerializable()}
|
||||
|
||||
# Attempting to add frontmatter with non-serializable content
|
||||
with patch('basic_memory.utils.file_utils.add_frontmatter') as mock_add:
|
||||
mock_add.side_effect = FileOperationError("Failed to serialize metadata")
|
||||
with pytest.raises(FileOperationError):
|
||||
await file_service.add_frontmatter(
|
||||
"content",
|
||||
id=123,
|
||||
metadata=bad_metadata
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_unicode_content(tmp_path: Path, file_service: FileService):
|
||||
"""Test handling of unicode content."""
|
||||
test_path = tmp_path / "test.md"
|
||||
test_content = """
|
||||
# Test Unicode
|
||||
- Emoji: 🚀 ⭐️ 🔥
|
||||
- Chinese: 你好世界
|
||||
- Arabic: مرحبا بالعالم
|
||||
- Russian: Привет, мир
|
||||
"""
|
||||
|
||||
# Write and read back
|
||||
await file_service.write_file(test_path, test_content)
|
||||
content, _ = await file_service.read_file(test_path)
|
||||
|
||||
assert content == test_content
|
||||
Reference in New Issue
Block a user