From a3bfc3c42d9f6ecfa04621282736faf05ff7d19c Mon Sep 17 00:00:00 2001 From: phernandez Date: Thu, 5 Dec 2024 15:49:36 -0600 Subject: [PATCH] Implement relations --- src/basic_memory/fileio.py | 51 ++++++++++++++++++++++++++-- src/basic_memory/schemas.py | 60 ++++++++++++++++++++++++++++----- src/basic_memory/services.py | 35 +++++++------------ tests/.coverage | Bin 0 -> 53248 bytes tests/test_relation_service.py | 37 +++++++++++++------- 5 files changed, 137 insertions(+), 46 deletions(-) create mode 100644 tests/.coverage diff --git a/src/basic_memory/fileio.py b/src/basic_memory/fileio.py index 4b29fbc6..39a434fe 100644 --- a/src/basic_memory/fileio.py +++ b/src/basic_memory/fileio.py @@ -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 ) diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 009b6cba..0e1c187d 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -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 \ No newline at end of file + +# Update forward refs +Entity.model_rebuild() +Relation.model_rebuild() \ No newline at end of file diff --git a/src/basic_memory/services.py b/src/basic_memory/services.py index c99c2d41..c96ce4dc 100644 --- a/src/basic_memory/services.py +++ b/src/basic_memory/services.py @@ -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)}") diff --git a/tests/.coverage b/tests/.coverage new file mode 100644 index 0000000000000000000000000000000000000000..27d02386a7d0b5198e2632233f77f9b2f0007a14 GIT binary patch literal 53248 zcmeI)%WoUU9S87P?vg7?q(+Y6u&NMy0N17k>k%P9k|uy-r*!}$Mr0>$F9k|3$&s`P zxl8XZWhp^W(nHQJT^0|}8-AJniKWICVdPs~3!w`p+5<(Q|QKm<>*tB6~JM?U< z+umriC}v*ll;y8QsqiNue_NiH)1`~${UcXP?cyIs{%SkL2A#kH0SG_<0zYnn-pP_( z86P*!zZ^OBmW~3ap@U>y{>8ILkDNNHP91sj_)(P{QwN3>Ez{HLhzk6(s-pw7;I_1K zy{6l6BG+3~(UQ(WcSGIeiH?rZRL22l;`qE(bDI<^(uU8X$WeSdL=y5qW4=hryKe4evK->Hb<+{?{^H@l|6fm3uoen!aTLQ8-;5d z#eqh%c6fuP!{;|{H{44G8wvNKR!jK{ahKuPDybZ5&FlNjY=+?k+|iVu2Be(+pskHF zOL8v$@+T7aaaBsX|S0*Nm^ToIsrL~nl*#k`|tI698quHspcEWCu4J6++NM`uR-~=m-8+MO0Pllp%G|o6^h@P0#X(m8Cz^Arnru?KG3_&Xuw* z#j{KLmgJJz=$B5UBe$&;2Z=KiH>|Txs2YK$lQfr3R?Y8f23%XUIbQQs@({n3>+Lb@ z%E5z1ubXzpWSpsmG^w|pEZgaBZlICebh6)G=4rsCcRTCb3EQpX;N9dEcyLQ zyEhAnTtff?5P$##AOHafKmY;|fB*y_kO`E`2aN1ffZVV-QA|Gs;P3zK@>@dwL!Op$ z^tI8&(OtBQ1p*L&00bZa0SG_<0uX=z1nxqCdD{}by|cdx>3;dy&XNv1y2oC>INQ|A zvvntQ8+5Dr68{`XvNzq!i;_IJ`6L6~@#$xCNQ4rdH@Ake?KgFcbNO_R5}eq4E}L~$G`2>u zR&zZ!s?{h-{{G)CUlH6EX?vC-E@?;l;d3mZAI1OW&@00Izz00bZa0SG_<0z#ay zdV7t{Ut_La|Jx_6-o7nJuz&quc-HDww<5*<^?&}9)jPNaNe*29TU$iAcKvT2w|Wy> zkjwb`KS%XaEpEP|*IobP_x~_}KmY;|fB*y_009U<00Izz00g$RfN7Yb%-{bT@@qkV zSReoa2tWV=5P$##AOHafKmY;|*oFe8SuV!+|I4P3U&xQ-+wyO+CnMU$0s#m>00Izz z00bZa0SG_<0uX?}_Y@elEu*;G{`8b2E;(Ed*f=RYZ3KlH)wJnbCHUy|>R zKPk=`V(hRO5~nXuh_OFEF(l5N+i%hNXRJ#lak(-ka&&UV60{bEx@=cXqrhKD>V8(tPQ&H!nW;<$LeEl)EoS=l19B7h+6^z2e+qLk#5&t5_Tz zj_?1MSB3maekT7bKbG&yf6^`%2tWV=5P$##AOHafKmY;|fB*#UW&!?e!LaT4sZfZY z^7;75vf?MxjGuBji#`B|zyIe?1+YK>0uX=z1Rwwb2tWV=5P$##Ah3M}`1^lc|8L*= zMadul0SG_<0uX=z1Rwwb2tWV=ERfv)e~thD|Nrua{Ez%Z{##yy4Fn(n0SG_<0uX=z j1Rwwb2tWV=_l$sP(=Ys+1zzQOW%0`7RgPEu`~Tkpi(Ltp literal 0 HcmV?d00001 diff --git a/tests/test_relation_service.py b/tests/test_relation_service.py index 764887fa..0e797414 100644 --- a/tests/test_relation_service.py +++ b/tests/test_relation_service.py @@ -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 \ No newline at end of file