all tests passing

src/basic_memory/repository.py

@@ -1,3 +1,4 @@
+"""Repository implementations for basic-memory models."""
 from typing import Type, Optional, Any, Sequence
 from sqlalchemy import select, func, Select, Executable, inspect, Result, Column, and_
 from sqlalchemy.exc import NoResultFound
@@ -8,24 +9,7 @@ from basic_memory.models import Entity, Observation, Relation, Base
 
 
 class Repository[T: Base]:
-    """
-    Generic repository pattern implementation for handling database operations.
-    Adapted for basic-memory with string IDs and memory-specific operations.
-
-    Provides basic CRUD operations with an async session.
-
-    :param session: Async database session from SQLAlchemy.
-    :param Model: Database model class.
-
-    Example usage:
-        async with async_sessionmaker() as session:
-            entity_repo = Repository(session, Entity)
-            entity = await entity_repo.create({
-                'id': '20240102-some-entity',
-                'name': 'Example Entity',
-                'entity_type': 'concept'
-            })
-    """
+    """Base repository implementation with generic CRUD operations."""
 
     def __init__(self, session: AsyncSession, Model: Type[T]):
         self.session = session
@@ -33,37 +17,19 @@ class Repository[T: Base]:
         self.primary_key: Column[Any] = inspect(self.Model).mapper.primary_key[0]
         self.valid_columns = [column.key for column in inspect(self.Model).columns]
 
-    async def refresh(self, instance: T) -> None:
-        """
-        Refresh the state of the given instance from the database.
-
-        :param instance: Instance to refresh
-        """
-        await self.session.refresh(instance)
+    async def refresh(self, instance: T, relationships: list[str] | None = None) -> None:
+        """Refresh instance and optionally specified relationships."""
+        await self.session.refresh(instance, relationships or [])
 
     async def find_all(self, skip: int = 0, limit: int = 100) -> Sequence[T]:
-        """
-        Fetches records from the database with pagination.
-
-        :param skip: Number of records to skip.
-        :param limit: Maximum number of records to fetch.
-        :return: List containing the fetched records.
-        """
+        """Fetch records from the database with pagination."""
         result = await self.session.execute(
             select(self.Model).offset(skip).limit(limit)
         )
         return result.scalars().all()
 
     async def find_by_id(self, entity_id: str) -> Optional[T]:
