MemoryService.create_relation implemented

This commit is contained in:
phernandez
2024-12-07 17:51:52 -06:00
parent b285b4619f
commit 122dbd7d31
6 changed files with 137 additions and 74 deletions
+2 -2
View File
@@ -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
))
+33 -1
View File
@@ -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.
+30 -2
View File
@@ -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"
+10 -21
View File
@@ -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)
+18 -30
View File
@@ -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)