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]