mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
memory_service.create_entities
This commit is contained in:
phernandez
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)}")
|
||||
@@ -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',
|
||||
]
|
||||
@@ -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:
|
||||
|
||||
@@ -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]
|
||||
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()
|
||||
}
|
||||
+40
-27
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user