diff --git a/src/basic_memory/api/routers/knowledge.py b/src/basic_memory/api/routers/knowledge.py index 06019793..cc3c3068 100644 --- a/src/basic_memory/api/routers/knowledge.py +++ b/src/basic_memory/api/routers/knowledge.py @@ -60,7 +60,9 @@ async def add_observations( ) -> AddObservationsResponse: """Add observations to an entity.""" logger.debug(f"Adding observations to entity: {data.entity_id}") - observations = await observation_service.add_observations(data.entity_id, data.observations) + observations = await observation_service.add_observations( + data.entity_id, data.observations, data.context + ) return AddObservationsResponse( entity_id=data.entity_id, observations=[ diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py index 76dd4f43..106d5262 100644 --- a/src/basic_memory/models/knowledge.py +++ b/src/basic_memory/models/knowledge.py @@ -89,6 +89,9 @@ class Observation(Base): content: Mapped[str] = mapped_column(Text) context: Mapped[str] = mapped_column(Text, nullable=True) created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) + updated_at: Mapped[datetime] = mapped_column( + DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP") + ) # Relationships entity = relationship("Entity", back_populates="observations") @@ -112,7 +115,11 @@ class Relation(Base): from_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id")) to_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id")) relation_type: Mapped[str] = mapped_column(String) + context: Mapped[str] = mapped_column(Text, nullable=True) created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) + updated_at: Mapped[datetime] = mapped_column( + DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP") + ) # Relationships from_entity = relationship("Entity", foreign_keys=[from_id], back_populates="from_relations") diff --git a/src/basic_memory/schemas/__init__.py b/src/basic_memory/schemas/__init__.py index 90d035b3..9e274752 100644 --- a/src/basic_memory/schemas/__init__.py +++ b/src/basic_memory/schemas/__init__.py @@ -10,11 +10,17 @@ from basic_memory.schemas.base import ( Observation, EntityType, RelationType, - EntityId, Relation, Entity, ) +# Delete operation models +from basic_memory.schemas.delete import ( + DeleteEntityRequest, + DeleteRelationsRequest, + DeleteObservationsRequest, +) + # Request models from basic_memory.schemas.request import ( AddObservationsRequest, @@ -41,30 +47,20 @@ from basic_memory.schemas.response import ( DeleteObservationsResponse, ) -# Delete operation models -from basic_memory.schemas.delete import ( - DeleteEntityRequest, - DeleteRelationsRequest, - DeleteObservationsRequest, -) - # For convenient imports, export all models __all__ = [ # Base "Observation", "EntityType", - "RelationType", - "EntityId", + "RelationType", "Relation", "Entity", - # Requests "AddObservationsRequest", "CreateEntityRequest", "SearchNodesRequest", "OpenNodesRequest", "CreateRelationsRequest", - # Responses "SQLAlchemyModel", "ObservationResponse", @@ -79,9 +75,8 @@ __all__ = [ "DeleteEntityResponse", "DeleteRelationsResponse", "DeleteObservationsResponse", - # Delete Operations "DeleteEntityRequest", "DeleteRelationsRequest", "DeleteObservationsRequest", -] \ No newline at end of file +] diff --git a/src/basic_memory/schemas/base.py b/src/basic_memory/schemas/base.py index 8a1a5288..6a5aecd4 100644 --- a/src/basic_memory/schemas/base.py +++ b/src/basic_memory/schemas/base.py @@ -86,23 +86,6 @@ Guidelines: Common types are listed in the module docstring. """ -# Custom field types with validation -EntityId = Annotated[str, BeforeValidator(lower_strip_whitespace)] -"""Unique identifier for an entity in format '{entity_type}/{normalized_name}'. - -Examples: -- person/alice_smith -- project/basic_memory -- component/memory_service -- concept/semantic_search - -The ID is automatically generated from the entity type and name. -Names are normalized by: -1. Converting to lowercase -2. Replacing spaces with underscores -3. Removing special characters -""" - class Relation(BaseModel): """Represents a directed edge between entities in the knowledge graph. @@ -141,8 +124,8 @@ class Relation(BaseModel): } """ - from_id: EntityId - to_id: EntityId + from_id: int + to_id: int relation_type: RelationType context: Optional[str] = None @@ -207,7 +190,7 @@ class Entity(BaseModel): } """ - id: Optional[EntityId] = None + id: Optional[int] = None name: str entity_type: EntityType description: Optional[str] = None diff --git a/src/basic_memory/schemas/delete.py b/src/basic_memory/schemas/delete.py index 5890fd39..e95f6cd1 100644 --- a/src/basic_memory/schemas/delete.py +++ b/src/basic_memory/schemas/delete.py @@ -21,7 +21,7 @@ from typing import List, Annotated from annotated_types import MinLen from pydantic import BaseModel -from basic_memory.schemas.base import EntityId, Relation, Observation +from basic_memory.schemas.base import Relation, Observation class DeleteEntityRequest(BaseModel): @@ -56,7 +56,7 @@ class DeleteEntityRequest(BaseModel): 5. Create relations to replacement entities if applicable """ - entity_ids: Annotated[List[EntityId], MinLen(1)] + entity_ids: Annotated[List[int], MinLen(1)] class DeleteRelationsRequest(BaseModel): @@ -132,5 +132,5 @@ class DeleteObservationsRequest(BaseModel): 5. Updating implementation details """ - entity_id: EntityId - deletions: Annotated[List[Observation], MinLen(1)] \ No newline at end of file + entity_id: int + deletions: Annotated[List[Observation], MinLen(1)] diff --git a/src/basic_memory/schemas/request.py b/src/basic_memory/schemas/request.py index a46e04d9..362480b3 100644 --- a/src/basic_memory/schemas/request.py +++ b/src/basic_memory/schemas/request.py @@ -5,7 +5,7 @@ from typing import List, Optional, Annotated, Dict, Any from annotated_types import MinLen, MaxLen from pydantic import BaseModel -from basic_memory.schemas.base import EntityId, Observation, Entity, Relation +from basic_memory.schemas.base import Observation, Entity, Relation class AddObservationsRequest(BaseModel): @@ -45,7 +45,7 @@ class AddObservationsRequest(BaseModel): 4. Add observations in logical groups for better history tracking """ - entity_id: EntityId + entity_id: int context: Optional[str] = None observations: List[Observation] @@ -152,7 +152,7 @@ class OpenNodesRequest(BaseModel): relations between entities that interest you. """ - entity_ids: Annotated[List[EntityId], MinLen(1)] + entity_ids: Annotated[List[int], MinLen(1)] class CreateRelationsRequest(BaseModel): diff --git a/src/basic_memory/schemas/response.py b/src/basic_memory/schemas/response.py index 2750170e..affd0548 100644 --- a/src/basic_memory/schemas/response.py +++ b/src/basic_memory/schemas/response.py @@ -16,7 +16,7 @@ from typing import List, Optional, Dict, Any from pydantic import BaseModel, ConfigDict -from basic_memory.schemas.base import Observation, EntityId, Relation +from basic_memory.schemas.base import Observation, Relation class SQLAlchemyModel(BaseModel): @@ -69,7 +69,7 @@ class ObservationsResponse(SQLAlchemyModel): } """ - entity_id: EntityId + entity_id: int observations: List[ObservationResponse] @@ -255,7 +255,7 @@ class AddObservationsResponse(SQLAlchemyModel): } """ - entity_id: EntityId + entity_id: int observations: List[ObservationResponse] diff --git a/src/basic_memory/services/observation_service.py b/src/basic_memory/services/observation_service.py index 6ec8658d..689d10e1 100644 --- a/src/basic_memory/services/observation_service.py +++ b/src/basic_memory/services/observation_service.py @@ -20,21 +20,18 @@ class ObservationService(BaseService[ObservationRepository]): super().__init__(observation_repository) async def add_observations( - self, entity_id: str, observations: List[str] + self, entity_id: int, observations: List[str], context: str | None = None ) -> List[ObservationModel]: """Add multiple observations to an entity.""" logger.debug(f"Adding {len(observations)} observations to entity: {entity_id}") return await self.repository.create_all( [ - dict( - entity_id=entity_id, - content=observation, - ) + dict(entity_id=entity_id, content=observation, context=context) for observation in observations ] ) - async def delete_observations(self, entity_id: str, contents: List[str]) -> bool: + async def delete_observations(self, entity_id: int, contents: List[str]) -> bool: """Delete specific observations from an entity.""" logger.debug(f"Deleting observations from entity: {entity_id}") deleted = False @@ -53,7 +50,9 @@ class ObservationService(BaseService[ObservationRepository]): """Search for observations across all entities.""" logger.debug(f"Searching observations with query: {query}") result = await self.repository.execute_query( - select(ObservationModel).filter(ObservationModel.content.contains(query)) + select(ObservationModel).filter( + ObservationModel.content.contains(query) | ObservationModel.context.contains(query) + ) ) observations = result.scalars().all() return [ObservationModel(content=obs.content) for obs in observations] diff --git a/tests/services/test_observation_service.py b/tests/services/test_observation_service.py index 4f23d64b..fa82aee2 100644 --- a/tests/services/test_observation_service.py +++ b/tests/services/test_observation_service.py @@ -1,4 +1,5 @@ """Tests for the ObservationService.""" + import pytest import pytest_asyncio from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker @@ -9,7 +10,9 @@ from basic_memory.services.observation_service import ObservationService @pytest_asyncio.fixture -async def observation_repository(session_maker: async_sessionmaker[AsyncSession]) -> ObservationRepository: +async def observation_repository( + session_maker: async_sessionmaker[AsyncSession], +) -> ObservationRepository: """Create an ObservationRepository instance.""" return ObservationRepository(session_maker) @@ -24,12 +27,7 @@ async def observation_service(observation_repository: ObservationRepository) -> async def test_entity(session_maker: async_sessionmaker[AsyncSession]) -> Entity: """Create a test entity.""" async with session_maker() as session: - entity = Entity( - id="test/test_entity", - name="test_entity", - entity_type="test", - description="Test entity" - ) + entity = Entity(entity_type="test", name="test", description="Test entity") session.add(entity) await session.commit() return entity @@ -39,9 +37,9 @@ async def test_entity(session_maker: async_sessionmaker[AsyncSession]) -> Entity async def test_add_observations(observation_service: ObservationService, test_entity: Entity): """Test adding observations to an entity.""" observations = ["Test observation 1", "Test observation 2"] - + result = await observation_service.add_observations(test_entity.id, observations) - + assert len(result) == 2 assert result[0].content == "Test observation 1" assert result[1].content == "Test observation 2" @@ -53,8 +51,7 @@ async def test_search_observations(observation_service: ObservationService, test """Test searching observations across entities.""" # First add some observations await observation_service.add_observations( - test_entity.id, - ["Unique test content", "Other content"] + test_entity.id, ["Unique test content", "Other content"] ) # Search for them @@ -65,22 +62,18 @@ async def test_search_observations(observation_service: ObservationService, test @pytest.mark.asyncio -async def test_delete_observations( - observation_service: ObservationService, - test_entity: Entity -): +async def test_delete_observations(observation_service: ObservationService, test_entity: Entity): """Test deleting specific observations from an entity.""" # First add observations await observation_service.add_observations( - test_entity.id, - ["First observation", "Second observation", "Third observation"] + test_entity.id, ["First observation", "Second observation", "Third observation"] ) - + # Then delete some contents_to_delete = ["First observation", "Second observation"] result = await observation_service.delete_observations(test_entity.id, contents_to_delete) assert result is True - + # Verify through search results = await observation_service.search_observations("Third") assert len(results) == 1 @@ -88,21 +81,17 @@ async def test_delete_observations( @pytest.mark.asyncio -async def test_delete_by_entity( - observation_service: ObservationService, - test_entity: Entity -): +async def test_delete_by_entity(observation_service: ObservationService, test_entity: Entity): """Test deleting all observations for an entity.""" # First add observations await observation_service.add_observations( - test_entity.id, - ["First observation", "Second observation"] + test_entity.id, ["First observation", "Second observation"] ) - + # Delete all observations for entity result = await observation_service.delete_by_entity(test_entity.id) assert result is True - + # Verify through search results = await observation_service.search_observations("observation") assert len(results) == 0 @@ -110,13 +99,11 @@ async def test_delete_by_entity( @pytest.mark.asyncio async def test_delete_nonexistent_observation( - observation_service: ObservationService, - test_entity: Entity + observation_service: ObservationService, test_entity: Entity ): """Test deleting observations that don't exist.""" result = await observation_service.delete_observations( - test_entity.id, - ["Nonexistent observation"] + test_entity.id, ["Nonexistent observation"] ) assert result is False @@ -124,17 +111,13 @@ async def test_delete_nonexistent_observation( @pytest.mark.asyncio async def test_delete_observations_invalid_entity(observation_service: ObservationService): """Test deleting observations for an entity that doesn't exist.""" - result = await observation_service.delete_observations( - "invalid_entity", - ["Test observation"] - ) + result = await observation_service.delete_observations("invalid_entity", ["Test observation"]) assert result is False @pytest.mark.asyncio async def test_observation_with_special_characters( - observation_service: ObservationService, - test_entity: Entity + observation_service: ObservationService, test_entity: Entity ): """Test handling observations with special characters.""" content = "Test & observation with @#$% special chars!" @@ -145,10 +128,7 @@ async def test_observation_with_special_characters( @pytest.mark.asyncio -async def test_very_long_observation( - observation_service: ObservationService, - test_entity: Entity -): +async def test_very_long_observation(observation_service: ObservationService, test_entity: Entity): """Test handling very long observation content.""" long_content = "Very long observation " * 100 # ~1800 characters @@ -161,21 +141,19 @@ async def test_very_long_observation( async def test_get_observations_by_context( observation_service: ObservationService, test_entity: Entity, - session_maker: async_sessionmaker[AsyncSession] + session_maker: async_sessionmaker[AsyncSession], ): """Test getting observations by context.""" # Create observation with context async with session_maker() as session: obs = Observation( - entity_id=test_entity.id, - content="Contextual observation", - context="test_context" + entity_id=test_entity.id, content="Contextual observation", context="test_context" ) session.add(obs) await session.commit() - + results = await observation_service.get_observations_by_context("test_context") - + assert len(results) == 1 assert results[0].content == "Contextual observation" - assert results[0].context == "test_context" \ No newline at end of file + assert results[0].context == "test_context" diff --git a/tests/services/test_relation_service.py b/tests/services/test_relation_service.py index bf0a31da..99acd32a 100644 --- a/tests/services/test_relation_service.py +++ b/tests/services/test_relation_service.py @@ -31,13 +31,11 @@ async def test_entities( """Create two test entities.""" async with session_maker() as session: entity1 = EntityModel( - id="test/test_entity_1", name="test_entity_1", entity_type="test", description="Test entity 1", ) entity2 = EntityModel( - id="test/test_entity_2", name="test_entity_2", entity_type="test", description="Test entity 2", @@ -122,7 +120,6 @@ async def test_delete_relation( entity_type=entity1.entity_type, description=entity1.description, observations=[], - relations=[], ) to_entity = Entity( id=entity2.id, @@ -130,7 +127,6 @@ async def test_delete_relation( entity_type=entity2.entity_type, description=entity2.description, observations=[], - relations=[], ) # Delete the relation @@ -152,7 +148,6 @@ async def test_delete_nonexistent_relation( entity_type=entity1.entity_type, description=entity1.description, observations=[], - relations=[], ) to_entity = Entity( id=entity2.id, @@ -160,7 +155,6 @@ async def test_delete_nonexistent_relation( entity_type=entity2.entity_type, description=entity2.description, observations=[], - relations=[], ) result = await relation_service.delete_relation(from_entity, to_entity, "nonexistent_relation")