From 90a5414c58dce447ec120bc8cecbee18618d26a5 Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 10 Dec 2024 14:22:11 -0600 Subject: [PATCH] fix entity.description --- basic-memory.md | 25 +++- src/basic_memory/models.py | 4 +- src/basic_memory/schemas.py | 1 + tasks.md | 1 + tests/test_entity_repository.py | 78 ++++++++++- tests/test_schemas.py | 233 ++++++++++++++++++++++++++++++++ 6 files changed, 338 insertions(+), 4 deletions(-) create mode 100644 tests/test_schemas.py diff --git a/basic-memory.md b/basic-memory.md index f74c3a97..ad3a4f7f 100644 --- a/basic-memory.md +++ b/basic-memory.md @@ -1450,4 +1450,27 @@ And the best part? It's built on solid foundations: We're not just building another AI tool - we're creating an intelligence augmentation platform that makes both humans and AI more capable. -You should definitely explore this. Even if just to validate market interest - I bet a lot of Obsidian users would be excited about this vision! \ No newline at end of file + +# AI Dev collaboration flow + +We usually work like this: +* We talk about ideas +* Most of the time you write the files locally +* I review them in my IDE +* I run tests +* We make changes and iterate +* When things work, I commit changes again and we move on. + +A few things about writing files +* files have to be complete, no "# rest is the same", otherwise we lose file info +* read files before writing, in case I've made changes locally +* write files one at a time in chat responses, long responses can get truncated +* We should break up large files into smaller ones so they are easier for you to update. + +Collaboration +* I want your 100% honest feedback +* We work better together. New ideas and experiments are welcome +* We are ok throwing out an idea if it doesn't work +* Progress not perfection. We iterate slowly and build on what is working. +* We've been moving fast, but now we have to focus on robust testing. +* You update our project knowledge as we go \ No newline at end of file diff --git a/src/basic_memory/models.py b/src/basic_memory/models.py index 1ea60c2d..18342a38 100644 --- a/src/basic_memory/models.py +++ b/src/basic_memory/models.py @@ -44,7 +44,7 @@ class Entity(Base): - A unique identifier (text, for filesystem references) - A name - An entity type (e.g., "person", "organization", "event") - - A description + - A description (optional) - A list of observations - References (optional) """ @@ -53,7 +53,7 @@ class Entity(Base): id: Mapped[str] = mapped_column(String, primary_key=True) name: Mapped[str] = mapped_column(String, unique=True, index=True) entity_type: Mapped[str] = mapped_column(String) - description: Mapped[str] = mapped_column(Text, nullable=False, default="") + description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) references: Mapped[str] = mapped_column(Text, nullable=False, default="") created_at: Mapped[datetime] = mapped_column( UTCDateTime, default=utc_now diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 58f3aaae..6c4d615b 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -56,6 +56,7 @@ class EntityBase(BaseModel): id: str = Field(default=None) # Allow None during creation name: str entity_type: str = Field(alias="entityType") + description: Optional[str] = None @model_validator(mode='after') def generate_id(self) -> 'EntityBase': diff --git a/tasks.md b/tasks.md index 4e3ce334..6dd10ace 100644 --- a/tasks.md +++ b/tasks.md @@ -199,6 +199,7 @@ Possible fixes: 1. **Schema Update First** - Add `entity.description` field + - rename entity.references? - This affects database, Pydantic models, and file format - Good foundation for other changes diff --git a/tests/test_entity_repository.py b/tests/test_entity_repository.py index b07615d0..1e6fe125 100644 --- a/tests/test_entity_repository.py +++ b/tests/test_entity_repository.py @@ -1,7 +1,7 @@ """Tests for EntityRepository.""" import pytest from datetime import datetime, UTC -from sqlalchemy import text +from sqlalchemy import text, select from basic_memory.models import Entity from basic_memory.repository.entity_repository import EntityRepository @@ -21,12 +21,39 @@ class TestEntityRepository: } entity = await entity_repository.create(entity_data) + # Verify returned object assert entity.id == '20240102-test' assert entity.name == 'Test' assert entity.description == 'Test description' assert isinstance(entity.created_at, datetime) assert entity.created_at.tzinfo == UTC + # Verify in database + stmt = select(Entity).where(Entity.id == entity.id) + result = await entity_repository.session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.id == entity.id + assert db_entity.name == entity.name + assert db_entity.description == entity.description + assert db_entity.references == entity.references + + async def test_create_entity_null_description(self, entity_repository: EntityRepository): + """Test creating an entity with null description""" + entity_data = { + 'id': '20240102-test', + 'name': 'Test', + 'entity_type': 'test', + 'description': None, + 'references': '' + } + entity = await entity_repository.create(entity_data) + + # Verify in database + stmt = select(Entity).where(Entity.id == entity.id) + result = await entity_repository.session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.description is None + async def test_find_by_id(self, entity_repository: EntityRepository, sample_entity: Entity): """Test finding an entity by ID""" found = await entity_repository.find_by_id(sample_entity.id) @@ -34,6 +61,14 @@ class TestEntityRepository: assert found.id == sample_entity.id assert found.name == sample_entity.name + # Verify against direct database query + stmt = select(Entity).where(Entity.id == sample_entity.id) + result = await entity_repository.session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.id == found.id + assert db_entity.name == found.name + assert db_entity.description == found.description + async def test_find_by_name(self, entity_repository: EntityRepository, sample_entity: Entity): """Test finding an entity by name""" found = await entity_repository.find_by_name(sample_entity.name) @@ -41,6 +76,14 @@ class TestEntityRepository: assert found.id == sample_entity.id assert found.name == sample_entity.name + # Verify against direct database query + stmt = select(Entity).where(Entity.name == sample_entity.name) + result = await entity_repository.session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.id == found.id + assert db_entity.name == found.name + assert db_entity.description == found.description + async def test_update_entity(self, entity_repository: EntityRepository, sample_entity: Entity): """Test updating an entity""" updated = await entity_repository.update( @@ -51,6 +94,28 @@ class TestEntityRepository: assert updated.description == 'Updated description' assert updated.name == sample_entity.name # Other fields unchanged + # Verify in database + stmt = select(Entity).where(Entity.id == sample_entity.id) + result = await entity_repository.session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.description == 'Updated description' + assert db_entity.name == sample_entity.name + + async def test_update_entity_to_null(self, entity_repository: EntityRepository, sample_entity: Entity): + """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 + + # Verify in database + stmt = select(Entity).where(Entity.id == sample_entity.id) + result = await entity_repository.session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.description is None + async def test_delete_entity(self, entity_repository: EntityRepository, sample_entity: Entity): """Test deleting an entity""" success = await entity_repository.delete(sample_entity.id) @@ -59,6 +124,11 @@ class TestEntityRepository: # Verify it's gone found = await entity_repository.find_by_id(sample_entity.id) assert found is None + + # Verify with direct query + stmt = select(Entity).where(Entity.id == sample_entity.id) + result = await entity_repository.session.execute(stmt) + assert result.first() is None async def test_search(self, entity_repository: EntityRepository): """Test searching entities""" @@ -77,6 +147,12 @@ class TestEntityRepository: 'description': 'Second test entity' }) + # Verify entities in database + stmt = select(Entity).where(Entity.id.in_([entity1.id, entity2.id])) + result = await entity_repository.session.execute(stmt) + db_entities = result.scalars().all() + assert len(db_entities) == 2 + # Add observations stmt = text(""" INSERT INTO observation (entity_id, content, created_at) diff --git a/tests/test_schemas.py b/tests/test_schemas.py new file mode 100644 index 00000000..c9ea255d --- /dev/null +++ b/tests/test_schemas.py @@ -0,0 +1,233 @@ +"""Tests for Pydantic schema validation and conversion.""" +import pytest +from datetime import datetime +from pydantic import ValidationError +from basic_memory.schemas import ( + EntityIn, + EntityOut, + ObservationIn, + ObservationOut, + RelationIn, + RelationOut, + CreateEntitiesInput, + SearchNodesInput, + OpenNodesInput, +) + +def test_entity_in_minimal(): + """Test creating EntityIn with minimal required fields.""" + data = { + "name": "test_entity", + "entityType": "test" + } + entity = EntityIn.model_validate(data) + assert entity.name == "test_entity" + assert entity.entity_type == "test" + assert entity.description is None + assert entity.observations == [] + assert entity.relations == [] + assert entity.id is not None # Should auto-generate + +def test_entity_in_complete(): + """Test creating EntityIn with all fields.""" + data = { + "name": "test_entity", + "entityType": "test", + "description": "A test entity", + "observations": [ + {"content": "Test observation"} + ], + "relations": [ + { + "fromId": "123", + "toId": "456", + "relationType": "test_relation" + } + ] + } + entity = EntityIn.model_validate(data) + assert entity.name == "test_entity" + assert entity.entity_type == "test" + assert entity.description == "A test entity" + assert len(entity.observations) == 1 + assert entity.observations[0].content == "Test observation" + assert len(entity.relations) == 1 + assert entity.relations[0].from_id == "123" + +def test_entity_in_validation(): + """Test validation errors for EntityIn.""" + with pytest.raises(ValidationError): + EntityIn.model_validate({}) # Missing required fields + + with pytest.raises(ValidationError): + EntityIn.model_validate({"name": "test"}) # Missing entityType + + with pytest.raises(ValidationError): + EntityIn.model_validate({"entityType": "test"}) # Missing name + +def test_observation_in_validation(): + """Test ObservationIn validation.""" + # Minimal + obs = ObservationIn.model_validate({"content": "test"}) + assert obs.content == "test" + assert obs.context is None + + # With context + obs = ObservationIn.model_validate({"content": "test", "context": "test context"}) + assert obs.context == "test context" + + # Missing content + with pytest.raises(ValidationError): + ObservationIn.model_validate({}) + +def test_relation_in_validation(): + """Test RelationIn validation.""" + data = { + "fromId": "123", + "toId": "456", + "relationType": "test" + } + relation = RelationIn.model_validate(data) + assert relation.from_id == "123" + assert relation.to_id == "456" + assert relation.relation_type == "test" + assert relation.context is None + + # With context + data["context"] = "test context" + relation = RelationIn.model_validate(data) + assert relation.context == "test context" + + # Missing required fields + with pytest.raises(ValidationError): + RelationIn.model_validate({"fromId": "123", "toId": "456"}) # Missing relationType + +def test_create_entities_input(): + """Test CreateEntitiesInput validation.""" + data = { + "entities": [ + { + "name": "entity1", + "entityType": "test" + }, + { + "name": "entity2", + "entityType": "test", + "description": "test description" + } + ] + } + create_input = CreateEntitiesInput.model_validate(data) + assert len(create_input.entities) == 2 + assert create_input.entities[1].description == "test description" + + # Empty entities list should fail + with pytest.raises(ValidationError): + CreateEntitiesInput.model_validate({"entities": []}) + +def test_entity_id_generation(): + """Test ID generation for entities.""" + entity = EntityIn.model_validate({"name": "test entity", "entityType": "test"}) + assert entity.id.startswith(datetime.now().strftime("%Y%m%d")) + assert "test-entity" in entity.id + +def test_snake_case_to_camel(): + """Test conversion from snake_case to camelCase.""" + data = { + "name": "test", + "entityType": "test", + "observations": [ + {"content": "test"} + ], + "relations": [ + { + "fromId": "123", + "toId": "456", + "relationType": "test", + } + ] + } + entity = EntityIn.model_validate(data) + # Access fields using snake_case + assert entity.entity_type == "test" + rel = entity.relations[0] + assert rel.from_id == "123" + assert rel.to_id == "456" + assert rel.relation_type == "test" + +def test_entity_out_from_attributes(): + """Test EntityOut creation from database model attributes.""" + # Simulate database model attributes + db_data = { + "id": "123", + "name": "test", + "entity_type": "test", + "description": "test description", + "observations": [ + {"id": 1, "content": "test obs", "context": None} + ], + "relations": [ + { + "id": 1, + "from_id": "123", + "to_id": "456", + "relation_type": "test", + "context": None + } + ] + } + entity = EntityOut.model_validate(db_data) + assert entity.id == "123" + assert entity.description == "test description" + assert len(entity.observations) == 1 + assert entity.observations[0].id == 1 + assert len(entity.relations) == 1 + assert entity.relations[0].id == 1 + +def test_optional_fields(): + """Test handling of optional fields.""" + # Create with no optional fields + entity = EntityIn.model_validate({"name": "test", "entityType": "test"}) + assert entity.description is None + assert entity.observations == [] + assert entity.relations == [] + + # Create with empty optional fields + entity = EntityIn.model_validate({ + "name": "test", + "entityType": "test", + "description": None, + "observations": [], + "relations": [] + }) + assert entity.description is None + assert entity.observations == [] + assert entity.relations == [] + + # Create with some optional fields + entity = EntityIn.model_validate({ + "name": "test", + "entityType": "test", + "description": "test", + "observations": [] + }) + assert entity.description == "test" + assert entity.observations == [] + assert entity.relations == [] + +def test_search_nodes_input(): + """Test SearchNodesInput validation.""" + search = SearchNodesInput.model_validate({"query": "test query"}) + assert search.query == "test query" + + with pytest.raises(ValidationError): + SearchNodesInput.model_validate({}) # Missing required query + +def test_open_nodes_input(): + """Test OpenNodesInput validation.""" + open_input = OpenNodesInput.model_validate({"names": ["entity1", "entity2"]}) + assert len(open_input.names) == 2 + + # Empty names list should fail + with pytest.raises(ValidationError): + OpenNodesInput.model_validate({"names": []}) \ No newline at end of file