fix observation service tests, remove EntityId from schemas.py

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=[

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")

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",
-]
+]

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

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)]
+    entity_id: int
+    deletions: Annotated[List[Observation], MinLen(1)]

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):

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]
 
 

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]

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"
+    assert results[0].context == "test_context"

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")