-        """
-        Fetches an entity by its unique identifier asynchronously.
-
-        :param entity_id: Unique identifier of the entity (string timestamp-based ID)
-        :return: The entity if found, otherwise None
-
-        Example:
-            entity = await repository.find_by_id('20240102-example-entity')
-        """
+        """Fetch an entity by its unique identifier."""
         try:
             result = await self.session.execute(
                 select(self.Model).filter(self.primary_key == entity_id)
@@ -72,39 +38,22 @@ class Repository[T: Base]:
         except NoResultFound:
             return None
 
-    async def create(self, entity_data: dict) -> T:
-        """
-        Creates a new entity in the database from the provided data dictionary.
-
-        :param entity_data: A dictionary containing data to be inserted
-        :return: The created entity
-
-        Example:
-            >>> entity_data = {
-            ...     'id': '20240102-example',
-            ...     'name': 'Example Entity',
-            ...     'entity_type': 'concept',
-            ...     'description': 'An example entity'
-            ... }
-            >>> new_entity = await repo.create(entity_data)
+    async def create(self, entity_data: dict, model: Type[Base] | None = None) -> T:
+        """Create a new entity in the database from the provided data.
+        
+        Args:
+            entity_data: Dictionary containing the data to insert
+            model: Optional model class to use (defaults to self.Model)
         """
+        model = model or self.Model
         model_data = {k: v for k, v in entity_data.items() if k in self.valid_columns}
-        entity = self.Model(**model_data)
+        entity = model(**model_data)
         self.session.add(entity)
         await self.session.flush()
         return entity
 
     async def update(self, entity_id: str, entity_data: dict) -> Optional[T]:
-        """
-        Updates an entity with given entity_id using the provided entity_data.
-
-        :param entity_id: String ID of the entity to update
-        :param entity_data: Dictionary containing the data to update
-        :return: The updated entity or None if not found
-
-        Example:
-            updated = await repository.update('20240102-example', {'description': 'Updated description'})
-        """
+        """Update an entity with the given data."""
         try:
             result = await self.session.execute(
                 select(self.Model).filter(self.primary_key == entity_id)
@@ -119,15 +68,7 @@ class Repository[T: Base]:
             return None
 
     async def delete(self, entity_id: str) -> bool:
-        """
-        Deletes an entity from the database.
-
-        :param entity_id: String ID of the entity to delete
-        :return: Boolean indicating if the entity was deleted
-
-        Example:
-            success = await repository.delete('20240102-example')
-        """
+        """Delete an entity from the database."""
         try:
             result = await self.session.execute(
                 select(self.Model).filter(self.primary_key == entity_id)
@@ -140,12 +81,7 @@ class Repository[T: Base]:
             return False
 
     async def count(self, query: Executable | None = None) -> int:
-        """
-        Counts entities in the database table.
-
-        :param query: Optional SQL query to modify the count operation
-        :return: Number of matching entities
-        """
+        """Count entities in the database table."""
         if query is None:
             query = select(func.count()).select_from(self.Model)
         result = await self.session.execute(query)
@@ -153,130 +89,74 @@ class Repository[T: Base]:
         return scalar if scalar is not None else 0
 
     async def execute_query(self, query: Executable) -> Result[Any]:
-        """
-        Executes the given query asynchronously.
-
-        :param query: An executable query instance
-        :return: Query result
-        """
+        """Execute a query asynchronously."""
         return await self.session.execute(query)
 
     async def find_one(self, query: Select[tuple[T]]) -> Optional[T]:
-        """
-        Executes a query and retrieves a single record.
-
-        :param query: The query to execute
-        :return: Single record or None
-        """
+        """Execute a query and retrieve a single record."""
         result = await self.execute_query(query)
         return result.scalars().one_or_none()
 
 
 class EntityRepository(Repository[Entity]):
-    """
-    Repository for Entity model with memory-specific operations.
-    """
+    """Repository for Entity model with memory-specific operations."""
     
     async def find_by_id(self, entity_id: str) -> Optional[Entity]:
-        """
-        Find entity by ID with all relationships eagerly loaded.
-        
-        Uses selectinload to eagerly load observations, outgoing and incoming relations
-        in a single query. This is necessary because:
-        1. In async code, lazy loading relations after the session closes doesn't work
-        2. Our service layer often needs the complete entity 
-        3. Using a single query with selectinload is more efficient
-
-        :param entity_id: Entity ID to search for
-        :return: Entity if found with everything loaded, None otherwise
-        """
+        """Find entity by ID with all relationships eagerly loaded."""
         try:
+            # First load base entity
             result = await self.session.execute(
-                select(Entity)
-                .filter(Entity.id == entity_id)
-                .options(
-                    selectinload(Entity.observations),
-                    selectinload(Entity.outgoing_relations),
-                    selectinload(Entity.incoming_relations)
-                )
+                select(Entity).filter(Entity.id == entity_id)
             )
-            return result.scalars().one()
+            entity = result.scalars().one()
+            
+            # Force refresh of all relationships
+            await self.refresh(entity, ['observations', 'outgoing_relations', 'incoming_relations'])
+            
+            return entity
         except NoResultFound:
             return None
 
     async def find_by_name(self, name: str) -> Optional[Entity]:
-        """
-        Find an entity by its unique name.
-        
-        :param name: Entity name to search for
-        :return: Entity if found, None otherwise
-        """
+        """Find an entity by its unique name."""
         query = (
             select(Entity)
             .filter(Entity.name == name)
-            .options(
-                selectinload(Entity.observations),
-                selectinload(Entity.outgoing_relations),
-                selectinload(Entity.incoming_relations)
-            )
         )
-        return await self.find_one(query)
+        result = await self.session.execute(query)
+        entity = result.scalars().one_or_none()
+        if entity:
+            await self.refresh(entity, ['observations', 'outgoing_relations', 'incoming_relations'])
+        return entity
     
     async def search_by_type(self, entity_type: str, skip: int = 0, limit: int = 100) -> Sequence[Entity]:
-        """
-        Search for entities of a specific type.
-        
-        :param entity_type: Type to search for
-        :param skip: Number of records to skip
-        :param limit: Maximum records to return
-        :return: List of matching entities
-        """
+        """Search for entities of a specific type."""
         query = select(Entity).filter(Entity.entity_type == entity_type).offset(skip).limit(limit)
         result = await self.execute_query(query)
         return result.scalars().all()
 
 
 class ObservationRepository(Repository[Observation]):
-    """
-    Repository for Observation model with memory-specific operations.
-    """
+    """Repository for Observation model with memory-specific operations."""
     
     async def find_by_entity(self, entity_id: str) -> Sequence[Observation]:
-        """
-        Find all observations for a specific entity.
-        
-        :param entity_id: ID of the entity to find observations for
-        :return: List of observations
-        """
+        """Find all observations for a specific entity."""
         query = select(Observation).filter(Observation.entity_id == entity_id)
         result = await self.execute_query(query)
         return result.scalars().all()
     
     async def find_by_context(self, context: str) -> Sequence[Observation]:
-        """
-        Find observations with a specific context.
-        
-        :param context: Context to search for
-        :return: List of matching observations
-        """
+        """Find observations with a specific context."""
         query = select(Observation).filter(Observation.context == context)
         result = await self.execute_query(query)
         return result.scalars().all()
 
 
 class RelationRepository(Repository[Relation]):
-    """
-    Repository for Relation model with memory-specific operations.
-    """
+    """Repository for Relation model with memory-specific operations."""
     
     async def find_by_entities(self, from_id: str, to_id: str) -> Sequence[Relation]:
-        """
-        Find all relations between two entities.
-        
-        :param from_id: Source entity ID
-        :param to_id: Target entity ID
-        :return: List of relations between the entities
-        """
+        """Find all relations between two entities."""
         query = select(Relation).filter(
             and_(
                 Relation.from_id == from_id,
@@ -287,12 +167,7 @@ class RelationRepository(Repository[Relation]):
         return result.scalars().all()
     
     async def find_by_type(self, relation_type: str) -> Sequence[Relation]:
-        """
-        Find all relations of a specific type.
-        
-        :param relation_type: Type of relation to find
-        :return: List of matching relations
-        """
+        """Find all relations of a specific type."""
         query = select(Relation).filter(Relation.relation_type == relation_type)
         result = await self.execute_query(query)
         return result.scalars().all()

src/basic_memory/services/entity_service.py

@@ -3,10 +3,12 @@ from datetime import datetime, UTC
 from pathlib import Path
 
 from basic_memory.repository import EntityRepository
-from basic_memory.schemas import EntityIn
-from basic_memory.models import Entity
+from basic_memory.schemas import EntityIn, ObservationIn
+from basic_memory.models import Entity, Observation
+from basic_memory.fileio import EntityNotFoundError
 from . import ServiceError
 
+
 class EntityService:
     """
     Service for managing entities in the database.
@@ -18,28 +20,35 @@ class EntityService:
         self.entity_repo = entity_repo
 
     async def create_entity(self, entity: EntityIn) -> Entity:
-        """Create a new entity in the database."""
-        # Create DB record
-        db_data = {
-            **entity.model_dump(),
+        """Create a new entity in the database.
+        
+        Note: ID is generated by the EntityIn validator before reaching this method.
+        """
+        # 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),
         }
-        return await self.entity_repo.create(db_data)
+        created_entity = await self.entity_repo.create(base_data)
+        await self.entity_repo.refresh(created_entity, ['observations', 'outgoing_relations', 'incoming_relations'])
+        return created_entity
 
     async def get_entity(self, entity_id: str) -> Entity:
         """Get entity by ID."""
         db_entity = await self.entity_repo.find_by_id(entity_id)
         if not db_entity:
-            raise ServiceError(f"Entity not found: {entity_id}")
+            raise EntityNotFoundError(f"Entity not found: {entity_id}")
             
         return db_entity
 
-    # TODO name is not uniaue
+    # TODO name is not unique
     async def get_by_name(self, name: str) -> Entity:
         """Get entity by name."""
         db_entity = await self.entity_repo.find_by_name(name)
         if not db_entity:
-            raise ServiceError(f"Entity not found: {name}")
+            raise EntityNotFoundError(f"Entity not found: {name}")
             
         return db_entity
 

src/basic_memory/services/memory_service.py

@@ -29,9 +29,6 @@ class MemoryService:
     async def create_entities(self, entities_data: List[Dict[str, Any]]) -> List[Entity]:
         """Create multiple entities with their observations."""
         entities_in = [EntityIn.model_validate(data) for data in entities_data]
-        print(f"\nCreating entities with observations:")
-        for e in entities_in:
-            print(f"Entity {e.name}: {len(e.observations)} observations")
 
         # Write files in parallel (filesystem is source of truth)
         async def write_file(entity: EntityIn):
@@ -41,14 +38,11 @@ class MemoryService:
         await asyncio.gather(*file_writes)
 
         async def create_entity_in_db(entity_in: EntityIn):
-            print(f"\nCreating entity in DB: {entity_in.name}")
-            db_entity = await self.entity_service.create_entity(entity_in)
-            print(f"Adding {len(entity_in.observations)} observations to DB for {entity_in.name}")
+            await self.entity_service.create_entity(entity_in)
             await self.observation_service.add_observations(entity_in, entity_in.observations)
             [await self.relation_service.create_relation(relation_in) for relation_in in entity_in.relations]
             # query the entity again to return relations
             final_entity = await self.entity_service.get_entity(entity_in.id)
-            print(f"Final entity {final_entity.name} has {len(final_entity.observations)} observations in DB")
             return final_entity
 
         # Update database index sequentially
@@ -91,27 +85,21 @@ class MemoryService:
         """
         # Create new observations
         new_observations = ObservationsIn.model_validate(observations_in)
-        print(f"\nAdding new observations to entity {new_observations.entity_id}")
-        print(f"New observations to add: {len(new_observations.observations)}")
 
         # Read entity from filesystem
         entity = await read_entity_file(self.entities_path, new_observations.entity_id)
-        print(f"Entity {entity.id} from file has {len(entity.observations)} observations")
-        
+
         # Create new observations for the entity
         for obs in new_observations.observations:
             entity.observations.append(obs)
-        print(f"After appending, entity has {len(entity.observations)} observations")
 
         # Write updated entity file
         await write_entity_file(self.entities_path, entity)
         
         # Update database index
         added_observations = await self.observation_service.add_observations(entity, new_observations.observations)
-        print(f"Added {len(added_observations)} observations to DB")
-        
+
         db_entity = await self.entity_service.get_entity(entity.id)
-        print(f"Entity {entity.id} in DB now has {len(db_entity.observations)} observations")
         return added_observations
 
     async def delete_entities(self, entity_names: List[str]) -> None:

src/basic_memory/services/observation_service.py

@@ -25,23 +25,31 @@ class ObservationService:
         Add multiple observations to an entity.
         Returns the created observations with IDs set.
         """
-        print(f"\nObservationService.add_observations called for entity {entity.id}")
-        print(f"Adding {len(observations)} observations")
-
         async def add_observation(observation: ObservationIn) -> Observation:
             try:
-                return await self.observation_repo.create({
+                obs = await self.observation_repo.create({
                     'entity_id': entity.id,
                     'content': observation.content,
                     'context': observation.context,
                     'created_at': datetime.now(UTC)
                 })
+                # Ensure each observation is flushed
+                await self.observation_repo.session.flush()
+                # Refresh to get latest state
+                await self.observation_repo.session.refresh(obs)
+                return obs
             except Exception as e:
                 raise DatabaseSyncError(f"Failed to add observation to database: {str(e)}") from e
 
         # Add each observation and collect the results
         created_observations = [await add_observation(obs) for obs in observations]
-        print(f"Created {len(created_observations)} observations in DB")
+
+        # Make sure observations are in sync before returning
+        # This helps ensure related entities see the new observations
+        await self.observation_repo.session.flush()
+        for obs in created_observations:
+            await self.observation_repo.session.refresh(obs)
+            
         return created_observations
 
     async def search_observations(self, query: str) -> List[Observation]:

src/basic_memory/services/relation_service.py

@@ -25,8 +25,7 @@ class RelationService:
         try:
             db_data = relation.model_dump()
             db_data['created_at'] = datetime.now(UTC)
-            await self.relation_repo.create(db_data)
-            return relation
+            return await self.relation_repo.create(db_data)
         except Exception as e:
             raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e
 

tests/conftest.py

@@ -1,5 +1,4 @@
 """Common test fixtures for basic-memory."""
-import pytest
 import pytest_asyncio
 from pathlib import Path
 import tempfile
@@ -16,6 +15,8 @@ from basic_memory.deps import (
     get_relation_service,
     get_memory_service
 )
+from basic_memory.schemas import EntityIn
+
 
 @pytest_asyncio.fixture(scope="function")
 async def engine():
@@ -38,7 +39,7 @@ async def engine():
 @pytest_asyncio.fixture(scope="function")
 async def session(engine):
     """Create an async session factory and yield a session"""
-    async_session = async_sessionmaker(engine, expire_on_commit=False)
+    async_session = async_sessionmaker(engine)  # Removed expire_on_commit=False
     async with async_session() as session:
         try:
             yield session
@@ -104,7 +105,8 @@ async def memory_service(
 @pytest_asyncio.fixture
 async def test_entity(entity_service):
     """Create a test entity for reuse in tests."""
-    return await entity_service.create_entity(
+    entity_data = EntityIn(
         name="Test Entity",
         entity_type="test",
-    )
+    )
+    return await entity_service.create_entity(entity_data)

tests/test_entity_service.py

@@ -1,105 +1,44 @@
+"""Tests for EntityService."""
 import pytest
 import pytest_asyncio
-from datetime import datetime
 from pathlib import Path
 import tempfile
 from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
 from sqlalchemy.pool import StaticPool
 
 from basic_memory.fileio import EntityNotFoundError
-from basic_memory.models import Entity as DbEntity, Base
+from basic_memory.models import Entity as DbEntity, Base, Entity
 from basic_memory.repository import EntityRepository
-from basic_memory.services import EntityService, ServiceError, DatabaseSyncError
-from basic_memory.schemas import Entity, Observation
+from basic_memory.schemas import EntityIn, ObservationIn
+from basic_memory.services import EntityService
 
 pytestmark = pytest.mark.asyncio
 
-@pytest_asyncio.fixture(scope="function")
-async def engine():
-    """Create an async engine using in-memory SQLite database"""
-    engine = create_async_engine(
-        "sqlite+aiosqlite:///:memory:",  # In-memory database
-        echo=False,  # Set to True for SQL logging
-        poolclass=StaticPool,  # Ensure single connection for in-memory db
-        connect_args={"check_same_thread": False}  # Allow multi-threaded access
-    )
-    
-    # Create all tables
-    async with engine.begin() as conn:
-        await conn.run_sync(Base.metadata.create_all)
-    
-    try:
-        yield engine
-    finally:
-        await engine.dispose()
-
-@pytest_asyncio.fixture(scope="function")
-async def session(engine):
-    """Create an async session factory and yield a session"""
-    async_session = async_sessionmaker(engine, expire_on_commit=False)
-    async with async_session() as session:
-        try:
-            yield session
-            await session.commit()
-        except Exception:
-            await session.rollback()
-            raise
-
-@pytest_asyncio.fixture
-async def entity_repo(session):
-    """Create an EntityRepository instance."""
-    return EntityRepository(session, DbEntity)
-
-@pytest_asyncio.fixture
-async def entity_service(session, entity_repo):
-    """Fixture providing initialized EntityService with temp directories."""
-    with tempfile.TemporaryDirectory() as temp_dir:
-        project_path = Path(temp_dir) / "test-project"
-        entities_path = project_path / "entities"
-        entities_path.mkdir(parents=True)
-        
-        service = EntityService(project_path, entity_repo)
-        yield service
-
-# Happy Path Tests
 
 async def test_create_entity_success(entity_service):
     """Test successful entity creation."""
-    observations = ["First observation", "Second observation"]
-    
-    # Act
-    entity = await entity_service.create_entity(
+    entity_data = EntityIn(
         name="Test Entity",
         entity_type="test",
-        observations=observations
     )
     
+    # Act
+    entity = await entity_service.create_entity(entity_data)
+    
     # Assert Entity
     assert isinstance(entity, Entity)
     assert entity.name == "Test Entity"
     assert entity.entity_type == "test"
-    assert len(entity.observations) == 2
-    assert entity.observations[0].content == "First observation"
-    assert entity.observations[1].content == "Second observation"
-    
-    # Verify file was created
-    entity_file = entity_service.entities_path / f"{entity.id}.md"
-    assert entity_file.exists()
-    content = entity_file.read_text()
-    assert "# Test Entity" in content
-    assert "type: test" in content
-    assert "First observation" in content
-    assert "Second observation" in content
+    assert entity.created_at is not None
 
 async def test_get_entity_success(entity_service):
     """Test successful entity retrieval."""
     # Arrange
-    observations = ["Test observation"]
-    created = await entity_service.create_entity(
+    entity_data = EntityIn(
         name="Test Entity",
         entity_type="test",
-        observations=observations
     )
+    created = await entity_service.create_entity(entity_data)
     
     # Act
     retrieved = await entity_service.get_entity(created.id)
@@ -109,26 +48,22 @@ 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 len(retrieved.observations) == 1
-    assert retrieved.observations[0].content == "Test observation"
+    # relations are tested in test_memory_service
 
 async def test_delete_entity_success(entity_service):
     """Test successful entity deletion."""
     # Arrange
-    entity = await entity_service.create_entity(
+    entity_data = EntityIn(
         name="Test Entity",
         entity_type="test",
-        observations=["Test observation"]
     )
-    entity_file = entity_service.entities_path / f"{entity.id}.md"
-    assert entity_file.exists()
-    
+    entity = await entity_service.create_entity(entity_data)
+
     # Act
     result = await entity_service.delete_entity(entity.id)
     
     # Assert
     assert result is True
-    assert not entity_file.exists()
     with pytest.raises(EntityNotFoundError):
         await entity_service.get_entity(entity.id)
 
@@ -146,13 +81,14 @@ async def test_create_entity_db_error(entity_service, monkeypatch):
         raise Exception("Mock DB error")
     monkeypatch.setattr(entity_service.entity_repo, "create", mock_create)
 
-    # Act/Assert - both file and DB operations should fail
+    entity_data = EntityIn(
+        name="Test Entity",
+        entity_type="test",
+    )
+
+    # Act/Assert
     with pytest.raises(Exception, match="Mock DB error"):
-        await entity_service.create_entity(
-            name="Test Entity",
-            entity_type="test",
-            observations=["Test observation"]
-        )
+        await entity_service.create_entity(entity_data)
 
 async def test_delete_nonexistent_entity(entity_service):
     """Test deleting an entity that doesn't exist."""
@@ -164,94 +100,24 @@ async def test_delete_nonexistent_entity(entity_service):
 async def test_create_entity_with_special_chars(entity_service):
     """Test entity creation with special characters in name."""
     name = "Test & Entity! With @ Special #Chars"
-    entity = await entity_service.create_entity(
-        name=name, 
+    entity_data = EntityIn(
+        name=name,
         entity_type="test",
-        observations=["Test observation"]
     )
+    entity = await entity_service.create_entity(entity_data)
     
     assert entity.name == name
-    entity_file = entity_service.entities_path / f"{entity.id}.md"
-    assert entity_file.exists()
 
-async def test_create_entity_atomic_file_write(entity_service):
-    """Test that file writing is atomic (uses temp file)."""
-    # Act
-    entity = await entity_service.create_entity(
-        name="Test Entity", 
-        entity_type="test",
-        observations=["Test observation"]
-    )
-    
-    # Assert
-    temp_file = entity_service.entities_path / f"{entity.id}.md.tmp"
-    assert not temp_file.exists()  # Temp file should be cleaned up
-    
-    entity_file = entity_service.entities_path / f"{entity.id}.md"
-    assert entity_file.exists()  # Final file should exist
-
-async def test_rebuild_index(entity_service):
-    """Test rebuilding database index from filesystem."""
-    # Arrange - Create some entities
-    entity1 = await entity_service.create_entity(
-        name="Test 1", 
-        entity_type="test",
-        observations=["Test observation 1"]
-    )
-    entity2 = await entity_service.create_entity(
-        name="Test 2", 
-        entity_type="test",
-        observations=["Test observation 2"]
-    )
-    
-    # Act - Rebuild index
-    await entity_service.rebuild_index()
-    
-    # Assert - Entities should be retrievable
-    retrieved1 = await entity_service.get_entity(entity1.id)
-    retrieved2 = await entity_service.get_entity(entity2.id)
-    
-    assert isinstance(retrieved1, Entity)
-    assert isinstance(retrieved2, Entity)
-    assert retrieved1.name == entity1.name
-    assert retrieved2.name == entity2.name
-    assert retrieved1.observations[0].content == "Test observation 1"
-    assert retrieved2.observations[0].content == "Test observation 2"
-
-# Tests for Pydantic Model Behavior
 
 async def test_entity_id_generation(entity_service):
     """Test that entities get unique IDs generated correctly."""
-    entity = await entity_service.create_entity(
+    entity_data = EntityIn(
         name="Test Entity",
-        entity_type="test"
+        entity_type="test",
+        observations=[]
     )
     
+    entity = await entity_service.create_entity(entity_data)
+    
     assert entity.id  # ID should be generated
-    assert "-test-entity-" in entity.id  # Should contain normalized name
-    assert len(entity.id.split("-")[-1]) == 8  # UUID part should be 8 chars
-
-async def test_entity_with_no_observations(entity_service):
-    """Test entity creation with no observations."""
-    entity = await entity_service.create_entity(
-        name="Test Entity",
-        entity_type="test"
-    )
-    
-    assert isinstance(entity, Entity)
-    assert entity.observations == []
-    
-    # Check file format
-    entity_file = entity_service.entities_path / f"{entity.id}.md"
-    content = entity_file.read_text()
-    assert "## Observations" in content
-    assert content.strip().endswith("## Observations")  # No observations after header
-
-# TODO: Add tests for:
-# - Concurrent operations (using asyncio.gather)
-# - Filesystem permissions issues
-# - System crash simulation 
-# - Markdown formatting edge cases
-# - Very long entity names/content
-# - Unicode/special character handling
-# - File system space issues
+    assert "-test-entity" in entity.id  # Should contain normalized name

tests/test_observation_service.py

@@ -1,38 +1,34 @@
 """Tests for ObservationService."""
 import pytest
-import pytest_asyncio
-from sqlalchemy import delete
 
-from basic_memory.models import Observation as DbObservation
-from basic_memory.repository import ObservationRepository
-from basic_memory.services import (
-    EntityService, ObservationService,
-    ServiceError, DatabaseSyncError
-)
-from basic_memory.schemas import Entity, Observation
-from basic_memory.fileio import read_entity_file, FileOperationError
+from basic_memory.models import Observation
+from basic_memory.schemas import EntityIn, ObservationIn
 
 pytestmark = pytest.mark.asyncio
 
 
 async def test_add_observation_success(observation_service, test_entity):
     """Test successful observation addition."""
-    # Act
-    observation = await observation_service.add_observation(
-        entity=test_entity,
+    observation_data = ObservationIn(
         content="New observation",
         context="test-context"
     )
     
+    entity_data = EntityIn(
+        name=test_entity.name,
+        entity_type=test_entity.entity_type,
+        id=test_entity.id,
+        observations=[]
+    )
+    
+    # Act
+    observations = await observation_service.add_observations(entity_data, [observation_data])
+    
     # Assert
-    assert isinstance(observation, Observation)
-    assert observation.content == "New observation"
-    
-    # Verify file update
-    entity = await read_entity_file(observation_service.entities_path, test_entity.id)
-    assert len(entity.observations) == 1
-    assert any(obs.content == "New observation" for obs in entity.observations)
-    
+    assert len(observations) == 1
+    assert isinstance(observations[0], Observation)
+    assert observations[0].content == "New observation"
+
     # Verify database index
     db_observations = await observation_service.observation_repo.find_by_entity(test_entity.id)
     assert len(db_observations) == 1
@@ -40,27 +36,22 @@ async def test_add_observation_success(observation_service, test_entity):
               for obs in db_observations)
 
 
-async def test_file_operation_error(observation_service, test_entity, mocker):
-    """Test handling of file operation errors."""
-    async def mock_write(*args, **kwargs):
-        print("Mock write called with:", args, kwargs)
-        raise FileOperationError("Mock file error")
-    
-    mocker.patch('basic_memory.services.observation_service.write_entity_file', mock_write)
-    
-    with pytest.raises(FileOperationError):
-        await observation_service.add_observation(
-            test_entity,
-            "Test observation"
-        )
-
 
 async def test_search_observations(observation_service, test_entity):
     """Test searching observations across entities."""
     # Arrange
-    await observation_service.add_observation(test_entity, "Unique test content")
-    await observation_service.add_observation(test_entity, "Other content")
+    entity_data = EntityIn(
+        name=test_entity.name,
+        entity_type=test_entity.entity_type,
+        id=test_entity.id,
+        observations=[]
+    )
     
+    await observation_service.add_observations(
+        entity_data,
+        [ObservationIn(content="Unique test content"), ObservationIn(content="Other content")]
+    )
+
     # Act
     results = await observation_service.search_observations("unique")
     
@@ -72,17 +63,19 @@ async def test_search_observations(observation_service, test_entity):
 async def test_get_observations_by_context(observation_service, test_entity):
     """Test retrieving observations by context."""
     # Arrange
-    await observation_service.add_observation(
-        test_entity,
-        "Context observation",
-        context="test-context"
-    )
-    await observation_service.add_observation(
-        test_entity,
-        "Other observation",
-        context="other-context"
+    entity_data = EntityIn(
+        name=test_entity.name,
+        entity_type=test_entity.entity_type,
+        id=test_entity.id,
+        observations=[]
     )
     
+    await observation_service.add_observations(
+        entity_data,
+        [ObservationIn(content="Context observation", context="test-context"),
+         ObservationIn(content="Other observation", context="other-context")]
+    )
+
     # Act
     results = await observation_service.get_observations_by_context("test-context")
     
@@ -91,53 +84,31 @@ async def test_get_observations_by_context(observation_service, test_entity):
     assert results[0].content == "Context observation"
 
 
-async def test_rebuild_observation_index(observation_service, test_entity):
-    """Test rebuilding observation index from filesystem."""
-    # Arrange - Add observations and clear database
-    await observation_service.add_observation(test_entity, "Test observation 1")
-    await observation_service.add_observation(test_entity, "Test observation 2")
-    
-    # Clear database but keep files
-    await observation_service.observation_repo.execute_query(delete(DbObservation))
-    
-    # Act
-    await observation_service.rebuild_observation_index()
-    
-    # Assert
-    db_observations = await observation_service.observation_repo.find_by_entity(test_entity.id)
-    assert len(db_observations) == 2
-    observation_contents = {obs.content for obs in db_observations}
-    assert observation_contents == {
-        "Test observation 1",
-        "Test observation 2"
-    }
-
-
 # Edge Cases
 
 async def test_observation_with_special_characters(observation_service, test_entity):
     """Test handling observations with special characters."""
     content = "Test & observation with @#$% special chars!"
-    observation = await observation_service.add_observation(
-        test_entity,
-        content
+    entity_data = EntityIn(
+        name=test_entity.name,
+        entity_type=test_entity.entity_type,
+        id=test_entity.id,
     )
-    assert observation.content == content
-    
-    # Verify file content
-    entity = await read_entity_file(observation_service.entities_path, test_entity.id)
-    assert any(obs.content == content for obs in entity.observations)
+
+    observations = await observation_service.add_observations(entity_data, [ObservationIn(content=content)])
+    assert observations[0].content == content
 
 
 async def test_very_long_observation(observation_service, test_entity):
     """Test handling very long observation content."""
     long_content = "Very long observation " * 100  # ~1800 characters
-    observation = await observation_service.add_observation(
-        test_entity,
-        long_content
+    entity_data = EntityIn(
+        name=test_entity.name,
+        entity_type=test_entity.entity_type,
+        id=test_entity.id,
+        observations=[]
     )
-    assert observation.content == long_content
-    
-    # Verify file content
-    entity = await read_entity_file(observation_service.entities_path, test_entity.id)
-    assert any(obs.content.rstrip() == long_content.rstrip() for obs in entity.observations)
+
+    observations = await observation_service.add_observations(entity_data, [ObservationIn(content=long_content)])
+    assert observations[0].content == long_content
+    

tests/test_relation_service.py

@@ -1,16 +1,9 @@
 """Tests for RelationService."""
 import pytest
 import pytest_asyncio
-from sqlalchemy import delete, select
 
-from basic_memory.models import Relation as DbRelation
-from basic_memory.repository import EntityRepository, RelationRepository
-from basic_memory.services import (
-    EntityService, RelationService,
-    ServiceError, DatabaseSyncError, RelationError
-)
-from basic_memory.schemas import Entity, Relation
-from basic_memory.fileio import read_entity_file, FileOperationError, write_entity_file
+from basic_memory.fileio import FileOperationError
+from basic_memory.schemas import EntityIn, RelationIn
 
 pytestmark = pytest.mark.asyncio
 
@@ -18,14 +11,20 @@ pytestmark = pytest.mark.asyncio
 @pytest_asyncio.fixture
 async def sample_entities(entity_service):
     """Create two sample entities for testing relations"""
-    entity1 = await entity_service.create_entity(
+    entity1_data = EntityIn(
         name="test_entity_1",
-        entity_type="test_type"
+        entity_type="test_type",
+        observations=[],
+        relations=[]
     )
-    entity2 = await entity_service.create_entity(
+    entity2_data = EntityIn(
         name="test_entity_2",
-        entity_type="test_type"
+        entity_type="test_type",
+        observations=[],
+        relations=[]
     )
+    entity1 = await entity_service.create_entity(entity1_data)
+    entity2 = await entity_service.create_entity(entity2_data)
     return entity1, entity2
 
 
@@ -38,29 +37,19 @@ async def test_create_relation(relation_service, sample_entities):
     """Test creating a basic relation between two entities"""
     entity1, entity2 = sample_entities
     
-    relation = await relation_service.create_relation(
-        from_entity=entity1,
-        to_entity=entity2,
+    relation_data = RelationIn(
+        from_id=entity1.id,
+        to_id=entity2.id,
         relation_type="test_relation"
     )
     
-    # Check Entity objects in relation
-    assert relation.from_entity.id == entity1.id
-    assert relation.to_entity.id == entity2.id
+    relation = await relation_service.create_relation(relation_data)
+    
+    # Check relation was created correctly
+    assert relation.from_id == entity1.id
+    assert relation.to_id == entity2.id
     assert relation.relation_type == "test_relation"
-    
-    # Verify relation was added to source entity's relations
-    assert hasattr(entity1, 'relations')
-    assert len(entity1.relations) == 1
-    assert entity1.relations[0].id == relation.id
-    
-    # Verify file was written with relation
-    entity_file = relation_service.entities_path / f"{entity1.id}.md"
-    assert entity_file.exists()
-    content = entity_file.read_text()
-    assert "## Relations" in content
-    assert f"[{entity2.id}] test_relation" in content
-    
+
     # Verify database was updated with correct IDs
     db_relation = await relation_service.relation_repo.find_by_id(relation.id)
     assert db_relation is not None
@@ -69,44 +58,22 @@ async def test_create_relation(relation_service, sample_entities):
     assert db_relation.relation_type == "test_relation"
 
 
-async def test_file_operation_error(relation_service, sample_entities, mocker):
-    """Test handling of file operation errors."""
-    entity1, entity2 = sample_entities
-    
-    # Add debug to see if mock is being called
-    async def mock_write(*args, **kwargs):
-        print("Mock write called with:", args, kwargs)
-        raise FileOperationError("Mock file error")
-    
-    # Patch where the function is used, not where it's imported from
-    mocker.patch('basic_memory.services.relation_service.write_entity_file', mock_write)
-    
-    with pytest.raises(FileOperationError):
-        await relation_service.create_relation(
-            from_entity=entity1,
-            to_entity=entity2,
-            relation_type="test_relation"
-        )
-
 
 async def test_create_relation_with_context(relation_service, sample_entities):
     """Test creating a relation with context information"""
     entity1, entity2 = sample_entities
     
-    relation = await relation_service.create_relation(
-        from_entity=entity1,
-        to_entity=entity2,
+    relation_data = RelationIn(
+        from_id=entity1.id,
+        to_id=entity2.id,
         relation_type="test_relation",
         context="test context"
     )
     
+    relation = await relation_service.create_relation(relation_data)
+    
     assert relation.context == "test context"
-    
-    # Verify context in file
-    entity_file = relation_service.entities_path / f"{entity1.id}.md"
-    content = entity_file.read_text()
-    assert f"[{entity2.id}] test_relation | test context" in content
-    
+
     # Verify context in database
     db_relation = await relation_service.relation_repo.find_by_id(relation.id)
     assert db_relation.context == "test context"

tests/test_repository.py

@@ -123,7 +123,6 @@ class TestObservationRepository:
     async def sample_observation(self, observation_repository: ObservationRepository, sample_entity: Entity):
         """Create a sample observation for testing"""
         observation_data = {
-            'id': '20240102-test-obs',
             'entity_id': sample_entity.id,
             'content': 'Test observation',
             'context': 'test-context'
@@ -137,16 +136,15 @@ class TestObservationRepository:
     ):
         """Test creating a new observation"""
         observation_data = {
-            'id': '20240102-obs',
             'entity_id': sample_entity.id,
             'content': 'Test content',
             'context': 'test-context'
         }
         observation = await observation_repository.create(observation_data)
         
-        assert observation.id == '20240102-obs'
         assert observation.entity_id == sample_entity.id
         assert observation.content == 'Test content'
+        assert observation.id is not None  # Should be auto-generated
 
     async def test_find_by_entity(
         self,
@@ -193,7 +191,6 @@ class TestRelationRepository:
     ):
         """Create a sample relation for testing"""
         relation_data = {
-            'id': '20240102-test-rel',
             'from_id': sample_entity.id,
             'to_id': related_entity.id,
             'relation_type': 'test_relation',
@@ -209,7 +206,6 @@ class TestRelationRepository:
     ):
         """Test creating a new relation"""
         relation_data = {
-            'id': '20240102-rel',
             'from_id': sample_entity.id,
             'to_id': related_entity.id,
             'relation_type': 'test_relation',
@@ -217,10 +213,10 @@ class TestRelationRepository:
         }
         relation = await relation_repository.create(relation_data)
         
-        assert relation.id == '20240102-rel'
         assert relation.from_id == sample_entity.id
         assert relation.to_id == related_entity.id
         assert relation.relation_type == 'test_relation'
+        assert relation.id is not None  # Should be auto-generated
 
     async def test_find_by_entities(
         self,