diff --git a/basic-memory.bundle b/basic-memory.bundle new file mode 100644 index 00000000..651122eb Binary files /dev/null and b/basic-memory.bundle differ diff --git a/db/schema.sql b/db/schema.sql index b980064b..4c545f3a 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -1,39 +1 @@ -CREATE TABLE IF NOT EXISTS "schema_migrations" (version varchar(128) primary key); -CREATE TABLE IF NOT EXISTS "entity" ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL, - entity_type TEXT NOT NULL, - description TEXT NULL, - created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP -); -CREATE UNIQUE INDEX idx_entity_type_name ON entity(entity_type, name); -CREATE TABLE IF NOT EXISTS "observation" ( - id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT, - entity_id VARCHAR NOT NULL, - content VARCHAR NOT NULL, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - context VARCHAR, - FOREIGN KEY(entity_id) REFERENCES entity (id) ON DELETE CASCADE -); -CREATE INDEX ix_observation_entity_id ON observation (entity_id); -CREATE TABLE IF NOT EXISTS "relation" ( - id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT, - from_id VARCHAR NOT NULL, - to_id VARCHAR NOT NULL, - relation_type VARCHAR NOT NULL, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - context VARCHAR, - FOREIGN KEY(from_id) REFERENCES entity (id) ON DELETE CASCADE, - FOREIGN KEY(to_id) REFERENCES entity (id) ON DELETE CASCADE -); -CREATE INDEX ix_relation_from_id ON relation (from_id); -CREATE INDEX ix_relation_to_id ON relation (to_id); --- Dbmate schema migrations -INSERT INTO "schema_migrations" (version) VALUES - ('20240101000000'), - ('20241210213454'), - ('20241211034719'), - ('20241211052101'), - ('20241211190000'), - ('20241213022126'); +gi \ No newline at end of file diff --git a/reponame.bundle b/reponame.bundle new file mode 100644 index 00000000..16f027d5 Binary files /dev/null and b/reponame.bundle differ diff --git a/src/basic_memory/api/app.py b/src/basic_memory/api/app.py index bd8b25cb..88317b03 100644 --- a/src/basic_memory/api/app.py +++ b/src/basic_memory/api/app.py @@ -5,17 +5,15 @@ from loguru import logger from .routers import knowledge - # Initialize FastAPI app app = FastAPI( - title="Basic Memory API", - description="Knowledge graph API for basic-memory", - version="0.1.0" + title="Basic Memory API", description="Knowledge graph API for basic-memory", version="0.1.0" ) # Include routers app.include_router(knowledge.router) + # Add startup event @app.on_event("startup") async def startup_event(): diff --git a/src/basic_memory/api/routers/__init__.py b/src/basic_memory/api/routers/__init__.py index 02ddf061..8b10c988 100644 --- a/src/basic_memory/api/routers/__init__.py +++ b/src/basic_memory/api/routers/__init__.py @@ -1,4 +1,5 @@ """API routers.""" + from . import knowledge -__all__ = ["knowledge"] \ No newline at end of file +__all__ = ["knowledge"] diff --git a/src/basic_memory/api/routers/knowledge.py b/src/basic_memory/api/routers/knowledge.py index e9f04d5d..4eef2d27 100644 --- a/src/basic_memory/api/routers/knowledge.py +++ b/src/basic_memory/api/routers/knowledge.py @@ -1,8 +1,9 @@ """Router for knowledge graph operations.""" from fastapi import APIRouter, HTTPException +from loguru import logger -from basic_memory.deps import MemoryServiceDep +from basic_memory.deps import EntityServiceDep, RelationServiceDep, ObservationServiceDep from basic_memory.fileio import EntityNotFoundError from basic_memory.schemas import ( CreateEntityRequest, @@ -24,7 +25,6 @@ from basic_memory.schemas import ( AddObservationsResponse, RelationResponse, DeleteEntityRequest, - Entity, ) router = APIRouter(prefix="/knowledge", tags=["knowledge"]) @@ -34,10 +34,10 @@ router = APIRouter(prefix="/knowledge", tags=["knowledge"]) @router.post("/entities", response_model=CreateEntityResponse) async def create_entities( - data: CreateEntityRequest, memory_service: MemoryServiceDep + data: CreateEntityRequest, entity_service: EntityServiceDep ) -> CreateEntityResponse: """Create new entities in the knowledge graph.""" - entities = await memory_service.create_entities(data.entities) + entities = await entity_service.create_entities(data.entities) return CreateEntityResponse( entities=[EntityResponse.model_validate(entity) for entity in entities] ) @@ -45,10 +45,10 @@ async def create_entities( @router.post("/relations", response_model=CreateRelationsResponse) async def create_relations( - data: CreateRelationsRequest, memory_service: MemoryServiceDep + data: CreateRelationsRequest, relation_service: RelationServiceDep ) -> CreateRelationsResponse: """Create relations between entities.""" - relations = await memory_service.create_relations(data.relations) + relations = await relation_service.create_relations(data.relations) return CreateRelationsResponse( relations=[RelationResponse.model_validate(relation) for relation in relations] ) @@ -56,10 +56,11 @@ async def create_relations( @router.post("/observations", response_model=AddObservationsResponse) async def add_observations( - data: AddObservationsRequest, memory_service: MemoryServiceDep + data: AddObservationsRequest, observation_service: ObservationServiceDep ) -> AddObservationsResponse: """Add observations to an entity.""" - observations = await memory_service.add_observations(data) + logger.debug(f"Adding observations to entity: {data.entity_id}") + observations = await observation_service.add_observations(data.entity_id, data.observations) return AddObservationsResponse( entity_id=data.entity_id, observations=[ @@ -72,10 +73,10 @@ async def add_observations( @router.get("/entities/{entity_id:path}", response_model=EntityResponse) -async def get_entity(entity_id: str, memory_service: MemoryServiceDep) -> EntityResponse: +async def get_entity(entity_id: str, entity_service: EntityServiceDep) -> EntityResponse: """Get a specific entity by ID.""" try: - entity = await memory_service.get_entity(entity_id) + entity = await entity_service.get_entity(entity_id) return EntityResponse.model_validate(entity) except EntityNotFoundError: raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found") @@ -83,20 +84,25 @@ async def get_entity(entity_id: str, memory_service: MemoryServiceDep) -> Entity @router.post("/search", response_model=SearchNodesResponse) async def search_nodes( - data: SearchNodesRequest, memory_service: MemoryServiceDep + data: SearchNodesRequest, entity_service: EntityServiceDep ) -> SearchNodesResponse: """Search for entities in the knowledge graph.""" - matches = await memory_service.search_nodes(data.query) + logger.debug(f"Searching nodes with query: {data.query}") + matches = await entity_service.search(data.query) + logger.debug(f"Found {len(matches)} matches for '{data.query}'") + return SearchNodesResponse( matches=[EntityResponse.model_validate(entity) for entity in matches], query=data.query ) @router.post("/nodes", response_model=OpenNodesResponse) -async def open_nodes(data: OpenNodesRequest, memory_service: MemoryServiceDep) -> OpenNodesResponse: +async def open_nodes(data: OpenNodesRequest, entity_service: EntityServiceDep) -> OpenNodesResponse: """Open specific nodes by their names.""" - entities = await memory_service.open_nodes(data.entity_ids) - return OpenNodesResponse(entities=[Entity.model_validate(entity) for entity in entities]) + entities = await entity_service.open_nodes(data.entity_ids) + return OpenNodesResponse( + entities=[EntityResponse.model_validate(entity) for entity in entities] + ) ## Delete endpoints @@ -104,26 +110,26 @@ async def open_nodes(data: OpenNodesRequest, memory_service: MemoryServiceDep) - @router.post("/entities/delete", response_model=DeleteEntityResponse) async def delete_entity( - data: DeleteEntityRequest, memory_service: MemoryServiceDep + data: DeleteEntityRequest, entity_service: EntityServiceDep ) -> DeleteEntityResponse: """Delete a specific entity by ID.""" - deleted = await memory_service.delete_entities(data.entity_ids) + deleted = await entity_service.delete_entities(data.entity_ids) return DeleteEntityResponse(deleted=deleted) @router.post("/observations/delete", response_model=DeleteObservationsResponse) async def delete_observations( - data: DeleteObservationsRequest, memory_service: MemoryServiceDep + data: DeleteObservationsRequest, observation_service: ObservationServiceDep ) -> DeleteObservationsResponse: """Delete observations from an entity.""" entity_id = data.entity_id - deleted = await memory_service.delete_observations(entity_id, data.deletions) + deleted = await observation_service.delete_observations(entity_id, data.deletions) return DeleteObservationsResponse(deleted=deleted) @router.post("/relations/delete", response_model=DeleteRelationsResponse) async def delete_relations( - data: DeleteRelationsRequest, memory_service: MemoryServiceDep + data: DeleteRelationsRequest, relation_service: RelationServiceDep ) -> DeleteRelationsResponse: """Delete relations between entities.""" to_delete = [ @@ -134,5 +140,5 @@ async def delete_relations( } for relation in data.relations ] - deleted = await memory_service.delete_relations(to_delete) + deleted = await relation_service.delete_relations(to_delete) return DeleteRelationsResponse(deleted=deleted) diff --git a/src/basic_memory/db.py b/src/basic_memory/db.py index f4bb8abd..09c0de7c 100644 --- a/src/basic_memory/db.py +++ b/src/basic_memory/db.py @@ -1,128 +1,87 @@ -"""Database configuration and initialization for basic-memory.""" - +import asyncio from contextlib import asynccontextmanager -from enum import Enum +from enum import Enum, auto from pathlib import Path from typing import AsyncGenerator -from loguru import logger from sqlalchemy import text from sqlalchemy.ext.asyncio import ( create_async_engine, async_sessionmaker, - AsyncEngine, AsyncSession, + AsyncEngine, + async_scoped_session, ) -from sqlalchemy.pool import StaticPool - -from basic_memory.models import Base class DatabaseType(Enum): - """Types of database configurations.""" + """Types of supported databases.""" - MEMORY = "memory" # In-memory SQLite for testing - FILESYSTEM = "file" # File-based SQLite for projects + MEMORY = auto() + FILESYSTEM = auto() + + @classmethod + def get_db_path(cls, project_path: Path, db_type: "DatabaseType") -> Path: + """Get database path based on type.""" + if db_type == cls.MEMORY: + return Path(":memory:") + else: + return project_path / "data" / "memory.db" + + @classmethod + def get_db_url(cls, db_path: Path) -> str: + """Get SQLAlchemy URL for database path.""" + if str(db_path) == ":memory:": + return "sqlite+aiosqlite://" + return f"sqlite+aiosqlite:///{db_path}" -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}" +def get_scoped_session_factory( + session_maker: async_sessionmaker[AsyncSession], +) -> async_scoped_session: + """Create a scoped session factory scoped to current task.""" + return async_scoped_session(session_maker, scopefunc=asyncio.current_task) -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, - ) - else: - 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) - - return engine, session_factory +# @asynccontextmanager +# 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 +# """ +# session = session_factory() +# try: +# await session.execute(text("PRAGMA foreign_keys=ON")) +# yield session +# await session.commit() +# except Exception: +# await session.rollback() +# raise +# finally: +# await session.close() +# @asynccontextmanager -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) - logger.debug(f"engine url: {engine.url}") - try: - yield engine, session_factory - finally: - await engine.dispose() - - -@asynccontextmanager -async def session( - session_factory: async_sessionmaker[AsyncSession], +async def scoped_session( + session_maker: async_sessionmaker[AsyncSession], ) -> AsyncGenerator[AsyncSession, None]: """ - Get database session with proper lifecycle management. + Get a scoped session with proper lifecycle management. Args: - session_factory: Async session factory to create session from - - Yields: - AsyncSession configured for engine + session_maker: Session maker to create scoped sessions from """ - # Create and yield session - session = session_factory() + factory = get_scoped_session_factory(session_maker) + session = factory() try: - # Ensure foreign keys enabled for this session await session.execute(text("PRAGMA foreign_keys=ON")) yield session await session.commit() @@ -131,13 +90,24 @@ async def session( raise finally: await session.close() + await factory.remove() -async def dispose_database(engine: AsyncEngine): - """ - Clean up database engine. +@asynccontextmanager +async def engine_session_factory( + project_path: Path, + db_type: DatabaseType = DatabaseType.FILESYSTEM, +) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]: + """Create engine and session factory.""" + db_path = DatabaseType.get_db_path(project_path, db_type) + db_url = DatabaseType.get_db_url(db_path) + engine = create_async_engine(db_url, connect_args={"check_same_thread": False}) + try: + factory = async_sessionmaker(engine, expire_on_commit=False) + async with scoped_session(factory) as db_session: + # Initialize database + await db_session.execute(text("PRAGMA foreign_keys=ON")) - Args: - engine: Engine to dispose - """ - await engine.dispose() + yield engine, factory + finally: + await engine.dispose() diff --git a/src/basic_memory/debug_utils.py b/src/basic_memory/debug_utils.py deleted file mode 100644 index 463eea25..00000000 --- a/src/basic_memory/debug_utils.py +++ /dev/null @@ -1,118 +0,0 @@ -"""Debugging utilities for test suite.""" -from sqlalchemy import text, inspect, MetaData, Table -from sqlalchemy.ext.asyncio import AsyncEngine -from loguru import logger - -async def dump_sqlite_master(conn) -> None: - """Dump entire sqlite_master table to see what SQLite thinks exists.""" - try: - # Check if sqlite_master exists first - result = await conn.execute(text( - "SELECT COUNT(*) FROM sqlite_master WHERE type='table';" - )) - if result.scalar(): - result = await conn.execute(text("SELECT * FROM sqlite_master;")) - rows = result.all() - logger.info("SQLITE_MASTER TABLE CONTENTS:") - for row in rows: - logger.info(f"{row}") - else: - logger.info("No tables exist in sqlite_master") - except Exception as e: - logger.warning(f"Could not dump sqlite_master: {e}") - -async def dump_table_schema(conn, table_name: str) -> None: - """Dump detailed schema info for a specific table.""" - try: - logger.info(f"\nDETAILED SCHEMA FOR {table_name}:") - - # Get CREATE statement - result = await conn.execute( - text("SELECT sql FROM sqlite_master WHERE type='table' AND name=:name"), - {"name": table_name} - ) - create_sql = result.scalar() - if create_sql: - logger.info(f"CREATE statement:\n{create_sql}") - else: - logger.info(f"No CREATE statement found for {table_name}") - - # Get column info - result = await conn.execute(text(f"PRAGMA table_info('{table_name}');")) - columns = result.all() - if columns: - logger.info("Column definitions:") - for col in columns: - logger.info(f" {col}") - else: - logger.info("No column definitions found") - - except Exception as e: - logger.warning(f"Could not dump schema for {table_name}: {e}") - -async def dump_sqlalchemy_metadata(engine: AsyncEngine) -> None: - """Dump SQLAlchemy's view of the table metadata.""" - try: - inspector = inspect(engine) - - logger.info("\nSQLALCHEMY METADATA:") - # Get all tables - tables = await inspector.get_table_names() - if not tables: - logger.info("No tables found in SQLAlchemy metadata") - return - - for table in tables: - logger.info(f"\nTable: {table}") - # Get columns - columns = await inspector.get_columns(table) - if columns: - logger.info("Columns:") - for col in columns: - logger.info(f" {col}") - - # Get indexes - indexes = await inspector.get_indexes(table) - if indexes: - logger.info("Indexes:") - for idx in indexes: - logger.info(f" {idx}") - - # Get foreign keys - fks = await inspector.get_foreign_keys(table) - if fks: - logger.info("Foreign keys:") - for fk in fks: - logger.info(f" {fk}") - except Exception as e: - logger.warning(f"Could not dump SQLAlchemy metadata: {e}") - -async def dump_db_state(engine: AsyncEngine) -> None: - """Dump complete database state for debugging.""" - logger.info("\n=== BEGINNING DATABASE STATE DUMP ===\n") - - try: - async with engine.begin() as conn: - # First dump sqlite_master - await dump_sqlite_master(conn) - - # Try to get list of tables - try: - result = await conn.execute(text( - "SELECT name FROM sqlite_master WHERE type='table';" - )) - tables = [row[0] for row in result.all()] - - # Dump each table's schema - for table in tables: - await dump_table_schema(conn, table) - except Exception as e: - logger.warning(f"Could not list tables: {e}") - - # Dump SQLAlchemy's view - await dump_sqlalchemy_metadata(engine) - - except Exception as e: - logger.warning(f"Error during state dump: {e}") - - logger.info("\n=== END DATABASE STATE DUMP ===\n") \ No newline at end of file diff --git a/src/basic_memory/deps.py b/src/basic_memory/deps.py index 29883dd1..2899e589 100644 --- a/src/basic_memory/deps.py +++ b/src/basic_memory/deps.py @@ -1,109 +1,113 @@ """Dependency injection functions for basic-memory services.""" -from contextlib import asynccontextmanager + from pathlib import Path -from typing import AsyncGenerator, Annotated +from typing import Annotated, AsyncGenerator from fastapi import Depends -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession, AsyncEngine, async_sessionmaker +from sqlalchemy.ext.asyncio import ( + AsyncSession, + AsyncEngine, + async_sessionmaker, +) +from basic_memory import db from basic_memory.config import ProjectConfig, config +from basic_memory.db import DatabaseType from basic_memory.repository.entity_repository import EntityRepository from basic_memory.repository.observation_repository import ObservationRepository from basic_memory.repository.relation_repository import RelationRepository -from basic_memory.services import EntityService, ObservationService, RelationService, MemoryService -from basic_memory import db +from basic_memory.services import EntityService, ObservationService, RelationService + def get_project_config() -> ProjectConfig: return config + + ProjectConfigDep = Annotated[ProjectConfig, Depends(get_project_config)] + def get_project_path(project_config: ProjectConfigDep) -> Path: return Path(project_config.path) + + ProjectPathDep = Annotated[Path, Depends(get_project_path)] -async def get_engine_factory(project_path: ProjectPathDep, db_type=db.DatabaseType.FILESYSTEM) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]: - async with db.engine_session_factory(project_path, db_type) as (engine, factory): - yield engine, factory - -EngineFactoryDep = Annotated[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], Depends(get_engine_factory)] - -async def get_session(engine_factory: EngineFactoryDep) -> AsyncGenerator[AsyncSession, None]: - _, factory = engine_factory - async with db.session(factory) as session: - yield session - -AsyncSessionDep = Annotated[AsyncSession, Depends(get_session)] +async def get_engine_factory( + project_path: ProjectPathDep, db_type=DatabaseType.FILESYSTEM +) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]: + async with db.engine_session_factory(project_path=project_path, db_type=db_type) as ( + engine, + session_maker, + ): + yield engine, session_maker -async def get_entity_repo(session: AsyncSessionDep) -> EntityRepository: - """Get an EntityRepository instance.""" - return EntityRepository(session) # Entity type is handled in EntityRepository.__init__ +EngineFactoryDep = Annotated[ + tuple[AsyncEngine, async_sessionmaker[AsyncSession]], Depends(get_engine_factory) +] -EntityRepositoryDep = Annotated[EntityRepository, Depends(get_entity_repo)] -async def get_observation_repo(session: AsyncSessionDep) -> ObservationRepository: - """Get an ObservationRepository instance.""" - return ObservationRepository(session) +async def get_session_maker(engine_factory: EngineFactoryDep) -> async_sessionmaker[AsyncSession]: + """Get session maker for tests.""" + _, session_maker = engine_factory + return session_maker -ObservationRepositoryDep = Annotated[ObservationRepository, Depends(get_observation_repo)] -async def get_relation_repo(session: AsyncSessionDep) -> RelationRepository: - """Get a RelationRepository instance.""" - return RelationRepository(session) +SessionMakerDep = Annotated[async_sessionmaker, Depends(get_session_maker)] -RelationRepositoryDep = Annotated[RelationRepository, Depends(get_relation_repo)] -async def get_entity_service( - project_path: ProjectPathDep, - entity_repo: EntityRepositoryDep -) -> EntityService: - """Get an EntityService instance.""" - return EntityService(project_path, entity_repo) +async def get_entity_repository( + session_maker: SessionMakerDep, +) -> EntityRepository: + """Create an EntityRepository instance.""" + return EntityRepository(session_maker) + + +EntityRepositoryDep = Annotated[EntityRepository, Depends(get_entity_repository)] + + +async def get_observation_repository( + session_maker: SessionMakerDep, +) -> ObservationRepository: + """Create an ObservationRepository instance.""" + return ObservationRepository(session_maker) + + +ObservationRepositoryDep = Annotated[ObservationRepository, Depends(get_observation_repository)] + + +async def get_relation_repository( + session_maker: SessionMakerDep, +) -> RelationRepository: + """Create a RelationRepository instance.""" + return RelationRepository(session_maker) + + +RelationRepositoryDep = Annotated[RelationRepository, Depends(get_relation_repository)] + + +async def get_entity_service(entity_repository: EntityRepositoryDep) -> EntityService: + """Create EntityService with repository.""" + return EntityService(entity_repository) + EntityServiceDep = Annotated[EntityService, Depends(get_entity_service)] + async def get_observation_service( - project_path: ProjectPathDep, - observation_repo: ObservationRepositoryDep + observation_repository: ObservationRepositoryDep, ) -> ObservationService: - """Get an ObservationService instance.""" - return ObservationService(project_path, observation_repo) + """Create ObservationService with repository.""" + return ObservationService(observation_repository) + ObservationServiceDep = Annotated[ObservationService, Depends(get_observation_service)] -async def get_relation_service( - project_path: ProjectPathDep, - relation_repo: RelationRepositoryDep -) -> RelationService: - """Get a RelationService instance.""" - return RelationService(project_path, relation_repo) + +async def get_relation_service(relation_repository: RelationRepositoryDep) -> RelationService: + """Create RelationService with repository.""" + return RelationService(relation_repository) + RelationServiceDep = Annotated[RelationService, Depends(get_relation_service)] - -@asynccontextmanager -async def memory_service( - project_path: ProjectPathDep, - entity_service: EntityServiceDep, - relation_service: RelationServiceDep, - observation_service: ObservationServiceDep -) -> AsyncGenerator[MemoryService, None]: - """Get a fully configured MemoryService instance.""" - yield MemoryService( - project_path=project_path, - entity_service=entity_service, - relation_service=relation_service, - observation_service=observation_service - ) - -async def get_memory_service( - project_path: ProjectPathDep, - entity_service: EntityServiceDep, - relation_service: RelationServiceDep, - observation_service: ObservationServiceDep -) -> AsyncGenerator[MemoryService, None]: - async with memory_service(project_path, entity_service, relation_service, observation_service) as service: - yield service - -MemoryServiceDep = Annotated[MemoryService, Depends(get_memory_service)] \ No newline at end of file diff --git a/src/basic_memory/repository/__init__.py b/src/basic_memory/repository/__init__.py index b1fc32c4..e69de29b 100644 --- a/src/basic_memory/repository/__init__.py +++ b/src/basic_memory/repository/__init__.py @@ -1,235 +0,0 @@ -"""Base repository implementation.""" -from typing import Type, Optional, Any, Sequence, TypeVar, List, Dict -from sqlalchemy import select, func, Select, Executable, inspect, Result, Column, insert, and_, delete -from sqlalchemy.exc import NoResultFound -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import Mapped -from loguru import logger - -from basic_memory.models import Base, Entity - -T = TypeVar('T', bound=Base) - -class Repository[T: Base]: - """Base repository implementation with generic CRUD operations.""" - - def __init__(self, session: AsyncSession, Model: Type[T]): - self.session = session - self.Model = Model - self.mapper = inspect(self.Model).mapper - self.primary_key: Column[Any] = self.mapper.primary_key[0] - self.valid_columns = [column.key for column in self.mapper.columns] - - def select(self, *entities: Any) -> Select: - """Create a new SELECT statement. - - Returns: - A SQLAlchemy Select object configured with the provided entities - or this repository's model if no entities provided. - """ - if not entities: - entities = (self.Model,) - return select(*entities) - - async def refresh(self, instance: T, relationships: list[str] | None = None) -> None: - """Refresh instance and optionally specified relationships.""" - logger.debug(f"Refreshing {self.Model.__name__} instance: {getattr(instance, 'id', None)}") - try: - await self.session.refresh(instance, relationships or []) - logger.debug(f"Refreshed relationships: {relationships}") - except Exception: - logger.exception(f"Failed to refresh {self.Model.__name__} instance") - raise - - async def find_all(self, skip: int = 0, limit: int = 100) -> Sequence[T]: - """Fetch records from the database with pagination.""" - logger.debug(f"Finding all {self.Model.__name__} (skip={skip}, limit={limit})") - try: - result = await self.session.execute( - select(self.Model).offset(skip).limit(limit) - ) - items = result.scalars().all() - logger.debug(f"Found {len(items)} {self.Model.__name__} records") - return items - except Exception: - logger.exception(f"Failed to find all {self.Model.__name__}") - raise - - async def find_by_id(self, entity_id: str) -> Optional[T]: - """Fetch an entity by its unique identifier.""" - logger.debug(f"Finding {self.Model.__name__} by ID: {entity_id}") - try: - result = await self.session.execute( - select(self.Model).filter(self.primary_key == entity_id) - ) - entity = result.scalars().one() - logger.debug(f"Found {self.Model.__name__}: {entity_id}") - return entity - except NoResultFound: - logger.debug(f"No {self.Model.__name__} found with ID: {entity_id}") - return None - except Exception: - logger.exception(f"Failed to find {self.Model.__name__} by ID: {entity_id}") - raise - - async def create(self, entity_data: dict, model: Type[Base] | None = None) -> T: - """Create a new entity in the database from the provided data.""" - model = model or self.Model - logger.debug(f"Creating {model.__name__} with data: {entity_data}") - try: - # Only include valid columns that are provided in entity_data - model_data = { - k: v for k, v in entity_data.items() - if k in self.valid_columns and v is not None - } - - # Generate ID if this is an Entity model and no ID provided - if model is Entity and 'id' not in model_data: - model_data['id'] = Entity.generate_id( - model_data['entity_type'], - model_data['name'] - ) - - logger.debug(f"Filtered data for valid columns: {model_data}") - - # Create insert statement with only provided data - stmt = insert(model).values(**model_data).returning(model) - result = await self.session.execute(stmt) - entity: T = result.scalar_one() # pyright: ignore [reportAssignmentType] - - logger.debug(f"Created {model.__name__}: {getattr(entity, 'id', None)}") - return entity - - except Exception: - logger.exception(f"Failed to create {model.__name__}") - raise - - async def instance_create(self, instance: T) -> T: - """Create a new record from a model instance.""" - logger.debug(f"Creating {self.Model.__name__} from instance: {instance}") - try: - self.session.add(instance) - await self.session.flush() - return instance - except Exception: - logger.exception(f"Failed to create {self.Model.__name__}") - raise - - async def bulk_create(self, instances: List[T]) -> List[T]: - """Create multiple records in a single transaction.""" - logger.debug(f"Bulk creating {len(instances)} {self.Model.__name__} instances") - try: - for instance in instances: - self.session.add(instance) - await self.session.flush() - return instances - except Exception: - logger.exception(f"Failed to bulk create {self.Model.__name__}") - raise - - async def update(self, entity_id: str, entity_data: dict) -> Optional[T]: - """Update an entity with the given data.""" - logger.debug(f"Updating {self.Model.__name__} {entity_id} with data: {entity_data}") - try: - result = await self.session.execute( - select(self.Model).filter(self.primary_key == entity_id) - ) - entity = result.scalars().one() - - for key, value in entity_data.items(): - if key in self.valid_columns: - setattr(entity, key, value) - await self.session.flush() - - logger.debug(f"Updated {self.Model.__name__}: {entity_id}") - return entity - except NoResultFound: - logger.debug(f"No {self.Model.__name__} found to update: {entity_id}") - return None - except Exception: - logger.exception(f"Failed to update {self.Model.__name__}: {entity_id}") - raise - - async def delete(self, entity_id: str) -> bool: - """Delete an entity from the database.""" - logger.debug(f"Deleting {self.Model.__name__}: {entity_id}") - try: - result = await self.session.execute( - select(self.Model).filter(self.primary_key == entity_id) - ) - entity = result.scalars().one() - await self.session.delete(entity) - await self.session.flush() - - logger.debug(f"Deleted {self.Model.__name__}: {entity_id}") - return True - except NoResultFound: - logger.debug(f"No {self.Model.__name__} found to delete: {entity_id}") - return False - except Exception: - logger.exception(f"Failed to delete {self.Model.__name__}: {entity_id}") - raise - - async def delete_by_fields(self, **filters: Dict[str, Any]) -> bool: - """ - Delete records matching given field values. - - Args: - **filters: Field names and values to filter by - - Returns: - bool: True if any records were deleted - """ - logger.debug(f"Deleting {self.Model.__name__} by fields: {filters}") - try: - conditions = [getattr(self.Model, field) == value for field, value in filters.items()] - query = delete(self.Model).where(and_(*conditions)) - result = await self.execute_query(query) - await self.session.flush() - deleted = result.rowcount > 0 # pyright: ignore [reportAttributeAccessIssue] - logger.debug(f"Deleted {result.rowcount} records") # pyright: ignore [reportAttributeAccessIssue] - return deleted # pyright: ignore [reportAttributeAccessIssue] - except Exception: - logger.exception(f"Failed to delete {self.Model.__name__} by fields") - raise - - async def count(self, query: Executable | None = None) -> int: - """Count entities in the database table.""" - try: - if query is None: - query = select(func.count()).select_from(self.Model) - result = await self.session.execute(query) - scalar = result.scalar() - count = scalar if scalar is not None else 0 - - logger.debug(f"Counted {count} {self.Model.__name__} records") - return count - except Exception: - logger.exception(f"Failed to count {self.Model.__name__}") - raise - - async def execute_query(self, query: Executable) -> Result[Any]: - """Execute a query asynchronously.""" - logger.debug(f"Executing query: {query}") - try: - result = await self.session.execute(query) - logger.debug("Query executed successfully") - return result - except Exception: - logger.exception("Failed to execute query") - raise - - async def find_one(self, query: Select[tuple[T]]) -> Optional[T]: - """Execute a query and retrieve a single record.""" - logger.debug(f"Finding one {self.Model.__name__} with query: {query}") - try: - result = await self.execute_query(query) - entity = result.scalars().one_or_none() - if entity: - logger.debug(f"Found {self.Model.__name__}: {getattr(entity, 'id', None)}") - else: - logger.debug(f"No {self.Model.__name__} found") - return entity - except Exception: - logger.exception(f"Failed to find one {self.Model.__name__}") - raise \ No newline at end of file diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index 0d87cf9d..c4e6d3e4 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -1,132 +1,165 @@ """Repository for managing Entity objects.""" -from typing import Optional, Sequence, Type + +from typing import Optional, Sequence, List + +from loguru import logger from sqlalchemy import select, or_ from sqlalchemy.exc import NoResultFound +from sqlalchemy.ext.asyncio import async_sessionmaker from sqlalchemy.orm import selectinload +from basic_memory import db from basic_memory.models import Entity, Observation -from basic_memory.repository import Repository -from loguru import logger +from basic_memory.repository.repository import Repository class EntityRepository(Repository[Entity]): """Repository for Entity model with memory-specific operations.""" - - def __init__(self, session): - super().__init__(session, Entity) + + def __init__(self, session_maker: async_sessionmaker): + super().__init__(session_maker, Entity) logger.debug("Initialized EntityRepository") - async def create(self, entity_data: dict, model: Type[Entity] | None = None) -> Entity: + async def create(self, data: dict) -> Entity: # pyright: ignore [reportIncompatibleMethodOverride] """Create a new entity in the database from the provided data.""" - entity_id = Entity.generate_id( - entity_data['entity_type'], - entity_data['name'] - ) - return await super().create({**entity_data, 'id': entity_id}) + entity_id = Entity.generate_id(data["entity_type"], data["name"]) + await super().create({**data, "id": entity_id}) + + # we have to find to get relations + created = await self.find_by_id(entity_id) + assert created is not None, f"Created entity {entity_id} should not be None" + return created + + async def create_all(self, data_list: List[dict]) -> Sequence[Entity]: # pyright: ignore [reportIncompatibleMethodOverride] + """Create a new entity in the database from the provided data.""" + for data in data_list: + entity_id = Entity.generate_id(data["entity_type"], data["name"]) + data["id"] = entity_id + created = await super().create_all(data_list) + # we have to find to get relations + return await self.find_by_ids([e.id for e in created]) async def find_by_id(self, entity_id: str) -> Optional[Entity]: """Find entity by ID with all relationships eagerly loaded.""" logger.debug(f"Finding entity by ID: {entity_id}") - try: - # First load base entity - result = await self.session.execute( - select(Entity).filter(Entity.id == entity_id) + async with db.scoped_session(self.session_maker) as session: + try: + result = await session.execute( + select(Entity) + .filter(Entity.id == entity_id) + .options( + selectinload(Entity.observations), + selectinload(Entity.outgoing_relations), + selectinload(Entity.incoming_relations), + ) + ) + entity = result.scalars().one() + logger.debug(f"Found entity: {entity.id}") + return entity + except NoResultFound: + logger.debug(f"No entity found with ID: {entity_id}") + return None + + async def find_by_ids(self, ids: List[str]) -> Sequence[Entity]: + """Search for entities of a specific type.""" + logger.debug(f"Find entities by ids: {ids}") + async with db.scoped_session(self.session_maker) as session: + result = await session.execute( + select(Entity) + .where(self.primary_key.in_(ids)) + .options( + selectinload(Entity.observations), + selectinload(Entity.outgoing_relations), + selectinload(Entity.incoming_relations), + ) ) - entity = result.scalars().one() - logger.debug(f"Found base entity: {entity.id}") - - # Force refresh of all relationships - await self.refresh(entity, ['observations', 'outgoing_relations', 'incoming_relations']) - logger.debug(f"Refreshed entity relationships: {entity.id}") - - return entity - except NoResultFound: - logger.debug(f"No entity found with ID: {entity_id}") - return None - except Exception as e: - logger.exception(f"Error finding entity by ID: {entity_id}") - raise + entities = result.scalars().all() + logger.debug(f"Found {len(entities)}") + return entities async def find_by_name(self, name: str) -> Optional[Entity]: """Find an entity by its unique name.""" logger.debug(f"Finding entity by name: {name}") - try: - query = ( + async with db.scoped_session(self.session_maker) as session: + result = await session.execute( select(Entity) .filter(Entity.name == name) + .options( + selectinload(Entity.observations), + selectinload(Entity.outgoing_relations), + selectinload(Entity.incoming_relations), + ) ) - result = await self.session.execute(query) entity = result.scalars().one_or_none() if entity: logger.debug(f"Found entity: {entity.id}") - await self.refresh(entity, ['observations', 'outgoing_relations', 'incoming_relations']) - logger.debug(f"Refreshed entity relationships: {entity.id}") else: logger.debug(f"No entity found with name: {name}") return entity - except Exception as e: - logger.exception(f"Error finding entity by name: {name}") - raise async def find_by_type_and_name(self, entity_type: str, name: str) -> Optional[Entity]: """Find an entity by its type and name combination.""" logger.debug(f"Finding entity by type and name: {entity_type}/{name}") - try: - query = ( + async with db.scoped_session(self.session_maker) as session: + result = await session.execute( select(Entity) .filter(Entity.entity_type == entity_type) .filter(Entity.name == name) .options( selectinload(Entity.observations), selectinload(Entity.outgoing_relations), - selectinload(Entity.incoming_relations) + selectinload(Entity.incoming_relations), ) ) - result = await self.session.execute(query) entity = result.scalars().one_or_none() if entity: logger.debug(f"Found entity: {entity.id}") else: logger.debug(f"No entity found with type/name: {entity_type}/{name}") return entity - except Exception as e: - logger.exception(f"Error finding entity by type/name: {entity_type}/{name}") - raise - - async def search_by_type(self, entity_type: str, skip: int = 0, limit: int = 100) -> Sequence[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.""" logger.debug(f"Searching entities by type: {entity_type} (skip={skip}, limit={limit})") - try: - query = select(Entity).filter(Entity.entity_type == entity_type).offset(skip).limit(limit) - result = await self.execute_query(query) + async with db.scoped_session(self.session_maker) as session: + result = await session.execute( + select(Entity) + .filter(Entity.entity_type == entity_type) + .options( + selectinload(Entity.observations), + selectinload(Entity.outgoing_relations), + selectinload(Entity.incoming_relations), + ) + .offset(skip) + .limit(limit) + ) entities = result.scalars().all() logger.debug(f"Found {len(entities)} entities of type {entity_type}") return entities - except Exception as e: - logger.exception(f"Error searching entities by type: {entity_type}") - raise async def search(self, query: str) -> Sequence[Entity]: """Search entities using LIKE pattern matching.""" logger.debug(f"Searching entities with query: {query}") - try: - stmt = select(Entity).distinct().where( - or_( - Entity.name.ilike(f"%{query}%"), - Entity.entity_type.ilike(f"%{query}%"), - Entity.observations.any( - Observation.content.ilike(f"%{query}%") + async with db.scoped_session(self.session_maker) as session: + stmt = ( + select(Entity) + .distinct() + .where( + or_( + Entity.name.ilike(f"%{query}%"), + Entity.entity_type.ilike(f"%{query}%"), + Entity.observations.any(Observation.content.ilike(f"%{query}%")), ) ) - ).options( - selectinload(Entity.observations), - selectinload(Entity.outgoing_relations), - selectinload(Entity.incoming_relations) + .options( + selectinload(Entity.observations), + selectinload(Entity.outgoing_relations), + selectinload(Entity.incoming_relations), + ) ) - result = await self.session.execute(stmt) + result = await session.execute(stmt) entities = list(result.scalars()) logger.debug(f"Found {len(entities)} matching entities") return entities - except Exception as e: - logger.exception(f"Error searching entities: {query}") - raise \ No newline at end of file diff --git a/src/basic_memory/repository/observation_repository.py b/src/basic_memory/repository/observation_repository.py index 831829d7..a62a50dd 100644 --- a/src/basic_memory/repository/observation_repository.py +++ b/src/basic_memory/repository/observation_repository.py @@ -1,25 +1,30 @@ """Repository for managing Observation objects.""" + from typing import Sequence + +from loguru import logger from sqlalchemy import select +from sqlalchemy.ext.asyncio import async_sessionmaker from basic_memory.models import Observation -from basic_memory.repository import Repository +from basic_memory.repository.repository import Repository class ObservationRepository(Repository[Observation]): """Repository for Observation model with memory-specific operations.""" - - def __init__(self, session): - super().__init__(session, 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]: """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.""" query = select(Observation).filter(Observation.context == context) result = await self.execute_query(query) - return result.scalars().all() \ No newline at end of file + return result.scalars().all() diff --git a/src/basic_memory/repository/relation_repository.py b/src/basic_memory/repository/relation_repository.py index f17ccb5f..f7c2c3c9 100644 --- a/src/basic_memory/repository/relation_repository.py +++ b/src/basic_memory/repository/relation_repository.py @@ -1,16 +1,19 @@ """Repository for managing Relation objects.""" + from typing import Sequence + from sqlalchemy import select, and_ +from sqlalchemy.ext.asyncio import async_sessionmaker from basic_memory.models import Relation -from basic_memory.repository import Repository +from basic_memory.repository.repository import Repository class RelationRepository(Repository[Relation]): """Repository for Relation model with memory-specific operations.""" - - def __init__(self, session): - super().__init__(session, 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]: """Find all relations from a specific entity.""" @@ -20,17 +23,12 @@ class RelationRepository(Repository[Relation]): async def find_by_entities(self, from_id: str, to_id: str) -> Sequence[Relation]: """Find all relations between two entities.""" - query = select(Relation).filter( - and_( - Relation.from_id == from_id, - Relation.to_id == to_id - ) - ) + query = select(Relation).filter(and_(Relation.from_id == from_id, Relation.to_id == to_id)) result = await self.execute_query(query) return result.scalars().all() - + async def find_by_type(self, relation_type: str) -> Sequence[Relation]: """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() \ No newline at end of file + return result.scalars().all() diff --git a/src/basic_memory/repository/repository.py b/src/basic_memory/repository/repository.py new file mode 100644 index 00000000..855ce836 --- /dev/null +++ b/src/basic_memory/repository/repository.py @@ -0,0 +1,223 @@ +"""Base repository implementation.""" + +from typing import Type, Optional, Any, Sequence, TypeVar, List + +from loguru import logger +from sqlalchemy import ( + select, + func, + Select, + Executable, + inspect, + Result, + Column, + and_, + delete, +) +from sqlalchemy.exc import NoResultFound +from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession + +from basic_memory import db +from basic_memory.models import Base + +T = TypeVar("T", bound=Base) + + +class Repository[T: Base]: + """Base repository implementation with generic CRUD operations.""" + + def __init__(self, session_maker: async_sessionmaker[AsyncSession], Model: Type[T]): + self.session_maker = session_maker + self.Model = Model + self.mapper = inspect(self.Model).mapper + self.primary_key: Column[Any] = self.mapper.primary_key[0] + self.valid_columns = [column.key for column in self.mapper.columns] + + def get_model_data(self, entity_data): + model_data = { + k: v for k, v in entity_data.items() if k in self.valid_columns and v is not None + } + return model_data + + async def add(self, model: T) -> T: + """ + Add a model to the repository. This will also add related objects + :param model: the model to add + :return: the added model instance + """ + async with db.scoped_session(self.session_maker) as session: + session.add(model) + await session.flush() + found = await self.find_by_id(model.id) # pyright: ignore [reportAttributeAccessIssue] + assert found is not None, "can't find model after session.add" + return found + + async def add_all(self, models: List[T]) -> Sequence[T]: + """ + Add a list of models to the repository. This will also add related objects + :param model: the models to add + :return: the added models instances + """ + async with db.scoped_session(self.session_maker) as session: + session.add_all(models) + await session.flush() + # we have to find to get relations + return await self.find_by_ids([m.id for m in models]) # pyright: ignore [reportAttributeAccessIssue] + + def select(self, *entities: Any) -> Select: + """Create a new SELECT statement. + + Returns: + A SQLAlchemy Select object configured with the provided entities + or this repository's model if no entities provided. + """ + if not entities: + entities = (self.Model,) + return select(*entities) + + async def refresh(self, instance: T, relationships: list[str] | None = None) -> None: + """Refresh instance and optionally specified relationships.""" + logger.debug(f"Refreshing {self.Model.__name__} instance: {getattr(instance, 'id', None)}") + async with db.scoped_session(self.session_maker) as session: + await session.refresh(instance, relationships or []) + logger.debug(f"Refreshed relationships: {relationships}") + + async def find_all(self, skip: int = 0, limit: int = 100) -> Sequence[T]: + """Fetch records from the database with pagination.""" + logger.debug(f"Finding all {self.Model.__name__} (skip={skip}, limit={limit})") + async with db.scoped_session(self.session_maker) as session: + result = await session.execute(select(self.Model).offset(skip).limit(limit)) + items = result.scalars().all() + logger.debug(f"Found {len(items)} {self.Model.__name__} records") + return items + + async def find_by_id(self, entity_id: str) -> Optional[T]: + """Fetch an entity by its unique identifier.""" + logger.debug(f"Finding {self.Model.__name__} by ID: {entity_id}") + async with db.scoped_session(self.session_maker) as session: + try: + result = await session.execute( + select(self.Model).filter(self.primary_key == entity_id) + ) + entity = result.scalars().one() + logger.debug(f"Found {self.Model.__name__}: {entity_id}") + return entity + except NoResultFound: + logger.debug(f"No {self.Model.__name__} found with ID: {entity_id}") + return None + + async def find_one(self, query: Select[tuple[T]]) -> Optional[T]: + """Execute a query and retrieve a single record.""" + logger.debug(f"Finding one {self.Model.__name__} with query: {query}") + result = await self.execute_query(query) + entity = result.scalars().one_or_none() + if entity: + logger.debug(f"Found {self.Model.__name__}: {getattr(entity, 'id', None)}") + else: + logger.debug(f"No {self.Model.__name__} found") + return entity + + async def find_by_ids(self, ids: List[str]) -> Sequence[T]: + """Fetch multiple entities by their identifiers in a single query.""" + logger.debug(f"Finding {self.Model.__name__} by IDs: {ids}") + async with db.scoped_session(self.session_maker) as session: + result = await session.execute(select(self.Model).where(self.primary_key.in_(ids))) + entities = result.scalars().all() + logger.debug(f"Found {len(entities)} {self.Model.__name__} records") + return entities + + async def create(self, data: dict) -> T: + """Create a new record from a model instance.""" + logger.debug(f"Creating {self.Model.__name__} from entity_data: {data}") + async with db.scoped_session(self.session_maker) as session: + # Only include valid columns that are provided in entity_data + model_data = self.get_model_data(data) + model = self.Model(**model_data) + session.add(model) + await session.flush() + return model + + async def create_all(self, data_list: List[dict]) -> List[T]: + """Create multiple records in a single transaction.""" + logger.debug(f"Bulk creating {len(data_list)} {self.Model.__name__} instances") + async with db.scoped_session(self.session_maker) as session: + # Only include valid columns that are provided in entity_data + model_list = [self.Model(**self.get_model_data(d)) for d in data_list] + session.add_all(model_list) + return model_list + + async def update(self, entity_id: str, entity_data: dict) -> Optional[T]: + """Update an entity with the given data.""" + logger.debug(f"Updating {self.Model.__name__} {entity_id} with data: {entity_data}") + async with db.scoped_session(self.session_maker) as session: + try: + result = await session.execute( + select(self.Model).filter(self.primary_key == entity_id) + ) + entity = result.scalars().one() + + for key, value in entity_data.items(): + if key in self.valid_columns: + setattr(entity, key, value) + + logger.debug(f"Updated {self.Model.__name__}: {entity_id}") + return entity + except NoResultFound: + logger.debug(f"No {self.Model.__name__} found to update: {entity_id}") + return None + + async def delete(self, entity_id: str) -> bool: + """Delete an entity from the database.""" + logger.debug(f"Deleting {self.Model.__name__}: {entity_id}") + async with db.scoped_session(self.session_maker) as session: + try: + result = await session.execute( + select(self.Model).filter(self.primary_key == entity_id) + ) + entity = result.scalars().one() + await session.delete(entity) + + logger.debug(f"Deleted {self.Model.__name__}: {entity_id}") + return True + except NoResultFound: + logger.debug(f"No {self.Model.__name__} found to delete: {entity_id}") + return False + + async def delete_by_ids(self, ids: List[str]) -> int: + """Delete records matching given field values.""" + logger.debug(f"Deleting {self.Model.__name__} by ids: {ids}") + async with db.scoped_session(self.session_maker) as session: + query = delete(self.Model).where(self.primary_key.in_(ids)) + result = await session.execute(query) + logger.debug(f"Deleted {result.rowcount} records") + return result.rowcount + + async def delete_by_fields(self, **filters: Any) -> bool: + """Delete records matching given field values.""" + logger.debug(f"Deleting {self.Model.__name__} by fields: {filters}") + async with db.scoped_session(self.session_maker) as session: + conditions = [getattr(self.Model, field) == value for field, value in filters.items()] + query = delete(self.Model).where(and_(*conditions)) + result = await session.execute(query) + deleted = result.rowcount > 0 + logger.debug(f"Deleted {result.rowcount} records") + return deleted + + async def count(self, query: Executable | None = None) -> int: + """Count entities in the database table.""" + async with db.scoped_session(self.session_maker) as session: + if query is None: + query = select(func.count()).select_from(self.Model) + result = await session.execute(query) + scalar = result.scalar() + count = scalar if scalar is not None else 0 + logger.debug(f"Counted {count} {self.Model.__name__} records") + return count + + async def execute_query(self, query: Executable) -> Result[Any]: + """Execute a query asynchronously.""" + logger.debug(f"Executing query: {query}") + async with db.scoped_session(self.session_maker) as session: + result = await session.execute(query) + logger.debug("Query executed successfully") + return result diff --git a/src/basic_memory/schemas/base.py b/src/basic_memory/schemas/base.py index 752e51b0..8a1a5288 100644 --- a/src/basic_memory/schemas/base.py +++ b/src/basic_memory/schemas/base.py @@ -1,7 +1,7 @@ """Core pydantic models for basic-memory entities, observations, and relations. This module defines the foundational data structures for the knowledge graph system. -The graph consists of entities (nodes) connected by relations (edges), where each +The graph consists of entities (nodes) connected by relations (edges), where each entity can have multiple observations (facts) attached to it. Key Concepts: @@ -106,7 +106,7 @@ Names are normalized by: class Relation(BaseModel): """Represents a directed edge between entities in the knowledge graph. - + Relations are directed connections stored in active voice (e.g., "created", "depends_on"). The from_id represents the source or actor entity, while to_id represents the target or recipient entity. @@ -212,9 +212,8 @@ class Entity(BaseModel): entity_type: EntityType description: Optional[str] = None observations: List[Observation] = [] - relations: List[Relation] = [] @property def file_path(self) -> str: """The relative file path for this entity.""" - return f"{id}.md" \ No newline at end of file + return f"{id}.md" diff --git a/src/basic_memory/schemas/response.py b/src/basic_memory/schemas/response.py index 5a56f677..69ed5af6 100644 --- a/src/basic_memory/schemas/response.py +++ b/src/basic_memory/schemas/response.py @@ -15,12 +15,12 @@ from typing import List, Optional from pydantic import BaseModel, ConfigDict -from basic_memory.schemas.base import Observation, EntityId, Relation, Entity +from basic_memory.schemas.base import Observation, EntityId, Relation class SQLAlchemyModel(BaseModel): """Base class for models that read from SQLAlchemy attributes. - + This base class handles conversion of SQLAlchemy model attributes to Pydantic model fields. All response models extend this to ensure proper handling of database results. @@ -229,7 +229,7 @@ class OpenNodesResponse(SQLAlchemyModel): } """ - entities: List[Entity] + entities: List[EntityResponse] class AddObservationsResponse(SQLAlchemyModel): @@ -323,4 +323,4 @@ class DeleteObservationsResponse(SQLAlchemyModel): } """ - deleted: bool \ No newline at end of file + deleted: bool diff --git a/src/basic_memory/services/__init__.py b/src/basic_memory/services/__init__.py index cc9f43a2..b4a87349 100644 --- a/src/basic_memory/services/__init__.py +++ b/src/basic_memory/services/__init__.py @@ -1,31 +1,33 @@ """Service layer exceptions and imports.""" + class ServiceError(Exception): """Base exception for service errors""" + pass class DatabaseSyncError(ServiceError): """Raised when database sync fails""" + pass class RelationError(ServiceError): """Base exception for relation-specific errors""" + pass from .entity_service import EntityService from .observation_service import ObservationService from .relation_service import RelationService -from .memory_service import MemoryService __all__ = [ - 'ServiceError', - 'DatabaseSyncError', - 'RelationError', - 'EntityService', - 'ObservationService', - 'RelationService', - 'MemoryService', -] \ No newline at end of file + "ServiceError", + "DatabaseSyncError", + "RelationError", + "EntityService", + "ObservationService", + "RelationService", +] diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index e0c04dbd..6497b0f1 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -1,112 +1,90 @@ """Service for managing entities in the database.""" -from pathlib import Path -from typing import Dict, Any, Sequence -from basic_memory.repository.entity_repository import EntityRepository -from basic_memory.schemas import Entity -from basic_memory.models import Entity as EntityModel -from basic_memory.fileio import EntityNotFoundError +from typing import Dict, Any, Sequence, List + from loguru import logger +from basic_memory.fileio import EntityNotFoundError +from basic_memory.models import Entity as EntityModel, Observation +from basic_memory.repository.entity_repository import EntityRepository +from basic_memory.schemas import Entity as EntitySchema +from .service import BaseService -class EntityService: - """ - Service for managing entities in the database. - File operations are handled by MemoryService. - """ - - def __init__(self, project_path: Path, entity_repo: EntityRepository): - self.project_path = project_path - self.entity_repo = entity_repo - logger.debug(f"Initialized EntityService with path: {project_path}") + +def entity_model(entity): + model = EntityModel( + id=EntityModel.generate_id(entity.entity_type, entity.name), + name=entity.name, + entity_type=entity.entity_type, + description=entity.description, + observations=[Observation(content=observation) for observation in entity.observations], + ) + return model + + +class EntityService(BaseService[EntityRepository]): + """Service for managing entities in the database.""" + + def __init__(self, entity_repository: EntityRepository): + super().__init__(entity_repository) async def search(self, query: str) -> Sequence[EntityModel]: """Search entities using LIKE pattern matching.""" logger.debug(f"Searching entities with query: {query}") - try: - results = await self.entity_repo.search(query) - logger.debug(f"Found {len(results)} matches") - return results - except Exception: - logger.exception(f"Failed to search entities with query: {query}") - raise + return await self.repository.search(query) - async def create_entity(self, entity: Entity) -> EntityModel: + async def create_entity(self, entity: EntitySchema) -> EntityModel: """Create a new entity in the database.""" logger.debug(f"Creating entity in DB: {entity}") - try: - created_entity = await self.entity_repo.create(entity.model_dump()) - logger.debug(f"Created base entity: {created_entity.id}") + model = entity_model(entity) + return await self.repository.add(model) - await self.entity_repo.refresh(created_entity, ['observations', 'outgoing_relations', 'incoming_relations']) - logger.debug(f"Refreshed entity relationships: {created_entity.id}") - - return created_entity - except Exception: - logger.exception(f"Failed to create entity: {entity}") - raise + async def create_entities(self, entities_in: List[EntitySchema]) -> Sequence[EntityModel]: + """Create multiple entities with their observations.""" + logger.debug(f"Creating {len(entities_in)} entities") + created = await self.repository.add_all([entity_model(entity) for entity in entities_in]) + return created async def update_entity(self, entity_id: str, update_data: Dict[str, Any]) -> EntityModel: """Update an entity's fields.""" logger.debug(f"Updating entity {entity_id} with data: {update_data}") - try: - updated = await self.entity_repo.update(entity_id, update_data) - if not updated: - raise EntityNotFoundError(f"Entity not found: {entity_id}") - - logger.debug(f"Updated entity: {updated.id}") - return updated - except EntityNotFoundError: - raise - except Exception: - logger.exception(f"Failed to update entity: {entity_id}") - raise + updated = await self.repository.update(entity_id, update_data) + if not updated: + raise EntityNotFoundError(f"Entity not found: {entity_id}") + return updated async def get_entity(self, entity_id: str) -> EntityModel: """Get entity by ID.""" logger.debug(f"Getting entity by ID: {entity_id}") - try: - db_entity = await self.entity_repo.find_by_id(entity_id) - if not db_entity: - logger.error(f"Entity not found: {entity_id}") - raise EntityNotFoundError(f"Entity not found: {entity_id}") - - logger.debug(f"Found entity: {db_entity.id}") - return db_entity - except EntityNotFoundError: - raise - except Exception: - logger.exception(f"Failed to get entity: {entity_id}") - raise + db_entity = await self.repository.find_by_id(entity_id) + if not db_entity: + raise EntityNotFoundError(f"Entity not found: {entity_id}") + return db_entity async def get_by_type_and_name(self, entity_type: str, name: str) -> EntityModel: """Get entity by type and name combination.""" logger.debug(f"Getting entity by type/name: {entity_type}/{name}") - try: - db_entity = await self.entity_repo.find_by_type_and_name(entity_type, name) - if not db_entity: - logger.error(f"Entity not found: {entity_type}/{name}") - raise EntityNotFoundError(f"Entity not found: {entity_type}/{name}") - - logger.debug(f"Found entity: {db_entity.id}") - return db_entity - except EntityNotFoundError: - raise - except Exception: - logger.exception(f"Failed to get entity by type/name: {entity_type}/{name}") - raise + db_entity = await self.repository.find_by_type_and_name(entity_type, name) + if not db_entity: + raise EntityNotFoundError(f"Entity not found: {entity_type}/{name}") + return db_entity async def get_all(self) -> Sequence[EntityModel]: - return await self.entity_repo.find_all() + """Get all entities.""" + return await self.repository.find_all() async def delete_entity(self, entity_id: str) -> bool: """Delete entity from database.""" logger.debug(f"Deleting entity: {entity_id}") - try: - result = await self.entity_repo.delete(entity_id) - logger.debug(f"Entity deleted: {entity_id}") - return result - except Exception: - logger.exception(f"Failed to delete entity: {entity_id}") - raise + return await self.repository.delete(entity_id) + async def open_nodes(self, entity_ids: List[str]) -> Sequence[EntityModel]: + """Get specific nodes and their relationships.""" + logger.debug(f"Opening nodes entity_ids: {entity_ids}") + return await self.repository.find_by_ids(entity_ids) + + async def delete_entities(self, entity_ids: List[str]) -> bool: + """Delete entities and their files.""" + logger.debug(f"Deleting entities: {entity_ids}") + deleted_count = await self.repository.delete_by_ids(entity_ids) + return deleted_count > 0 diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py deleted file mode 100644 index 2364665f..00000000 --- a/src/basic_memory/services/memory_service.py +++ /dev/null @@ -1,371 +0,0 @@ -"""Service for orchestrating entity, relation, and observation operations.""" - -import asyncio -from pathlib import Path -from typing import List, Dict, Any, Optional, Sequence - -from loguru import logger - -from basic_memory.fileio import ( - write_entity_file, - read_entity_file, - EntityNotFoundError, - get_entity_path, -) -from basic_memory.models import ( - Entity as EntityModel, - Observation as ObservationModel, - Relation as RelationModel, -) -from basic_memory.schemas import AddObservationsRequest, Entity, Relation -from basic_memory.services import EntityService, RelationService, ObservationService - - -class MemoryService: - """Orchestrates entity, relation, and observation operations with filesystem handling.""" - - def __init__( - self, - project_path: Optional[Path], - entity_service: EntityService, - relation_service: RelationService, - observation_service: ObservationService, - ): - if project_path: - assert ( - project_path.is_dir() - ), "Path does not exist or is not a directory: {project_path}" - self.project_path = project_path - self.entities_path = project_path / "entities" - - self.entity_service = entity_service - self.relation_service = relation_service - self.observation_service = observation_service - logger.debug(f"Initialized MemoryService with path: {project_path}") - - def get_entity_file_path(self, entity_id: str) -> Path: - return get_entity_path(self.entities_path, entity_id) - - async def create_entities(self, entities_in: List[Entity]) -> List[EntityModel]: - """Create multiple entities with their observations.""" - logger.debug(f"Creating {len(entities_in)} entities") - - # TODO this could be better - for e in entities_in: - # Generate ID and write file - try: - existing = await self.entity_service.get_entity( - EntityModel.generate_id(e.entity_type, e.name) - ) - if existing: - raise ValueError( - f"Entity {e.entity_type}/{e.name} already exists, id: {EntityModel.generate_id(e.entity_type, e.name)}" - ) - except EntityNotFoundError: - # Good - entity doesn't exist yet - pass - - # Write files in parallel (filesystem is source of truth) - async def write_file(entity: Entity): - # Generate ID and write file - entity_id = EntityModel.generate_id(entity.entity_type, entity.name) - await write_entity_file(self.entities_path, entity_id, entity) - - file_writes = [write_file(entity) for entity in entities_in] - logger.debug("Starting parallel file writes") - await asyncio.gather(*file_writes) - logger.debug("Completed all file writes") - - async def create_entity_in_db(entity_create: Entity): - logger.debug(f"Creating entity in DB: {entity_create}") - try: - # Create base entity - created_entity = await self.entity_service.create_entity(entity_create) - logger.debug(f"Created base entity: {created_entity.id}") - - # Add observations - await self.observation_service.add_observations( - created_entity.id, entity_create.observations - ) - logger.debug( - f"Added {len(entity_create.observations)} observations to {created_entity.id}" - ) - - # Add relations - for relation in entity_create.relations: - await self.relation_service.create_relation(relation) - logger.debug( - f"Added {len(entity_create.relations)} relations for {created_entity.id}" - ) - - # Query final state - final_entity = await self.entity_service.get_entity(created_entity.id) - logger.debug(f"Retrieved final entity state: {final_entity}") - return final_entity - except Exception: - logger.exception(f"Failed to create entity in DB: {entity_create}") - raise - - # Update database index sequentially - logger.debug("Starting DB updates") - try: - entities = [] - for entity_in in entities_in: - entity = await create_entity_in_db(entity_in) - entities.append(entity) - logger.debug(f"Successfully created {len(entities)} entities in DB") - return entities - except Exception: - # On failure, we should try to clean up any files we wrote - logger.exception("Failed to create entities in DB") - for entity in entities_in: - try: - entity_id = EntityModel.generate_id(entity.entity_type, entity.name) - path = self.entities_path / entity_id - if path.exists(): - path.unlink() - except Exception as cleanup_error: - logger.error(f"Failed to clean up file for {entity.id}: {cleanup_error}") - raise - - async def get_entity(self, entity_id: str): - logger.debug(f"Get entity {entity_id} entities") - entity = await self.entity_service.get_entity(entity_id) - logger.debug(f"Found entity {entity}") - return entity - - async def create_relations(self, relations_data: List[Relation]) -> List[RelationModel]: - """Create multiple relations between entities.""" - logger.debug(f"Creating {len(relations_data)} relations") - - relations = [] - for relation in relations_data: - logger.debug(f"Processing relation: {relation.from_id} -> {relation.to_id}") - try: - # First read complete entities from filesystem - from_entity = await read_entity_file(self.entities_path, relation.from_id) - to_entity = await read_entity_file(self.entities_path, relation.to_id) - logger.debug(f"Read entities for relation: {from_entity.id}, {to_entity.id}") - - # Add the new relation to the source entity - if not hasattr(from_entity, "relations"): - from_entity.relations = [] - from_entity.relations.append(relation) - logger.debug(f"Added relation to source entity: {from_entity.id}") - - # Write updated entity files (filesystem is source of truth) - logger.debug("Writing updated entity files") - assert from_entity.id is not None - assert to_entity.id is not None - - await asyncio.gather( - *[ - write_entity_file(self.entities_path, from_entity.id, from_entity), - write_entity_file(self.entities_path, to_entity.id, to_entity), - ] - ) - logger.debug("Wrote updated entity files") - - # Now update the database index - relation = await self.relation_service.create_relation(relation) - relations.append(relation) - logger.debug(f"Created relation in DB: {relation.id}") - except Exception: - logger.exception(f"Failed to create relation: {relation}") - raise - - logger.debug(f"Successfully created {len(relations)} relations") - return relations - - async def add_observations( - self, observations_in: AddObservationsRequest - ) -> List[ObservationModel]: - """Add observations to an existing entity.""" - logger.debug(f"Adding observations to entity: {observations_in.entity_id}") - try: - # First get the entity from DB to get its ID - db_entity = await self.entity_service.get_entity(observations_in.entity_id) - logger.debug(f"Found entity in DB: {db_entity.id}") - - # Read entity from filesystem using the ID - entity = await read_entity_file(self.entities_path, db_entity.id) - logger.debug(f"Read entity from filesystem: {db_entity.id}") - - # Create new observations for the entity - entity.observations += observations_in.observations - logger.debug(f"Added {len(observations_in.observations)} observations to entity") - - # Write updated entity file - logger.debug("Writing updated entity file") - await write_entity_file(self.entities_path, db_entity.id, entity) - logger.debug("Wrote updated entity file") - - # Update database index - added_observations = await self.observation_service.add_observations( - db_entity.id, observations_in.observations - ) - logger.debug(f"Added {len(added_observations)} observations to DB") - - return added_observations - except Exception: - logger.exception(f"Failed to add observations to entity: {observations_in.entity_id}") - raise - - async def delete_entities(self, entity_ids: List[str]) -> bool: - """Delete entities and their files.""" - logger.debug(f"Deleting entities: {entity_ids}") - try: - deleted = False - for entity_id in entity_ids: - # First read the entity to make sure it exists - try: - entity = await read_entity_file(self.entities_path, entity_id) - except EntityNotFoundError: - logger.debug(f"Entity file not found: {entity_id}") - continue - - # Delete the file first since it's source of truth - file_path = self.get_entity_file_path(entity_id) - if file_path.exists(): - file_path.unlink() - logger.debug(f"Deleted entity {entity_id} file: {file_path}") - - # Then update database - result = await self.entity_service.delete_entity(entity_id) - if result: - deleted = True - logger.debug(f"Deleted entity from database: {entity_id}") - - return deleted - except Exception: - logger.exception("Failed to delete entities") - raise - - async def delete_observations(self, entity_id: str, contents: List[str]) -> bool: - """Delete specific observations from an entity.""" - logger.debug(f"Deleting observations from entity: {entity_id}") - try: - # First read the entity - entity = await read_entity_file(self.entities_path, entity_id) - - # Remove observations from entity - original_count = len(entity.observations) - entity.observations = [obs for obs in entity.observations if obs not in contents] - - # Only write file if we actually removed anything - if len(entity.observations) < original_count: - # Write updated entity file first (source of truth) - await write_entity_file(self.entities_path, entity_id, entity) - logger.debug(f"Updated entity file: {entity_id}") - - # Then update database - result = await self.observation_service.delete_observations(entity_id, contents) - logger.debug(f"Deleted observations from database: {result}") - return True - - return False - except Exception: - logger.exception("Failed to delete observations") - raise - - async def delete_relations(self, relations: List[Dict[str, Any]]) -> bool: - """Delete relations between entities.""" - logger.debug(f"Deleting relations: {relations}") - try: - deleted = False - for relation in relations: - # First read the source entity - try: - from_entity = await read_entity_file(self.entities_path, relation["from_id"]) - except EntityNotFoundError: - logger.debug(f"Source entity not found: {relation['from_id']}") - continue - - # Update the entity's relations - if hasattr(from_entity, "relations"): - original_count = len(from_entity.relations) - relation_type = relation.get("relation_type") - - if relation_type: - from_entity.relations = [ - r - for r in from_entity.relations - if not ( - r.to_id == relation["to_id"] and r.relation_type == relation_type - ) - ] - else: - from_entity.relations = [ - r for r in from_entity.relations if r.to_id != relation["to_id"] - ] - - # Only write file if we actually removed any relations - if len(from_entity.relations) < original_count: - # Write updated entity file first (source of truth) - await write_entity_file(self.entities_path, from_entity.id, from_entity) # pyright: ignore [reportArgumentType] - logger.debug(f"Updated source entity file: {from_entity.id}") - - # Then update database - result = await self.relation_service.delete_relations([relation]) - if result: - deleted = True - logger.debug("Deleted relations from database") - - return deleted - except Exception: - logger.exception("Failed to delete relations") - raise - - async def read_graph(self) -> Sequence[EntityModel]: - """Read the entire knowledge graph.""" - logger.debug("Reading entire knowledge graph") - try: - entities = await self.entity_service.get_all() - logger.debug(f"Read {len(entities)} entities from graph") - return entities - except Exception: - logger.exception("Failed to read graph") - raise - - async def search_nodes(self, query: str) -> Sequence[EntityModel]: - """Search for nodes in the knowledge graph.""" - logger.debug(f"Searching nodes with query: {query}") - try: - results = await self.entity_service.search(query) - logger.debug(f"Found {len(results)} matches for '{query}'") - return results - except Exception: - logger.exception(f"Failed to search nodes with query: {query}") - raise - - async def open_nodes(self, entity_ids: List[str]) -> List[Entity]: - """Get specific nodes and their relationships.""" - logger.debug(f"Opening nodes entity_ids: {entity_ids}") - - async def read_node(id: str) -> Optional[Entity]: - try: - # Get ID from name first - logger.debug(f"Looking up entity: {id}") - db_entity = await self.entity_service.get_entity(id) - if db_entity: - logger.debug(f"Found entity in DB: {db_entity.id}") - entity = await read_entity_file(self.entities_path, db_entity.id) - logger.debug(f"Read entity from filesystem: {entity.id}") - return entity - logger.debug(f"Entity not found: {id}") - return None - except Exception: - logger.exception(f"Failed to read node: {id}") - return None - - try: - entities = [ - entity - for entity in await asyncio.gather(*(read_node(id) for id in entity_ids)) - if entity is not None - ] - logger.debug(f"Opened {len(entities)} entities") - return entities - except Exception: - logger.exception("Failed to open nodes") - raise diff --git a/src/basic_memory/services/observation_service.py b/src/basic_memory/services/observation_service.py index f0a5ca9a..6ec8658d 100644 --- a/src/basic_memory/services/observation_service.py +++ b/src/basic_memory/services/observation_service.py @@ -1,99 +1,64 @@ -"""Service for managing observations in both filesystem and database.""" -from pathlib import Path -from typing import List, Sequence, Dict, Any +"""Service for managing observations in the database.""" + +from typing import List, Sequence + +from loguru import logger from sqlalchemy import select from basic_memory.models import Observation as ObservationModel from basic_memory.repository.observation_repository import ObservationRepository -from . import DatabaseSyncError -from basic_memory.schemas import Observation +from .service import BaseService -class ObservationService: +class ObservationService(BaseService[ObservationRepository]): """ Service for managing observations in the database. File operations are handled by MemoryService. """ - - def __init__(self, project_path: Path, observation_repo: ObservationRepository): - self.project_path = project_path - self.observation_repo = observation_repo - - async def add_observations(self, entity_id: str, observations: List[Observation]) -> List[ObservationModel]: - """ - Add multiple observations to an entity. - Returns the created observations with IDs set. - """ - try: - return await self.observation_repo.bulk_create([ - ObservationModel( + + def __init__(self, observation_repository: ObservationRepository): + super().__init__(observation_repository) + + async def add_observations( + self, entity_id: str, observations: List[str] + ) -> List[ObservationModel]: + """Add multiple observations to an entity.""" + logger.debug(f"Adding {len(observations)} observations to entity: {entity_id}") + return await self.repository.create_all( + [ + dict( entity_id=entity_id, content=observation, ) for observation in observations - ]) - except Exception as e: - raise DatabaseSyncError(f"Failed to add observations to database: {str(e)}") from e + ] + ) - async def delete_observations(self, entity_id: str, contents: List[str]) -> int: - """ - Delete specific observations from an entity. - - Args: - entity_id: ID of the entity - contents: List of observation contents to delete - - Returns: - Number of observations deleted - """ - try: - deleted = False - for content in contents: - result = await self.observation_repo.delete_by_fields( - entity_id=entity_id, - content=content - ) - if result: - deleted = True - return deleted - except Exception as e: - raise DatabaseSyncError(f"Failed to delete observations from database: {str(e)}") from e + async def delete_observations(self, entity_id: str, contents: List[str]) -> bool: + """Delete specific observations from an entity.""" + logger.debug(f"Deleting observations from entity: {entity_id}") + deleted = False + for content in contents: + result = await self.repository.delete_by_fields(entity_id=entity_id, content=content) + if result: + deleted = True + return deleted async def delete_by_entity(self, entity_id: str) -> bool: - """ - Delete all observations for an entity. - - Args: - entity_id: ID of the entity - - Returns: - True if any observations were deleted - """ - try: - return await self.observation_repo.delete_by_fields(entity_id=entity_id) - except Exception as e: - raise DatabaseSyncError(f"Failed to delete observations from database: {str(e)}") from e + """Delete all observations for an entity.""" + logger.debug(f"Deleting all observations for entity: {entity_id}") + return await self.repository.delete_by_fields(entity_id=entity_id) async def search_observations(self, query: str) -> List[ObservationModel]: - """ - Search for observations across all entities. - - Args: - query: Text to search for in observation content - - Returns: - List of matching observations with their entity contexts - """ - result = await self.observation_repo.execute_query( - select(ObservationModel).filter( - ObservationModel.content.contains(query) - ) + """Search for observations across all entities.""" + logger.debug(f"Searching observations with query: {query}") + result = await self.repository.execute_query( + select(ObservationModel).filter(ObservationModel.content.contains(query)) ) - return [ - ObservationModel(content=obs.content) - for obs in result.scalars().all() - ] - + observations = result.scalars().all() + return [ObservationModel(content=obs.content) for obs in observations] + async def get_observations_by_context(self, context: str) -> Sequence[ObservationModel]: """Get all observations with a specific context.""" - return await self.observation_repo.find_by_context(context) + logger.debug(f"Getting observations for context: {context}") + return await self.repository.find_by_context(context) diff --git a/src/basic_memory/services/relation_service.py b/src/basic_memory/services/relation_service.py index dbe44ec9..2d2e646b 100644 --- a/src/basic_memory/services/relation_service.py +++ b/src/basic_memory/services/relation_service.py @@ -1,78 +1,61 @@ """Service for managing relations in the database.""" -from pathlib import Path + from typing import List, Dict, Any -from basic_memory.repository.relation_repository import RelationRepository -from basic_memory.schemas import Entity, Relation -from . import DatabaseSyncError -from basic_memory.models import Relation as RelationModel +from loguru import logger -class RelationService: +from basic_memory.models import Relation as RelationModel +from basic_memory.repository.relation_repository import RelationRepository +from basic_memory.schemas import Entity as EntitySchema, Relation +from .service import BaseService + + +class RelationService(BaseService[RelationRepository]): """ Service for managing relations in the database. File operations are handled by MemoryService. """ - - def __init__(self, project_path: Path, relation_repo: RelationRepository): - self.project_path = project_path - self.relation_repo = relation_repo + + def __init__(self, relation_repository: RelationRepository): + super().__init__(relation_repository) async def create_relation(self, relation: Relation) -> RelationModel: """Create a new relation in the database.""" - try: - return await self.relation_repo.create(relation.model_dump()) - except Exception as e: - raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e + logger.debug(f"Creating relation: {relation}") + return await self.repository.create(relation.model_dump()) - async def delete_relation(self, from_entity: Entity, to_entity: Entity, relation_type: str) -> bool: + async def delete_relation( + self, from_entity: EntitySchema, to_entity: EntitySchema, relation_type: str + ) -> bool: """Delete a specific relation between entities.""" - try: - # Use repository to find and delete the relation - filters = { - 'from_id': from_entity.id, - 'to_id': to_entity.id, - 'relation_type': relation_type - } - - # Delete in database - result = await self.relation_repo.delete_by_fields(**filters) - - # Update in-memory relations if present - if hasattr(from_entity, 'relations'): - from_entity.relations = [ - r for r in from_entity.relations - if not (r.to_id == to_entity.id and r.relation_type == relation_type) - ] - - return result - - except Exception as e: - raise DatabaseSyncError(f"Failed to delete relation: {str(e)}") from e + logger.debug(f"Deleting relation between {from_entity.id} and {to_entity.id}") + + assert from_entity.id is not None, "from_entity.id must not be None" + result = await self.repository.delete_by_fields( + from_id=from_entity.id, + to_id=to_entity.id, + relation_type=relation_type, + ) + return result async def delete_relations(self, relations: List[Dict[str, Any]]) -> bool: - """ - Delete relations matching specified criteria. - - Args: - relations: List of dicts with from_id, to_id, and optional relation_type - - Returns: - True if any relations were deleted - """ - try: - deleted = False - for relation in relations: - filters = { - 'from_id': relation['from_id'], - 'to_id': relation['to_id'] - } - if 'relation_type' in relation: - filters['relation_type'] = relation['relation_type'] - - result = await self.relation_repo.delete_by_fields(**filters) - if result: - deleted = True - - return deleted - except Exception as e: - raise DatabaseSyncError(f"Failed to delete relations: {str(e)}") from e \ No newline at end of file + """Delete relations matching specified criteria.""" + logger.debug(f"Deleting {len(relations)} relations") + deleted = False + for relation in relations: + filters = {"from_id": relation["from_id"], "to_id": relation["to_id"]} + if "relation_type" in relation: + filters["relation_type"] = relation["relation_type"] + + result = await self.repository.delete_by_fields(**filters) + if result: + deleted = True + + return deleted + + async def create_relations(self, relations_data: List[Relation]) -> List[RelationModel]: + """Create multiple relations between entities.""" + logger.debug(f"Creating {len(relations_data)} relations") + return await self.repository.create_all( + [Relation.model_dump(relation) for relation in relations_data] + ) diff --git a/src/basic_memory/services/service.py b/src/basic_memory/services/service.py new file mode 100644 index 00000000..763db145 --- /dev/null +++ b/src/basic_memory/services/service.py @@ -0,0 +1,23 @@ +"""Base service class.""" + +from typing import TypeVar, Generic, List, Sequence + +from basic_memory.repository.repository import Repository + +T = TypeVar("T", bound=Repository) + + +class BaseService(Generic[T]): + """Base service that takes a repository.""" + + def __init__(self, repository: T): + """Initialize service with repository.""" + self.repository = repository + + async def add(self, model: T) -> T: + """Add model to repository.""" + return await self.repository.add(model) + + async def add_all(self, models: List[T]) -> Sequence[T]: + """Add a List of models to repository.""" + return await self.repository.add_all(models) diff --git a/tests/api/test_knowledge_router.py b/tests/api/test_knowledge_router.py index 2835dbbf..f405bd6a 100644 --- a/tests/api/test_knowledge_router.py +++ b/tests/api/test_knowledge_router.py @@ -20,7 +20,7 @@ from basic_memory.schemas import ( @pytest_asyncio.fixture -def app(test_config, engine_session_factory) -> FastAPI: +def app(test_config, engine_factory) -> FastAPI: """Create FastAPI test application.""" # Lazy import router to avoid app startup issues from basic_memory.api.routers.knowledge import router @@ -29,7 +29,7 @@ def app(test_config, engine_session_factory) -> FastAPI: app.include_router(router) app.dependency_overrides[get_project_config] = lambda: test_config - app.dependency_overrides[get_engine_factory] = lambda: engine_session_factory + app.dependency_overrides[get_engine_factory] = lambda: engine_factory return app @@ -40,9 +40,12 @@ async def client(app: FastAPI) -> AsyncGenerator[AsyncClient, None]: yield client -async def create_entity( - client, data={"name": "Test Entity", "entity_type": "test"} -) -> EntityResponse: +async def create_entity(client) -> EntityResponse: + data = { + "name": "Test Entity", + "entity_type": "test", + "observations": ["First observation", "Second observation"], + } # Create an entity response = await client.post("/knowledge/entities", json={"entities": [data]}) # Verify creation @@ -56,6 +59,8 @@ async def create_entity( entity_type = entity.get("entity_type") assert entity_type == data["entity_type"] + assert len(entity["observations"]) == 2 + create_response = CreateEntityResponse.model_validate(response_data) return create_response.entities[0] @@ -193,6 +198,35 @@ async def test_search_nodes(client: AsyncClient): assert "Gamma Production" in names +@pytest.mark.asyncio +async def test_open_nodes(client: AsyncClient): + """Should search for entities in the knowledge graph.""" + # Create a few entities with different names + entities = [ + {"name": "Alpha Test", "entity_type": "test"}, + {"name": "Beta Test", "entity_type": "test"}, + ] + await client.post("/knowledge/entities", json={"entities": entities}) + + # open nodes + response = await client.post( + "/knowledge/nodes", + json={ + "entity_ids": [ + "test/alpha_test", + ] + }, + ) + + # Verify search results + assert response.status_code == 200 + data = response.json() + assert len(data["entities"]) == 1 + entity = data["entities"][0] + assert entity["name"] == "Alpha Test" + assert entity["entity_type"] == "test" + + @pytest.mark.asyncio async def test_delete_entity(client: AsyncClient): """Test DELETE /knowledge/entities/{entity_id}.""" @@ -214,7 +248,12 @@ async def test_delete_entity_bulk(client: AsyncClient): """Test DELETE /knowledge/entities/{entity_id}.""" # Create test entity entity1 = await create_entity(client) - entity2 = await create_entity(client, data={"name": "Test Entity2", "entity_type": "test"}) + + e2_response = await client.post( + "/knowledge/entities", json={"entities": [{"name": "Test Entity2", "entity_type": "test"}]} + ) + create_response = CreateEntityResponse.model_validate(e2_response.json()) + entity2 = create_response.entities[0] # Test deletion response = await client.post( @@ -253,7 +292,7 @@ async def test_delete_observations(client, observation_repository): """Test DELETE /knowledge/entities/{entity_id}/observations.""" # Create test data entity = await create_entity(client) - observations = await add_observations(client, entity.id) + observations = await add_observations(client, entity.id) # adds 2 # Delete specific observations request_data = {"entity_id": entity.id, "deletions": [observations[0].content]} @@ -263,7 +302,7 @@ async def test_delete_observations(client, observation_repository): # Verify only specified observations were deleted remaining = await observation_repository.find_by_entity(entity.id) - assert len(remaining) == 1 + assert len(remaining) == 2 # because entity originally had 1 assert remaining[0].content == observations[1].content @@ -332,27 +371,6 @@ async def test_delete_nonexistent_relations(client): assert response.json() == {"deleted": False} -# @pytest.mark.asyncio -# TODO should this be a 404? -# - should we check the file system first -# - should we fail if the entity is not found in the db -# async def test_delete_observations_entity_does_not_exist(client): -# """Test when entity_id in path doesn't match request body.""" -# # Create test entity -# entity = await create_entity(client) -# -# # Try to delete with mismatched IDs -# request_data = { -# "entity_id": "different/id", -# "deletions": ["Some observation"] -# } -# response = await client.post( -# f"/knowledge/entities/{entity.id}/observations/delete", -# json=request_data -# ) -# assert response.status_code == 404 - - @pytest.mark.asyncio async def test_full_knowledge_flow(client: AsyncClient): """Test a complete knowledge graph flow with multiple operations.""" @@ -367,9 +385,9 @@ async def test_full_knowledge_flow(client: AsyncClient): }, ) assert main_response.status_code == 200 - entity0_id = main_response.json()["entities"][0]["id"] - entity1_id = main_response.json()["entities"][1]["id"] - assert entity0_id is not None + main_entity_id = "test/main_entity" + non_entity_id = "n_a/non_entity" + assert main_entity_id is not None # 2. Create related entities related_response = await client.post( @@ -391,8 +409,16 @@ async def test_full_knowledge_flow(client: AsyncClient): "/knowledge/relations", json={ "relations": [ - {"from_id": entity0_id, "to_id": related_ids[0], "relation_type": "connects_to"}, - {"from_id": entity0_id, "to_id": related_ids[1], "relation_type": "connects_to"}, + { + "from_id": main_entity_id, + "to_id": related_ids[0], + "relation_type": "connects_to", + }, + { + "from_id": main_entity_id, + "to_id": related_ids[1], + "relation_type": "connects_to", + }, ] }, ) @@ -403,7 +429,7 @@ async def test_full_knowledge_flow(client: AsyncClient): await client.post( "/knowledge/observations", json={ - "entity_id": entity0_id, + "entity_id": main_entity_id, "observations": [ "Connected to first related entity", "Connected to second related entity", @@ -412,7 +438,7 @@ async def test_full_knowledge_flow(client: AsyncClient): ) # 5. Verify full graph structure - main_get = await client.get(f"/knowledge/entities/{entity0_id}") + main_get = await client.get(f"/knowledge/entities/{main_entity_id}") main_entity = main_get.json() # Check entity structure @@ -427,7 +453,7 @@ async def test_full_knowledge_flow(client: AsyncClient): # 7. delete entities response = await client.post( - "/knowledge/entities/delete", json={"entity_ids": [entity0_id, entity1_id]} + "/knowledge/entities/delete", json={"entity_ids": [main_entity_id, non_entity_id]} ) assert response.status_code == 200 diff --git a/tests/conftest.py b/tests/conftest.py index 001d7450..c7a64cf2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,22 +5,26 @@ from pathlib import Path from typing import AsyncGenerator import pytest_asyncio -from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession, AsyncEngine +from sqlalchemy import text +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.deps import ( - get_entity_service, - get_observation_service, - get_relation_service, - get_relation_repo, - get_observation_repo, - get_entity_repo, -) +from basic_memory.models import Base, Entity as EntityModel from basic_memory.repository.entity_repository import EntityRepository +from basic_memory.repository.observation_repository import ObservationRepository +from basic_memory.repository.relation_repository import RelationRepository from basic_memory.schemas import Entity -from basic_memory.services import MemoryService +from basic_memory.services import ( + EntityService, + ObservationService, + RelationService, +) @pytest_asyncio.fixture @@ -39,22 +43,27 @@ def test_config(tmp_path): @pytest_asyncio.fixture(scope="function") -async def engine_session_factory( +async def engine_factory( test_config, ) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]: - """Create an async engine using in-memory SQLite database""" + """Create engine and session factory using in-memory SQLite database.""" async with db.engine_session_factory( project_path=test_config.path, db_type=DatabaseType.MEMORY - ) as (engine, session_factory): - yield engine, session_factory + ) as (engine, session_maker): + # Initialize database + async with db.scoped_session(session_maker) as session: + await session.execute(text("PRAGMA foreign_keys=ON")) + conn = await session.connection() + await conn.run_sync(Base.metadata.create_all) + + yield engine, session_maker -@pytest_asyncio.fixture(scope="function") -async def session(engine_session_factory): - """Create an async session factory and yield a session""" - engine, session_factory = engine_session_factory - async with db.session(session_factory) as session: - yield session +@pytest_asyncio.fixture +async def session_maker(engine_factory) -> async_sessionmaker[AsyncSession]: + """Get session maker for tests.""" + _, session_maker = engine_factory + return session_maker @pytest_asyncio.fixture @@ -68,52 +77,50 @@ async def test_project_path(): @pytest_asyncio.fixture(scope="function") -async def entity_repository(session: AsyncSession): - """Create an EntityRepository instance""" - return await get_entity_repo(session) +async def entity_repository(session_maker: async_sessionmaker[AsyncSession]) -> EntityRepository: + """Create an EntityRepository instance.""" + return EntityRepository(session_maker) @pytest_asyncio.fixture(scope="function") -async def observation_repository(session: AsyncSession): - """Create an ObservationRepository instance""" - return await get_observation_repo(session) +async def observation_repository( + session_maker: async_sessionmaker[AsyncSession], +) -> ObservationRepository: + """Create an ObservationRepository instance.""" + return ObservationRepository(session_maker) @pytest_asyncio.fixture(scope="function") -async def relation_repository(session: AsyncSession): - """Create a RelationRepository instance""" - return await get_relation_repo(session) +async def relation_repository( + session_maker: async_sessionmaker[AsyncSession], +) -> RelationRepository: + """Create a RelationRepository instance.""" + return RelationRepository(session_maker) @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) +async def entity_service(entity_repository: EntityRepository) -> EntityService: + """Create EntityService with repository.""" + return EntityService(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) +async def relation_service(relation_repository: RelationRepository) -> RelationService: + """Create RelationService with repository.""" + return RelationService(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): - """Fixture providing initialized MemoryService.""" - return MemoryService(test_project_path, entity_service, relation_service, observation_service) +async def observation_service(observation_repository: ObservationRepository) -> ObservationService: + """Create ObservationService with repository.""" + return ObservationService(observation_repository) @pytest_asyncio.fixture(scope="function") -async def sample_entity(entity_repository: EntityRepository): - """Create a sample entity for testing""" +async def sample_entity(entity_repository: EntityRepository) -> EntityModel: + """Create a sample entity for testing.""" entity_data = { - "id": "20240102-test-entity", + "id": "test/test_entity", "name": "Test Entity", "entity_type": "test", "description": "A test entity", @@ -122,10 +129,7 @@ async def sample_entity(entity_repository: EntityRepository): @pytest_asyncio.fixture -async def test_entity(entity_service): +async def test_entity(entity_service: EntityService) -> EntityModel: """Create a test entity for reuse in tests.""" - entity_data = Entity( # pyright: ignore [reportCallIssue] - name="Test Entity", - entity_type="test", # pyright: ignore [reportCallIssue] - ) + entity_data = Entity(name="Test Entity", entity_type="test", observations=[]) return await entity_service.create_entity(entity_data) diff --git a/tests/mcp/conftest.py b/tests/mcp/conftest.py index c7febde1..adb3d59d 100644 --- a/tests/mcp/conftest.py +++ b/tests/mcp/conftest.py @@ -9,11 +9,11 @@ from basic_memory.deps import get_project_config, get_engine_factory @pytest_asyncio.fixture -def app(test_config, engine_session_factory) -> FastAPI: +def app(test_config, engine_factory) -> FastAPI: """Create test FastAPI application.""" app = fastapi_app app.dependency_overrides[get_project_config] = lambda: test_config - app.dependency_overrides[get_engine_factory] = lambda: engine_session_factory + app.dependency_overrides[get_engine_factory] = lambda: engine_factory return app diff --git a/tests/mcp/test_open_nodes.py b/tests/mcp/test_open_nodes.py index a2ac5d52..ec802673 100644 --- a/tests/mcp/test_open_nodes.py +++ b/tests/mcp/test_open_nodes.py @@ -57,4 +57,4 @@ async def test_open_nodes(app): assert entity.id == "test/opentesta" assert entity.entity_type == "test" assert len(entity.observations) == 1 - assert entity.observations[0] == "First test entity" + assert entity.observations[0].content == "First test entity" diff --git a/tests/mcp/test_tool_docs.py b/tests/mcp/test_tool_docs.py new file mode 100644 index 00000000..cef7268c --- /dev/null +++ b/tests/mcp/test_tool_docs.py @@ -0,0 +1,25 @@ +"""Test to show MCP tool documentation.""" + +import json +import pytest +from mcp.types import Tool +from basic_memory.mcp.server import handle_list_tools + + +@pytest.mark.asyncio +async def test_list_tools(): + """List available tools and their documentation.""" + tools = await handle_list_tools() + assert isinstance(tools, list) + assert all(isinstance(t, Tool) for t in tools) + + print("\nAvailable MCP Tools:\n") + + # Print each tool's documentation + for tool in tools: + print(f"Tool: {tool.name}") + print(f"Description: {tool.description}") + print("Required fields:", tool.inputSchema.get("required", [])) + print() + print("Schema:", json.dumps(tool.inputSchema, indent=2)) + print("-" * 80 + "\n") \ No newline at end of file diff --git a/tests/repository/test_entity_repository.py b/tests/repository/test_entity_repository.py index a5e57919..02bcce1f 100644 --- a/tests/repository/test_entity_repository.py +++ b/tests/repository/test_entity_repository.py @@ -1,4 +1,5 @@ """Tests for the EntityRepository.""" + from datetime import datetime, UTC import pytest @@ -6,135 +7,164 @@ import pytest_asyncio from sqlalchemy import select, text from sqlalchemy.exc import IntegrityError +from basic_memory import db from basic_memory.models import Entity, Observation, Relation from basic_memory.repository.entity_repository import EntityRepository @pytest_asyncio.fixture -async def entity_repo(session): - """Create an EntityRepository with test DB session.""" - return EntityRepository(session) - - -@pytest_asyncio.fixture -async def test_entity(session): +async def test_entity(session_maker): """Create a test entity.""" - entity = Entity( - id="test/test_entity", - name="test_entity", - entity_type="test", - description="Test entity" - ) - session.add(entity) - await session.flush() - return 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" + ) + session.add(entity) + return entity @pytest_asyncio.fixture -async def entity_with_observations(session, test_entity): +async def entity_with_observations(session_maker, test_entity): """Create an entity with observations.""" - observations = [ - Observation(entity_id=test_entity.id, content="First observation"), - Observation(entity_id=test_entity.id, content="Second observation") - ] - session.add_all(observations) - await session.flush() - return test_entity + async with db.scoped_session(session_maker) as session: + observations = [ + Observation(entity_id=test_entity.id, content="First observation"), + Observation(entity_id=test_entity.id, content="Second observation"), + ] + session.add_all(observations) + return test_entity @pytest_asyncio.fixture -async def related_entities(session): +async def related_entities(session_maker): """Create entities with relations between them.""" - source = Entity( - id="source/test_entity", - name="source", - entity_type="source", - description="Source entity" - ) - target = Entity( - id="target/test_entity", - name="target", - entity_type="target", - description="Target entity" - ) - session.add_all([source, target]) - await session.flush() + async with db.scoped_session(session_maker) as session: + source = Entity( + id="source/test_entity", + name="source", + entity_type="source", + description="Source entity", + ) + target = Entity( + id="target/test_entity", + name="target", + entity_type="target", + description="Target entity", + ) + session.add_all([source, target]) + + relation = Relation(from_id=source.id, to_id=target.id, relation_type="connects_to") + session.add(relation) + + return source, target, relation - relation = Relation( - from_id=source.id, - to_id=target.id, - relation_type="connects_to" - ) - session.add(relation) - await session.flush() - - return source, target, relation @pytest.mark.asyncio async def test_create_entity(entity_repository: EntityRepository): """Test creating a new entity""" entity_data = { - 'name': 'Test', - 'entity_type': 'test', - 'description': 'Test description', + "name": "Test", + "entity_type": "test", + "description": "Test description", } entity = await entity_repository.create(entity_data) # Verify returned object - assert entity.id == 'test/test' - assert entity.name == 'Test' - assert entity.description == 'Test description' + assert entity.id == "test/test" + assert entity.name == "Test" + assert entity.description == "Test description" assert isinstance(entity.created_at, datetime) # Verify in database - stmt = select(Entity).where(Entity.id == entity.id) - result = await entity_repository.session.execute(stmt) - db_entity = result.scalar_one() - assert db_entity.id == entity.id - assert db_entity.name == entity.name - assert db_entity.description == entity.description + found = await entity_repository.find_by_id(entity.id) + assert found is not None + assert found.id is not None + assert found.id == entity.id + assert found.name == entity.name + assert found.description == entity.description + + # assert relations are eagerly loaded + assert len(entity.observations) == 0 + assert len(entity.relations) == 0 + + +@pytest.mark.asyncio +async def test_create_all(entity_repository: EntityRepository): + """Test creating a new entity""" + entity_data = [ + { + "name": "Test_1", + "entity_type": "test", + "description": "Test description", + }, + { + "name": "Test-2", + "entity_type": "test", + "description": "Test description", + }, + ] + entities = await entity_repository.create_all(entity_data) + + assert len(entities) == 2 + entity = entities[0] + + # Verify in database + found = await entity_repository.find_by_id(entity.id) + assert found is not None + assert found.id is not None + assert found.id == entity.id + assert found.name == entity.name + assert found.description == entity.description + + # assert relations are eagerly loaded + assert len(entity.observations) == 0 + assert len(entity.relations) == 0 + @pytest.mark.asyncio async def test_entity_type_name_unique_constraint(entity_repository: EntityRepository): """Test the unique constraint on entity_type + name combination.""" # Create first entity entity1_data = { - 'id': '20240102-test1', - 'name': 'Test Entity', - 'entity_type': 'type1', - 'description': 'First entity' + "id": "20240102-test1", + "name": "Test Entity", + "entity_type": "type1", + "description": "First entity", } await entity_repository.create(entity1_data) # Try to create another entity with same type and name entity2_data = { - 'id': '20240102-test2', - 'name': 'Test Entity', # Same name - 'entity_type': 'type1', # Same type - 'description': 'Second entity' + "id": "20240102-test2", + "name": "Test Entity", # Same name + "entity_type": "type1", # Same type + "description": "Second entity", } # Should raise IntegrityError with pytest.raises(IntegrityError) as exc_info: await entity_repository.create(entity2_data) - assert 'UNIQUE constraint failed: entity.entity_type, entity.name' in str(exc_info.value) + assert "UNIQUE constraint failed: entity.entity_type, entity.name" in str(exc_info.value) + @pytest.mark.asyncio -async def test_create_entity_null_description(entity_repository: EntityRepository): +async def test_create_entity_null_description(session_maker, entity_repository: EntityRepository): """Test creating an entity with null description""" entity_data = { - 'id': '20240102-test', - 'name': 'Test', - 'entity_type': 'test', - 'description': None, + "id": "20240102-test", + "name": "Test", + "entity_type": "test", + "description": None, } entity = await entity_repository.create(entity_data) # Verify in database - stmt = select(Entity).where(Entity.id == entity.id) - result = await entity_repository.session.execute(stmt) - db_entity = result.scalar_one() - assert db_entity.description is None + async with db.scoped_session(session_maker) as session: + stmt = select(Entity).where(Entity.id == entity.id) + result = await session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.description is None + @pytest.mark.asyncio async def test_find_by_id(entity_repository: EntityRepository, sample_entity: Entity): @@ -145,12 +175,14 @@ async def test_find_by_id(entity_repository: EntityRepository, sample_entity: En assert found.name == sample_entity.name # Verify against direct database query - stmt = select(Entity).where(Entity.id == sample_entity.id) - result = await entity_repository.session.execute(stmt) - db_entity = result.scalar_one() - assert db_entity.id == found.id - assert db_entity.name == found.name - assert db_entity.description == found.description + async with db.scoped_session(entity_repository.session_maker) as session: + stmt = select(Entity).where(Entity.id == sample_entity.id) + result = await session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.id == found.id + assert db_entity.name == found.name + assert db_entity.description == found.description + @pytest.mark.asyncio async def test_find_by_name(entity_repository: EntityRepository, sample_entity: Entity): @@ -161,250 +193,174 @@ async def test_find_by_name(entity_repository: EntityRepository, sample_entity: assert found.name == sample_entity.name # Verify against direct database query - stmt = select(Entity).where(Entity.name == sample_entity.name) - result = await entity_repository.session.execute(stmt) - db_entity = result.scalar_one() - assert db_entity.id == found.id - assert db_entity.name == found.name - assert db_entity.description == found.description + async with db.scoped_session(entity_repository.session_maker) as session: + stmt = select(Entity).where(Entity.name == sample_entity.name) + result = await session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.id == found.id + assert db_entity.name == found.name + assert db_entity.description == found.description + @pytest.mark.asyncio async def test_update_entity(entity_repository: EntityRepository, sample_entity: Entity): """Test updating an entity""" updated = await entity_repository.update( - sample_entity.id, - {'description': 'Updated description'} + sample_entity.id, {"description": "Updated description"} ) assert updated is not None - assert updated.description == 'Updated description' + assert updated.description == "Updated description" assert updated.name == sample_entity.name # Other fields unchanged # Verify in database - stmt = select(Entity).where(Entity.id == sample_entity.id) - result = await entity_repository.session.execute(stmt) - db_entity = result.scalar_one() - assert db_entity.description == 'Updated description' - assert db_entity.name == sample_entity.name + async with db.scoped_session(entity_repository.session_maker) as session: + stmt = select(Entity).where(Entity.id == sample_entity.id) + result = await session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.description == "Updated description" + assert db_entity.name == sample_entity.name + @pytest.mark.asyncio async def test_update_entity_to_null(entity_repository: EntityRepository, sample_entity: Entity): """Test updating an entity's description to null""" - updated = await entity_repository.update( - sample_entity.id, - {'description': None} - ) + updated = await entity_repository.update(sample_entity.id, {"description": None}) assert updated is not None assert updated.description is None # Verify in database - stmt = select(Entity).where(Entity.id == sample_entity.id) - result = await entity_repository.session.execute(stmt) - db_entity = result.scalar_one() - assert db_entity.description is None + async with db.scoped_session(entity_repository.session_maker) as session: + stmt = select(Entity).where(Entity.id == sample_entity.id) + result = await session.execute(stmt) + db_entity = result.scalar_one() + assert db_entity.description is None + @pytest.mark.asyncio -async def test_delete_entity_find_by_id(entity_repository: EntityRepository, sample_entity: Entity): - """Test deleting an entity""" - success = await entity_repository.delete(sample_entity.id) - assert success is True - - # Verify it's gone - found = await entity_repository.find_by_id(sample_entity.id) - assert found is None - - # Verify with direct query - stmt = select(Entity).where(Entity.id == sample_entity.id) - result = await entity_repository.session.execute(stmt) - assert result.first() is None - -@pytest.mark.asyncio -async def test_search(entity_repository: EntityRepository): - """Test searching entities""" - # Create test entities with observations - entity1 = await entity_repository.create({ - 'id': '20240102-test1', - 'name': 'Search Test 1', - 'entity_type': 'test', - 'description': 'First test entity' - }) - - entity2 = await entity_repository.create({ - 'id': '20240102-test2', - 'name': 'Search Test 2', - 'entity_type': 'other', - 'description': 'Second test entity' - }) - - # Verify entities in database - stmt = select(Entity).where(Entity.id.in_([entity1.id, entity2.id])) - result = await entity_repository.session.execute(stmt) - db_entities = result.scalars().all() - assert len(db_entities) == 2 - - # Add observations - stmt = text(""" - INSERT INTO observation (entity_id, content, created_at) - VALUES (:e1_id, :e1_obs, :ts), (:e2_id, :e2_obs, :ts) - """) - ts = datetime.now(UTC) - await entity_repository.session.execute(stmt, { - "e1_id": entity1.id, - "e1_obs": "First observation with searchable content", - "e2_id": entity2.id, - "e2_obs": "Another observation to find", - "ts": ts - }) - await entity_repository.session.commit() - - # Test search by name - results = await entity_repository.search('Search Test') - assert len(results) == 2 - names = {e.name for e in results} - assert 'Search Test 1' in names - assert 'Search Test 2' in names - - # Test search by type - results = await entity_repository.search('other') - assert len(results) == 1 - assert results[0].entity_type == 'other' - - # Test search by observation content - results = await entity_repository.search('searchable') - assert len(results) == 1 - assert results[0].id == entity1.id - -@pytest.mark.asyncio -async def test_find_by_type_and_name(entity_repository: EntityRepository): - """Test finding an entity by type and name combination.""" - # Create two entities with same name but different types - entity1 = await entity_repository.create({ - 'id': '20240102-test1', - 'name': 'Test Entity', - 'entity_type': 'type1', - 'description': 'First test entity' - }) - - entity2 = await entity_repository.create({ - 'id': '20240102-test2', - 'name': 'Test Entity', - 'entity_type': 'type2', - 'description': 'Second test entity' - }) - - # Should find correct entity when both type and name match - found = await entity_repository.find_by_type_and_name('type1', 'Test Entity') - assert found is not None - assert found.id == entity1.id - assert found.entity_type == 'type1' - assert found.name == 'Test Entity' - - # Should find other entity with same name but different type - found = await entity_repository.find_by_type_and_name('type2', 'Test Entity') - assert found is not None - assert found.id == entity2.id - assert found.entity_type == 'type2' - assert found.name == 'Test Entity' - - # Should return None when type doesn't match - found = await entity_repository.find_by_type_and_name('nonexistent', 'Test Entity') - assert found is None - - # Should return None when name doesn't match - found = await entity_repository.find_by_type_and_name('type1', 'Nonexistent') - assert found is None - - # Verify relationships are loaded - entity3 = await entity_repository.create({ - 'id': '20240102-test3', - 'name': 'Entity With Relations', - 'entity_type': 'type3', - 'description': 'Entity with observations and relations' - }) - - # Add an observation - stmt = text(""" - INSERT INTO observation (entity_id, content, created_at) - VALUES (:entity_id, :content, :ts) - """) - ts = datetime.now(UTC) - await entity_repository.session.execute(stmt, { - "entity_id": entity3.id, - "content": "Test observation", - "ts": ts - }) - await entity_repository.session.commit() - - # Find entity and verify relationships are loaded - found = await entity_repository.find_by_type_and_name('type3', 'Entity With Relations') - assert found is not None - assert len(found.observations) == 1 - assert found.observations[0].content == "Test observation" - -@pytest.mark.asyncio -async def test_delete_entity(entity_repo, test_entity): +async def test_delete_entity(entity_repository: EntityRepository, test_entity): """Test deleting an entity.""" - result = await entity_repo.delete(test_entity.id) + result = await entity_repository.delete(test_entity.id) assert result is True - + # Verify deletion - deleted = await entity_repo.find_by_id(test_entity.id) + deleted = await entity_repository.find_by_id(test_entity.id) assert deleted is None @pytest.mark.asyncio -async def test_delete_entity_with_observations(entity_repo, entity_with_observations): +async def test_delete_entity_with_observations( + entity_repository: EntityRepository, entity_with_observations +): """Test deleting an entity cascades to its observations.""" entity = entity_with_observations - - result = await entity_repo.delete(entity.id) + + result = await entity_repository.delete(entity.id) assert result is True - + # Verify entity deletion - deleted = await entity_repo.find_by_id(entity.id) + deleted = await entity_repository.find_by_id(entity.id) assert deleted is None - + # Verify observations were cascaded - query = select(Observation).filter(Observation.entity_id == entity.id) - result = await entity_repo.execute_query(query) - remaining_observations = result.scalars().all() - assert len(remaining_observations) == 0 + async with db.scoped_session(entity_repository.session_maker) as session: + query = select(Observation).filter(Observation.entity_id == entity.id) + result = await session.execute(query) + remaining_observations = result.scalars().all() + assert len(remaining_observations) == 0 @pytest.mark.asyncio -async def test_delete_entities_by_type(entity_repo, test_entity): +async def test_delete_entities_by_type(entity_repository: EntityRepository, test_entity): """Test deleting entities by type.""" - result = await entity_repo.delete_by_fields(entity_type=test_entity.entity_type) + result = await entity_repository.delete_by_fields(entity_type=test_entity.entity_type) assert result is True - + # Verify deletion - query = select(Entity).filter(Entity.entity_type == test_entity.entity_type) - result = await entity_repo.execute_query(query) - remaining = result.scalars().all() - assert len(remaining) == 0 + async with db.scoped_session(entity_repository.session_maker) as session: + query = select(Entity).filter(Entity.entity_type == test_entity.entity_type) + result = await session.execute(query) + remaining = result.scalars().all() + assert len(remaining) == 0 @pytest.mark.asyncio -async def test_delete_entity_with_relations(entity_repo, related_entities): +async def test_delete_entity_with_relations(entity_repository: EntityRepository, related_entities): """Test deleting an entity cascades to its relations.""" source, target, relation = related_entities - + # Delete source entity - result = await entity_repo.delete(source.id) + result = await entity_repository.delete(source.id) assert result is True - + # Verify relation was cascaded - query = select(Relation).filter(Relation.from_id == source.id) - result = await entity_repo.execute_query(query) - remaining_relations = result.scalars().all() - assert len(remaining_relations) == 0 - - # Verify target entity still exists - target_exists = await entity_repo.find_by_id(target.id) - assert target_exists is not None + async with db.scoped_session(entity_repository.session_maker) as session: + query = select(Relation).filter(Relation.from_id == source.id) + result = await session.execute(query) + remaining_relations = result.scalars().all() + assert len(remaining_relations) == 0 + + # Verify target entity still exists + target_exists = await entity_repository.find_by_id(target.id) + assert target_exists is not None @pytest.mark.asyncio -async def test_delete_nonexistent_entity(entity_repo): +async def test_delete_nonexistent_entity(entity_repository: EntityRepository): """Test deleting an entity that doesn't exist.""" - result = await entity_repo.delete("nonexistent/id") - assert result is False \ No newline at end of file + result = await entity_repository.delete("nonexistent/id") + assert result is False + + +@pytest.mark.asyncio +async def test_search(session_maker, entity_repository: EntityRepository): + """Test searching entities""" + # First create and commit the entities + async with db.scoped_session(session_maker) as session: + entity1 = Entity( + id="20240102-test1", + name="Search Test 1", + entity_type="test", + description="First test entity", + ) + entity2 = Entity( + id="20240102-test2", + name="Search Test 2", + entity_type="other", + description="Second test entity", + ) + session.add_all([entity1, entity2]) + + # Then add observations in a new transaction + async with db.scoped_session(session_maker) as session: + ts = datetime.now(UTC) + stmt = text(""" + INSERT INTO observation (entity_id, content, created_at) + VALUES (:e1_id, :e1_obs, :ts), (:e2_id, :e2_obs, :ts) + """) + await session.execute( + stmt, + { + "e1_id": entity1.id, + "e1_obs": "First observation with searchable content", + "e2_id": entity2.id, + "e2_obs": "Another observation to find", + "ts": ts, + }, + ) + + # Test search by name + results = await entity_repository.search("Search Test") + assert len(results) == 2 + names = {e.name for e in results} + assert "Search Test 1" in names + assert "Search Test 2" in names + + # Test search by type + results = await entity_repository.search("other") + assert len(results) == 1 + assert results[0].entity_type == "other" + + # Test search by observation content + results = await entity_repository.search("searchable") + assert len(results) == 1 + assert results[0].id == entity1.id diff --git a/tests/repository/test_observation_repository.py b/tests/repository/test_observation_repository.py index ed2e010b..1c4ff99b 100644 --- a/tests/repository/test_observation_repository.py +++ b/tests/repository/test_observation_repository.py @@ -3,8 +3,9 @@ import pytest import pytest_asyncio import sqlalchemy -from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.ext.asyncio import async_sessionmaker +from basic_memory import db from basic_memory.models import Entity, Observation from basic_memory.repository.observation_repository import ObservationRepository @@ -82,20 +83,19 @@ async def test_find_by_context( @pytest.mark.asyncio -async def test_delete_observations(session: AsyncSession, repo): +async def test_delete_observations(session_maker: async_sessionmaker, 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" - ) - session.add(entity) - await session.flush() + async with db.scoped_session(session_maker) as session: + entity = Entity( + id="test/test_entity", name="test_entity", entity_type="test", description="Test entity" + ) + session.add(entity) - # Create test observations - obs1 = Observation(entity_id=entity.id, content="Test observation 1") - obs2 = Observation(entity_id=entity.id, content="Test observation 2") - session.add_all([obs1, obs2]) - await session.flush() + # Create test observations + obs1 = Observation(entity_id=entity.id, content="Test observation 1") + obs2 = Observation(entity_id=entity.id, content="Test observation 2") + session.add_all([obs1, obs2]) # Test deletion by entity_id deleted = await repo.delete_by_fields(entity_id=entity.id) @@ -107,19 +107,18 @@ async def test_delete_observations(session: AsyncSession, repo): @pytest.mark.asyncio -async def test_delete_observation_by_id(session: AsyncSession, repo): +async def test_delete_observation_by_id(session_maker: async_sessionmaker, 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" - ) - session.add(entity) - await session.flush() + async with db.scoped_session(session_maker) as session: + entity = Entity( + id="test/test_entity", name="test_entity", entity_type="test", description="Test entity" + ) + session.add(entity) - # Create test observation - obs = Observation(entity_id=entity.id, content="Test observation") - session.add(obs) - await session.flush() + # Create test observation + obs = Observation(entity_id=entity.id, content="Test observation") + session.add(obs) # Test deletion by ID deleted = await repo.delete(obs.id) @@ -131,20 +130,19 @@ async def test_delete_observation_by_id(session: AsyncSession, repo): @pytest.mark.asyncio -async def test_delete_observation_by_content(session: AsyncSession, repo): +async def test_delete_observation_by_content(session_maker: async_sessionmaker, repo): """Test deleting observations by content.""" # Create test entity - entity = Entity( - id="test/test_entity", name="test_entity", entity_type="test", description="Test entity" - ) - session.add(entity) - await session.flush() + async with db.scoped_session(session_maker) as session: + entity = Entity( + id="test/test_entity", name="test_entity", entity_type="test", description="Test entity" + ) + session.add(entity) - # Create test observations - obs1 = Observation(entity_id=entity.id, content="Delete this observation") - obs2 = Observation(entity_id=entity.id, content="Keep this observation") - session.add_all([obs1, obs2]) - await session.flush() + # Create test observations + obs1 = Observation(entity_id=entity.id, content="Delete this observation") + obs2 = Observation(entity_id=entity.id, content="Keep this observation") + session.add_all([obs1, obs2]) # Test deletion by content deleted = await repo.delete_by_fields(content="Delete this observation") diff --git a/tests/repository/test_relation_repository.py b/tests/repository/test_relation_repository.py index 5532624f..455631cc 100644 --- a/tests/repository/test_relation_repository.py +++ b/tests/repository/test_relation_repository.py @@ -4,18 +4,13 @@ import pytest import pytest_asyncio import sqlalchemy +from basic_memory import db 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.""" - return RelationRepository(session) - - -@pytest_asyncio.fixture -async def source_entity(session): +async def source_entity(session_maker): """Create a source entity for testing relations.""" entity = Entity( id="source/test_entity", @@ -23,13 +18,14 @@ async def source_entity(session): entity_type="source", description="Source entity", ) - session.add(entity) - await session.flush() - return entity + async with db.scoped_session(session_maker) as session: + session.add(entity) + await session.flush() + return entity @pytest_asyncio.fixture -async def target_entity(session): +async def target_entity(session_maker): """Create a target entity for testing relations.""" entity = Entity( id="target/test_entity", @@ -37,21 +33,23 @@ async def target_entity(session): entity_type="target", description="Target entity", ) - session.add(entity) - await session.flush() - return entity + async with db.scoped_session(session_maker) as session: + session.add(entity) + await session.flush() + return entity @pytest_asyncio.fixture -async def test_relations(session, source_entity, target_entity): +async def test_relations(session_maker, 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"), ] - session.add_all(relations) - await session.flush() - return relations + async with db.scoped_session(session_maker) as session: + session.add_all(relations) + await session.flush() + return relations @pytest_asyncio.fixture(scope="function") @@ -242,50 +240,52 @@ async def test_delete_by_fields_all_fields( @pytest.mark.asyncio -async def test_delete_relation_by_id(relation_repo, test_relations): +async def test_delete_relation_by_id(relation_repository, test_relations): """Test deleting a relation by ID.""" relation = test_relations[0] - result = await relation_repo.delete(relation.id) + result = await relation_repository.delete(relation.id) assert result is True # Verify deletion - remaining = await relation_repo.find_one( - relation_repo.select(Relation).filter(Relation.id == relation.id) + remaining = await relation_repository.find_one( + relation_repository.select(Relation).filter(Relation.id == relation.id) ) assert remaining is None @pytest.mark.asyncio -async def test_delete_relations_by_type(relation_repo, test_relations): +async def test_delete_relations_by_type(relation_repository, test_relations): """Test deleting relations by type.""" - result = await relation_repo.delete_by_fields(relation_type="connects_to") + result = await relation_repository.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") + remaining = await relation_repository.find_by_type("connects_to") assert len(remaining) == 0 # Verify other type still exists - others = await relation_repo.find_by_type("depends_on") + others = await relation_repository.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 + relation_repository, 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_repository.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) + remaining = await relation_repository.find_by_entities(source_entity.id, target_entity.id) assert len(remaining) == 0 @pytest.mark.asyncio -async def test_delete_nonexistent_relation(relation_repo): +async def test_delete_nonexistent_relation(relation_repository): """Test deleting a relation that doesn't exist.""" - result = await relation_repo.delete_by_fields(relation_type="nonexistent") + result = await relation_repository.delete_by_fields(relation_type="nonexistent") assert result is False diff --git a/tests/repository/test_repository.py b/tests/repository/test_repository.py new file mode 100644 index 00000000..e9f7fad6 --- /dev/null +++ b/tests/repository/test_repository.py @@ -0,0 +1,116 @@ +"""Test repository implementation.""" + +import pytest +from sqlalchemy import String +from sqlalchemy.orm import Mapped, mapped_column + +from basic_memory.models import Base +from basic_memory.repository.repository import Repository + + +class TestModel(Base): + """Test model for repository tests.""" + + __tablename__ = "test_model" + + id: Mapped[str] = mapped_column(String(255), primary_key=True) + name: Mapped[str] = mapped_column(String(255)) + description: Mapped[str | None] = mapped_column(String(255), nullable=True) + + +@pytest.fixture +def repository(session_maker): + """Create a test repository.""" + return Repository(session_maker, TestModel) + + +@pytest.mark.asyncio +async def test_add(repository): + """Test bulk creation of entities.""" + # Create test instances + instance = TestModel(id="test_add", name="Test Add") + await repository.add(instance) + + # Verify we can find in db + found = await repository.find_by_id("test_add") + assert found is not None + assert found.name == "Test Add" + + +@pytest.mark.asyncio +async def test_add_all(repository): + """Test bulk creation of entities.""" + # Create test instances + instances = [TestModel(id=f"test_{i}", name=f"Test {i}") for i in range(3)] + await repository.add_all(instances) + + # Verify we can find them in db + found = await repository.find_by_id("test_0") + assert found is not None + assert found.name == "Test 0" + + +@pytest.mark.asyncio +async def test_bulk_create(repository): + """Test bulk creation of entities.""" + # Create test instances + instances = [TestModel(id=f"test_{i}", name=f"Test {i}") for i in range(3)] + + # Bulk create + await repository.create_all([instance.__dict__ for instance in instances]) + + # Verify we can find them in db + found = await repository.find_by_id("test_0") + assert found is not None + assert found.name == "Test 0" + + +@pytest.mark.asyncio +async def test_find_by_ids(repository): + """Test finding multiple entities by IDs.""" + # Create test data + instances = [TestModel(id=f"test_{i}", name=f"Test {i}") for i in range(5)] + await repository.create_all([instance.__dict__ for instance in instances]) + + # Test finding subset of entities + ids_to_find = ["test_0", "test_2", "test_4"] + found = await repository.find_by_ids(ids_to_find) + assert len(found) == 3 + assert sorted([e.id for e in found]) == sorted(ids_to_find) + + # Test finding with some non-existent IDs + mixed_ids = ["test_0", "nonexistent", "test_4"] + partial_found = await repository.find_by_ids(mixed_ids) + assert len(partial_found) == 2 + assert sorted([e.id for e in partial_found]) == ["test_0", "test_4"] + + # Test with empty list + empty_found = await repository.find_by_ids([]) + assert len(empty_found) == 0 + + # Test with all non-existent IDs + not_found = await repository.find_by_ids(["fake1", "fake2"]) + assert len(not_found) == 0 + + +@pytest.mark.asyncio +async def test_delete_by_ids(repository): + """Test finding multiple entities by IDs.""" + # Create test data + instances = [TestModel(id=f"test_{i}", name=f"Test {i}") for i in range(5)] + await repository.create_all([instance.__dict__ for instance in instances]) + + # Test delete subset of entities + ids_to_delete = ["test_0", "test_2", "test_4"] + deleted_count = await repository.delete_by_ids(ids_to_delete) + assert deleted_count == 3 + + # Test finding subset of entities + ids_to_find = ["test_1", "test_3"] + found = await repository.find_by_ids(ids_to_find) + assert len(found) == 2 + assert sorted([e.id for e in found]) == sorted(ids_to_find) + + assert await repository.find_by_id(ids_to_delete[0]) is None + assert await repository.find_by_id(ids_to_delete[1]) is None + assert await repository.find_by_id(ids_to_delete[2]) is None diff --git a/tests/schemas/test_schemas.py b/tests/schemas/test_schemas.py index 8962b9b6..e3a3f81c 100644 --- a/tests/schemas/test_schemas.py +++ b/tests/schemas/test_schemas.py @@ -21,7 +21,6 @@ def test_entity_in_minimal(): assert entity.entity_type == "test" assert entity.description is None assert entity.observations == [] - assert entity.relations == [] def test_entity_in_complete(): @@ -31,7 +30,6 @@ def test_entity_in_complete(): "entity_type": "test", "description": "A test entity", "observations": ["Test observation"], - "relations": [{"from_id": "123", "to_id": "456", "relation_type": "test_relation"}], } entity = Entity.model_validate(data) assert entity.name == "test_entity" @@ -39,8 +37,6 @@ def test_entity_in_complete(): assert entity.description == "A test entity" assert len(entity.observations) == 1 assert entity.observations[0] == "Test observation" - assert len(entity.relations) == 1 - assert entity.relations[0].from_id == "123" def test_entity_in_validation(): @@ -118,7 +114,6 @@ def test_optional_fields(): entity = Entity.model_validate({"name": "test", "entity_type": "test"}) assert entity.description is None assert entity.observations == [] - assert entity.relations == [] # Create with empty optional fields entity = Entity.model_validate( @@ -127,12 +122,10 @@ def test_optional_fields(): "entity_type": "test", "description": None, "observations": [], - "relations": [], } ) assert entity.description is None assert entity.observations == [] - assert entity.relations == [] # Create with some optional fields entity = Entity.model_validate( @@ -140,7 +133,6 @@ def test_optional_fields(): ) assert entity.description == "test" assert entity.observations == [] - assert entity.relations == [] def test_search_nodes_input(): diff --git a/tests/services/test_entity_service.py b/tests/services/test_entity_service.py index 9660fe75..2b836c3a 100644 --- a/tests/services/test_entity_service.py +++ b/tests/services/test_entity_service.py @@ -1,49 +1,124 @@ """Tests for EntityService.""" + import pytest +import pytest_asyncio +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from basic_memory.fileio import EntityNotFoundError from basic_memory.models import Entity as EntityModel +from basic_memory.repository.entity_repository import EntityRepository from basic_memory.schemas import Entity +from basic_memory.services.entity_service import EntityService pytestmark = pytest.mark.asyncio -async def test_create_entity_success(entity_service): +@pytest_asyncio.fixture +async def entity_repository(session_maker: async_sessionmaker[AsyncSession]) -> EntityRepository: + """Create an EntityRepository instance.""" + return EntityRepository(session_maker) + + +@pytest_asyncio.fixture +async def entity_service(entity_repository: EntityRepository) -> EntityService: + """Create EntityService with repository.""" + return EntityService(entity_repository) + + +async def test_create_entity(entity_service: EntityService): """Test successful entity creation.""" entity_data = Entity( name="Test Entity", entity_type="test", - description="A test entity description" + description="A test entity description", + observations=["this is a test observation"], ) - + # Act entity = await entity_service.create_entity(entity_data) - + # Assert Entity assert isinstance(entity, EntityModel) assert entity.name == "Test Entity" assert entity.entity_type == "test" assert entity.description == "A test entity description" assert entity.created_at is not None + assert entity.observations[0].content == "this is a test observation" + assert len(entity.relations) == 0 # Verify we can retrieve it retrieved = await entity_service.get_entity(entity.id) assert retrieved.description == "A test entity description" + assert retrieved.name == "Test Entity" + assert retrieved.entity_type == "test" + assert retrieved.description == "A test entity description" + assert retrieved.created_at is not None + assert retrieved.observations[0].content == "this is a test observation" -async def test_get_by_type_and_name(entity_service): + +async def test_create_entities(entity_service: EntityService): + """Test successful entity creation.""" + entity_data = [ + Entity( + name="Test Entity_1", + entity_type="test", + description="A test entity description", + observations=["this is a test observation"], + ), + Entity( + name="Test Entity_2", + entity_type="test", + description="A test entity description", + observations=["this is a test observation"], + ), + ] + + # Act + entities = await entity_service.create_entities(entity_data) + + # Assert Entity + assert len(entities) == 2 + entity1 = entities[0] + assert isinstance(entity1, EntityModel) + assert entity1.name == "Test Entity_1" + assert entity1.entity_type == "test" + assert entity1.description == "A test entity description" + assert entity1.created_at is not None + assert entity1.observations[0].content == "this is a test observation" + assert len(entity1.relations) == 0 + + entity2 = entities[1] + assert isinstance(entity1, EntityModel) + assert entity2.name == "Test Entity_2" + assert entity2.entity_type == "test" + assert entity2.description == "A test entity description" + assert entity2.created_at is not None + assert entity2.observations[0].content == "this is a test observation" + + # Verify we can retrieve them + retrieved1 = await entity_service.get_entity(entity1.id) + assert retrieved1.description == "A test entity description" + + retrieved2 = await entity_service.get_entity(entity2.id) + assert retrieved2.description == "A test entity description" + + +async def test_get_by_type_and_name(entity_service: EntityService): """Test finding entity by type and name combination.""" # Create two entities with same name but different types entity1_data = Entity( name="Test Entity", entity_type="type1", - description="First test entity" + description="First test entity", + observations=[], ) entity1 = await entity_service.create_entity(entity1_data) entity2_data = Entity( name="Test Entity", # Same name entity_type="type2", # Different type - description="Second test entity" + description="Second test entity", + observations=[], ) entity2 = await entity_service.create_entity(entity2_data) @@ -65,13 +140,11 @@ async def test_get_by_type_and_name(entity_service): with pytest.raises(EntityNotFoundError): await entity_service.get_by_type_and_name("nonexistent", "Test Entity") -async def test_create_entity_no_description(entity_service): + +async def test_create_entity_no_description(entity_service: EntityService): """Test creating entity without description (should be None).""" - entity_data = Entity( - name="Test Entity", - entity_type="test", - ) - + entity_data = Entity(name="Test Entity", entity_type="test", observations=[], relations=[]) + entity = await entity_service.create_entity(entity_data) assert entity.description is None @@ -79,37 +152,38 @@ async def test_create_entity_no_description(entity_service): retrieved = await entity_service.get_entity(entity.id) assert retrieved.description is None -async def test_get_entity_success(entity_service): + +async def test_get_entity_success(entity_service: EntityService): """Test successful entity retrieval.""" - # Arrange entity_data = Entity( name="Test Entity", entity_type="test", - description="Test description" + description="Test description", + observations=[], ) created = await entity_service.create_entity(entity_data) - + # Act retrieved = await entity_service.get_entity(created.id) - + # Assert assert isinstance(retrieved, EntityModel) assert retrieved.id == created.id assert retrieved.name == created.name assert retrieved.entity_type == created.entity_type assert retrieved.description == "Test description" - # relations are tested in test_memory_service -async def test_update_entity_description(entity_service): + +async def test_update_entity_description(entity_service: EntityService): """Test updating an entity's description.""" - # Create entity with description entity_data = Entity( name="Test Entity", entity_type="test", - description="Initial description" + description="Initial description", + observations=[], ) entity = await entity_service.create_entity(entity_data) - + # Update description updated = await entity_service.update_entity(entity.id, {"description": "Updated description"}) assert updated.description == "Updated description" @@ -118,16 +192,17 @@ async def test_update_entity_description(entity_service): retrieved = await entity_service.get_entity(entity.id) assert retrieved.description == "Updated description" -async def test_update_entity_description_to_none(entity_service): + +async def test_update_entity_description_to_none(entity_service: EntityService): """Test updating an entity's description to None.""" - # Create entity with description entity_data = Entity( name="Test Entity", entity_type="test", - description="Initial description" + description="Initial description", + observations=[], ) entity = await entity_service.create_entity(entity_data) - + # Update description to None updated = await entity_service.update_entity(entity.id, {"description": None}) assert updated.description is None @@ -136,65 +211,42 @@ async def test_update_entity_description_to_none(entity_service): retrieved = await entity_service.get_entity(entity.id) assert retrieved.description is None -async def test_delete_entity_success(entity_service): + +async def test_delete_entity_success(entity_service: EntityService): """Test successful entity deletion.""" - # Arrange - entity_data = Entity( - name="Test Entity", - entity_type="test", - ) + entity_data = Entity(name="Test Entity", entity_type="test", observations=[], relations=[]) entity = await entity_service.create_entity(entity_data) # Act result = await entity_service.delete_entity(entity.id) - + # Assert assert result is True with pytest.raises(EntityNotFoundError): await entity_service.get_entity(entity.id) -# Error Path Tests -async def test_get_entity_not_found(entity_service): +async def test_get_entity_not_found(entity_service: EntityService): """Test handling of non-existent entity retrieval.""" with pytest.raises(EntityNotFoundError): await entity_service.get_entity("nonexistent-id") -async def test_create_entity_db_error(entity_service, monkeypatch): - """Test handling of database errors during creation.""" - # Arrange - make db operations fail - async def mock_create(*args, **kwargs): - raise Exception("Mock DB error") - monkeypatch.setattr(entity_service.entity_repo, "create", mock_create) - entity_data = Entity( - name="Test Entity", - entity_type="test", - description="Test description" - ) - - # Act/Assert - with pytest.raises(Exception, match="Mock DB error"): - await entity_service.create_entity(entity_data) - -async def test_delete_nonexistent_entity(entity_service): +async def test_delete_nonexistent_entity(entity_service: EntityService): """Test deleting an entity that doesn't exist.""" - await entity_service.delete_entity("nonexistent-id") - # If we get here, the deletion succeeded or failed silently as expected + result = await entity_service.delete_entity("nonexistent-id") + assert result is False -# Edge Cases -async def test_create_entity_with_special_chars(entity_service): +async def test_create_entity_with_special_chars(entity_service: EntityService): """Test entity creation with special characters in name and description.""" name = "Test & Entity! With @ Special #Chars" description = "Description with $pecial chars & symbols!" entity_data = Entity( - name=name, - entity_type="test", - description=description + name=name, entity_type="test", description=description, observations=[], relations=[] ) entity = await entity_service.create_entity(entity_data) - + assert entity.name == name assert entity.description == description @@ -202,32 +254,35 @@ async def test_create_entity_with_special_chars(entity_service): retrieved = await entity_service.get_entity(entity.id) assert retrieved.description == description -async def test_entity_id_generation(entity_service): + +async def test_entity_id_generation(entity_service: EntityService): """Test that entities get unique IDs generated correctly.""" entity_data = Entity( name="Test Entity", entity_type="test", description="Test description", - observations=[] + observations=[], ) - - entity = await entity_service.create_entity(entity_data) - - assert entity.id # ID should be generated - assert "test/test_entity" == entity.id # Should contain normalized name -async def test_create_entity_long_description(entity_service): + entity = await entity_service.create_entity(entity_data) + + assert entity.id # ID should be generated + assert "test/test_entity" == entity.id # Should contain normalized name + + +async def test_create_entity_long_description(entity_service: EntityService): """Test creating entity with a long description.""" long_description = "A" * 1000 # 1000 character description entity_data = Entity( name="Test Entity", entity_type="test", - description=long_description + description=long_description, + observations=[], ) - + entity = await entity_service.create_entity(entity_data) assert entity.description == long_description # Verify after retrieval retrieved = await entity_service.get_entity(entity.id) - assert retrieved.description == long_description \ No newline at end of file + assert retrieved.description == long_description diff --git a/tests/services/test_memory_service.py b/tests/services/test_memory_service.py deleted file mode 100644 index f4fc00b3..00000000 --- a/tests/services/test_memory_service.py +++ /dev/null @@ -1,395 +0,0 @@ -"""Tests for MemoryService delete operations.""" - -import pytest -from basic_memory.services import MemoryService -from basic_memory.fileio import read_entity_file -from basic_memory.schemas import CreateEntityRequest, CreateRelationsRequest, AddObservationsRequest, Relation - -test_create_entity_input = [ - { - "name": "Test_Entity_1", - "entity_type": "test", - "observations": ["Observation 1.1", "Observation 1.2"] - }, - { - "name": "Test_Entity_2", - "entity_type": "test", - "observations": ["Observation 2.1", "Observation 2.2"] - } -] - - -@pytest.mark.asyncio -async def test_create_entities(memory_service: MemoryService): - """Should create multiple entities in parallel with their observations.""" - - entity_input = CreateEntityRequest.model_validate({"entities": test_create_entity_input}) - entities = await memory_service.create_entities(entity_input.entities) - - # Verify the SQLAlchemy models were created - assert len(entities) == 2 - - # Check first entity - assert entities[0].name == "Test_Entity_1" - assert entities[0].entity_type == "test" - assert len(entities[0].observations) == 2 - assert entities[0].observations[0].content == "Observation 1.1" - assert entities[0].observations[1].content == "Observation 1.2" - - # Check second entity - assert entities[1].name == "Test_Entity_2" - assert entities[1].entity_type == "test" - assert len(entities[1].observations) == 2 - assert entities[1].observations[0].content == "Observation 2.1" - assert entities[1].observations[1].content == "Observation 2.2" - - # Verify files were created (returns Pydantic Entity) - file_entity1 = await read_entity_file(memory_service.entities_path, entities[0].id) - file_entity2 = await read_entity_file(memory_service.entities_path, entities[1].id) - - assert file_entity1.name == "Test_Entity_1" - assert file_entity2.name == "Test_Entity_2" - - # Verify files are present - for entity in entities: - entity_file_path = memory_service.get_entity_file_path(entity.id) - assert entity_file_path.exists() - - -@pytest.mark.asyncio -async def test_add_observations(memory_service: MemoryService): - """Should add observations to an existing entity.""" - entity_input = CreateEntityRequest.model_validate({"entities": test_create_entity_input}) - entities = await memory_service.create_entities([entity_input.entities[0]]) - entity = entities[0] - - # Create observations input - observations_data = { - "entity_id": entity.id, - "observations": [ - "New observation 1", - "New observation 2" - ] - } - - # Add observations - returns List[models.Observation] - observation_input = AddObservationsRequest.model_validate(observations_data) - added_observations = await memory_service.add_observations(observation_input) - - # Check the SQLAlchemy model results - assert len(added_observations) == 2 - assert added_observations[0].content == "New observation 1" - assert added_observations[0].context is None - assert added_observations[1].content == "New observation 2" - assert added_observations[1].context is None - - # Verify file was updated - returns Pydantic Entity - updated_entity = await read_entity_file(memory_service.entities_path, entity.id) - assert len(updated_entity.observations) == 4 # 2 original + 2 new - # assert updated_entity.observations[2] == "New observation 1" - # assert updated_entity.observations[3] == "New observation 2" - - # Verify database - returns SQLAlchemy Entity - db_entity = await memory_service.entity_service.get_entity(entity.id) - assert len(db_entity.observations) == 4 - - -@pytest.mark.asyncio -async def test_add_observations_nonexistent_entity(memory_service: MemoryService): - """Should raise an appropriate error when adding observations to a non-existent entity.""" - observations_data = { - "entity_id": "nonexistent-id", - "observations": ["Test observation"] - } - - with pytest.raises(Exception) as exc: # We might want to define a specific error type - observation_input = AddObservationsRequest.model_validate(observations_data) - await memory_service.add_observations(observation_input) - - -@pytest.mark.asyncio -async def test_create_relations(memory_service: MemoryService): - """Should create relations between entities and update both filesystem and database.""" - entity_input = CreateEntityRequest.model_validate({"entities": test_create_entity_input}) - entities = await memory_service.create_entities(entity_input.entities) - entity1, entity2 = entities - - # Create test relations data using actual entity IDs - test_relations_data = [ - { - "from_id": entity1.id, - "to_id": entity2.id, - "relation_type": "connects_to" - }, - { - "from_id": entity2.id, - "to_id": entity1.id, - "relation_type": "references", - "context": "test context" - } - ] - - # Create relations - returns List[models.Relation] - input_args = CreateRelationsRequest.model_validate({"relations": test_relations_data}) - relations = await memory_service.create_relations(input_args.relations) - - # Verify SQLAlchemy Relation models were created - assert len(relations) == 2 - - # Check first relation - assert relations[0].from_id == entity1.id - assert relations[0].to_id == entity2.id - assert relations[0].relation_type == "connects_to" - assert relations[0].context is None - - # Check second relation - assert relations[1].from_id == entity2.id - assert relations[1].to_id == entity1.id - assert relations[1].relation_type == "references" - assert relations[1].context == "test context" - - # Read updated entities from filesystem - returns Pydantic Entities - updated_entity1 = await read_entity_file(memory_service.entities_path, entity1.id) - updated_entity2 = await read_entity_file(memory_service.entities_path, entity2.id) - - # Verify relations were added to entity1 in filesystem - assert len(updated_entity1.relations) == 1 - assert updated_entity1.relations[0].from_id == entity1.id - assert updated_entity1.relations[0].to_id == entity2.id - assert updated_entity1.relations[0].relation_type == "connects_to" - - # Verify relations were added to entity2 in filesystem - assert len(updated_entity2.relations) == 1 - assert updated_entity2.relations[0].from_id == entity2.id - assert updated_entity2.relations[0].to_id == entity1.id - assert updated_entity2.relations[0].relation_type == "references" - assert updated_entity2.relations[0].context == "test context" - - # Now verify database state - get SQLAlchemy Entity models - db_entity1 = await memory_service.entity_service.get_entity(entity1.id) - db_entity2 = await memory_service.entity_service.get_entity(entity2.id) - - # Entity 1 should have one outgoing relation to entity 2 - assert len(db_entity1.outgoing_relations) == 1 - outgoing = db_entity1.outgoing_relations[0] - assert outgoing.from_id == entity1.id - assert outgoing.to_id == entity2.id - assert outgoing.relation_type == "connects_to" - assert outgoing.context is None - - # Entity 1 should have one incoming relation from entity 2 - assert len(db_entity1.incoming_relations) == 1 - incoming = db_entity1.incoming_relations[0] - assert incoming.from_id == entity2.id - assert incoming.to_id == entity1.id - assert incoming.relation_type == "references" - assert incoming.context == "test context" - - # Verify the same for entity 2 (reversed) - assert len(db_entity2.outgoing_relations) == 1 - assert len(db_entity2.incoming_relations) == 1 - - -@pytest.mark.asyncio -async def test_create_relations_with_invalid_entity_id(memory_service: MemoryService): - """Should raise an appropriate error when trying to create relations with non-existent entity IDs.""" - # Create one entity - returns SQLAlchemy Entity - entity_input = CreateEntityRequest.model_validate({"entities": test_create_entity_input}) - entities = await memory_service.create_entities([entity_input.entities[0]]) - entity1 = entities[0] - - # Try to create relation with non-existent entity ID - bad_relation = { - "from_id": entity1.id, - "to_id": "nonexistent-id", - "relation_type": "connects_to" - } - - with pytest.raises(Exception) as exc: # We might want to define a specific error type - await memory_service.create_relations([Relation.model_validate(bad_relation)]) - - -@pytest.mark.asyncio -async def test_delete_entities(memory_service: MemoryService): - """Test deleting an entity deletes file and database record.""" - # Write the entity files - entity_input = CreateEntityRequest.model_validate({"entities": test_create_entity_input}) - entities = await memory_service.create_entities(entity_input.entities) - - # Verify files are present - for entity in entities: - entity_file_path = memory_service.get_entity_file_path(entity.id) - assert entity_file_path.exists() - - # Delete the entities - result = await memory_service.delete_entities([entity.id for entity in entities]) - assert result is True - - # Verify files are gone - for entity in entities: - assert not memory_service.get_entity_file_path(entity.id).exists() - - # Verify database records are deleted - for entity in entities: - deleted = await memory_service.entity_service.entity_repo.find_by_id(entity.id) - assert deleted is None - -@pytest.mark.asyncio -async def test_delete_entity_cascades(memory_service): - """Test deleting an entity cascades to observations.""" - - # Write the entity files - create_entity_input = test_create_entity_input[0] - entity_input = CreateEntityRequest.model_validate({"entities": [create_entity_input]}) - entities = await memory_service.create_entities(entity_input.entities) - assert len(entities) == 1 - test_entity = entities[0] - - # Delete the entity - result = await memory_service.delete_entities([test_entity.id]) - assert result is True - - # Verify file is gone - assert not memory_service.get_entity_file_path(test_entity.id).exists() - - # Verify observations are gone from database - obsv = await memory_service.observation_service.observation_repo.find_by_entity(test_entity.id) - assert len(obsv) == 0 - - # Verify observations are gone from database - rels = await memory_service.relation_service.relation_repo.find_by_entity(test_entity.id) - assert len(rels) == 0 - -@pytest.mark.asyncio -async def test_delete_observations(memory_service): - """Test deleting specific observations.""" - - # Set up entity file with observations - create_entity_input = test_create_entity_input[0] - create_entity_input['observations'] = ["First observation", "Second observation", "Third observation"] - create_entity = CreateEntityRequest.model_validate({"entities": [create_entity_input]}) - - entities = await memory_service.create_entities(create_entity.entities) - assert len(entities) == 1 - test_entity = entities[0] - - # Delete two observations - to_delete = ["First observation", "Second observation"] - result = await memory_service.delete_observations(test_entity.id, to_delete) - assert result is True - - # Verify file updated - updated_entities = await memory_service.open_nodes([test_entity.id]) - assert len(updated_entities) == 1 - updated_entity = updated_entities[0] - - assert len(updated_entity.observations) == 1 - assert updated_entity.observations[0] == "Third observation" - - # Verify database updated - remaining_db = await memory_service.observation_service.observation_repo.find_by_entity(test_entity.id) - assert len(remaining_db) == 1 - assert remaining_db[0].content == "Third observation" - -@pytest.mark.asyncio -async def test_delete_relations(memory_service): - """Test deleting relations between entities.""" - entity_input = CreateEntityRequest.model_validate({"entities": test_create_entity_input}) - entities = await memory_service.create_entities(entity_input.entities) - entity1, entity2 = entities - - # Create test relations data using actual entity IDs - test_relations_data = [ - { - "from_id": entity1.id, - "to_id": entity2.id, - "relation_type": "connects_to" - }, - { - "from_id": entity2.id, - "to_id": entity1.id, - "relation_type": "references", - "context": "test context" - } - ] - - # Create relations - returns List[models.Relation] - input_args = CreateRelationsRequest.model_validate({"relations": test_relations_data}) - relations = await memory_service.create_relations(input_args.relations) - - # Delete the relation - to_delete = [{ - "from_id": entity1.id, - "to_id": entity2.id, - "relation_type": "connects_to" - }] - result = await memory_service.delete_relations(to_delete) - assert result is True - - # Verify relation removed from source entity file - # Verify file updated - updated_entities = await memory_service.open_nodes([entity1.id]) - assert len(updated_entities) == 1 - updated_entity = updated_entities[0] - - assert len(updated_entity.relations) == 0 - - # Verify relation removed from database - relations = await memory_service.relation_service.relation_repo.find_by_entities(entity1.id, entity2.id) - assert len(relations) == 0 - -@pytest.mark.asyncio -async def test_delete_nonexistent_entity(memory_service): - """Test deleting an entity that doesn't exist.""" - result = await memory_service.delete_entities(["nonexistent/id"]) - assert result is False - -@pytest.mark.asyncio -async def test_delete_nonexistent_observations(memory_service): - """Test deleting observations that don't exist.""" - - # Set up entity file with observations - create_entity_input = test_create_entity_input[0] - create_entity_input['observations'] = ["First observation", "Second observation", "Third observation"] - create_entity = CreateEntityRequest.model_validate({"entities": [create_entity_input]}) - - entities = await memory_service.create_entities(create_entity.entities) - assert len(entities) == 1 - test_entity = entities[0] - - result = await memory_service.delete_observations(test_entity.id, ["Nonexistent"]) - assert result is False - -@pytest.mark.asyncio -async def test_delete_nonexistent_relations(memory_service): - """Test deleting relations that don't exist.""" - entity_input = CreateEntityRequest.model_validate({"entities": test_create_entity_input}) - entities = await memory_service.create_entities(entity_input.entities) - entity1, entity2 = entities - - # Create test relations data using actual entity IDs - test_relations_data = [ - { - "from_id": entity1.id, - "to_id": entity2.id, - "relation_type": "connects_to" - }, - { - "from_id": entity2.id, - "to_id": entity1.id, - "relation_type": "references", - "context": "test context" - } - ] - - # Create relations - returns List[models.Relation] - input_args = CreateRelationsRequest.model_validate({"relations": test_relations_data}) - relations = await memory_service.create_relations(input_args.relations) - - to_delete = [{ - "from_id": entity1.id, - "to_id": entity2.id, - "relation_type": "nonexistent" - }] - result = await memory_service.delete_relations(to_delete) - assert result is False \ No newline at end of file diff --git a/tests/services/test_observation_service.py b/tests/services/test_observation_service.py index 21e99057..4f23d64b 100644 --- a/tests/services/test_observation_service.py +++ b/tests/services/test_observation_service.py @@ -1,7 +1,7 @@ """Tests for the ObservationService.""" import pytest import pytest_asyncio -from pathlib import Path +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from basic_memory.models import Entity, Observation from basic_memory.repository.observation_repository import ObservationRepository @@ -9,95 +9,34 @@ from basic_memory.services.observation_service import ObservationService @pytest_asyncio.fixture -async def observation_service(session): - """Create a test ObservationService.""" - repo = ObservationRepository(session) - return ObservationService(Path("/test"), repo) +async def observation_repository(session_maker: async_sessionmaker[AsyncSession]) -> ObservationRepository: + """Create an ObservationRepository instance.""" + return ObservationRepository(session_maker) @pytest_asyncio.fixture -async def test_entity(session): +async def observation_service(observation_repository: ObservationRepository) -> ObservationService: + """Create ObservationService with repository.""" + return ObservationService(observation_repository) + + +@pytest_asyncio.fixture +async def test_entity(session_maker: async_sessionmaker[AsyncSession]) -> Entity: """Create a test entity.""" - entity = Entity( - id="test/test_entity", - name="test_entity", - entity_type="test", - description="Test entity" - ) - session.add(entity) - await session.flush() - return entity - - -@pytest_asyncio.fixture -async def test_add_observation_success(observation_service, test_entity): - """Test successful observation addition.""" - - # Act - observations = await observation_service.add_observations(test_entity.id, ["New observation"]) - - # Assert - 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 - assert any(obs.content == "New observation" - for obs in db_observations) - - -@pytest_asyncio.fixture -async def test_search_observations(observation_service, test_entity): - """Test searching observations across entities.""" - # Arrange - await observation_service.add_observations( - test_entity.id, - ["Unique test content", "Other content"] - ) - - # Act - results = await observation_service.search_observations("unique") - - # Assert - assert len(results) == 1 - assert results[0].content == "Unique test content" - - -@pytest_asyncio.fixture -async def test_observation_with_special_characters(observation_service, test_entity): - """Test handling observations with special characters.""" - content = "Test & observation with @#$% special chars!" - - observations = await observation_service.add_observations(test_entity.id, [content]) - assert observations[0].content == content - - -@pytest_asyncio.fixture -async def test_very_long_observation(observation_service, test_entity): - """Test handling very long observation content.""" - long_content = "Very long observation " * 100 # ~1800 characters - - observations = await observation_service.add_observations(test_entity.id, [long_content]) - assert observations[0].content == long_content - - -@pytest_asyncio.fixture -async def test_observations(session, test_entity): - """Create test observations.""" - observations = [ - Observation(entity_id=test_entity.id, content="First observation"), - Observation(entity_id=test_entity.id, content="Second observation"), - Observation(entity_id=test_entity.id, content="Third observation") - ] - session.add_all(observations) - await session.flush() - return observations + async with session_maker() as session: + entity = Entity( + id="test/test_entity", + name="test_entity", + entity_type="test", + description="Test entity" + ) + session.add(entity) + await session.commit() + return entity @pytest.mark.asyncio -async def test_add_observations(observation_service, test_entity): +async def test_add_observations(observation_service: ObservationService, test_entity: Entity): """Test adding observations to an entity.""" observations = ["Test observation 1", "Test observation 2"] @@ -110,59 +49,130 @@ async def test_add_observations(observation_service, test_entity): @pytest.mark.asyncio -async def test_delete_observations(observation_service, test_entity, test_observations): - """Test deleting specific observations from an entity.""" - contents_to_delete = ["First observation", "Second observation"] - - result = await observation_service.delete_observations(test_entity.id, contents_to_delete) - - assert result is True - - # Verify observations were deleted - remaining = await observation_service.observation_repo.find_by_entity(test_entity.id) - assert len(remaining) == 1 - assert remaining[0].content == "Third observation" - - -@pytest.mark.asyncio -async def test_delete_by_entity(observation_service, test_entity, test_observations): - """Test deleting all observations for an entity.""" - result = await observation_service.delete_by_entity(test_entity.id) - - assert result is True - - # Verify all observations were deleted - remaining = await observation_service.observation_repo.find_by_entity(test_entity.id) - assert len(remaining) == 0 - - -@pytest.mark.asyncio -async def test_delete_nonexistent_observation(observation_service, test_entity): - """Test deleting observations that don't exist.""" - result = await observation_service.delete_observations(test_entity.id, ["Nonexistent observation"]) - - assert result is False - - -@pytest.mark.asyncio -async def test_delete_observations_invalid_entity(observation_service): - """Test deleting observations for an entity that doesn't exist.""" - result = await observation_service.delete_observations("invalid_entity", ["Test observation"]) - - # Should return False since there were no observations to delete - assert result is False - - -@pytest.mark.asyncio -async def test_get_observations_by_context(observation_service, session, test_entity): - """Test getting observations by context.""" - obs = Observation( - entity_id=test_entity.id, - content="Contextual observation", - context="test_context" +async def test_search_observations(observation_service: ObservationService, test_entity: Entity): + """Test searching observations across entities.""" + # First add some observations + await observation_service.add_observations( + test_entity.id, + ["Unique test content", "Other content"] ) - session.add(obs) - await session.flush() + + # Search for them + results = await observation_service.search_observations("unique") + + assert len(results) == 1 + assert results[0].content == "Unique test content" + + +@pytest.mark.asyncio +async def test_delete_observations( + observation_service: ObservationService, + test_entity: Entity +): + """Test deleting specific observations from an entity.""" + # First add observations + await observation_service.add_observations( + test_entity.id, + ["First observation", "Second observation", "Third observation"] + ) + + # Then delete some + contents_to_delete = ["First observation", "Second observation"] + result = await observation_service.delete_observations(test_entity.id, contents_to_delete) + assert result is True + + # Verify through search + results = await observation_service.search_observations("Third") + assert len(results) == 1 + assert results[0].content == "Third observation" + + +@pytest.mark.asyncio +async def test_delete_by_entity( + observation_service: ObservationService, + test_entity: Entity +): + """Test deleting all observations for an entity.""" + # First add observations + await observation_service.add_observations( + test_entity.id, + ["First observation", "Second observation"] + ) + + # Delete all observations for entity + result = await observation_service.delete_by_entity(test_entity.id) + assert result is True + + # Verify through search + results = await observation_service.search_observations("observation") + assert len(results) == 0 + + +@pytest.mark.asyncio +async def test_delete_nonexistent_observation( + observation_service: ObservationService, + test_entity: Entity +): + """Test deleting observations that don't exist.""" + result = await observation_service.delete_observations( + test_entity.id, + ["Nonexistent observation"] + ) + assert result is False + + +@pytest.mark.asyncio +async def test_delete_observations_invalid_entity(observation_service: ObservationService): + """Test deleting observations for an entity that doesn't exist.""" + result = await observation_service.delete_observations( + "invalid_entity", + ["Test observation"] + ) + assert result is False + + +@pytest.mark.asyncio +async def test_observation_with_special_characters( + observation_service: ObservationService, + test_entity: Entity +): + """Test handling observations with special characters.""" + content = "Test & observation with @#$% special chars!" + + results = await observation_service.add_observations(test_entity.id, [content]) + assert len(results) == 1 + assert results[0].content == content + + +@pytest.mark.asyncio +async def test_very_long_observation( + observation_service: ObservationService, + test_entity: Entity +): + """Test handling very long observation content.""" + long_content = "Very long observation " * 100 # ~1800 characters + + results = await observation_service.add_observations(test_entity.id, [long_content]) + assert len(results) == 1 + assert results[0].content == long_content + + +@pytest.mark.asyncio +async def test_get_observations_by_context( + observation_service: ObservationService, + test_entity: Entity, + session_maker: async_sessionmaker[AsyncSession] +): + """Test getting observations by context.""" + # Create observation with context + async with session_maker() as session: + obs = Observation( + entity_id=test_entity.id, + content="Contextual observation", + context="test_context" + ) + session.add(obs) + await session.commit() results = await observation_service.get_observations_by_context("test_context") diff --git a/tests/services/test_relation_service.py b/tests/services/test_relation_service.py index 04b9ed65..bf0a31da 100644 --- a/tests/services/test_relation_service.py +++ b/tests/services/test_relation_service.py @@ -1,78 +1,191 @@ """Tests for RelationService.""" + import pytest import pytest_asyncio +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from basic_memory.models import Entity as EntityModel +from basic_memory.repository.relation_repository import RelationRepository from basic_memory.schemas import Entity, Relation - -pytestmark = pytest.mark.asyncio +from basic_memory.services.relation_service import RelationService @pytest_asyncio.fixture -async def sample_entities(entity_service): - """Create two sample entities for testing relations""" - entity1_data = Entity( - name="test_entity_1", - entity_type="test_type", - observations=[], - relations=[] - ) - entity2_data = Entity( - name="test_entity_2", - 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 +async def relation_repository( + session_maker: async_sessionmaker[AsyncSession], +) -> RelationRepository: + """Create a RelationRepository instance.""" + return RelationRepository(session_maker) -def normalize_whitespace(s: str) -> str: - """Normalize whitespace in a string for comparison.""" - return ' '.join(s.split()) +@pytest_asyncio.fixture +async def relation_service(relation_repository: RelationRepository) -> RelationService: + """Create RelationService with repository.""" + return RelationService(relation_repository) -async def test_create_relation(relation_service, sample_entities): - """Test creating a basic relation between two entities""" - entity1, entity2 = sample_entities - - relation_data = Relation( - from_id=entity1.id, - to_id=entity2.id, - relation_type="test_relation" - ) - +@pytest_asyncio.fixture +async def test_entities( + session_maker: async_sessionmaker[AsyncSession], +) -> tuple[EntityModel, EntityModel]: + """Create two test entities.""" + async with session_maker() as session: + entity1 = EntityModel( + id="test/test_entity_1", + name="test_entity_1", + entity_type="test", + description="Test entity 1", + ) + entity2 = EntityModel( + id="test/test_entity_2", + name="test_entity_2", + entity_type="test", + description="Test entity 2", + ) + session.add_all([entity1, entity2]) + await session.commit() + return entity1, entity2 + + +@pytest.mark.asyncio +async def test_create_relation( + relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel] +): + """Test creating a basic relation between two entities.""" + entity1, entity2 = test_entities + + relation_data = Relation(from_id=entity1.id, to_id=entity2.id, relation_type="test_relation") + 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 database was updated with correct IDs - db_relation = await relation_service.relation_repo.find_by_id(relation.id) - assert db_relation is not None - assert db_relation.from_id == entity1.id - assert db_relation.to_id == entity2.id - assert db_relation.relation_type == "test_relation" + +@pytest.mark.asyncio +async def test_create_relations( + relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel] +): + """Test creating a basic relation between two entities.""" + entity1, entity2 = test_entities + + relation_data = [ + Relation(from_id=entity1.id, to_id=entity2.id, relation_type="type_0"), + Relation(from_id=entity1.id, to_id=entity2.id, relation_type="type_1"), + ] + + relations = await relation_service.create_relations(relation_data) + assert len(relations) == 2 + relation0 = relations[0] + assert relation0.from_id == entity1.id + assert relation0.to_id == entity2.id + assert relation0.relation_type == "type_0" + + relation1 = relations[1] + assert relation1.from_id == entity1.id + assert relation1.to_id == entity2.id + assert relation1.relation_type == "type_1" +@pytest.mark.asyncio +async def test_create_relation_with_context( + relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel] +): + """Test creating a relation with context information.""" + entity1, entity2 = test_entities -async def test_create_relation_with_context(relation_service, sample_entities): - """Test creating a relation with context information""" - entity1, entity2 = sample_entities - relation_data = Relation( - from_id=entity1.id, - to_id=entity2.id, - relation_type="test_relation", - context="test context" + 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 database - db_relation = await relation_service.relation_repo.find_by_id(relation.id) - assert db_relation.context == "test context" \ No newline at end of file + +@pytest.mark.asyncio +async def test_delete_relation( + relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel] +): + """Test deleting a relation between entities.""" + entity1, entity2 = test_entities + + # Create a relation first + relation_data = Relation(from_id=entity1.id, to_id=entity2.id, relation_type="test_relation") + await relation_service.create_relation(relation_data) + + # Create Entity schema instances for delete_relation call + from_entity = Entity( + id=entity1.id, + name=entity1.name, + entity_type=entity1.entity_type, + description=entity1.description, + observations=[], + relations=[], + ) + to_entity = Entity( + id=entity2.id, + name=entity2.name, + entity_type=entity2.entity_type, + description=entity2.description, + observations=[], + relations=[], + ) + + # Delete the relation + result = await relation_service.delete_relation(from_entity, to_entity, "test_relation") + + assert result is True + + +@pytest.mark.asyncio +async def test_delete_nonexistent_relation( + relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel] +): + """Test trying to delete a relation that doesn't exist.""" + entity1, entity2 = test_entities + + from_entity = Entity( + id=entity1.id, + name=entity1.name, + entity_type=entity1.entity_type, + description=entity1.description, + observations=[], + relations=[], + ) + to_entity = Entity( + id=entity2.id, + name=entity2.name, + entity_type=entity2.entity_type, + description=entity2.description, + observations=[], + relations=[], + ) + + result = await relation_service.delete_relation(from_entity, to_entity, "nonexistent_relation") + + assert result is False + + +@pytest.mark.asyncio +async def test_delete_relations_by_criteria( + relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel] +): + """Test deleting relations by criteria.""" + entity1, entity2 = test_entities + + # Create test relations + await relation_service.create_relation( + Relation(from_id=entity1.id, to_id=entity2.id, relation_type="relation1") + ) + await relation_service.create_relation( + Relation(from_id=entity1.id, to_id=entity2.id, relation_type="relation2") + ) + + # Delete relations matching criteria + result = await relation_service.delete_relations( + [{"from_id": entity1.id, "to_id": entity2.id, "relation_type": "relation1"}] + ) + + assert result is True