mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
relation service wip
This commit is contained in:
phernandez
+124
-64
@@ -1,13 +1,14 @@
|
||||
from datetime import datetime, UTC
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from typing import Optional, List
|
||||
from uuid import uuid4
|
||||
from sqlalchemy import and_, select, delete
|
||||
|
||||
from basic_memory.models import Entity as DbEntity # Rename to avoid confusion
|
||||
from basic_memory.models import Observation as DbObservation
|
||||
from basic_memory.repository import EntityRepository, ObservationRepository
|
||||
from basic_memory.schemas import Entity, Observation
|
||||
from basic_memory.models import Relation as DbRelation
|
||||
from basic_memory.repository import EntityRepository, ObservationRepository, RelationRepository
|
||||
from basic_memory.schemas import Entity, Observation, Relation
|
||||
from basic_memory.fileio import (
|
||||
read_entity_file, write_entity_file, delete_entity_file,
|
||||
FileOperationError, EntityNotFoundError
|
||||
@@ -24,12 +25,14 @@ class DatabaseSyncError(ServiceError):
|
||||
pass
|
||||
|
||||
|
||||
class RelationError(ServiceError):
|
||||
"""Base exception for relation-specific errors"""
|
||||
pass
|
||||
|
||||
|
||||
class EntityService:
|
||||
"""
|
||||
Service for managing entities in the filesystem and database.
|
||||
Follows the "filesystem is source of truth" principle.
|
||||
"""
|
||||
|
||||
"""Service for managing entities in the filesystem and database."""
|
||||
# [Previous EntityService implementation remains unchanged]
|
||||
def __init__(self, project_path: Path, entity_repo: EntityRepository):
|
||||
self.project_path = project_path
|
||||
self.entity_repo = entity_repo
|
||||
@@ -111,14 +114,8 @@ class EntityService:
|
||||
|
||||
|
||||
class ObservationService:
|
||||
"""
|
||||
Service for managing observations in the filesystem and database.
|
||||
Follows the "filesystem is source of truth" principle.
|
||||
|
||||
Observations are stored in entity markdown files and indexed in the database
|
||||
for efficient querying.
|
||||
"""
|
||||
|
||||
"""Service for managing observations in the filesystem and database."""
|
||||
# [Previous ObservationService implementation remains unchanged]
|
||||
def __init__(self, project_path: Path, observation_repo: ObservationRepository):
|
||||
self.project_path = project_path
|
||||
self.entities_path = project_path / "entities"
|
||||
@@ -126,22 +123,7 @@ class ObservationService:
|
||||
|
||||
async def add_observation(self, entity: Entity, content: str,
|
||||
context: Optional[str] = None) -> Observation:
|
||||
"""
|
||||
Add a new observation to an entity.
|
||||
|
||||
Args:
|
||||
entity: Entity to add observation to
|
||||
content: Content of the observation
|
||||
context: Optional context for the observation
|
||||
|
||||
Returns:
|
||||
The created Observation
|
||||
|
||||
Raises:
|
||||
FileOperationError: If file operations fail
|
||||
DatabaseSyncError: If database sync fails
|
||||
"""
|
||||
# Create new observation
|
||||
"""Add a new observation to an entity."""
|
||||
observation = Observation(content=content)
|
||||
entity.observations.append(observation)
|
||||
|
||||
@@ -160,41 +142,117 @@ class ObservationService:
|
||||
return observation
|
||||
except Exception as e:
|
||||
raise DatabaseSyncError(f"Failed to sync observation to database: {str(e)}") from e
|
||||
|
||||
async def search_observations(self, query: str) -> list[Observation]:
|
||||
|
||||
|
||||
class RelationService:
|
||||
"""
|
||||
Service for managing relations between entities.
|
||||
Follows the "filesystem is source of truth" principle.
|
||||
|
||||
Relations are stored in entity markdown files and indexed in the database
|
||||
for efficient querying.
|
||||
"""
|
||||
|
||||
def __init__(self, project_path: Path, relation_repo: RelationRepository):
|
||||
self.project_path = project_path
|
||||
self.entities_path = project_path / "entities"
|
||||
self.relation_repo = relation_repo
|
||||
|
||||
async def create_relation(self, from_entity: Entity, to_entity: Entity, relation_type: str,
|
||||
context: Optional[str] = None) -> Relation:
|
||||
"""
|
||||
Search for observations across all entities.
|
||||
Create a new relation between two entities.
|
||||
|
||||
Args:
|
||||
query: Text to search for in observation content
|
||||
from_entity: Source entity
|
||||
to_entity: Target entity
|
||||
relation_type: Type of relation
|
||||
context: Optional context for the relation
|
||||
|
||||
Returns:
|
||||
List of matching observations with their entity contexts
|
||||
The created Relation
|
||||
|
||||
Raises:
|
||||
FileOperationError: If file operations fail
|
||||
DatabaseSyncError: If database sync fails
|
||||
"""
|
||||
result = await self.observation_repo.execute_query(
|
||||
select(DbObservation).filter(
|
||||
DbObservation.content.contains(query)
|
||||
)
|
||||
# Create new relation
|
||||
relation = Relation(
|
||||
id=f"rel-{uuid4().hex[:8]}",
|
||||
from_id=from_entity.id,
|
||||
to_id=to_entity.id,
|
||||
relation_type=relation_type,
|
||||
context=context
|
||||
)
|
||||
return [
|
||||
Observation(content=obs.content)
|
||||
for obs in result.scalars().all()
|
||||
]
|
||||
|
||||
async def get_observations_by_context(self, context: str) -> list[Observation]:
|
||||
"""Get all observations with a specific context."""
|
||||
db_observations = await self.observation_repo.find_by_context(context)
|
||||
return [
|
||||
Observation(content=obs.content)
|
||||
for obs in db_observations
|
||||
]
|
||||
# Add relation to source entity's relations list
|
||||
if not hasattr(from_entity, 'relations'):
|
||||
from_entity.relations = []
|
||||
from_entity.relations.append(relation)
|
||||
|
||||
# Update filesystem first (source of truth)
|
||||
await write_entity_file(self.entities_path, from_entity)
|
||||
|
||||
# Update database index
|
||||
try:
|
||||
await self.relation_repo.create({
|
||||
'id': relation.id,
|
||||
'from_id': from_entity.id,
|
||||
'to_id': to_entity.id,
|
||||
'relation_type': relation_type,
|
||||
'context': context,
|
||||
'created_at': datetime.now(UTC)
|
||||
})
|
||||
return relation
|
||||
except Exception as e:
|
||||
raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e
|
||||
|
||||
async def rebuild_observation_index(self) -> None:
|
||||
async def get_entity_relations(self, entity: Entity) -> List[Relation]:
|
||||
"""
|
||||
Rebuild the observation database index from filesystem contents.
|
||||
Get all relations for an entity (both outgoing and incoming).
|
||||
|
||||
Args:
|
||||
entity: Entity to get relations for
|
||||
|
||||
Returns:
|
||||
List of relations where the entity is either source or target
|
||||
"""
|
||||
# Relations are stored in the entity object already
|
||||
return getattr(entity, 'relations', [])
|
||||
|
||||
async def delete_relation(self, from_entity: Entity, relation_id: str) -> bool:
|
||||
"""
|
||||
Delete a relation from both filesystem and database.
|
||||
|
||||
Args:
|
||||
from_entity: Source entity containing the relation
|
||||
relation_id: ID of the relation to delete
|
||||
|
||||
Returns:
|
||||
True if deletion was successful
|
||||
|
||||
Raises:
|
||||
RelationError: If relation cannot be found or deleted
|
||||
"""
|
||||
# Remove relation from entity's relations
|
||||
if hasattr(from_entity, 'relations'):
|
||||
from_entity.relations = [
|
||||
r for r in from_entity.relations
|
||||
if r.id != relation_id
|
||||
]
|
||||
|
||||
# Update filesystem first (source of truth)
|
||||
await write_entity_file(self.entities_path, from_entity)
|
||||
|
||||
# Remove from database index
|
||||
await self.relation_repo.delete(relation_id)
|
||||
return True
|
||||
|
||||
async def rebuild_relation_index(self) -> None:
|
||||
"""
|
||||
Rebuild the relation database index from filesystem contents.
|
||||
Used for recovery or ensuring sync.
|
||||
"""
|
||||
# List all entity files
|
||||
if not self.entities_path.exists():
|
||||
return
|
||||
|
||||
@@ -202,20 +260,22 @@ class ObservationService:
|
||||
entity_files = list(self.entities_path.glob("*.md"))
|
||||
except Exception as e:
|
||||
raise FileOperationError(f"Failed to read entities directory: {str(e)}") from e
|
||||
|
||||
# Clear existing observation index
|
||||
await self.observation_repo.execute_query(delete(DbObservation))
|
||||
|
||||
# Clear existing relation index
|
||||
await self.relation_repo.execute_query(delete(DbRelation))
|
||||
|
||||
# Rebuild from each entity file
|
||||
for entity_file in entity_files:
|
||||
try:
|
||||
entity = await read_entity_file(self.entities_path, entity_file.stem)
|
||||
for obs in entity.observations:
|
||||
await self.observation_repo.create({
|
||||
'id': f"{entity.id}-obs-{uuid4().hex[:8]}",
|
||||
'entity_id': entity.id,
|
||||
'content': obs.content,
|
||||
for relation in getattr(entity, 'relations', []):
|
||||
await self.relation_repo.create({
|
||||
'id': relation.id,
|
||||
'from_id': relation.from_id,
|
||||
'to_id': relation.to_id,
|
||||
'relation_type': relation.relation_type,
|
||||
'context': relation.context,
|
||||
'created_at': datetime.now(UTC)
|
||||
})
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to reindex observations for {entity_file}: {str(e)}")
|
||||
print(f"Warning: Failed to reindex relations for {entity_file}: {str(e)}")
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
import pytest
|
||||
from datetime import datetime, UTC
|
||||
from pathlib import Path
|
||||
|
||||
from basic_memory.models import Entity as DbEntity
|
||||
from basic_memory.models import Relation as DbRelation
|
||||
from basic_memory.schemas import Entity, Relation
|
||||
from basic_memory.services import RelationService, EntityService
|
||||
from basic_memory.repository import EntityRepository, RelationRepository
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def entity_repo(db_session):
|
||||
return EntityRepository(db_session)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def relation_repo(db_session):
|
||||
return RelationRepository(db_session)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def entity_service(tmp_path, entity_repo):
|
||||
return EntityService(tmp_path, entity_repo)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def relation_service(tmp_path, relation_repo):
|
||||
return RelationService(tmp_path, relation_repo)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def sample_entities(entity_service):
|
||||
"""Create two sample entities for testing relations"""
|
||||
entity1 = await entity_service.create_entity(
|
||||
name="test_entity_1",
|
||||
entity_type="test_type"
|
||||
)
|
||||
entity2 = await entity_service.create_entity(
|
||||
name="test_entity_2",
|
||||
entity_type="test_type"
|
||||
)
|
||||
return entity1, entity2
|
||||
|
||||
|
||||
async def test_create_relation(relation_service, sample_entities):
|
||||
"""Test creating a basic relation between two entities"""
|
||||
entity1, entity2 = sample_entities
|
||||
|
||||
relation = await relation_service.create_relation(
|
||||
from_entity=entity1,
|
||||
to_entity=entity2,
|
||||
relation_type="test_relation"
|
||||
)
|
||||
|
||||
assert relation.from_id == entity1.id
|
||||
assert relation.to_id == entity2.id
|
||||
assert relation.relation_type == "test_relation"
|
||||
|
||||
# Verify relation was added to source entity's relations
|
||||
assert hasattr(entity1, 'relations')
|
||||
assert len(entity1.relations) == 1
|
||||
assert entity1.relations[0].id == relation.id
|
||||
|
||||
# Verify file was written with relation
|
||||
entity_file = relation_service.entities_path / f"{entity1.id}.md"
|
||||
assert entity_file.exists()
|
||||
content = entity_file.read_text()
|
||||
assert "## Relations" in content
|
||||
assert f"[{entity2.id}] test_relation" in content
|
||||
|
||||
# Verify database was updated
|
||||
db_relation = await relation_service.relation_repo.find_by_id(relation.id)
|
||||
assert db_relation is not None
|
||||
assert db_relation.from_id == entity1.id
|
||||
assert db_relation.to_id == entity2.id
|
||||
assert db_relation.relation_type == "test_relation"
|
||||
|
||||
|
||||
async def test_create_relation_with_context(relation_service, sample_entities):
|
||||
"""Test creating a relation with context information"""
|
||||
entity1, entity2 = sample_entities
|
||||
|
||||
relation = await relation_service.create_relation(
|
||||
from_entity=entity1,
|
||||
to_entity=entity2,
|
||||
relation_type="test_relation",
|
||||
context="test context"
|
||||
)
|
||||
|
||||
assert relation.context == "test context"
|
||||
|
||||
# Verify context in file
|
||||
entity_file = relation_service.entities_path / f"{entity1.id}.md"
|
||||
content = entity_file.read_text()
|
||||
assert f"[{entity2.id}] test_relation | test context" in content
|
||||
|
||||
# Verify context in database
|
||||
db_relation = await relation_service.relation_repo.find_by_id(relation.id)
|
||||
assert db_relation.context == "test context"
|
||||
|
||||
|
||||
async def test_get_entity_relations(relation_service, sample_entities):
|
||||
"""Test retrieving relations for an entity"""
|
||||
entity1, entity2 = sample_entities
|
||||
|
||||
# Create test relation
|
||||
relation = await relation_service.create_relation(
|
||||
from_entity=entity1,
|
||||
to_entity=entity2,
|
||||
relation_type="test_relation"
|
||||
)
|
||||
|
||||
# Get relations from entity
|
||||
relations = await relation_service.get_entity_relations(entity1)
|
||||
|
||||
assert len(relations) == 1
|
||||
assert relations[0].from_id == entity1.id
|
||||
assert relations[0].to_id == entity2.id
|
||||
assert relations[0].relation_type == "test_relation"
|
||||
|
||||
|
||||
async def test_delete_relation(relation_service, sample_entities):
|
||||
"""Test deleting a relation"""
|
||||
entity1, entity2 = sample_entities
|
||||
|
||||
# Create then delete a relation
|
||||
relation = await relation_service.create_relation(
|
||||
from_entity=entity1,
|
||||
to_entity=entity2,
|
||||
relation_type="test_relation"
|
||||
)
|
||||
|
||||
success = await relation_service.delete_relation(entity1, relation.id)
|
||||
assert success is True
|
||||
|
||||
# Verify removed from entity relations
|
||||
assert not entity1.relations or relation.id not in [r.id for r in entity1.relations]
|
||||
|
||||
# Verify removed from file
|
||||
entity_file = relation_service.entities_path / f"{entity1.id}.md"
|
||||
content = entity_file.read_text()
|
||||
assert f"[{entity2.id}] test_relation" not in content
|
||||
|
||||
# Verify removed from database
|
||||
db_relation = await relation_service.relation_repo.find_by_id(relation.id)
|
||||
assert db_relation is None
|
||||
|
||||
|
||||
async def test_rebuild_relation_index(relation_service, sample_entities):
|
||||
"""Test rebuilding the relation index from files"""
|
||||
entity1, entity2 = sample_entities
|
||||
|
||||
# Create some test relations
|
||||
relation1 = await relation_service.create_relation(
|
||||
from_entity=entity1,
|
||||
to_entity=entity2,
|
||||
relation_type="test_relation_1"
|
||||
)
|
||||
relation2 = await relation_service.create_relation(
|
||||
from_entity=entity2,
|
||||
to_entity=entity1,
|
||||
relation_type="test_relation_2"
|
||||
)
|
||||
|
||||
# Clear the database relations
|
||||
await relation_service.relation_repo.execute_query(
|
||||
'DELETE FROM relation'
|
||||
)
|
||||
|
||||
# Rebuild index
|
||||
await relation_service.rebuild_relation_index()
|
||||
|
||||
# Verify relations were restored
|
||||
db_relations = await relation_service.relation_repo.execute_query(
|
||||
'SELECT * FROM relation'
|
||||
)
|
||||
relations = db_relations.scalars().all()
|
||||
|
||||
assert len(relations) == 2
|
||||
relation_types = {r.relation_type for r in relations}
|
||||
assert relation_types == {"test_relation_1", "test_relation_2"}
|
||||
Reference in New Issue
Block a user