From 61a8cfcf0ae01895530a2267560905239cfc5271 Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 10 Dec 2024 14:53:51 -0600 Subject: [PATCH] use db defaults for created_at --- src/basic_memory/schemas.py | 45 ++++++----- src/basic_memory/services/entity_service.py | 40 +++++----- tests/test_entity_service.py | 84 ++++++++++++++++++++- 3 files changed, 125 insertions(+), 44 deletions(-) diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 6c4d615b..4665cd08 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -9,6 +9,11 @@ from uuid import uuid4 from annotated_types import Gt, Len from pydantic import BaseModel, Field, model_validator, ConfigDict +# Base output model for SQLAlchemy attribute conversion +class SQLAlchemyOut(BaseModel): + """Base class for models that read from SQLAlchemy attributes.""" + model_config = ConfigDict(from_attributes=True) + # Base Models class ObservationIn(BaseModel): """Schema for creating a single observation.""" @@ -21,16 +26,15 @@ class ObservationsIn(BaseModel): observations: List[ObservationIn] model_config = ConfigDict(populate_by_name=True) -class ObservationOut(ObservationIn): +class ObservationOut(ObservationIn, SQLAlchemyOut): """Schema for observation data returned from the service.""" id: int - model_config = ConfigDict(from_attributes=True) -class ObservationsOut(BaseModel): +class ObservationsOut(SQLAlchemyOut): """Schema for bulk observation operation results.""" entity_id: str = Field(alias="entityId") observations: List[ObservationOut] - model_config = ConfigDict(populate_by_name=True, from_attributes=True) + model_config = ConfigDict(populate_by_name=True) class RelationIn(BaseModel): """ @@ -44,13 +48,13 @@ class RelationIn(BaseModel): model_config = ConfigDict(populate_by_name=True) -class RelationOut(BaseModel): +class RelationOut(SQLAlchemyOut): id: int from_id: str = Field(alias="fromId") to_id: str = Field(alias="toId") relation_type: str = Field(alias="relationType") context: Optional[str] = None - model_config = ConfigDict(from_attributes=True, populate_by_name=True) + model_config = ConfigDict(populate_by_name=True) class EntityBase(BaseModel): id: str = Field(default=None) # Allow None during creation @@ -81,13 +85,13 @@ class EntityIn(EntityBase): """ observations: List[ObservationIn] = [] relations: List[RelationIn] = [] - model_config = ConfigDict(populate_by_name=True, from_attributes=True) + model_config = ConfigDict(populate_by_name=True) -class EntityOut(EntityBase): +class EntityOut(EntityBase, SQLAlchemyOut): """Schema for entity data returned from the service.""" observations: List[ObservationOut] = [] relations: List[RelationOut] = [] - model_config = ConfigDict(populate_by_name=True, from_attributes=True) + model_config = ConfigDict(populate_by_name=True) # Tool Input Schemas class CreateEntitiesInput(BaseModel): @@ -121,40 +125,33 @@ class DeleteObservationsInput(BaseModel): deletions: List[Dict[str, Any]] # TODO: Make this more specific # Tool Response Schemas -class CreateEntitiesResponse(BaseModel): +class CreateEntitiesResponse(SQLAlchemyOut): """Response for create_entities tool.""" entities: List[EntityOut] - model_config = ConfigDict(from_attributes=True) -class SearchNodesResponse(BaseModel): +class SearchNodesResponse(SQLAlchemyOut): """Response for search_nodes tool.""" matches: List[EntityOut] query: str - model_config = ConfigDict(from_attributes=True) -class OpenNodesResponse(BaseModel): +class OpenNodesResponse(SQLAlchemyOut): """Response for open_nodes tool.""" entities: List[EntityOut] - model_config = ConfigDict(from_attributes=True) -class AddObservationsResponse(BaseModel): +class AddObservationsResponse(SQLAlchemyOut): """Response for add_observations tool.""" entity_id: str added_observations: List[ObservationOut] - model_config = ConfigDict(from_attributes=True) -class CreateRelationsResponse(BaseModel): +class CreateRelationsResponse(SQLAlchemyOut): """Response for create_relations tool.""" relations: List[RelationOut] - model_config = ConfigDict(from_attributes=True) -class DeleteEntitiesResponse(BaseModel): +class DeleteEntitiesResponse(SQLAlchemyOut): """Response for delete_entities tool.""" deleted: List[str] - model_config = ConfigDict(from_attributes=True) -class DeleteObservationsResponse(BaseModel): +class DeleteObservationsResponse(SQLAlchemyOut): """Response for delete_observations tool.""" entity_id: str - deleted: List[str] - model_config = ConfigDict(from_attributes=True) \ No newline at end of file + deleted: List[str] \ No newline at end of file diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index 564ba2d1..dbef6258 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -1,11 +1,10 @@ """Service for managing entities in the database.""" -from datetime import datetime, UTC from pathlib import Path -from typing import List +from typing import List, Dict, Any from basic_memory.repository.entity_repository import EntityRepository -from basic_memory.schemas import EntityIn, ObservationIn -from basic_memory.models import Entity, Observation +from basic_memory.schemas import EntityIn +from basic_memory.models import Entity from basic_memory.fileio import EntityNotFoundError from loguru import logger from . import ServiceError @@ -34,22 +33,10 @@ class EntityService: raise async def create_entity(self, entity: EntityIn) -> Entity: - """Create a new entity in the database. - - Note: ID is generated by the EntityIn validator before reaching this method. - """ + """Create a new entity in the database.""" logger.debug(f"Creating entity in DB: {entity.id}") try: - # Create base entity first - base_data = { - "id": entity.id, # Include the generated ID - "name": entity.name, - "entity_type": entity.entity_type, - "created_at": datetime.now(UTC), - } - logger.debug(f"Base entity data: {base_data}") - - created_entity = await self.entity_repo.create(base_data) + created_entity = await self.entity_repo.create(entity.model_dump()) logger.debug(f"Created base entity: {created_entity.id}") await self.entity_repo.refresh(created_entity, ['observations', 'outgoing_relations', 'incoming_relations']) @@ -60,6 +47,22 @@ class EntityService: logger.exception(f"Failed to create entity: {entity.id}") raise + async def update_entity(self, entity_id: str, update_data: Dict[str, Any]) -> Entity: + """Update an entity's fields.""" + logger.debug(f"Updating entity {entity_id} with data: {update_data}") + try: + updated = await self.entity_repo.update(entity_id, update_data) + if not updated: + raise EntityNotFoundError(f"Entity not found: {entity_id}") + + logger.debug(f"Updated entity: {updated.id}") + return updated + except EntityNotFoundError: + raise + except Exception as e: + logger.exception(f"Failed to update entity: {entity_id}") + raise + async def get_entity(self, entity_id: str) -> Entity: """Get entity by ID.""" logger.debug(f"Getting entity by ID: {entity_id}") @@ -77,7 +80,6 @@ class EntityService: logger.exception(f"Failed to get entity: {entity_id}") raise - # TODO name is not unique async def get_by_name(self, name: str) -> Entity: """Get entity by name.""" logger.debug(f"Getting entity by name: {name}") diff --git a/tests/test_entity_service.py b/tests/test_entity_service.py index 1cb5914f..6f381a84 100644 --- a/tests/test_entity_service.py +++ b/tests/test_entity_service.py @@ -13,6 +13,7 @@ async def test_create_entity_success(entity_service): entity_data = EntityIn( name="Test Entity", entity_type="test", + description="A test entity description" ) # Act @@ -22,14 +23,34 @@ async def test_create_entity_success(entity_service): assert isinstance(entity, Entity) assert entity.name == "Test Entity" assert entity.entity_type == "test" + assert entity.description == "A test entity description" assert entity.created_at is not None + # Verify we can retrieve it + retrieved = await entity_service.get_entity(entity.id) + assert retrieved.description == "A test entity description" + +async def test_create_entity_no_description(entity_service): + """Test creating entity without description (should be None).""" + entity_data = EntityIn( + name="Test Entity", + entity_type="test", + ) + + entity = await entity_service.create_entity(entity_data) + assert entity.description is None + + # Verify after retrieval + retrieved = await entity_service.get_entity(entity.id) + assert retrieved.description is None + async def test_get_entity_success(entity_service): """Test successful entity retrieval.""" # Arrange entity_data = EntityIn( name="Test Entity", entity_type="test", + description="Test description" ) created = await entity_service.create_entity(entity_data) @@ -41,8 +62,45 @@ async def test_get_entity_success(entity_service): assert retrieved.id == created.id assert retrieved.name == created.name assert retrieved.entity_type == created.entity_type + assert retrieved.description == "Test description" # relations are tested in test_memory_service +async def test_update_entity_description(entity_service): + """Test updating an entity's description.""" + # Create entity with description + entity_data = EntityIn( + name="Test Entity", + entity_type="test", + description="Initial description" + ) + entity = await entity_service.create_entity(entity_data) + + # Update description + updated = await entity_service.update_entity(entity.id, {"description": "Updated description"}) + assert updated.description == "Updated description" + + # Verify after retrieval + retrieved = await entity_service.get_entity(entity.id) + assert retrieved.description == "Updated description" + +async def test_update_entity_description_to_none(entity_service): + """Test updating an entity's description to None.""" + # Create entity with description + entity_data = EntityIn( + name="Test Entity", + entity_type="test", + description="Initial description" + ) + entity = await entity_service.create_entity(entity_data) + + # Update description to None + updated = await entity_service.update_entity(entity.id, {"description": None}) + assert updated.description is None + + # Verify after retrieval + retrieved = await entity_service.get_entity(entity.id) + assert retrieved.description is None + async def test_delete_entity_success(entity_service): """Test successful entity deletion.""" # Arrange @@ -77,6 +135,7 @@ async def test_create_entity_db_error(entity_service, monkeypatch): entity_data = EntityIn( name="Test Entity", entity_type="test", + description="Test description" ) # Act/Assert @@ -91,22 +150,29 @@ async def test_delete_nonexistent_entity(entity_service): # Edge Cases async def test_create_entity_with_special_chars(entity_service): - """Test entity creation with special characters in name.""" + """Test entity creation with special characters in name and description.""" name = "Test & Entity! With @ Special #Chars" + description = "Description with $pecial chars & symbols!" entity_data = EntityIn( name=name, entity_type="test", + description=description ) entity = await entity_service.create_entity(entity_data) assert entity.name == name + assert entity.description == description + # Verify after retrieval + retrieved = await entity_service.get_entity(entity.id) + assert retrieved.description == description async def test_entity_id_generation(entity_service): """Test that entities get unique IDs generated correctly.""" entity_data = EntityIn( name="Test Entity", entity_type="test", + description="Test description", observations=[] ) @@ -114,3 +180,19 @@ async def test_entity_id_generation(entity_service): assert entity.id # ID should be generated assert "-test-entity" in entity.id # Should contain normalized name + +async def test_create_entity_long_description(entity_service): + """Test creating entity with a long description.""" + long_description = "A" * 1000 # 1000 character description + entity_data = EntityIn( + name="Test Entity", + entity_type="test", + description=long_description + ) + + entity = await entity_service.create_entity(entity_data) + assert entity.description == long_description + + # Verify after retrieval + retrieved = await entity_service.get_entity(entity.id) + assert retrieved.description == long_description \ No newline at end of file