mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
MemoryService.create_relation implemented
This commit is contained in:
@@ -152,8 +152,8 @@ async def read_entity_file(entities_path: Path, entity_id: str) -> Entity:
|
||||
source_entity = Entity(id=entity_id, name=name, entity_type=entity_type)
|
||||
|
||||
relations.append(Relation(
|
||||
from_entity=source_entity,
|
||||
to_entity=target_entity,
|
||||
from_id=source_entity.id,
|
||||
to_id=target_entity.id,
|
||||
relation_type=relation_type,
|
||||
context=context
|
||||
))
|
||||
|
||||
@@ -2,7 +2,7 @@ from typing import Type, Optional, Any, Sequence
|
||||
from sqlalchemy import select, func, Select, Executable, inspect, Result, Column, and_
|
||||
from sqlalchemy.exc import NoResultFound
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm import Mapped
|
||||
from sqlalchemy.orm import Mapped, selectinload
|
||||
|
||||
from basic_memory.models import Entity, Observation, Relation, Base
|
||||
|
||||
@@ -177,6 +177,38 @@ class EntityRepository(Repository[Entity]):
|
||||
Repository for Entity model with memory-specific operations.
|
||||
"""
|
||||
|
||||
async def find_by_id(self, entity_id: str) -> Optional[Entity]:
|
||||
"""
|
||||
Find entity by ID with relations eagerly loaded.
|
||||
|
||||
Uses selectinload to eagerly load outgoing and incoming relations in a single query.
|
||||
This is necessary because:
|
||||
1. In async code, lazy loading relations after the session closes doesn't work
|
||||
2. Our service layer often needs the complete entity with its relations
|
||||
3. Using a single query with selectinload is more efficient than multiple lazy-loaded queries
|
||||
|
||||
:param entity_id: Entity ID to search for
|
||||
:return: Entity if found with all relations loaded, None otherwise
|
||||
|
||||
Example:
|
||||
entity = await repo.find_by_id('20240102-entity-123')
|
||||
# Relations are already loaded - no additional queries needed
|
||||
for relation in entity.outgoing_relations:
|
||||
print(f"Related to {relation.to_id} via {relation.relation_type}")
|
||||
"""
|
||||
try:
|
||||
result = await self.session.execute(
|
||||
select(Entity)
|
||||
.filter(Entity.id == entity_id)
|
||||
.options(
|
||||
selectinload(Entity.outgoing_relations),
|
||||
selectinload(Entity.incoming_relations)
|
||||
)
|
||||
)
|
||||
return result.scalars().one()
|
||||
except NoResultFound:
|
||||
return None
|
||||
|
||||
async def find_by_name(self, name: str) -> Optional[Entity]:
|
||||
"""
|
||||
Find an entity by its unique name.
|
||||
|
||||
@@ -5,8 +5,10 @@ independent from storage/persistence concerns.
|
||||
"""
|
||||
|
||||
from datetime import datetime, UTC
|
||||
from typing import List, Optional
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Optional, Dict, Any
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel, model_validator
|
||||
|
||||
|
||||
class Observation(BaseModel):
|
||||
@@ -47,6 +49,32 @@ class Entity(BaseModel):
|
||||
observations: List[Observation] = []
|
||||
relations: List[Relation] = []
|
||||
|
||||
@model_validator(mode='before')
|
||||
@classmethod
|
||||
def generate_id(cls, data: dict) -> dict:
|
||||
"""Generate an ID if one wasn't provided during instantiation"""
|
||||
if not data.get('id') and data.get('name'):
|
||||
timestamp = datetime.now(UTC).strftime("%Y%m%d")
|
||||
normalized_name = data['name'].lower().replace(" ", "-")
|
||||
data['id'] = f"{timestamp}-{normalized_name}-{uuid4().hex[:8]}"
|
||||
return data
|
||||
|
||||
def model_dump(self, **kwargs) -> Dict[str, Any]:
|
||||
"""Serialize entity, handling relations to prevent circular references"""
|
||||
# Get basic data without relations
|
||||
exclude = kwargs.pop('exclude', set())
|
||||
exclude.add('relations')
|
||||
basic_data = super().model_dump(exclude=exclude, **kwargs)
|
||||
|
||||
# Add serialized relations if we have any
|
||||
if 'relations' not in exclude and self.relations:
|
||||
basic_data['relations'] = [
|
||||
relation.model_dump(**kwargs)
|
||||
for relation in self.relations
|
||||
]
|
||||
|
||||
return basic_data
|
||||
|
||||
def file_name(self) -> str:
|
||||
"""Get the markdown file name for this entity."""
|
||||
return f"{self.id}.md"
|
||||
@@ -1,12 +1,11 @@
|
||||
"""Service for managing entities in the database."""
|
||||
from datetime import datetime, UTC
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from basic_memory.models import Entity as DbEntity
|
||||
from basic_memory.repository import EntityRepository
|
||||
from basic_memory.schemas import Entity
|
||||
from . import ServiceError, DatabaseSyncError
|
||||
from basic_memory.models import Entity as EntityModel
|
||||
from . import ServiceError
|
||||
|
||||
class EntityService:
|
||||
"""
|
||||
@@ -18,42 +17,32 @@ class EntityService:
|
||||
self.project_path = project_path
|
||||
self.entity_repo = entity_repo
|
||||
|
||||
async def create_entity(self, entity: Entity) -> Entity:
|
||||
async def create_entity(self, entity: Entity) -> EntityModel:
|
||||
"""Create a new entity in the database."""
|
||||
# Create DB record
|
||||
db_data = {
|
||||
**entity.model_dump(),
|
||||
"created_at": datetime.now(UTC),
|
||||
"updated_at": datetime.now(UTC)
|
||||
}
|
||||
await self.entity_repo.create(db_data)
|
||||
return entity
|
||||
return await self.entity_repo.create(db_data)
|
||||
|
||||
async def get_entity(self, entity_id: str) -> Entity:
|
||||
async def get_entity(self, entity_id: str) -> EntityModel:
|
||||
"""Get entity by ID."""
|
||||
db_entity = await self.entity_repo.find_by_id(entity_id)
|
||||
if not db_entity:
|
||||
raise ServiceError(f"Entity not found: {entity_id}")
|
||||
|
||||
return Entity(
|
||||
id=db_entity.id,
|
||||
name=db_entity.name,
|
||||
entity_type=db_entity.entity_type
|
||||
)
|
||||
return db_entity
|
||||
|
||||
async def get_by_name(self, name: str) -> Entity:
|
||||
# TODO name is not uniaue
|
||||
async def get_by_name(self, name: str) -> EntityModel:
|
||||
"""Get entity by name."""
|
||||
db_entity = await self.entity_repo.find_by_name(name)
|
||||
if not db_entity:
|
||||
raise ServiceError(f"Entity not found: {name}")
|
||||
|
||||
return Entity(
|
||||
id=db_entity.id,
|
||||
name=db_entity.name,
|
||||
entity_type=db_entity.entity_type
|
||||
)
|
||||
return db_entity
|
||||
|
||||
async def delete_entity(self, entity_id: str) -> bool:
|
||||
"""Delete entity from database."""
|
||||
await self.entity_repo.delete(entity_id)
|
||||
return True
|
||||
return await self.entity_repo.delete(entity_id)
|
||||
@@ -42,40 +42,28 @@ class MemoryService:
|
||||
|
||||
async def create_relations(self, relations_data: List[Dict[str, Any]]) -> List[Relation]:
|
||||
"""Create multiple relations between entities."""
|
||||
relations = []
|
||||
|
||||
for data in relations_data:
|
||||
# Resolve entities by ID
|
||||
from_entity, to_entity = await asyncio.gather(
|
||||
self.entity_service.get_entity(data["from_id"]),
|
||||
self.entity_service.get_entity(data["to_id"])
|
||||
)
|
||||
|
||||
# Create relation from schema data
|
||||
relation = Relation(
|
||||
from_id=from_entity.id,
|
||||
to_id=to_entity.id,
|
||||
relation_type=data["relation_type"],
|
||||
context=data.get("context")
|
||||
)
|
||||
|
||||
# Create in database
|
||||
stored_relation = await self.relation_service.create_relation(relation)
|
||||
|
||||
# Add to source entity's relations list
|
||||
relations = [Relation.model_validate(data) for data in relations_data]
|
||||
|
||||
for relation in relations:
|
||||
# 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)
|
||||
|
||||
# Add the new relation to the source entity
|
||||
if not hasattr(from_entity, 'relations'):
|
||||
from_entity.relations = []
|
||||
from_entity.relations.append(stored_relation)
|
||||
relations.append((stored_relation, from_entity))
|
||||
from_entity.relations.append(relation)
|
||||
|
||||
# Write updated entities in parallel
|
||||
async def write_file(entity: Entity):
|
||||
await write_entity_file(self.entities_path, entity)
|
||||
# Write updated entity files (filesystem is source of truth)
|
||||
await asyncio.gather(
|
||||
write_entity_file(self.entities_path, from_entity),
|
||||
write_entity_file(self.entities_path, to_entity)
|
||||
)
|
||||
|
||||
file_writes = [write_file(entity) for _, entity in relations]
|
||||
await asyncio.gather(*file_writes)
|
||||
# Now update the database index
|
||||
await self.relation_service.create_relation(relation)
|
||||
|
||||
return [relation for relation, _ in relations]
|
||||
return relations
|
||||
|
||||
async def add_observations(self, observations_data: List[Dict[str, Any]]) -> None:
|
||||
"""Add observations to existing entities."""
|
||||
@@ -129,7 +117,7 @@ class MemoryService:
|
||||
]
|
||||
entity_updates.append(entity)
|
||||
|
||||
# Write updated files in parallel
|
||||
# Write updated entities in parallel
|
||||
async def write_file(entity: Entity):
|
||||
await write_entity_file(self.entities_path, entity)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user