add FKs to db connections

src/basic_memory/db.py

@@ -1,90 +1,101 @@
 """Database configuration and initialization for basic-memory."""
+
+from contextlib import asynccontextmanager
 from enum import Enum
 from pathlib import Path
 from typing import AsyncGenerator
-from contextlib import asynccontextmanager
 
 from loguru import logger
-from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncEngine, AsyncSession
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import (
+    create_async_engine,
+    async_sessionmaker,
+    AsyncEngine,
+    AsyncSession,
+)
 from sqlalchemy.pool import StaticPool
 
 from basic_memory.models import Base
 
+
 class DatabaseType(Enum):
     """Types of database configurations."""
-    MEMORY = "memory"      # In-memory SQLite for testing
-    FILESYSTEM = "file"    # File-based SQLite for projects
 
-def get_database_url(project_path: Path, db_type: DatabaseType, ) -> str:
+    MEMORY = "memory"  # In-memory SQLite for testing
+    FILESYSTEM = "file"  # File-based SQLite for projects
+
+
+def get_database_url(
+    project_path: Path,
+    db_type: DatabaseType,
+) -> str:
     """
     Get database URL based on type and optional project path.
-    
+
     Args:
         db_type: Type of database to configure
         project_path: Project directory for file-based DBs (required if type is FILESYSTEM)
-        
+
     Returns:
         Database URL string
-    
+
     Raises:
         ValueError: If project_path is required but not provided
     """
     match db_type:
         case DatabaseType.MEMORY:
             return "sqlite+aiosqlite:///:memory:"
-            
+
         case DatabaseType.FILESYSTEM:
             if not project_path:
                 raise ValueError("project_path required for filesystem database")
-                
+
             # Ensure data directory exists
             data_dir = project_path / "data"
             data_dir.mkdir(parents=True, exist_ok=True)
-            
+
             db_path = data_dir / "memory.db"
             return f"sqlite+aiosqlite:///{db_path}"
 
+
 async def init_database(url: str, echo: bool = False) -> tuple[AsyncEngine, async_sessionmaker]:
     """
     Initialize database with schema.
-    
+
     Args:
         url: Database URL
         echo: Whether to echo SQL statements
-        
+
     Returns:
         Configured async engine and session factory
     """
     # Configure engine based on URL
     connect_args = {"check_same_thread": False}
+
     if url == "sqlite+aiosqlite:///:memory:":
         engine = create_async_engine(
             url,
             echo=echo,
             poolclass=StaticPool,  # Single connection for in-memory
-            connect_args=connect_args
+            connect_args=connect_args,
         )
     else:
-        engine = create_async_engine(
-            url,
-            echo=echo,
-            connect_args=connect_args
-        )
+        engine = create_async_engine(url, echo=echo, connect_args=connect_args)
 
     # Create tables
     async with engine.begin() as conn:
         await conn.run_sync(Base.metadata.create_all)
 
     # Create session factory for this engine
-    session_factory = async_sessionmaker(
-        engine,
-        expire_on_commit=False
-    )
-        
+    session_factory = async_sessionmaker(engine, expire_on_commit=False)
+
     return engine, session_factory
 
+
 @asynccontextmanager
-async def engine_session_factory(project_path: Path, db_type=DatabaseType.FILESYSTEM) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]:
+async def engine_session_factory(
+    project_path: Path, db_type=DatabaseType.FILESYSTEM
+) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]:
     """Get database engine and session factory with proper lifecycle management."""
     url = get_database_url(project_path, db_type=db_type)
     engine, session_factory = await init_database(url)
