From f2c1e9390037e77a9a334da7e1818bc599f07de4 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 7 Dec 2024 15:22:14 -0600 Subject: [PATCH] memory_service.create_entities --- src/basic_memory/deps.py | 54 ++++ src/basic_memory/repository.py | 10 +- src/basic_memory/schemas.py | 32 +- src/basic_memory/services.py | 329 -------------------- src/basic_memory/services/__init__.py | 2 + src/basic_memory/services/entity_service.py | 8 +- src/basic_memory/services/memory_service.py | 177 +++++++++-- tests/conftest.py | 67 ++-- tests/test_memory_service.py | 46 +++ 9 files changed, 305 insertions(+), 420 deletions(-) create mode 100644 src/basic_memory/deps.py delete mode 100644 src/basic_memory/services.py create mode 100644 tests/test_memory_service.py diff --git a/src/basic_memory/deps.py b/src/basic_memory/deps.py new file mode 100644 index 00000000..a475e147 --- /dev/null +++ b/src/basic_memory/deps.py @@ -0,0 +1,54 @@ +"""Dependency injection functions for basic-memory services.""" +from pathlib import Path +from sqlalchemy.ext.asyncio import AsyncSession + +from basic_memory.models import Entity as DbEntity, Observation as DbObservation, Relation as DbRelation +from basic_memory.repository import EntityRepository, ObservationRepository, RelationRepository +from basic_memory.services import EntityService, ObservationService, RelationService, MemoryService + +async def get_entity_repo(session: AsyncSession) -> EntityRepository: + """Get an EntityRepository instance.""" + return EntityRepository(session, DbEntity) + +async def get_observation_repo(session: AsyncSession) -> ObservationRepository: + """Get an ObservationRepository instance.""" + return ObservationRepository(session, DbObservation) + +async def get_relation_repo(session: AsyncSession) -> RelationRepository: + """Get a RelationRepository instance.""" + return RelationRepository(session, DbRelation) + +async def get_entity_service( + project_path: Path, + entity_repo: EntityRepository +) -> EntityService: + """Get an EntityService instance.""" + return EntityService(project_path, entity_repo) + +async def get_observation_service( + project_path: Path, + observation_repo: ObservationRepository +) -> ObservationService: + """Get an ObservationService instance.""" + return ObservationService(project_path, observation_repo) + +async def get_relation_service( + project_path: Path, + relation_repo: RelationRepository +) -> RelationService: + """Get a RelationService instance.""" + return RelationService(project_path, relation_repo) + +async def get_memory_service( + project_path: Path, + entity_service: EntityService, + relation_service: RelationService, + observation_service: ObservationService +) -> MemoryService: + """Get a fully configured MemoryService instance.""" + return MemoryService( + project_path=project_path, + entity_service=entity_service, + relation_service=relation_service, + observation_service=observation_service + ) \ No newline at end of file diff --git a/src/basic_memory/repository.py b/src/basic_memory/repository.py index 357ed041..bfc44da3 100644 --- a/src/basic_memory/repository.py +++ b/src/basic_memory/repository.py @@ -20,7 +20,7 @@ class Repository[T: Base]: Example usage: async with async_sessionmaker() as session: entity_repo = Repository(session, Entity) - await entity_repo.create({ + entity = await entity_repo.create({ 'id': '20240102-some-entity', 'name': 'Example Entity', 'entity_type': 'concept' @@ -91,8 +91,7 @@ class Repository[T: Base]: model_data = {k: v for k, v in entity_data.items() if k in self.valid_columns} entity = self.Model(**model_data) self.session.add(entity) - await self.session.commit() - await self.session.refresh(entity) + await self.session.flush() return entity async def update(self, entity_id: str, entity_data: dict) -> Optional[T]: @@ -114,8 +113,7 @@ class Repository[T: Base]: for key, value in entity_data.items(): if key in self.valid_columns: setattr(entity, key, value) - await self.session.commit() - await self.session.refresh(entity) + await self.session.flush() return entity except NoResultFound: return None @@ -136,7 +134,7 @@ class Repository[T: Base]: ) entity = result.scalars().one() await self.session.delete(entity) - await self.session.commit() + await self.session.flush() return True except NoResultFound: return False diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 5aa79393..9e2cdf9c 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -20,13 +20,6 @@ class ObservationCreate(BaseModel): content: str -class EntityCreate(BaseModel): - """Schema for creating a new entity via the MCP tool interface.""" - name: str - entityType: str # Matches the JSON field name from MCP tool - observations: Optional[List[str]] = None - - class RelationCreate(BaseModel): """Schema for creating a new relation via the MCP tool interface.""" from_: str = None # Raw entity name from MCP tool @@ -79,15 +72,6 @@ class Relation(BaseModel): 'context': self.context } - @classmethod - def from_create(cls, create_data: RelationCreate, from_entity: 'Entity', to_entity: 'Entity') -> 'Relation': - """Create a Relation from a RelationCreate schema and actual entities.""" - return cls( - from_entity=from_entity, - to_entity=to_entity, - relation_type=create_data.relationType - ) - class Entity(BaseModel): """ @@ -102,7 +86,7 @@ class Entity(BaseModel): relations: List[Relation] = [] @model_validator(mode='before') - @classmethod + @classmethod def generate_id_if_needed(cls, data: dict) -> dict: """Generate an ID if one wasn't provided during instantiation""" if not data.get('id') and data.get('name'): @@ -111,16 +95,6 @@ class Entity(BaseModel): data['id'] = f"{timestamp}-{normalized_name}-{uuid4().hex[:8]}" return data - @classmethod - def from_create(cls, data: EntityCreate) -> 'Entity': - """Create an Entity from an EntityCreate schema.""" - observations = [Observation(content=obs) for obs in (data.observations or [])] - return cls( - name=data.name, - entity_type=data.entityType, - observations=observations - ) - def model_dump(self, **kwargs) -> Dict[str, Any]: """Serialize entity, handling relations to prevent circular references""" # Get basic data without relations @@ -137,6 +111,10 @@ class Entity(BaseModel): return basic_data + def file_name(self) -> str: + """Get the markdown file name for this entity.""" + return f"{self.id}.md" + # Update forward refs Entity.model_rebuild() diff --git a/src/basic_memory/services.py b/src/basic_memory/services.py deleted file mode 100644 index 04e912b8..00000000 --- a/src/basic_memory/services.py +++ /dev/null @@ -1,329 +0,0 @@ -from datetime import datetime, UTC -from pathlib import Path -from typing import Optional, List -from uuid import uuid4 -from sqlalchemy import and_, select, delete - -from basic_memory.models import Entity as DbEntity # Rename to avoid confusion -from basic_memory.models import Observation as DbObservation -from basic_memory.models import Relation as DbRelation -from basic_memory.repository import EntityRepository, ObservationRepository, RelationRepository -from basic_memory.schemas import Entity, Observation, Relation -from basic_memory.fileio import ( - read_entity_file, write_entity_file, delete_entity_file, - FileOperationError, EntityNotFoundError -) - - -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 - - -class EntityService: - """Service for managing entities in the filesystem and database.""" - def __init__(self, project_path: Path, entity_repo: EntityRepository): - self.project_path = project_path - self.entity_repo = entity_repo - self.entities_path = project_path / "entities" - - async def _update_db_index(self, entity: Entity) -> DbEntity: - """Update database index with entity data.""" - entity_data = { - **entity.model_dump(), - "created_at": datetime.now(UTC), - "updated_at": datetime.now(UTC) - } - - # Observations will be handled by ObservationService - entity_data.pop('observations', None) # Remove observations if present - entity_data.pop('relations', None) # Remove relations if present - - # Try to find existing entity first - if await self.entity_repo.find_by_id(entity.id): - return await self.entity_repo.update(entity.id, entity_data) - else: - return await self.entity_repo.create(entity_data) - - async def create_entity(self, name: str, entity_type: str, - observations: Optional[list[str]] = None) -> Entity: - """Create a new entity.""" - # Convert string observations to Observation objects if provided - obs_list = [Observation(content=obs) for obs in (observations or [])] - - # Create entity (ID will be auto-generated) - entity = Entity( - name=name, - entity_type=entity_type, - observations=obs_list - ) - - # Step 1: Write to filesystem (source of truth) - await write_entity_file(self.entities_path, entity) - - # Step 2: Update database index - await self._update_db_index(entity) - - return entity - - async def get_entity(self, entity_id: str) -> Entity: - """Get entity by ID, reading from filesystem first.""" - # Read from filesystem (source of truth) - entity = await read_entity_file(self.entities_path, entity_id) - - # Update database index - await self._update_db_index(entity) - - return entity - - async def delete_entity(self, entity_id: str) -> bool: - """Delete entity from filesystem and database.""" - # Delete from filesystem first (source of truth) - await delete_entity_file(self.entities_path, entity_id) - - # Delete from database index - await self.entity_repo.delete(entity_id) - return True - - async def rebuild_index(self) -> None: - """Rebuild database index from filesystem contents.""" - if not self.entities_path.exists(): - return - - try: - entity_files = list(self.entities_path.glob("*.md")) - except Exception as e: - raise FileOperationError(f"Failed to read entities directory: {str(e)}") from e - - for entity_file in entity_files: - try: - entity = await read_entity_file(self.entities_path, entity_file.stem) - await self._update_db_index(entity) - except Exception as e: - print(f"Warning: Failed to reindex {entity_file}: {str(e)}") - - -class ObservationService: - """Service for managing observations in the filesystem and database.""" - def __init__(self, project_path: Path, observation_repo: ObservationRepository): - self.project_path = project_path - self.entities_path = project_path / "entities" - self.observation_repo = observation_repo - - async def add_observation(self, entity: Entity, content: str, - context: Optional[str] = None) -> Observation: - """Add a new observation to an entity.""" - observation = Observation(content=content) - entity.observations.append(observation) - - # Update filesystem first (source of truth) - await write_entity_file(self.entities_path, entity) - - # Update database index - try: - db_observation = await self.observation_repo.create({ - 'id': f"{entity.id}-obs-{uuid4().hex[:8]}", - 'entity_id': entity.id, - 'content': content, - 'context': context, - 'created_at': datetime.now(UTC) - }) - return observation - except Exception as e: - raise DatabaseSyncError(f"Failed to sync observation to database: {str(e)}") from e - - async def search_observations(self, query: str) -> list[Observation]: - """ - 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(DbObservation).filter( - DbObservation.content.contains(query) - ) - ) - return [ - Observation(content=obs.content) - for obs in result.scalars().all() - ] - - async def get_observations_by_context(self, context: str) -> list[Observation]: - """Get all observations with a specific context.""" - db_observations = await self.observation_repo.find_by_context(context) - return [ - Observation(content=obs.content) - for obs in db_observations - ] - - async def rebuild_observation_index(self) -> None: - """ - Rebuild the observation database index from filesystem contents. - Used for recovery or ensuring sync. - """ - # List all entity files - if not self.entities_path.exists(): - return - - try: - entity_files = list(self.entities_path.glob("*.md")) - except Exception as e: - raise FileOperationError(f"Failed to read entities directory: {str(e)}") from e - - # Clear existing observation index - await self.observation_repo.execute_query(delete(DbObservation)) - - # Rebuild from each entity file - for entity_file in entity_files: - try: - entity = await read_entity_file(self.entities_path, entity_file.stem) - for obs in entity.observations: - await self.observation_repo.create({ - 'id': f"{entity.id}-obs-{uuid4().hex[:8]}", - 'entity_id': entity.id, - 'content': obs.content, - 'created_at': datetime.now(UTC) - }) - except Exception as e: - print(f"Warning: Failed to reindex observations for {entity_file}: {str(e)}") - - -class RelationService: - """ - Service for managing relations between entities. - Follows the "filesystem is source of truth" principle. - - Relations are stored in entity markdown files and indexed in the database - for efficient querying. - """ - - def __init__(self, project_path: Path, relation_repo: RelationRepository): - self.project_path = project_path - self.entities_path = project_path / "entities" - self.relation_repo = relation_repo - - async def create_relation(self, from_entity: Entity, to_entity: Entity, relation_type: str, - context: Optional[str] = None) -> Relation: - """ - Create a new relation between two entities. - - Args: - from_entity: Source entity - to_entity: Target entity - relation_type: Type of relation - context: Optional context for the relation - - Returns: - The created Relation - - Raises: - FileOperationError: If file operations fail - DatabaseSyncError: If database sync fails - """ - # Create new relation with actual Entity objects - relation = Relation( - from_entity=from_entity, - to_entity=to_entity, - relation_type=relation_type, - context=context - ) - - # Add relation to source entity's relations list - if not hasattr(from_entity, 'relations'): - from_entity.relations = [] - from_entity.relations.append(relation) - - # Update filesystem first (source of truth) - await write_entity_file(self.entities_path, from_entity) - - # Update database index - # model_dump will handle converting Entity refs to IDs - try: - db_data = relation.model_dump() - db_data['created_at'] = datetime.now(UTC) - await self.relation_repo.create(db_data) - return relation - except Exception as e: - raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e - - async def get_entity_relations(self, entity: Entity) -> List[Relation]: - """ - Get all relations for an entity (both outgoing and incoming). - - Args: - entity: Entity to get relations for - - Returns: - List of relations where the entity is either source or target - """ - # Relations are stored in the entity object - return getattr(entity, 'relations', []) - - async def delete_relation(self, from_entity: Entity, relation_id: str) -> bool: - """ - Delete a relation from both filesystem and database. - - Args: - from_entity: Source entity containing the relation - relation_id: ID of the relation to delete - - Returns: - True if deletion was successful - - Raises: - RelationError: If relation cannot be found or deleted - """ - # Remove relation from entity's relations - if hasattr(from_entity, 'relations'): - from_entity.relations = [ - r for r in from_entity.relations - if r.id != relation_id - ] - - # Update filesystem first (source of truth) - await write_entity_file(self.entities_path, from_entity) - - # Remove from database index - await self.relation_repo.delete(relation_id) - return True - - async def rebuild_relation_index(self) -> None: - """ - Rebuild the relation database index from filesystem contents. - Used for recovery or ensuring sync. - """ - if not self.entities_path.exists(): - return - - try: - entity_files = list(self.entities_path.glob("*.md")) - except Exception as e: - raise FileOperationError(f"Failed to read entities directory: {str(e)}") from e - - # Clear existing relation index - await self.relation_repo.execute_query(delete(DbRelation)) - - # Rebuild from each entity file - for entity_file in entity_files: - try: - entity = await read_entity_file(self.entities_path, entity_file.stem) - for relation in getattr(entity, 'relations', []): - db_data = relation.model_dump() - db_data['created_at'] = datetime.now(UTC) - await self.relation_repo.create(db_data) - except Exception as e: - print(f"Warning: Failed to reindex relations for {entity_file}: {str(e)}") diff --git a/src/basic_memory/services/__init__.py b/src/basic_memory/services/__init__.py index 5c0c469d..cc9f43a2 100644 --- a/src/basic_memory/services/__init__.py +++ b/src/basic_memory/services/__init__.py @@ -18,6 +18,7 @@ class RelationError(ServiceError): from .entity_service import EntityService from .observation_service import ObservationService from .relation_service import RelationService +from .memory_service import MemoryService __all__ = [ 'ServiceError', @@ -26,4 +27,5 @@ __all__ = [ 'EntityService', 'ObservationService', 'RelationService', + 'MemoryService', ] \ No newline at end of file diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index a0383102..55d715fe 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -5,7 +5,7 @@ from typing import Optional from basic_memory.models import Entity as DbEntity from basic_memory.repository import EntityRepository -from basic_memory.schemas import Entity, EntityCreate +from basic_memory.schemas import Entity from . import ServiceError, DatabaseSyncError class EntityService: @@ -18,11 +18,8 @@ class EntityService: self.project_path = project_path self.entity_repo = entity_repo - async def create_entity(self, create_data: EntityCreate) -> Entity: + async def create_entity(self, entity: Entity) -> Entity: """Create a new entity in the database.""" - # Create Entity from EntityCreate data - entity = Entity.from_create(create_data) - # Create DB record db_data = { **entity.model_dump(), @@ -30,7 +27,6 @@ class EntityService: "updated_at": datetime.now(UTC) } await self.entity_repo.create(db_data) - return entity async def get_entity(self, entity_id: str) -> Entity: diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index 26439b43..914000b3 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -1,50 +1,177 @@ """Service for orchestrating entity, relation, and observation operations.""" +import asyncio from typing import List, Dict, Any, Optional from pathlib import Path -from ..schemas import Entity, EntityCreate, Relation, RelationCreate, Observation -from ..fileio import write_entity_file, read_entity_file, delete_entity_file -from .entity_service import EntityService -from .relation_service import RelationService +from basic_memory.schemas import Entity, RelationCreate, Observation, Relation +from basic_memory.fileio import write_entity_file, read_entity_file, delete_entity_file +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] = None): + def __init__( + self, + project_path: Optional[Path], + entity_service: EntityService, + relation_service: RelationService, + observation_service: ObservationService + ): self.project_path = project_path self.entities_path = project_path / "entities" if project_path else None - # Initialize with repos when we add them - self.entity_service = EntityService() - self.relation_service = RelationService() + self.entity_service = entity_service + self.relation_service = relation_service + self.observation_service = observation_service async def create_entities(self, entities_data: List[Dict[str, Any]]) -> List[Entity]: """Create multiple entities with their observations.""" - async def create_and_write(data: Dict[str, Any]) -> Entity: - create_data = EntityCreate.model_validate(data) - entity = await self.entity_service.create_entity(create_data) + + entities = [Entity.model_validate(data) for data in entities_data] + + # Write files in parallel (filesystem is source of truth) + async def write_file(entity: Entity): await write_entity_file(self.entities_path, entity) - return entity - - return [await create_and_write(data) for data in entities_data] + + file_writes = [write_file(entity) for entity in entities] + await asyncio.gather(*file_writes) + + # Update database index sequentially + for entity in entities: + await self.entity_service.create_entity(entity) + + return entities async def create_relations(self, relations_data: List[Dict[str, Any]]) -> List[Relation]: """Create multiple relations between entities.""" - async def create_and_write_relation(data: Dict[str, Any]) -> Relation: + # First get all entities and create relations + relations = [] + for data in relations_data: create_data = RelationCreate.model_validate(data) - - # Get the actual entities from_entity = await self.entity_service.get_by_name(create_data.from_) to_entity = await self.entity_service.get_by_name(create_data.to) - - # Create relation relation = await self.relation_service.create_relation( create_data=create_data, from_entity=from_entity, to_entity=to_entity ) - - # Write updated from_entity to file - await write_entity_file(self.entities_path, from_entity) - return relation - - return [await create_and_write_relation(data) for data in relations_data] \ No newline at end of file + relations.append((relation, from_entity)) + + # Write updated entities in parallel + async def write_file(entity: Entity): + await write_entity_file(self.entities_path, entity) + + file_writes = [write_file(entity) for _, entity in relations] + await asyncio.gather(*file_writes) + + return [relation for relation, _ in relations] + + async def add_observations(self, observations_data: List[Dict[str, Any]]) -> None: + """Add observations to existing entities.""" + # First read all entities and create their observations + entity_updates = [] + for data in observations_data: + entity = await read_entity_file(self.entities_path, data["entityName"]) + new_observations = [Observation(content=content) for content in data["contents"]] + entity.observations.extend(new_observations) + entity_updates.append(entity) + + # Write updated entities in parallel + async def write_file(entity: Entity): + await write_entity_file(self.entities_path, entity) + + file_writes = [write_file(entity) for entity in entity_updates] + await asyncio.gather(*file_writes) + + # Update database indexes sequentially + for entity in entity_updates: + await self.entity_service.rebuild_index(entity) + + async def delete_entities(self, entity_names: List[str]) -> None: + """Delete multiple entities and their associated data.""" + # First get all entities to be deleted + entities = [] + for name in entity_names: + entity = await self.entity_service.get_by_name(name) + entities.append(entity) + + # Delete files in parallel + async def delete_file(entity: Entity): + await delete_entity_file(self.entities_path, entity.id) + + file_deletes = [delete_file(entity) for entity in entities] + await asyncio.gather(*file_deletes) + + # Update database sequentially + for entity in entities: + await self.entity_service.delete_entity(entity.id) + + async def delete_observations(self, deletions: List[Dict[str, Any]]) -> None: + """Delete specific observations from entities.""" + # First read and update all entities + entity_updates = [] + for deletion in deletions: + entity = await read_entity_file(self.entities_path, deletion["entityName"]) + entity.observations = [ + obs for obs in entity.observations + if obs.content not in deletion["observations"] + ] + entity_updates.append(entity) + + # Write updated files in parallel + async def write_file(entity: Entity): + await write_entity_file(self.entities_path, entity) + + file_writes = [write_file(entity) for entity in entity_updates] + await asyncio.gather(*file_writes) + + # Update database indexes sequentially + for entity in entity_updates: + await self.entity_service.rebuild_index(entity) + + async def delete_relations(self, relations: List[Dict[str, Any]]) -> None: + """Delete specific relations between entities.""" + # First get all entities and delete relations + updates = [] + for data in relations: + from_entity = await self.entity_service.get_by_name(data["from"]) + to_entity = await self.entity_service.get_by_name(data["to"]) + await self.relation_service.delete_relation(from_entity, to_entity, data["relationType"]) + updates.append(from_entity) + + # Write updated files in parallel + async def write_file(entity: Entity): + await write_entity_file(self.entities_path, entity) + + file_writes = [write_file(entity) for entity in updates] + await asyncio.gather(*file_writes) + + async def read_graph(self) -> Dict[str, Any]: + """Read the entire knowledge graph.""" + entities = await self.entity_service.get_all() + return { + "entities": [entity.model_dump() for entity in entities] + # Relations are included in entity.model_dump() + } + + async def search_nodes(self, query: str) -> Dict[str, Any]: + """Search for nodes in the knowledge graph.""" + results = await self.entity_service.search(query) + return { + "matches": [entity.model_dump() for entity in results], + "query": query + } + + async def open_nodes(self, names: List[str]) -> Dict[str, Any]: + """Get specific nodes and their relationships.""" + async def read_node(name: str) -> Optional[Entity]: + if entity := await read_entity_file(self.entities_path, name): + return entity + return None + + entities = [entity for entity in await asyncio.gather(*(read_node(name) for name in names)) + if entity is not None] + + return { + "entities": [entity.model_dump() for entity in entities] + # Relations between these entities are included in model_dump() + } \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index 8dcf74b1..f511ab10 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,9 +6,16 @@ import tempfile from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker from sqlalchemy.pool import StaticPool -from basic_memory.models import Base, Entity as DbEntity, Observation as DbObservation, Relation as DbRelation -from basic_memory.repository import EntityRepository, ObservationRepository, RelationRepository -from basic_memory.services import EntityService, ObservationService, RelationService +from basic_memory.models import Base +from basic_memory.deps import ( + get_entity_repo, + get_observation_repo, + get_relation_repo, + get_entity_service, + get_observation_service, + get_relation_service, + get_memory_service +) @pytest_asyncio.fixture(scope="function") async def engine(): @@ -43,50 +50,56 @@ async def session(engine): @pytest_asyncio.fixture async def entity_repo(session): """Create an EntityRepository instance.""" - return EntityRepository(session, DbEntity) + return await get_entity_repo(session) @pytest_asyncio.fixture async def observation_repo(session): """Create an ObservationRepository instance.""" - return ObservationRepository(session, DbObservation) + return await get_observation_repo(session) @pytest_asyncio.fixture async def relation_repo(session): """Create a RelationRepository instance.""" - return RelationRepository(session, DbRelation) + return await get_relation_repo(session) @pytest_asyncio.fixture -async def entity_service(session, entity_repo): - """Fixture providing initialized EntityService with temp directories.""" +async def test_project_path(): + """Create a temporary project directory.""" with tempfile.TemporaryDirectory() as temp_dir: project_path = Path(temp_dir) / "test-project" entities_path = project_path / "entities" entities_path.mkdir(parents=True) - - service = EntityService(project_path, entity_repo) - yield service + yield project_path @pytest_asyncio.fixture -async def observation_service(session, observation_repo): +async def entity_service(test_project_path, entity_repo): + """Fixture providing initialized EntityService.""" + return await get_entity_service(test_project_path, entity_repo) + +@pytest_asyncio.fixture +async def observation_service(test_project_path, observation_repo): """Fixture providing initialized ObservationService.""" - with tempfile.TemporaryDirectory() as temp_dir: - project_path = Path(temp_dir) / "test-project" - entities_path = project_path / "entities" - entities_path.mkdir(parents=True) - - service = ObservationService(project_path, observation_repo) - yield service + return await get_observation_service(test_project_path, observation_repo) @pytest_asyncio.fixture -async def relation_service(session, relation_repo): +async def relation_service(test_project_path, relation_repo): """Fixture providing initialized RelationService.""" - with tempfile.TemporaryDirectory() as temp_dir: - project_path = Path(temp_dir) / "test-project" - entities_path = project_path / "entities" - entities_path.mkdir(parents=True) - - service = RelationService(project_path, relation_repo) - yield service + return await get_relation_service(test_project_path, relation_repo) + +@pytest_asyncio.fixture +async def memory_service( + test_project_path, + entity_service, + relation_service, + observation_service +): + """Fixture providing initialized MemoryService.""" + return await get_memory_service( + test_project_path, + entity_service, + relation_service, + observation_service + ) @pytest_asyncio.fixture async def test_entity(entity_service): diff --git a/tests/test_memory_service.py b/tests/test_memory_service.py new file mode 100644 index 00000000..e8ae3460 --- /dev/null +++ b/tests/test_memory_service.py @@ -0,0 +1,46 @@ +"""Tests for the MemoryService class.""" +import pytest + +from basic_memory.services import MemoryService + +test_entities_data = [ + { + "name": "Test_Entity_1", + "entity_type": "test", + "observations": [{"content":"Observation 1.1"}, {"content":"Observation 1.2"}] + }, + { + "name": "Test_Entity_2", + "entity_type": "test", + "observations": [{"content":"Observation 2.1"}, {"content":"Observation 2.2"}] + } +] + +@pytest.mark.asyncio +async def test_create_entities(memory_service: MemoryService): + """Should create multiple entities in parallel with their observations.""" + # Create entities + entities = await memory_service.create_entities(test_entities_data) + + # Verify the entities 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 + entity1_path = memory_service.entities_path / entities[0].file_name() + entity2_path = memory_service.entities_path / entities[1].file_name() + assert entity1_path.exists() + assert entity2_path.exists() \ No newline at end of file