"""Service for orchestrating entity, relation, and observation operations.""" import asyncio from typing import List, Dict, Any, Optional from pathlib import Path from basic_memory.models import Entity, Observation, Relation from basic_memory.schemas import ( ObservationsIn, EntityIn, RelationIn, ObservationIn ) from basic_memory.fileio import write_entity_file, read_entity_file, EntityNotFoundError from basic_memory.services import EntityService, RelationService, ObservationService from loguru import logger 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}") async def create_entities(self, entities_in: List[EntityIn]) -> List[Entity]: """Create multiple entities with their observations.""" logger.debug(f"Creating {len(entities_in)} entities") # Write files in parallel (filesystem is source of truth) async def write_file(entity: EntityIn): try: existing = await self.entity_service.get_by_type_and_name( entity.entity_type, entity.name ) if existing: raise ValueError( f"Entity already exists: {entity.entity_type}/{entity.name}" ) except EntityNotFoundError: # Good - entity doesn't exist yet pass # Generate ID and write file entity_id = Entity.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_in: EntityIn): logger.debug(f"Creating entity in DB: {entity_in}") try: # Create base entity created_entity = await self.entity_service.create_entity(entity_in) logger.debug(f"Created base entity: {created_entity.id}") # Add observations await self.observation_service.add_observations(created_entity.id, entity_in.observations) logger.debug(f"Added {len(entity_in.observations)} observations to {created_entity.id}") # Add relations for relation in entity_in.relations: await self.relation_service.create_relation(relation) logger.debug(f"Added {len(entity_in.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 as e: logger.exception(f"Failed to create entity in DB: {entity_in}") 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 as e: # 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 = Entity.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 create_relations(self, relations_data: List[RelationIn]) -> List[Relation]: """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") 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 as e: 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: ObservationsIn) -> List[Observation]: """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 for obs in observations_in.observations: entity.observations.append(obs) 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 as e: logger.exception(f"Failed to add observations to entity: {observations_in.entity_id}") raise async def delete_entities(self, entity_names: List[str]) -> None: pass async def delete_observations(self, deletions: List[Dict[str, Any]]) -> None: pass async def delete_relations(self, relations: List[Dict[str, Any]]) -> None: pass async def read_graph(self) -> List[Entity]: """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 as e: logger.exception("Failed to read graph") raise async def search_nodes(self, query: str) -> List[Entity]: """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 as e: logger.exception(f"Failed to search nodes with query: {query}") raise async def open_nodes(self, names: List[str]) -> List[Entity]: """Get specific nodes and their relationships.""" logger.debug(f"Opening nodes: {names}") async def read_node(name: str) -> Optional[Entity]: try: # Get ID from name first logger.debug(f"Looking up entity: {name}") db_entity = await self.entity_service.get_entity(name) 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: {name}") return None except Exception as e: logger.exception(f"Failed to read node: {name}") return None try: entities = [entity for entity in await asyncio.gather(*(read_node(name) for name in names)) if entity is not None] logger.debug(f"Opened {len(entities)} entities") return entities except Exception as e: logger.exception("Failed to open nodes") raise