@@ -94,20 +105,25 @@ async def engine_session_factory(project_path: Path, db_type=DatabaseType.FILESY
     finally:
         await engine.dispose()
 
+
 @asynccontextmanager
-async def session(session_factory: async_sessionmaker[AsyncSession]) -> AsyncGenerator[AsyncSession, None]:
+async def session(
+    session_factory: async_sessionmaker[AsyncSession],
+) -> AsyncGenerator[AsyncSession, None]:
     """
     Get database session with proper lifecycle management.
-    
+
     Args:
         session_factory: Async session factory to create session from
-        
+
     Yields:
         AsyncSession configured for engine
     """
     # Create and yield session
     session = session_factory()
     try:
+        # Ensure foreign keys enabled for this session
+        await session.execute(text("PRAGMA foreign_keys=ON"))
         yield session
         await session.commit()
     except Exception:
@@ -116,11 +132,12 @@ async def session(session_factory: async_sessionmaker[AsyncSession]) -> AsyncGen
     finally:
         await session.close()
 
+
 async def dispose_database(engine: AsyncEngine):
     """
     Clean up database engine.
-    
+
     Args:
         engine: Engine to dispose
     """
-    await engine.dispose()
+    await engine.dispose()

src/basic_memory/models.py

@@ -5,6 +5,8 @@ from sqlalchemy import String, DateTime, ForeignKey, Text, Integer, text, Unique
 from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase
 from sqlalchemy.ext.asyncio import AsyncAttrs
 
+from basic_memory.utils import normalize_entity_id
+
 
 class Base(AsyncAttrs, DeclarativeBase):
     """Base class for all models"""
@@ -68,8 +70,8 @@ class Entity(Base):
     @classmethod
     def generate_id(cls, entity_type: str, name: str) -> str:
         """Generate a filesystem path-based ID for this entity."""
-        # Normalize name for filesystem (handle spaces, special chars etc)
-        safe_name = name.lower().replace(" ", "_")
+        # Use common normalization for filesystem safety
+        safe_name = normalize_entity_id(name)
         return f"{entity_type}/{safe_name}"
 
     def get_file_path(self) -> str:

src/basic_memory/schemas.py

@@ -3,11 +3,16 @@
 from typing import List, Optional, Annotated, TypeAlias
 
 from annotated_types import Len
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, BeforeValidator
+
+from basic_memory.utils import normalize_entity_id
 
 # Base Models
 Observation: TypeAlias = str
 
+# Custom field types with validation
+EntityId = Annotated[str, BeforeValidator(normalize_entity_id)]
+
 
 class Relation(BaseModel):
     """
@@ -15,8 +20,8 @@ class Relation(BaseModel):
     Relations are always stored in active voice (e.g. "created", "teaches", etc.)
     """
 
-    from_id: str
-    to_id: str
+    from_id: EntityId
+    to_id: EntityId
     relation_type: str
     context: Optional[str] = None
 
@@ -28,7 +33,7 @@ class Entity(BaseModel):
     associated observations.
     """
 
-    id: Optional[str] = None
+    id: Optional[EntityId] = None
     name: str
     entity_type: str
     description: Optional[str] = None
@@ -47,7 +52,7 @@ class Entity(BaseModel):
 class AddObservationsRequest(BaseModel):
     """Schema for adding observations to an entity."""
 
-    entity_id: str
+    entity_id: EntityId
     context: Optional[str] = None
     observations: List[Observation]
 
@@ -67,7 +72,7 @@ class SearchNodesRequest(BaseModel):
 class OpenNodesRequest(BaseModel):
     """Request schema for open_nodes tool."""
 
-    names: Annotated[List[str], Len(min_length=1)]
+    names: Annotated[List[EntityId], Len(min_length=1)]
 
 
 class CreateRelationsRequest(BaseModel):
@@ -82,7 +87,7 @@ class CreateRelationsRequest(BaseModel):
 class DeleteEntityRequest(BaseModel):
     """Request schema for delete_entities tool."""
 
-    entity_ids: List[str]
+    entity_ids: List[EntityId]
 
 
 class DeleteRelationsRequest(BaseModel):
@@ -94,7 +99,7 @@ class DeleteRelationsRequest(BaseModel):
 class DeleteObservationsRequest(BaseModel):
     """Request schema for delete_observations tool."""
 
-    entity_id: str
+    entity_id: EntityId
     deletions: List[Observation]
 
 
@@ -118,7 +123,7 @@ class ObservationResponse(SQLAlchemyModel):
 class ObservationsResponse(SQLAlchemyModel):
     """Schema for bulk observation operation results."""
 
-    entity_id: str
+    entity_id: EntityId
     observations: List[ObservationResponse]
 
 
@@ -129,7 +134,7 @@ class RelationResponse(Relation, SQLAlchemyModel):
 class EntityResponse(SQLAlchemyModel):
     """Schema for entity data returned from the service."""
 
-    id: str
+    id: EntityId
     name: str
     entity_type: str
     description: Optional[str] = None
@@ -159,7 +164,7 @@ class OpenNodesResponse(SQLAlchemyModel):
 class AddObservationsResponse(SQLAlchemyModel):
     """Response for add_observations tool."""
 
-    entity_id: str
+    entity_id: EntityId
     observations: List[ObservationResponse]
 
 

src/basic_memory/utils.py

@@ -0,0 +1,13 @@
+"""Utility functions for basic-memory."""
+
+def normalize_entity_id(entity_id: str) -> str:
+    """
+    Normalize an entity ID by converting to lowercase and replacing spaces with underscores.
+    
+    Args:
+        entity_id: Raw entity ID to normalize
+        
+    Returns:
+        Normalized entity ID suitable for filesystem and database use
+    """
+    return entity_id.lower().replace(" ", "_")

tests/conftest.py

@@ -1,4 +1,5 @@
 """Common test fixtures."""
+
 import tempfile
 from pathlib import Path
 from typing import AsyncGenerator
@@ -7,16 +8,18 @@ import pytest_asyncio
 from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession, AsyncEngine
 
 from basic_memory import db
+from basic_memory.config import ProjectConfig
 from basic_memory.db import DatabaseType
-from basic_memory.repository.entity_repository import EntityRepository
 from basic_memory.deps import (
     get_entity_service,
     get_observation_service,
     get_relation_service,
-    get_relation_repo, get_observation_repo, get_entity_repo
+    get_relation_repo,
+    get_observation_repo,
+    get_entity_repo,
 )
+from basic_memory.repository.entity_repository import EntityRepository
 from basic_memory.schemas import Entity
-from basic_memory.config import ProjectConfig
 from basic_memory.services import MemoryService
 
 
@@ -24,20 +27,25 @@ from basic_memory.services import MemoryService
 def anyio_backend():
     return "asyncio"
 
+
 @pytest_asyncio.fixture
 def test_config(tmp_path):
     """Test configuration using in-memory DB."""
     config = ProjectConfig(
         name="test",
     )
-    config.path=tmp_path
+    config.path = tmp_path
     return config
 
 
 @pytest_asyncio.fixture(scope="function")
-async def engine_session_factory(test_config)-> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]:
+async def engine_session_factory(
+    test_config,
+) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]:
     """Create an async engine using in-memory SQLite database"""
-    async with db.engine_session_factory(project_path=test_config.path, db_type=DatabaseType.MEMORY) as (engine, session_factory):
+    async with db.engine_session_factory(
+        project_path=test_config.path, db_type=DatabaseType.MEMORY
+    ) as (engine, session_factory):
         yield engine, session_factory
 
 
@@ -45,9 +53,10 @@ async def engine_session_factory(test_config)-> AsyncGenerator[tuple[AsyncEngine
 async def session(engine_session_factory):
     """Create an async session factory and yield a session"""
     engine, session_factory = engine_session_factory
-    async with session_factory() as session:
+    async with db.session(session_factory) as session:
         yield session
 
+
 @pytest_asyncio.fixture
 async def test_project_path():
     """Create a temporary project directory."""
@@ -57,62 +66,61 @@ async def test_project_path():
         entities_path.mkdir(parents=True)
         yield project_path
 
+
 @pytest_asyncio.fixture(scope="function")
 async def entity_repository(session: AsyncSession):
     """Create an EntityRepository instance"""
     return await get_entity_repo(session)
 
+
 @pytest_asyncio.fixture(scope="function")
 async def observation_repository(session: AsyncSession):
     """Create an ObservationRepository instance"""
     return await get_observation_repo(session)
 
+
 @pytest_asyncio.fixture(scope="function")
 async def relation_repository(session: AsyncSession):
     """Create a RelationRepository instance"""
     return await get_relation_repo(session)
 
+
 @pytest_asyncio.fixture
 async def entity_service(test_project_path, entity_repository):
     """Fixture providing initialized EntityService."""
     return await get_entity_service(test_project_path, entity_repository)
 
+
 @pytest_asyncio.fixture
 async def relation_service(test_project_path, relation_repository):
     """Fixture providing initialized RelationService."""
     return await get_relation_service(test_project_path, relation_repository)
 
+
 @pytest_asyncio.fixture
 async def observation_service(test_project_path, observation_repository):
     """Fixture providing initialized RelationService."""
     return await get_observation_service(test_project_path, observation_repository)
 
+
 @pytest_asyncio.fixture
-async def memory_service(
-    test_project_path,
-    entity_service,
-    relation_service,
-    observation_service
-):
+async def memory_service(test_project_path, entity_service, relation_service, observation_service):
     """Fixture providing initialized MemoryService."""
-    return MemoryService(
-        test_project_path,
-        entity_service,
-        relation_service,
-        observation_service
-    )
+    return MemoryService(test_project_path, entity_service, relation_service, observation_service)
+
 
 @pytest_asyncio.fixture(scope="function")
 async def sample_entity(entity_repository: EntityRepository):
     """Create a sample entity for testing"""
     entity_data = {
-        'id': '20240102-test-entity',
-        'name': 'Test Entity',
-        'entity_type': 'test',
-        'description': 'A test entity',
+        "id": "20240102-test-entity",
+        "name": "Test Entity",
+        "entity_type": "test",
+        "description": "A test entity",
     }
     return await entity_repository.create(entity_data)
 
+
 @pytest_asyncio.fixture
 async def test_entity(entity_service):
     """Create a test entity for reuse in tests."""
@@ -121,4 +129,3 @@ async def test_entity(entity_service):
         entity_type="test",  # pyright: ignore [reportCallIssue]
     )
     return await entity_service.create_entity(entity_data)
-

tests/mcp/test_mcp_server.py

@@ -11,6 +11,7 @@ from basic_memory.api.app import app as fastapi_app
 from basic_memory.deps import get_project_config, get_engine_factory
 from basic_memory.mcp.server import MemoryServer, MIME_TYPE, BASIC_MEMORY_URI
 from basic_memory.schemas import CreateEntityResponse, SearchNodesResponse, AddObservationsResponse
+from basic_memory.utils import normalize_entity_id
 
 
 @pytest_asyncio.fixture
@@ -70,6 +71,106 @@ def test_directory_entity_data():
     }
 
 
+@pytest.mark.asyncio
+async def test_create_single_entity(server):
+    """Test creating a single entity."""
+    entity_data = {
+        "entities": [
+            {"name": "SingleTest", "entity_type": "test", "observations": ["Test observation"]}
+        ]
+    }
+
+    result = await server.handle_call_tool("create_entities", entity_data)
+
+    # Verify response format
+    assert len(result) == 1
+    assert isinstance(result[0], EmbeddedResource)
+    assert result[0].type == "resource"
+    assert result[0].resource.mimeType == MIME_TYPE
+
+    # Verify entity creation
+    response = CreateEntityResponse.model_validate_json(result[0].resource.text)
+    assert len(response.entities) == 1
+    entity = response.entities[0]
+    assert entity.name == "SingleTest"
+    assert entity.entity_type == "test"
+    assert len(entity.observations) == 1
+    assert entity.observations[0].content == "Test observation"
+    assert entity.id == "test/singletest"
+
+    # Verify entity can be found via search
+    search_result = await server.handle_call_tool("search_nodes", {"query": "SingleTest"})
+    search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text)
+    assert len(search_response.matches) == 1
+    assert search_response.matches[0].name == "SingleTest"
+
+
+@pytest.mark.asyncio
+async def test_create_multiple_entities(server):
+    """Test creating multiple entities in one call."""
+    entity_data = {
+        "entities": [
+            {"name": "BulkTest1", "entity_type": "test", "observations": ["First bulk test"]},
+            {"name": "BulkTest2", "entity_type": "test", "observations": ["Second bulk test"]},
+            {"name": "BulkTest3", "entity_type": "demo", "observations": ["Third bulk test"]},
+        ]
+    }
+
+    result = await server.handle_call_tool("create_entities", entity_data)
+
+    # Verify response
+    assert len(result) == 1
+    response = CreateEntityResponse.model_validate_json(result[0].resource.text)
+
+    # Verify all entities were created
+    assert len(response.entities) == 3
+
+    # Check specific entities
+    entities = {e.name: e for e in response.entities}
+    assert "BulkTest1" in entities
+    assert "BulkTest2" in entities
+    assert "BulkTest3" in entities
+
+    # Verify IDs were generated correctly
+    assert entities["BulkTest1"].id == "test/bulktest1"
+    assert entities["BulkTest2"].id == "test/bulktest2"
+    assert entities["BulkTest3"].id == "demo/bulktest3"
+
+    # Verify observations were saved
+    assert len(entities["BulkTest1"].observations) == 1
+    assert entities["BulkTest1"].observations[0].content == "First bulk test"
+
+    # Verify entities can be found via search
+    search_result = await server.handle_call_tool("search_nodes", {"query": "BulkTest"})
+    search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text)
+    assert len(search_response.matches) == 3
+
+
+@pytest.mark.asyncio
+async def test_create_entity_with_all_fields(server):
+    """Test creating entity with all possible fields populated."""
+    entity_data = {
+        "entities": [
+            {
+                "name": "FullEntity",
+                "entity_type": "test",
+                "description": "A complete test entity",
+                "observations": ["First observation", "Second observation"],
+            }
+        ]
+    }
+
+    result = await server.handle_call_tool("create_entities", entity_data)
+    response = CreateEntityResponse.model_validate_json(result[0].resource.text)
+
+    entity = response.entities[0]
+    assert entity.name == "FullEntity"
+    assert entity.description == "A complete test entity"
+    assert len(entity.observations) == 2
+    assert entity.observations[0].content == "First observation"
+    assert entity.observations[1].content == "Second observation"
+
+
 @pytest.mark.asyncio
 async def test_list_tools(server):
     """Test that server exposes expected tools."""
@@ -116,7 +217,7 @@ async def test_search_nodes(test_entity_data, client, server):
     assert result[0].resource.mimeType == MIME_TYPE
 
     # Verify search results
-    response = SearchNodesResponse.model_validate_json(result[0].resource.text)  # pyright: ignore [reportAttributeAccessIssue]
+    response = SearchNodesResponse.model_validate_json(result[0].resource.text)
     assert len(response.matches) == 1
     assert response.matches[0].name == "Test Entity"
     assert response.query == "Test Entity"
@@ -135,7 +236,7 @@ async def test_add_observations(test_entity_data, client, server):
 
     # First create an entity
     create_result = await server.handle_call_tool("create_entities", test_entity_data)
-    create_response = CreateEntityResponse.model_validate_json(create_result[0].resource.text)  # pyright: ignore [reportAttributeAccessIssue]
+    create_response = CreateEntityResponse.model_validate_json(create_result[0].resource.text)
     entity_id = create_response.entities[0].id
 
     # Add new observation
@@ -162,6 +263,43 @@ async def test_add_observations(test_entity_data, client, server):
     assert "A new observation" in [o["content"] for o in entity["observations"]]
 
 
+@pytest.mark.asyncio
+async def test_create_relations(test_entity_data, client, server):
+    """Test creating relations between entities."""
+    # Create two test entities
+    entity_data = {
+        "entities": [
+            {"name": "TestEntityA", "entity_type": "test", "observations": ["Entity A"]},
+            {"name": "TestEntityB", "entity_type": "test", "observations": ["Entity B"]},
+        ]
+    }
+
+    await server.handle_call_tool("create_entities", entity_data)
+
+    # Create relation between them
+    relation_data = {
+        "relations": [
+            {
+                "from_id": "test/TestEntityA",
+                "to_id": "test/TestEntityB",
+                "relation_type": "relates_to",
+            }
+        ]
+    }
+
+    result = await server.handle_call_tool("create_relations", relation_data)
+
+    # Verify through search
+    search_result = await server.handle_call_tool("search_nodes", {"query": "TestEntityA"})
+    response = SearchNodesResponse.model_validate_json(search_result[0].resource.text)
+
+    assert len(response.matches) == 1
+    entity = response.matches[0]
+    assert len(entity.relations) == 1
+    assert entity.relations[0].to_id == normalize_entity_id("test/TestEntityB")
+    assert entity.relations[0].relation_type == "relates_to"
+
+
 @pytest.mark.asyncio
 async def test_invalid_tool_name(server):
     """Test calling a non-existent tool."""

tests/repository/test_observation_repository.py

@@ -1,6 +1,8 @@
 """Tests for the ObservationRepository."""
+
 import pytest
 import pytest_asyncio
+import sqlalchemy
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from basic_memory.models import Entity, Observation
@@ -12,38 +14,54 @@ async def repo(observation_repository):
     """Create an ObservationRepository instance"""
     return observation_repository
 
+
 @pytest_asyncio.fixture(scope="function")
 async def sample_observation(repo, sample_entity: Entity):
     """Create a sample observation for testing"""
     observation_data = {
-        'entity_id': sample_entity.id,
-        'content': 'Test observation',
-        'context': 'test-context'
+        "entity_id": sample_entity.id,
+        "content": "Test observation",
+        "context": "test-context",
     }
     return await repo.create(observation_data)
 
+
 @pytest.mark.asyncio
 async def test_create_observation(
-        observation_repository: ObservationRepository,
-        sample_entity: Entity
+    observation_repository: ObservationRepository, sample_entity: Entity
 ):
     """Test creating a new observation"""
     observation_data = {
-        'entity_id': sample_entity.id,
-        'content': 'Test content',
-        'context': 'test-context'
+        "entity_id": sample_entity.id,
+        "content": "Test content",
+        "context": "test-context",
     }
     observation = await observation_repository.create(observation_data)
 
     assert observation.entity_id == sample_entity.id
-    assert observation.content == 'Test content'
+    assert observation.content == "Test content"
     assert observation.id is not None  # Should be auto-generated
 
+
+@pytest.mark.asyncio
+async def test_create_observation_entity_does_not_exist(
+    observation_repository: ObservationRepository, sample_entity: Entity
+):
+    """Test creating a new observation"""
+    observation_data = {
+        "entity_id": "does-not-exist",
+        "content": "Test content",
+        "context": "test-context",
+    }
+    with pytest.raises(sqlalchemy.exc.IntegrityError):
+        await observation_repository.create(observation_data)
+
+
 @pytest.mark.asyncio
 async def test_find_by_entity(
-        observation_repository: ObservationRepository,
-        sample_observation: Observation,
-        sample_entity: Entity
+    observation_repository: ObservationRepository,
+    sample_observation: Observation,
+    sample_entity: Entity,
 ):
     """Test finding observations by entity"""
     observations = await observation_repository.find_by_entity(sample_entity.id)
@@ -51,13 +69,13 @@ async def test_find_by_entity(
     assert observations[0].id == sample_observation.id
     assert observations[0].content == sample_observation.content
 
+
 @pytest.mark.asyncio
 async def test_find_by_context(
-        observation_repository: ObservationRepository,
-        sample_observation: Observation
+    observation_repository: ObservationRepository, sample_observation: Observation
 ):
     """Test finding observations by context"""
-    observations = await observation_repository.find_by_context('test-context')
+    observations = await observation_repository.find_by_context("test-context")
     assert len(observations) == 1
     assert observations[0].id == sample_observation.id
     assert observations[0].content == sample_observation.content
@@ -68,10 +86,7 @@ async def test_delete_observations(session: AsyncSession, repo):
     """Test deleting observations by entity_id."""
     # Create test entity
     entity = Entity(
-        id="test/test_entity",
-        name="test_entity",
-        entity_type="test",
-        description="Test entity"
+        id="test/test_entity", name="test_entity", entity_type="test", description="Test entity"
     )
     session.add(entity)
     await session.flush()
@@ -90,15 +105,13 @@ async def test_delete_observations(session: AsyncSession, repo):
     remaining = await repo.find_by_entity(entity.id)
     assert len(remaining) == 0
 
+
 @pytest.mark.asyncio
 async def test_delete_observation_by_id(session: AsyncSession, repo):
     """Test deleting a single observation by its ID."""
     # Create test entity
     entity = Entity(
-        id="test/test_entity",
-        name="test_entity",
-        entity_type="test",
-        description="Test entity"
+        id="test/test_entity", name="test_entity", entity_type="test", description="Test entity"
     )
     session.add(entity)
     await session.flush()
@@ -116,15 +129,13 @@ async def test_delete_observation_by_id(session: AsyncSession, repo):
     remaining = await repo.find_by_id(obs.id)
     assert remaining is None
 
+
 @pytest.mark.asyncio
 async def test_delete_observation_by_content(session: AsyncSession, repo):
     """Test deleting observations by content."""
     # Create test entity
     entity = Entity(
-        id="test/test_entity",
-        name="test_entity",
-        entity_type="test",
-        description="Test entity"
+        id="test/test_entity", name="test_entity", entity_type="test", description="Test entity"
     )
     session.add(entity)
     await session.flush()
@@ -135,7 +146,6 @@ async def test_delete_observation_by_content(session: AsyncSession, repo):
     session.add_all([obs1, obs2])
     await session.flush()
 
-
     # Test deletion by content
     deleted = await repo.delete_by_fields(content="Delete this observation")
     assert deleted is True

tests/repository/test_relation_repository.py

@@ -1,11 +1,13 @@
 """Tests for the RelationRepository."""
+
 import pytest
 import pytest_asyncio
+import sqlalchemy
+
 from basic_memory.models import Entity, Relation
 from basic_memory.repository.relation_repository import RelationRepository
 
 
-
 @pytest_asyncio.fixture
 async def relation_repo(session):
     """Create a RelationRepository with test DB session."""
@@ -19,7 +21,7 @@ async def source_entity(session):
         id="source/test_entity",
         name="test_source",
         entity_type="source",
-        description="Source entity"
+        description="Source entity",
     )
     session.add(entity)
     await session.flush()
@@ -33,7 +35,7 @@ async def target_entity(session):
         id="target/test_entity",
         name="test_target",
         entity_type="target",
-        description="Target entity"
+        description="Target entity",
     )
     session.add(entity)
     await session.flush()
@@ -44,181 +46,181 @@ async def target_entity(session):
 async def test_relations(session, source_entity, target_entity):
     """Create test relations."""
     relations = [
-        Relation(
-            from_id=source_entity.id,
-            to_id=target_entity.id,
-            relation_type="connects_to"
-        ),
-        Relation(
-            from_id=source_entity.id,
-            to_id=target_entity.id,
-            relation_type="depends_on"
-        )
+        Relation(from_id=source_entity.id, to_id=target_entity.id, relation_type="connects_to"),
+        Relation(from_id=source_entity.id, to_id=target_entity.id, relation_type="depends_on"),
     ]
     session.add_all(relations)
     await session.flush()
     return relations
 
+
 @pytest_asyncio.fixture(scope="function")
 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',
-        'references': ''
+        "id": "20240102-related",
+        "name": "Related Entity",
+        "entity_type": "test",
+        "description": "A related test entity",
+        "references": "",
     }
     return await entity_repository.create(entity_data)
 
+
 @pytest_asyncio.fixture(scope="function")
 async def sample_relation(
-        relation_repository: RelationRepository,
-        sample_entity: Entity,
-        related_entity: Entity
+    relation_repository: RelationRepository, sample_entity: Entity, related_entity: Entity
 ):
     """Create a sample relation for testing"""
     relation_data = {
-        'from_id': sample_entity.id,
-        'to_id': related_entity.id,
-        'relation_type': 'test_relation',
-        'context': 'test-context'
+        "from_id": sample_entity.id,
+        "to_id": related_entity.id,
+        "relation_type": "test_relation",
+        "context": "test-context",
     }
     return await relation_repository.create(relation_data)
 
+
 @pytest_asyncio.fixture(scope="function")
 async def multiple_relations(
-        relation_repository: RelationRepository,
-        sample_entity: Entity,
-        related_entity: Entity
+    relation_repository: RelationRepository, sample_entity: Entity, related_entity: Entity
 ):
     """Create multiple relations for testing"""
     relations_data = [
         {
-            'from_id': sample_entity.id,
-            'to_id': related_entity.id,
-            'relation_type': 'relation_one',
-            'context': 'context_one'
+            "from_id": sample_entity.id,
+            "to_id": related_entity.id,
+            "relation_type": "relation_one",
+            "context": "context_one",
         },
         {
-            'from_id': sample_entity.id,
-            'to_id': related_entity.id,
-            'relation_type': 'relation_two',
-            'context': 'context_two'
+            "from_id": sample_entity.id,
+            "to_id": related_entity.id,
+            "relation_type": "relation_two",
+            "context": "context_two",
         },
         {
-            'from_id': related_entity.id,
-            'to_id': sample_entity.id,
-            'relation_type': 'relation_one',
-            'context': 'context_three'
-        }
+            "from_id": related_entity.id,
+            "to_id": sample_entity.id,
+            "relation_type": "relation_one",
+            "context": "context_three",
+        },
     ]
     return [await relation_repository.create(data) for data in relations_data]
 
+
 @pytest.mark.asyncio
 async def test_create_relation(
-        relation_repository: RelationRepository,
-        sample_entity: Entity,
-        related_entity: Entity
+    relation_repository: RelationRepository, sample_entity: Entity, related_entity: Entity
 ):
     """Test creating a new relation"""
     relation_data = {
-        'from_id': sample_entity.id,
-        'to_id': related_entity.id,
-        'relation_type': 'test_relation',
-        'context': 'test-context'
+        "from_id": sample_entity.id,
+        "to_id": related_entity.id,
+        "relation_type": "test_relation",
+        "context": "test-context",
     }
     relation = await relation_repository.create(relation_data)
 
     assert relation.from_id == sample_entity.id
     assert relation.to_id == related_entity.id
-    assert relation.relation_type == 'test_relation'
+    assert relation.relation_type == "test_relation"
     assert relation.id is not None  # Should be auto-generated
 
+
+@pytest.mark.asyncio
+async def test_create_relation_entity_does_not_exist(
+    relation_repository: RelationRepository, sample_entity: Entity, related_entity: Entity
+):
+    """Test creating a new relation"""
+    relation_data = {
+        "from_id": "not_exist",
+        "to_id": related_entity.id,
+        "relation_type": "test_relation",
+        "context": "test-context",
+    }
+    with pytest.raises(sqlalchemy.exc.IntegrityError):
+        await relation_repository.create(relation_data)
+
+
 @pytest.mark.asyncio
 async def test_find_by_entities(
-        relation_repository: RelationRepository,
-        sample_relation: Relation,
-        sample_entity: Entity,
-        related_entity: Entity
+    relation_repository: RelationRepository,
+    sample_relation: Relation,
+    sample_entity: Entity,
+    related_entity: Entity,
 ):
     """Test finding relations between specific entities"""
-    relations = await relation_repository.find_by_entities(
-        sample_entity.id,
-        related_entity.id
-    )
+    relations = await relation_repository.find_by_entities(sample_entity.id, related_entity.id)
     assert len(relations) == 1
     assert relations[0].id == sample_relation.id
     assert relations[0].relation_type == sample_relation.relation_type
 
+
 @pytest.mark.asyncio
-async def test_find_by_type(
-        relation_repository: RelationRepository,
-        sample_relation: Relation
-):
+async def test_find_by_type(relation_repository: RelationRepository, sample_relation: Relation):
     """Test finding relations by type"""
-    relations = await relation_repository.find_by_type('test_relation')
+    relations = await relation_repository.find_by_type("test_relation")
     assert len(relations) == 1
     assert relations[0].id == sample_relation.id
 
+
 @pytest.mark.asyncio
 async def test_delete_by_fields_single_field(
-        relation_repository: RelationRepository,
-        multiple_relations: list[Relation]
+    relation_repository: RelationRepository, multiple_relations: list[Relation]
 ):
     """Test deleting relations by a single field."""
     # Delete all relations of type 'relation_one'
-    result = await relation_repository.delete_by_fields(relation_type='relation_one')  # pyright: ignore [reportArgumentType]
+    result = await relation_repository.delete_by_fields(relation_type="relation_one")  # pyright: ignore [reportArgumentType]
     assert result is True
 
     # Verify deletion
-    remaining = await relation_repository.find_by_type('relation_one')
+    remaining = await relation_repository.find_by_type("relation_one")
     assert len(remaining) == 0
 
     # Other relations should still exist
-    others = await relation_repository.find_by_type('relation_two')
+    others = await relation_repository.find_by_type("relation_two")
     assert len(others) == 1
 
+
 @pytest.mark.asyncio
 async def test_delete_by_fields_multiple_fields(
-        relation_repository: RelationRepository,
-        multiple_relations: list[Relation],
-        sample_entity: Entity,
-        related_entity: Entity
+    relation_repository: RelationRepository,
+    multiple_relations: list[Relation],
+    sample_entity: Entity,
+    related_entity: Entity,
 ):
     """Test deleting relations by multiple fields."""
     # Delete specific relation matching both from_id and relation_type
     result = await relation_repository.delete_by_fields(
-        from_id=sample_entity.id, # pyright: ignore [reportArgumentType]
-        relation_type='relation_one' # pyright: ignore [reportArgumentType]
+        from_id=sample_entity.id,  # pyright: ignore [reportArgumentType]
+        relation_type="relation_one",  # pyright: ignore [reportArgumentType]
     )
     assert result is True
 
     # Verify correct relation was deleted
-    remaining = await relation_repository.find_by_entities(
-        sample_entity.id,
-        related_entity.id
-    )
+    remaining = await relation_repository.find_by_entities(sample_entity.id, related_entity.id)
     assert len(remaining) == 1  # Only relation_two should remain
-    assert remaining[0].relation_type == 'relation_two'
+    assert remaining[0].relation_type == "relation_two"
+
 
 @pytest.mark.asyncio
 async def test_delete_by_fields_no_match(
-        relation_repository: RelationRepository,
-        multiple_relations: list[Relation]
+    relation_repository: RelationRepository, multiple_relations: list[Relation]
 ):
     """Test delete_by_fields when no relations match."""
     result = await relation_repository.delete_by_fields(
-        relation_type='nonexistent_type'  # pyright: ignore [reportArgumentType]
+        relation_type="nonexistent_type"  # pyright: ignore [reportArgumentType]
     )
     assert result is False
 
+
 @pytest.mark.asyncio
 async def test_delete_by_fields_all_fields(
-        relation_repository: RelationRepository,
-        multiple_relations: list[Relation],
-        sample_entity: Entity,
-        related_entity: Entity
+    relation_repository: RelationRepository,
+    multiple_relations: list[Relation],
+    sample_entity: Entity,
+    related_entity: Entity,
 ):
     """Test deleting relation by matching all fields."""
     # Get first relation's data
@@ -226,10 +228,10 @@ async def test_delete_by_fields_all_fields(
 
     # Delete using all fields
     result = await relation_repository.delete_by_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]
+        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
 
@@ -239,15 +241,14 @@ async def test_delete_by_fields_all_fields(
     assert remaining[0].context != relation.context
 
 
-
 @pytest.mark.asyncio
 async def test_delete_relation_by_id(relation_repo, test_relations):
     """Test deleting a relation by ID."""
     relation = test_relations[0]
-    
+
     result = await relation_repo.delete(relation.id)
     assert result is True
-    
+
     # Verify deletion
     remaining = await relation_repo.find_one(
         relation_repo.select(Relation).filter(Relation.id == relation.id)
@@ -260,25 +261,24 @@ async def test_delete_relations_by_type(relation_repo, test_relations):
     """Test deleting relations by type."""
     result = await relation_repo.delete_by_fields(relation_type="connects_to")
     assert result is True
-    
+
     # Verify specific type was deleted
     remaining = await relation_repo.find_by_type("connects_to")
     assert len(remaining) == 0
-    
+
     # Verify other type still exists
     others = await relation_repo.find_by_type("depends_on")
     assert len(others) == 1
 
 
 @pytest.mark.asyncio
-async def test_delete_relations_by_entities(relation_repo, test_relations, source_entity, target_entity):
+async def test_delete_relations_by_entities(
+    relation_repo, test_relations, source_entity, target_entity
+):
     """Test deleting relations between specific entities."""
-    result = await relation_repo.delete_by_fields(
-        from_id=source_entity.id,
-        to_id=target_entity.id
-    )
+    result = await relation_repo.delete_by_fields(from_id=source_entity.id, to_id=target_entity.id)
     assert result is True
-    
+
     # Verify all relations between entities were deleted
     remaining = await relation_repo.find_by_entities(source_entity.id, target_entity.id)
     assert len(remaining) == 0
@@ -288,4 +288,4 @@ async def test_delete_relations_by_entities(relation_repo, test_relations, sourc
 async def test_delete_nonexistent_relation(relation_repo):
     """Test deleting a relation that doesn't exist."""
     result = await relation_repo.delete_by_fields(relation_type="nonexistent")
-    assert result is False
+    assert result is False