From 6686c331a518a5e8e72323de8dcf1b23f3597cc1 Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 7 Jan 2025 19:00:08 -0600 Subject: [PATCH] refactor entity.entity_type to be freeform, add summary and content_type --- src/basic_memory/markdown/knowledge_writer.py | 4 +- src/basic_memory/mcp/main.py | 4 +- src/basic_memory/models/knowledge.py | 18 ++------- .../repository/entity_repository.py | 2 +- src/basic_memory/services/activity_service.py | 2 +- .../services/knowledge/file_operations.py | 2 +- src/basic_memory/services/search_service.py | 2 +- .../sync/knowledge_sync_service.py | 4 +- tests/repository/test_entity_repository.py | 18 ++++----- tests/services/test_entity_service.py | 40 +++++++++---------- tests/services/test_knowledge_service.py | 2 +- tests/sync/test_knowledge_sync_service.py | 4 +- 12 files changed, 46 insertions(+), 56 deletions(-) diff --git a/src/basic_memory/markdown/knowledge_writer.py b/src/basic_memory/markdown/knowledge_writer.py index 6d18107e..0cc1f222 100644 --- a/src/basic_memory/markdown/knowledge_writer.py +++ b/src/basic_memory/markdown/knowledge_writer.py @@ -25,8 +25,8 @@ class KnowledgeWriter: "", # Empty line after name ] - if entity.description: - sections.extend([entity.description, ""]) + if entity.summary: + sections.extend([entity.summary, ""]) if entity.observations: sections.extend( diff --git a/src/basic_memory/mcp/main.py b/src/basic_memory/mcp/main.py index 43b110c0..2a27cf26 100644 --- a/src/basic_memory/mcp/main.py +++ b/src/basic_memory/mcp/main.py @@ -11,8 +11,8 @@ from basic_memory.config import config from basic_memory.mcp.server import mcp # Import tools to register them -from basic_memory.mcp.tools import knowledge, search, documents, discovery, help -__all__ = ["mcp", "knowledge", "search", "documents", "discovery", "help"] +from basic_memory.mcp.tools import knowledge, search, discovery, help +__all__ = ["mcp", "knowledge", "search", "discovery", "help"] def setup_logging(home_dir: str = config.home, log_file: str = "basic-memory.log"): diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py index 6675ba07..8ac9ed7c 100644 --- a/src/basic_memory/models/knowledge.py +++ b/src/basic_memory/models/knowledge.py @@ -21,13 +21,6 @@ from basic_memory.models.base import Base from enum import Enum -class EntityType(str, Enum): - """Types of knowledge nodes.""" - - KNOWLEDGE = "knowledge" - NOTE = "note" - - class Entity(Base): """ Core entity in the knowledge graph. @@ -45,17 +38,14 @@ class Entity(Base): Index("ix_entity_type", "entity_type"), Index("ix_entity_created_at", "created_at"), # For timeline queries Index("ix_entity_updated_at", "updated_at"), # For timeline queries - CheckConstraint( - f"entity_type IN {tuple(t.value for t in EntityType)}", name="check_entity_type" - ), ) # Core identity id: Mapped[int] = mapped_column(Integer, primary_key=True) name: Mapped[str] = mapped_column(String) - entity_type: Mapped[EntityType] = mapped_column(String, default=EntityType.KNOWLEDGE) + entity_type: Mapped[str] = mapped_column(String) entity_metadata: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True) - + content_type: Mapped[str] = mapped_column(String) # Normalized path for URIs path_id: Mapped[str] = mapped_column(String, unique=True, index=True) @@ -64,8 +54,8 @@ class Entity(Base): # checksum of file checksum: Mapped[Optional[str]] = mapped_column(String, nullable=True) - # Content for knowledge entity_type - description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + # Content summary + summary: Mapped[Optional[str]] = mapped_column(Text, nullable=True) # Metadata and tracking created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index 1701bde2..a9aa85ff 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -83,7 +83,7 @@ class EntityRepository(Repository[Entity]): or_( Entity.name.ilike(search_term), Entity.entity_type.ilike(search_term), - Entity.description.ilike(search_term), + Entity.summary.ilike(search_term), Entity.observations.any(Observation.content.ilike(search_term)), ) ) diff --git a/src/basic_memory/services/activity_service.py b/src/basic_memory/services/activity_service.py index 82f75029..32464e0e 100644 --- a/src/basic_memory/services/activity_service.py +++ b/src/basic_memory/services/activity_service.py @@ -97,7 +97,7 @@ class ActivityService: timestamp=updated_at, path_id=entity.path_id, summary=f"{change_type.value.title()} entity: {entity.name}", - content=entity.description + content=entity.summary ) ) diff --git a/src/basic_memory/services/knowledge/file_operations.py b/src/basic_memory/services/knowledge/file_operations.py index 87785f1b..493032fb 100644 --- a/src/basic_memory/services/knowledge/file_operations.py +++ b/src/basic_memory/services/knowledge/file_operations.py @@ -60,7 +60,7 @@ class FileOperations: frontmatter = await writer.format_frontmatter(entity) file_content = await writer.format_content( entity=entity, - content=content or entity.description or "", + content=content or entity.summary or "", ) # Add frontmatter and write diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index b02f3459..71c862df 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -62,7 +62,7 @@ class SearchService: content = "\n".join( [ entity.name, - entity.description or "", + entity.summary or "", # Add observations *[f"{obs.category}: {obs.content}" for obs in entity.observations], # Add relations diff --git a/src/basic_memory/sync/knowledge_sync_service.py b/src/basic_memory/sync/knowledge_sync_service.py index d35f5853..02611cbd 100644 --- a/src/basic_memory/sync/knowledge_sync_service.py +++ b/src/basic_memory/sync/knowledge_sync_service.py @@ -83,7 +83,7 @@ class KnowledgeSyncService: # Update fields from markdown db_entity.name = markdown.content.title db_entity.entity_type = markdown.frontmatter.type - db_entity.description = markdown.content.description + db_entity.summary = markdown.content.description # Clear and update observations await self.observation_service.delete_by_entity(db_entity.id) @@ -102,7 +102,7 @@ class KnowledgeSyncService: { "name": db_entity.name, "entity_type": db_entity.entity_type, - "description": db_entity.description, + "description": db_entity.summary, # Mark as incomplete "checksum": None, }, diff --git a/tests/repository/test_entity_repository.py b/tests/repository/test_entity_repository.py index e8297c3a..d218399f 100644 --- a/tests/repository/test_entity_repository.py +++ b/tests/repository/test_entity_repository.py @@ -68,7 +68,7 @@ async def test_create_entity(entity_repository: EntityRepository): # Verify returned object assert entity.id is not None assert entity.name == "Test" - assert entity.description == "Test description" + assert entity.summary == "Test description" assert isinstance(entity.created_at, datetime) assert isinstance(entity.updated_at, datetime) @@ -78,7 +78,7 @@ async def test_create_entity(entity_repository: EntityRepository): assert found.id is not None assert found.id == entity.id assert found.name == entity.name - assert found.description == entity.description + assert found.summary == entity.summary # assert relations are eagerly loaded assert len(entity.observations) == 0 @@ -115,7 +115,7 @@ async def test_create_all(entity_repository: EntityRepository): assert found.id is not None assert found.id == entity.id assert found.name == entity.name - assert found.description == entity.description + assert found.summary == entity.summary # assert relations are eagerly loaded assert len(entity.observations) == 0 @@ -139,7 +139,7 @@ async def test_create_entity_null_description(session_maker, entity_repository: stmt = select(Entity).where(Entity.id == entity.id) result = await session.execute(stmt) db_entity = result.scalar_one() - assert db_entity.description is None + assert db_entity.summary is None @pytest.mark.asyncio @@ -157,7 +157,7 @@ async def test_find_by_id(entity_repository: EntityRepository, sample_entity: En db_entity = result.scalar_one() assert db_entity.id == found.id assert db_entity.name == found.name - assert db_entity.description == found.description + assert db_entity.summary == found.summary @pytest.mark.asyncio @@ -167,7 +167,7 @@ async def test_update_entity(entity_repository: EntityRepository, sample_entity: sample_entity.id, {"description": "Updated description"} ) assert updated is not None - assert updated.description == "Updated description" + assert updated.summary == "Updated description" assert updated.name == sample_entity.name # Other fields unchanged # Verify in database @@ -175,7 +175,7 @@ async def test_update_entity(entity_repository: EntityRepository, sample_entity: stmt = select(Entity).where(Entity.id == sample_entity.id) result = await session.execute(stmt) db_entity = result.scalar_one() - assert db_entity.description == "Updated description" + assert db_entity.summary == "Updated description" assert db_entity.name == sample_entity.name @@ -184,14 +184,14 @@ async def test_update_entity_to_null(entity_repository: EntityRepository, sample """Test updating an entity's description to null""" updated = await entity_repository.update(sample_entity.id, {"description": None}) assert updated is not None - assert updated.description is None + assert updated.summary is None # Verify in database async with db.scoped_session(entity_repository.session_maker) as session: stmt = select(Entity).where(Entity.id == sample_entity.id) result = await session.execute(stmt) db_entity = result.scalar_one() - assert db_entity.description is None + assert db_entity.summary is None @pytest.mark.asyncio diff --git a/tests/services/test_entity_service.py b/tests/services/test_entity_service.py index 8d65a150..a46caaad 100644 --- a/tests/services/test_entity_service.py +++ b/tests/services/test_entity_service.py @@ -43,17 +43,17 @@ async def test_create_entity(entity_service: EntityService): assert entity.path_id == entity_data.path_id assert entity.file_path == entity_data.file_path assert entity.entity_type == EntityType.KNOWLEDGE - assert entity.description == "A test entity description" + assert entity.summary == "A test entity description" assert entity.created_at is not None assert entity.observations[0].content == "this is a test observation" assert len(entity.relations) == 0 # Verify we can retrieve it using path_id retrieved = await entity_service.get_by_path_id(entity_data.path_id) - assert retrieved.description == "A test entity description" + assert retrieved.summary == "A test entity description" assert retrieved.name == "TestEntity" assert retrieved.entity_type == EntityType.KNOWLEDGE - assert retrieved.description == "A test entity description" + assert retrieved.summary == "A test entity description" assert retrieved.created_at is not None assert retrieved.observations[0].content == "this is a test observation" @@ -84,7 +84,7 @@ async def test_create_entities(entity_service: EntityService): assert isinstance(entity1, EntityModel) assert entity1.name == "TestEntity1" assert entity1.entity_type == EntityType.KNOWLEDGE - assert entity1.description == "A test entity description" + assert entity1.summary == "A test entity description" assert entity1.created_at is not None assert entity1.observations[0].content == "this is a test observation" assert len(entity1.relations) == 0 @@ -93,16 +93,16 @@ async def test_create_entities(entity_service: EntityService): assert isinstance(entity1, EntityModel) assert entity2.name == "TestEntity2" assert entity2.entity_type == EntityType.KNOWLEDGE - assert entity2.description == "A test entity description" + assert entity2.summary == "A test entity description" assert entity2.created_at is not None assert entity2.observations[0].content == "this is a test observation" # Verify we can retrieve them using path_ids retrieved1 = await entity_service.get_by_path_id(entity_data[0].path_id) - assert retrieved1.description == "A test entity description" + assert retrieved1.summary == "A test entity description" retrieved2 = await entity_service.get_by_path_id(entity_data[1].path_id) - assert retrieved2.description == "A test entity description" + assert retrieved2.summary == "A test entity description" async def test_get_by_path_id(entity_service: EntityService): @@ -128,14 +128,14 @@ async def test_get_by_path_id(entity_service: EntityService): assert found is not None assert found.id == entity1.id assert found.entity_type == entity1.entity_type - assert found.description == "First test entity" + assert found.summary == "First test entity" # Find by type2 and name found = await entity_service.get_by_path_id(entity2_data.path_id) assert found is not None assert found.id == entity2.id assert found.entity_type == entity2.entity_type - assert found.description == "Second test entity" + assert found.summary == "Second test entity" # Test not found case with pytest.raises(EntityNotFoundError): @@ -147,11 +147,11 @@ async def test_create_entity_no_description(entity_service: EntityService): entity_data = EntitySchema(name="TestEntity", entity_type=EntityType.KNOWLEDGE, observations=[]) entity = await entity_service.create_entity(entity_data) - assert entity.description is None + assert entity.summary is None # Verify after retrieval retrieved = await entity_service.get_by_path_id(entity_data.path_id) - assert retrieved.description is None + assert retrieved.summary is None async def test_get_entity_success(entity_service: EntityService): @@ -170,7 +170,7 @@ async def test_get_entity_success(entity_service: EntityService): assert isinstance(retrieved, EntityModel) assert retrieved.name == "TestEntity" assert retrieved.entity_type == EntityType.KNOWLEDGE - assert retrieved.description == "Test description" + assert retrieved.summary == "Test description" async def test_update_entity_description(entity_service: EntityService): @@ -187,11 +187,11 @@ async def test_update_entity_description(entity_service: EntityService): updated = await entity_service.update_entity( entity_data.path_id, {"description": "Updated description"} ) - assert updated.description == "Updated description" + assert updated.summary == "Updated description" # Verify after retrieval retrieved = await entity_service.get_by_path_id(entity_data.path_id) - assert retrieved.description == "Updated description" + assert retrieved.summary == "Updated description" async def test_update_entity_description_to_none(entity_service: EntityService): @@ -206,11 +206,11 @@ async def test_update_entity_description_to_none(entity_service: EntityService): # Update description to None using path_id updated = await entity_service.update_entity(entity_data.path_id, {"description": None}) - assert updated.description is None + assert updated.summary is None # Verify after retrieval retrieved = await entity_service.get_by_path_id(entity_data.path_id) - assert retrieved.description is None + assert retrieved.summary is None async def test_delete_entity_success(entity_service: EntityService): @@ -255,11 +255,11 @@ async def test_create_entity_with_special_chars(entity_service: EntityService): entity = await entity_service.create_entity(entity_data) assert entity.name == name - assert entity.description == description + assert entity.summary == description # Verify after retrieval using path_id retrieved = await entity_service.get_by_path_id(entity_data.path_id) - assert retrieved.description == description + assert retrieved.summary == description async def test_create_entity_long_description(entity_service: EntityService): @@ -273,11 +273,11 @@ async def test_create_entity_long_description(entity_service: EntityService): ) entity = await entity_service.create_entity(entity_data) - assert entity.description == long_description + assert entity.summary == long_description # Verify after retrieval using path_id retrieved = await entity_service.get_by_path_id(entity_data.path_id) - assert retrieved.description == long_description + assert retrieved.summary == long_description async def test_open_nodes_by_path_ids(entity_service: EntityService): diff --git a/tests/services/test_knowledge_service.py b/tests/services/test_knowledge_service.py index 0dcdac86..28dbb46d 100644 --- a/tests/services/test_knowledge_service.py +++ b/tests/services/test_knowledge_service.py @@ -41,7 +41,7 @@ async def test_create_entity(knowledge_service: KnowledgeService): # Verify DB entity assert created.name == entity_schema.name assert created.entity_type == entity_schema.entity_type - assert created.description == entity_schema.description + assert created.summary == entity_schema.description assert created.checksum is not None assert created.path_id == "test_entity" assert created.file_path == "test_entity.md" diff --git a/tests/sync/test_knowledge_sync_service.py b/tests/sync/test_knowledge_sync_service.py index b7b1cd0f..4e993bbf 100644 --- a/tests/sync/test_knowledge_sync_service.py +++ b/tests/sync/test_knowledge_sync_service.py @@ -67,7 +67,7 @@ async def test_create_entity_without_relations( assert entity.name == "Test Entity" assert entity.entity_type == EntityType.KNOWLEDGE assert entity.path_id == "concept/test_entity" - assert entity.description == "A test entity description" + assert entity.summary == "A test entity description" # Check observations assert len(entity.observations) == 2 @@ -101,7 +101,7 @@ async def test_update_entity_without_relations( # Check fields updated assert updated.name == "Updated Title" - assert updated.description == "Updated description" + assert updated.summary == "Updated description" assert len(updated.observations) == 1 assert updated.observations[0].content == "Updated observation"