admin/basicmachines-co-basic-memory
1
0
Fork 0
mirror of https://github.com/basicmachines-co/basic-memory synced 2026-06-21 13:47:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Implement relations

Browse Source
This commit is contained in:
phernandez
2024-12-05 15:49:36 -06:00
parent 29ad6fcf4b
commit a3bfc3c42d
5 changed files with 137 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_memory/fileio.py
+49 -2
View File
@@ -4,7 +4,7 @@ Handles reading and writing entities and observations to the filesystem.
"""
from pathlib import Path
from basic_memory.schemas import Entity, Observation
from basic_memory.schemas import Entity, Observation, Relation
class FileOperationError(Exception):
@@ -50,6 +50,20 @@ async def write_entity_file(entities_path: Path, entity: Entity) -> bool:
# Add observations
for obs in entity.observations:
content.append(f"- {obs.content}\n")
# Add relations section if we have relations
if hasattr(entity, 'relations') and entity.relations:
content.extend([
"\n", # Blank line before relations
"## Relations\n"
])
# Use model_dump to get proper storage format
for rel in entity.relations:
rel_data = rel.model_dump()
relation_line = f"- [{rel_data['to_id']}] {rel_data['relation_type']}"
if rel_data.get('context'):
relation_line += f" | {rel_data['context']}"
content.append(f"{relation_line}\n")
# Handle atomic write operation
temp_path = entity_path.with_suffix('.tmp')
@@ -97,9 +111,12 @@ async def read_entity_file(entities_path: Path, entity_id: str) -> Entity:
# Parse metadata (type)
entity_type = ""
observations = []
relations = []
# Parse content sections
in_observations = False
in_relations = False
for line in content[1:]: # Skip the title line
line = line.strip()
if not line:
@@ -109,14 +126,44 @@ async def read_entity_file(entities_path: Path, entity_id: str) -> Entity:
entity_type = line.replace("type: ", "").strip()
elif line == "## Observations":
in_observations = True
in_relations = False
elif line == "## Relations":
in_observations = False
in_relations = True
elif in_observations and line.startswith("- "):
observations.append(Observation(content=line[2:]))
elif in_relations and line.startswith("- "):
# Parse relation line: - [target_id] relation_type | context
line = line[2:] # Remove the bullet point
if "] " not in line:
continue # Skip malformed lines
# Split on the first "] " to separate ID from relation_type
id_part, rest = line.split("] ", 1)
target_id = id_part[1:] # Remove leading [
# Split rest on " | " if there's a context
parts = rest.split(" | ", 1)
relation_type = parts[0]
context = parts[1] if len(parts) > 1 else None
# Create temporary entities for the relation
target_entity = Entity(id=target_id, name=target_id, entity_type="unknown")
source_entity = Entity(id=entity_id, name=name, entity_type=entity_type)
relations.append(Relation(
from_entity=source_entity,
to_entity=target_entity,
relation_type=relation_type,
context=context
))
return Entity(
id=entity_id,
name=name,
entity_type=entity_type,
observations=observations
observations=observations,
relations=relations
)
src/basic_memory/schemas.py
+51 -9
View File
@@ -6,7 +6,7 @@ independent from storage/persistence concerns.
from datetime import datetime, UTC
from uuid import uuid4
from typing import List
from typing import List, Optional, ForwardRef, Dict, Any
from pydantic import BaseModel, model_validator
@@ -15,6 +15,36 @@ class Observation(BaseModel):
content: str
class Relation(BaseModel):
"""
Represents a directed edge between entities in the knowledge graph.
Relations are always stored in active voice (e.g. "created", "teaches", etc.)
"""
id: str
from_entity: 'Entity'
to_entity: 'Entity'
relation_type: str
context: Optional[str] = None
@model_validator(mode='before')
@classmethod
def generate_id_if_needed(cls, data: dict) -> dict:
"""Generate an ID if one wasn't provided"""
if not data.get('id'):
data['id'] = f"rel-{uuid4().hex[:8]}"
return data
def model_dump(self, **kwargs) -> Dict[str, Any]:
"""Serialize to storage format with entity IDs"""
return {
'id': self.id,
'from_id': self.from_entity.id,
'to_id': self.to_entity.id,
'relation_type': self.relation_type,
'context': self.context
}
class Entity(BaseModel):
"""
Represents a node in our knowledge graph - could be a person, project,
@@ -25,6 +55,7 @@ class Entity(BaseModel):
name: str
entity_type: str
observations: List[Observation] = []
relations: List[Relation] = []
@model_validator(mode='before')
@classmethod
@@ -36,12 +67,23 @@ class Entity(BaseModel):
data['id'] = f"{timestamp}-{normalized_name}-{uuid4().hex[:8]}"
return data
def model_dump(self, **kwargs) -> Dict[str, Any]:
"""Serialize entity, handling relations to prevent circular references"""
# Get basic data without relations
exclude = kwargs.pop('exclude', set())
exclude.add('relations')
basic_data = super().model_dump(exclude=exclude, **kwargs)
# Add serialized relations if we have any
if 'relations' not in exclude and self.relations:
basic_data['relations'] = [
relation.model_dump(**kwargs)
for relation in self.relations
]
return basic_data
class Relation(BaseModel):
"""
Represents a directed edge between two entities in our knowledge graph.
Relations are always stored in active voice (e.g. "created", "teaches", etc.)
"""
from_entity: str
to_entity: str
relation_type: str
# Update forward refs
Entity.model_rebuild()
Relation.model_rebuild()
src/basic_memory/services.py
+12 -23
View File
@@ -32,7 +32,6 @@ class RelationError(ServiceError):
class EntityService:
"""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
@@ -48,6 +47,7 @@ class EntityService:
# Observations will be handled by ObservationService
entity_data.pop('observations', None) # Remove observations if present
entity_data.pop('relations', None) # Remove relations if present
# Try to find existing entity first
if await self.entity_repo.find_by_id(entity.id):
@@ -115,7 +115,6 @@ class EntityService:
class ObservationService:
"""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"
@@ -176,11 +175,10 @@ class RelationService:
FileOperationError: If file operations fail
DatabaseSyncError: If database sync fails
"""
# Create new relation
# Create new relation with actual Entity objects
relation = Relation(
id=f"rel-{uuid4().hex[:8]}",
from_id=from_entity.id,
to_id=to_entity.id,
from_entity=from_entity,
to_entity=to_entity,
relation_type=relation_type,
context=context
)
@@ -194,15 +192,11 @@ class RelationService:
await write_entity_file(self.entities_path, from_entity)
# Update database index
# model_dump will handle converting Entity refs to IDs
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)
})
db_data = relation.model_dump()
db_data['created_at'] = datetime.now(UTC)
await self.relation_repo.create(db_data)
return relation
except Exception as e:
raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e
@@ -217,7 +211,7 @@ class RelationService:
Returns:
List of relations where the entity is either source or target
"""
# Relations are stored in the entity object already
# Relations are stored in the entity object
return getattr(entity, 'relations', [])
async def delete_relation(self, from_entity: Entity, relation_id: str) -> bool:
@@ -269,13 +263,8 @@ class RelationService:
try:
entity = await read_entity_file(self.entities_path, entity_file.stem)
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)
})
db_data = relation.model_dump()
db_data['created_at'] = datetime.now(UTC)
await self.relation_repo.create(db_data)
except Exception as e:
print(f"Warning: Failed to reindex relations for {entity_file}: {str(e)}")
tests/.coverage
BIN
View File
Binary file not shown.
tests/test_relation_service.py
+25 -12
View File
@@ -1,6 +1,6 @@
import pytest
import pytest_asyncio
from sqlalchemy import delete
from sqlalchemy import delete, select
from basic_memory.models import Relation as DbRelation
from basic_memory.schemas import Relation
@@ -9,6 +9,7 @@ from basic_memory.fileio import read_entity_file
pytestmark = pytest.mark.asyncio
@pytest_asyncio.fixture
async def sample_entities(entity_service):
"""Create two sample entities for testing relations"""
@@ -22,6 +23,13 @@ async def sample_entities(entity_service):
)
return entity1, entity2
# Helper function for comparing strings with variable whitespace
def normalize_whitespace(s: str) -> str:
"""Normalize whitespace in a string for comparison."""
return ' '.join(s.split())
# Happy Path Tests
async def test_create_relation(relation_service, sample_entities):
@@ -34,8 +42,9 @@ async def test_create_relation(relation_service, sample_entities):
relation_type="test_relation"
)
assert relation.from_id == entity1.id
assert relation.to_id == entity2.id
# Check Entity objects in relation
assert relation.from_entity.id == entity1.id
assert relation.to_entity.id == entity2.id
assert relation.relation_type == "test_relation"
# Verify relation was added to source entity's relations
@@ -50,7 +59,7 @@ async def test_create_relation(relation_service, sample_entities):
assert "## Relations" in content
assert f"[{entity2.id}] test_relation" in content
# Verify database was updated
# Verify database was updated with correct IDs
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
@@ -96,8 +105,8 @@ async def test_get_entity_relations(relation_service, sample_entities):
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].from_entity.id == entity1.id
assert relations[0].to_entity.id == entity2.id
assert relations[0].relation_type == "test_relation"
@@ -150,16 +159,16 @@ async def test_rebuild_relation_index(relation_service, sample_entities):
# 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()
# Verify relations were restored using SQLAlchemy select
query = select(DbRelation)
result = await relation_service.relation_repo.execute_query(query)
relations = result.scalars().all()
assert len(relations) == 2
relation_types = {r.relation_type for r in relations}
assert relation_types == {"test_relation_1", "test_relation_2"}
# Error Path Tests
async def test_file_operation_error(relation_service, sample_entities, monkeypatch):
@@ -178,6 +187,7 @@ async def test_file_operation_error(relation_service, sample_entities, monkeypat
relation_type="test_relation"
)
async def test_database_sync_error(relation_service, sample_entities, monkeypatch):
"""Test handling of database sync errors."""
entity1, entity2 = sample_entities
@@ -194,6 +204,7 @@ async def test_database_sync_error(relation_service, sample_entities, monkeypatc
relation_type="test_relation"
)
# Edge Cases
async def test_relation_with_special_characters(relation_service, sample_entities):
@@ -229,4 +240,6 @@ async def test_very_long_relation_type(relation_service, sample_entities):
# Verify file content
entity = await read_entity_file(relation_service.entities_path, entity1.id)
assert any(r.relation_type == long_type for r in getattr(entity, 'relations', []))
# Compare with normalized whitespace
stored_types = {normalize_whitespace(r.relation_type) for r in getattr(entity, 'relations', [])}
assert normalize_whitespace(long_type) in stored_types