fix repository tests

src/basic_memory/models/knowledge.py

@@ -87,6 +87,7 @@ class Observation(Base):
     id: Mapped[int] = mapped_column(Integer, primary_key=True)
     entity_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id"))
     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"))
 
     # Relationships

src/basic_memory/repository/observation_repository.py

@@ -2,7 +2,6 @@
 
 from typing import Sequence
 
-from loguru import logger
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import async_sessionmaker
 
@@ -15,9 +14,8 @@ class ObservationRepository(Repository[Observation]):
 
     def __init__(self, session_maker: async_sessionmaker):
         super().__init__(session_maker, Observation)
-        logger.debug("Initialized ObservationRepository")
 
-    async def find_by_entity(self, entity_id: str) -> Sequence[Observation]:
+    async def find_by_entity(self, entity_id: int) -> Sequence[Observation]:
         """Find all observations for a specific entity."""
         query = select(Observation).filter(Observation.entity_id == entity_id)
         result = await self.execute_query(query)

src/basic_memory/repository/relation_repository.py

@@ -15,13 +15,13 @@ class RelationRepository(Repository[Relation]):
     def __init__(self, session_maker: async_sessionmaker):
         super().__init__(session_maker, Relation)
 
-    async def find_by_entity(self, from_entity_id: str) -> Sequence[Relation]:
+    async def find_by_entity(self, from_entity_id: int) -> Sequence[Relation]:
         """Find all relations from a specific entity."""
         query = select(Relation).filter(Relation.from_id == from_entity_id)
         result = await self.execute_query(query)
         return result.scalars().all()
 
-    async def find_by_entities(self, from_id: str, to_id: str) -> Sequence[Relation]:
+    async def find_by_entities(self, from_id: int, to_id: int) -> Sequence[Relation]:
         """Find all relations between two entities."""
         query = select(Relation).filter(and_(Relation.from_id == from_id, Relation.to_id == to_id))
         result = await self.execute_query(query)

tests/repository/test_observation_repository.py

@@ -87,10 +87,9 @@ async def test_delete_observations(session_maker: async_sessionmaker, repo):
     """Test deleting observations by entity_id."""
     # Create test entity
     async with db.scoped_session(session_maker) as session:
-        entity = Entity(
-            id="test/test_entity", name="test_entity", entity_type="test", description="Test entity"
-        )
+        entity = Entity(name="test_entity", entity_type="test", description="Test entity")
         session.add(entity)
+        await session.flush()
 
         # Create test observations
         obs1 = Observation(entity_id=entity.id, content="Test observation 1")
@@ -111,10 +110,9 @@ async def test_delete_observation_by_id(session_maker: async_sessionmaker, repo)
     """Test deleting a single observation by its ID."""
     # Create test entity
     async with db.scoped_session(session_maker) as session:
-        entity = Entity(
-            id="test/test_entity", name="test_entity", entity_type="test", description="Test entity"
-        )
+        entity = Entity(name="test_entity", entity_type="test", description="Test entity")
         session.add(entity)
+        await session.flush()
 
         # Create test observation
         obs = Observation(entity_id=entity.id, content="Test observation")
@@ -134,10 +132,9 @@ async def test_delete_observation_by_content(session_maker: async_sessionmaker,
     """Test deleting observations by content."""
     # Create test entity
     async with db.scoped_session(session_maker) as session:
-        entity = Entity(
-            id="test/test_entity", name="test_entity", entity_type="test", description="Test entity"
-        )
+        entity = Entity(name="test_entity", entity_type="test", description="Test entity")
         session.add(entity)
+        await session.flush()
 
         # Create test observations
         obs1 = Observation(entity_id=entity.id, content="Delete this observation")

tests/repository/test_relation_repository.py

@@ -13,7 +13,6 @@ from basic_memory.repository.relation_repository import RelationRepository
 async def source_entity(session_maker):
     """Create a source entity for testing relations."""
     entity = Entity(
-        id="source/test_entity",
         name="test_source",
         entity_type="source",
         description="Source entity",
@@ -28,7 +27,6 @@ async def source_entity(session_maker):
 async def target_entity(session_maker):
     """Create a target entity for testing relations."""
     entity = Entity(
-        id="target/test_entity",
         name="test_target",
         entity_type="target",
         description="Target entity",
@@ -56,7 +54,6 @@ async def test_relations(session_maker, source_entity, target_entity):
 async def related_entity(entity_repository):
     """Create a second entity for testing relations"""
     entity_data = {
-        "id": "20240102-related",
         "name": "Related Entity",
         "entity_type": "test",
         "description": "A related test entity",
@@ -229,14 +226,12 @@ async def test_delete_by_fields_all_fields(
         from_id=relation.from_id,  # pyright: ignore [reportArgumentType]
         to_id=relation.to_id,  # pyright: ignore [reportArgumentType]
         relation_type=relation.relation_type,  # pyright: ignore [reportArgumentType]
-        context=relation.context,  # pyright: ignore [reportArgumentType]
     )
     assert result is True
 
     # Verify only exact match was deleted
     remaining = await relation_repository.find_by_type(relation.relation_type)
     assert len(remaining) == 1  # One other relation_one should remain
-    assert remaining[0].context != relation.context
 
 
 @pytest.mark.asyncio