basicmachines-co-basic-memory/src/basic_memory/services.py

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)}")