mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
all tests passing
This commit is contained in:
phernandez
+43
-168
@@ -1,3 +1,4 @@
|
||||
"""Repository implementations for basic-memory models."""
|
||||
from typing import Type, Optional, Any, Sequence
|
||||
from sqlalchemy import select, func, Select, Executable, inspect, Result, Column, and_
|
||||
from sqlalchemy.exc import NoResultFound
|
||||
@@ -8,24 +9,7 @@ from basic_memory.models import Entity, Observation, Relation, Base
|
||||
|
||||
|
||||
class Repository[T: Base]:
|
||||
"""
|
||||
Generic repository pattern implementation for handling database operations.
|
||||
Adapted for basic-memory with string IDs and memory-specific operations.
|
||||
|
||||
Provides basic CRUD operations with an async session.
|
||||
|
||||
:param session: Async database session from SQLAlchemy.
|
||||
:param Model: Database model class.
|
||||
|
||||
Example usage:
|
||||
async with async_sessionmaker() as session:
|
||||
entity_repo = Repository(session, Entity)
|
||||
entity = await entity_repo.create({
|
||||
'id': '20240102-some-entity',
|
||||
'name': 'Example Entity',
|
||||
'entity_type': 'concept'
|
||||
})
|
||||
"""
|
||||
"""Base repository implementation with generic CRUD operations."""
|
||||
|
||||
def __init__(self, session: AsyncSession, Model: Type[T]):
|
||||
self.session = session
|
||||
@@ -33,37 +17,19 @@ class Repository[T: Base]:
|
||||
self.primary_key: Column[Any] = inspect(self.Model).mapper.primary_key[0]
|
||||
self.valid_columns = [column.key for column in inspect(self.Model).columns]
|
||||
|
||||
async def refresh(self, instance: T) -> None:
|
||||
"""
|
||||
Refresh the state of the given instance from the database.
|
||||
|
||||
:param instance: Instance to refresh
|
||||
"""
|
||||
await self.session.refresh(instance)
|
||||
async def refresh(self, instance: T, relationships: list[str] | None = None) -> None:
|
||||
"""Refresh instance and optionally specified relationships."""
|
||||
await self.session.refresh(instance, relationships or [])
|
||||
|
||||
async def find_all(self, skip: int = 0, limit: int = 100) -> Sequence[T]:
|
||||
"""
|
||||
Fetches records from the database with pagination.
|
||||
|
||||
:param skip: Number of records to skip.
|
||||
:param limit: Maximum number of records to fetch.
|
||||
:return: List containing the fetched records.
|
||||
"""
|
||||
"""Fetch records from the database with pagination."""
|
||||
result = await self.session.execute(
|
||||
select(self.Model).offset(skip).limit(limit)
|
||||
)
|
||||
return result.scalars().all()
|
||||
|
||||
async def find_by_id(self, entity_id: str) -> Optional[T]:
|
||||
"""
|
||||
Fetches an entity by its unique identifier asynchronously.
|
||||
|
||||
:param entity_id: Unique identifier of the entity (string timestamp-based ID)
|
||||
:return: The entity if found, otherwise None
|
||||
|
||||
Example:
|
||||
entity = await repository.find_by_id('20240102-example-entity')
|
||||
"""
|
||||
"""Fetch an entity by its unique identifier."""
|
||||
try:
|
||||
result = await self.session.execute(
|
||||
select(self.Model).filter(self.primary_key == entity_id)
|
||||
@@ -72,39 +38,22 @@ class Repository[T: Base]:
|
||||
except NoResultFound:
|
||||
return None
|
||||
|
||||
async def create(self, entity_data: dict) -> T:
|
||||
"""
|
||||
Creates a new entity in the database from the provided data dictionary.
|
||||
|
||||
:param entity_data: A dictionary containing data to be inserted
|
||||
:return: The created entity
|
||||
|
||||
Example:
|
||||
>>> entity_data = {
|
||||
... 'id': '20240102-example',
|
||||
... 'name': 'Example Entity',
|
||||
... 'entity_type': 'concept',
|
||||
... 'description': 'An example entity'
|
||||
... }
|
||||
>>> new_entity = await repo.create(entity_data)
|
||||
async def create(self, entity_data: dict, model: Type[Base] | None = None) -> T:
|
||||
"""Create a new entity in the database from the provided data.
|
||||
|
||||
Args:
|
||||
entity_data: Dictionary containing the data to insert
|
||||
model: Optional model class to use (defaults to self.Model)
|
||||
"""
|
||||
model = model or self.Model
|
||||
model_data = {k: v for k, v in entity_data.items() if k in self.valid_columns}
|
||||
entity = self.Model(**model_data)
|
||||
entity = model(**model_data)
|
||||
self.session.add(entity)
|
||||
await self.session.flush()
|
||||
return entity
|
||||
|
||||
async def update(self, entity_id: str, entity_data: dict) -> Optional[T]:
|
||||
"""
|
||||
Updates an entity with given entity_id using the provided entity_data.
|
||||
|
||||
:param entity_id: String ID of the entity to update
|
||||
:param entity_data: Dictionary containing the data to update
|
||||
:return: The updated entity or None if not found
|
||||
|
||||
Example:
|
||||
updated = await repository.update('20240102-example', {'description': 'Updated description'})
|
||||
"""
|
||||
"""Update an entity with the given data."""
|
||||
try:
|
||||
result = await self.session.execute(
|
||||
select(self.Model).filter(self.primary_key == entity_id)
|
||||
@@ -119,15 +68,7 @@ class Repository[T: Base]:
|
||||
return None
|
||||
|
||||
async def delete(self, entity_id: str) -> bool:
|
||||
"""
|
||||
Deletes an entity from the database.
|
||||
|
||||
:param entity_id: String ID of the entity to delete
|
||||
:return: Boolean indicating if the entity was deleted
|
||||
|
||||
Example:
|
||||
success = await repository.delete('20240102-example')
|
||||
"""
|
||||
"""Delete an entity from the database."""
|
||||
try:
|
||||
result = await self.session.execute(
|
||||
select(self.Model).filter(self.primary_key == entity_id)
|
||||
@@ -140,12 +81,7 @@ class Repository[T: Base]:
|
||||
return False
|
||||
|
||||
async def count(self, query: Executable | None = None) -> int:
|
||||
"""
|
||||
Counts entities in the database table.
|
||||
|
||||
:param query: Optional SQL query to modify the count operation
|
||||
:return: Number of matching entities
|
||||
"""
|
||||
"""Count entities in the database table."""
|
||||
if query is None:
|
||||
query = select(func.count()).select_from(self.Model)
|
||||
result = await self.session.execute(query)
|
||||
@@ -153,130 +89,74 @@ class Repository[T: Base]:
|
||||
return scalar if scalar is not None else 0
|
||||
|
||||
async def execute_query(self, query: Executable) -> Result[Any]:
|
||||
"""
|
||||
Executes the given query asynchronously.
|
||||
|
||||
:param query: An executable query instance
|
||||
:return: Query result
|
||||
"""
|
||||
"""Execute a query asynchronously."""
|
||||
return await self.session.execute(query)
|
||||
|
||||
async def find_one(self, query: Select[tuple[T]]) -> Optional[T]:
|
||||
"""
|
||||
Executes a query and retrieves a single record.
|
||||
|
||||
:param query: The query to execute
|
||||
:return: Single record or None
|
||||
"""
|
||||
"""Execute a query and retrieve a single record."""
|
||||
result = await self.execute_query(query)
|
||||
return result.scalars().one_or_none()
|
||||
|
||||
|
||||
class EntityRepository(Repository[Entity]):
|
||||
"""
|
||||
Repository for Entity model with memory-specific operations.
|
||||
"""
|
||||
"""Repository for Entity model with memory-specific operations."""
|
||||
|
||||
async def find_by_id(self, entity_id: str) -> Optional[Entity]:
|
||||
"""
|
||||
Find entity by ID with all relationships eagerly loaded.
|
||||
|
||||
Uses selectinload to eagerly load observations, 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
|
||||
3. Using a single query with selectinload is more efficient
|
||||
|
||||
:param entity_id: Entity ID to search for
|
||||
:return: Entity if found with everything loaded, None otherwise
|
||||
"""
|
||||
"""Find entity by ID with all relationships eagerly loaded."""
|
||||
try:
|
||||
# First load base entity
|
||||
result = await self.session.execute(
|
||||
select(Entity)
|
||||
.filter(Entity.id == entity_id)
|
||||
.options(
|
||||
selectinload(Entity.observations),
|
||||
selectinload(Entity.outgoing_relations),
|
||||
selectinload(Entity.incoming_relations)
|
||||
)
|
||||
select(Entity).filter(Entity.id == entity_id)
|
||||
)
|
||||
return result.scalars().one()
|
||||
entity = result.scalars().one()
|
||||
|
||||
# Force refresh of all relationships
|
||||
await self.refresh(entity, ['observations', 'outgoing_relations', 'incoming_relations'])
|
||||
|
||||
return entity
|
||||
except NoResultFound:
|
||||
return None
|
||||
|
||||
async def find_by_name(self, name: str) -> Optional[Entity]:
|
||||
"""
|
||||
Find an entity by its unique name.
|
||||
|
||||
:param name: Entity name to search for
|
||||
:return: Entity if found, None otherwise
|
||||
"""
|
||||
"""Find an entity by its unique name."""
|
||||
query = (
|
||||
select(Entity)
|
||||
.filter(Entity.name == name)
|
||||
.options(
|
||||
selectinload(Entity.observations),
|
||||
selectinload(Entity.outgoing_relations),
|
||||
selectinload(Entity.incoming_relations)
|
||||
)
|
||||
)
|
||||
return await self.find_one(query)
|
||||
result = await self.session.execute(query)
|
||||
entity = result.scalars().one_or_none()
|
||||
if entity:
|
||||
await self.refresh(entity, ['observations', 'outgoing_relations', 'incoming_relations'])
|
||||
return entity
|
||||
|
||||
async def search_by_type(self, entity_type: str, skip: int = 0, limit: int = 100) -> Sequence[Entity]:
|
||||
"""
|
||||
Search for entities of a specific type.
|
||||
|
||||
:param entity_type: Type to search for
|
||||
:param skip: Number of records to skip
|
||||
:param limit: Maximum records to return
|
||||
:return: List of matching entities
|
||||
"""
|
||||
"""Search for entities of a specific type."""
|
||||
query = select(Entity).filter(Entity.entity_type == entity_type).offset(skip).limit(limit)
|
||||
result = await self.execute_query(query)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
class ObservationRepository(Repository[Observation]):
|
||||
"""
|
||||
Repository for Observation model with memory-specific operations.
|
||||
"""
|
||||
"""Repository for Observation model with memory-specific operations."""
|
||||
|
||||
async def find_by_entity(self, entity_id: str) -> Sequence[Observation]:
|
||||
"""
|
||||
Find all observations for a specific entity.
|
||||
|
||||
:param entity_id: ID of the entity to find observations for
|
||||
:return: List of observations
|
||||
"""
|
||||
"""Find all observations for a specific entity."""
|
||||
query = select(Observation).filter(Observation.entity_id == entity_id)
|
||||
result = await self.execute_query(query)
|
||||
return result.scalars().all()
|
||||
|
||||
async def find_by_context(self, context: str) -> Sequence[Observation]:
|
||||
"""
|
||||
Find observations with a specific context.
|
||||
|
||||
:param context: Context to search for
|
||||
:return: List of matching observations
|
||||
"""
|
||||
"""Find observations with a specific context."""
|
||||
query = select(Observation).filter(Observation.context == context)
|
||||
result = await self.execute_query(query)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
class RelationRepository(Repository[Relation]):
|
||||
"""
|
||||
Repository for Relation model with memory-specific operations.
|
||||
"""
|
||||
"""Repository for Relation model with memory-specific operations."""
|
||||
|
||||
async def find_by_entities(self, from_id: str, to_id: str) -> Sequence[Relation]:
|
||||
"""
|
||||
Find all relations between two entities.
|
||||
|
||||
:param from_id: Source entity ID
|
||||
:param to_id: Target entity ID
|
||||
:return: List of relations between the entities
|
||||
"""
|
||||
"""Find all relations between two entities."""
|
||||
query = select(Relation).filter(
|
||||
and_(
|
||||
Relation.from_id == from_id,
|
||||
@@ -287,12 +167,7 @@ class RelationRepository(Repository[Relation]):
|
||||
return result.scalars().all()
|
||||
|
||||
async def find_by_type(self, relation_type: str) -> Sequence[Relation]:
|
||||
"""
|
||||
Find all relations of a specific type.
|
||||
|
||||
:param relation_type: Type of relation to find
|
||||
:return: List of matching relations
|
||||
"""
|
||||
"""Find all relations of a specific type."""
|
||||
query = select(Relation).filter(Relation.relation_type == relation_type)
|
||||
result = await self.execute_query(query)
|
||||
return result.scalars().all()
|
||||
@@ -3,10 +3,12 @@ from datetime import datetime, UTC
|
||||
from pathlib import Path
|
||||
|
||||
from basic_memory.repository import EntityRepository
|
||||
from basic_memory.schemas import EntityIn
|
||||
from basic_memory.models import Entity
|
||||
from basic_memory.schemas import EntityIn, ObservationIn
|
||||
from basic_memory.models import Entity, Observation
|
||||
from basic_memory.fileio import EntityNotFoundError
|
||||
from . import ServiceError
|
||||
|
||||
|
||||
class EntityService:
|
||||
"""
|
||||
Service for managing entities in the database.
|
||||
@@ -18,28 +20,35 @@ class EntityService:
|
||||
self.entity_repo = entity_repo
|
||||
|
||||
async def create_entity(self, entity: EntityIn) -> Entity:
|
||||
"""Create a new entity in the database."""
|
||||
# Create DB record
|
||||
db_data = {
|
||||
**entity.model_dump(),
|
||||
"""Create a new entity in the database.
|
||||
|
||||
Note: ID is generated by the EntityIn validator before reaching this method.
|
||||
"""
|
||||
# Create base entity first
|
||||
base_data = {
|
||||
"id": entity.id, # Include the generated ID
|
||||
"name": entity.name,
|
||||
"entity_type": entity.entity_type,
|
||||
"created_at": datetime.now(UTC),
|
||||
}
|
||||
return await self.entity_repo.create(db_data)
|
||||
created_entity = await self.entity_repo.create(base_data)
|
||||
await self.entity_repo.refresh(created_entity, ['observations', 'outgoing_relations', 'incoming_relations'])
|
||||
return created_entity
|
||||
|
||||
async def get_entity(self, entity_id: str) -> Entity:
|
||||
"""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}")
|
||||
raise EntityNotFoundError(f"Entity not found: {entity_id}")
|
||||
|
||||
return db_entity
|
||||
|
||||
# TODO name is not uniaue
|
||||
# TODO name is not unique
|
||||
async def get_by_name(self, name: str) -> Entity:
|
||||
"""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}")
|
||||
raise EntityNotFoundError(f"Entity not found: {name}")
|
||||
|
||||
return db_entity
|
||||
|
||||
|
||||
@@ -29,9 +29,6 @@ class MemoryService:
|
||||
async def create_entities(self, entities_data: List[Dict[str, Any]]) -> List[Entity]:
|
||||
"""Create multiple entities with their observations."""
|
||||
entities_in = [EntityIn.model_validate(data) for data in entities_data]
|
||||
print(f"\nCreating entities with observations:")
|
||||
for e in entities_in:
|
||||
print(f"Entity {e.name}: {len(e.observations)} observations")
|
||||
|
||||
# Write files in parallel (filesystem is source of truth)
|
||||
async def write_file(entity: EntityIn):
|
||||
@@ -41,14 +38,11 @@ class MemoryService:
|
||||
await asyncio.gather(*file_writes)
|
||||
|
||||
async def create_entity_in_db(entity_in: EntityIn):
|
||||
print(f"\nCreating entity in DB: {entity_in.name}")
|
||||
db_entity = await self.entity_service.create_entity(entity_in)
|
||||
print(f"Adding {len(entity_in.observations)} observations to DB for {entity_in.name}")
|
||||
await self.entity_service.create_entity(entity_in)
|
||||
await self.observation_service.add_observations(entity_in, entity_in.observations)
|
||||
[await self.relation_service.create_relation(relation_in) for relation_in in entity_in.relations]
|
||||
# query the entity again to return relations
|
||||
final_entity = await self.entity_service.get_entity(entity_in.id)
|
||||
print(f"Final entity {final_entity.name} has {len(final_entity.observations)} observations in DB")
|
||||
return final_entity
|
||||
|
||||
# Update database index sequentially
|
||||
@@ -91,27 +85,21 @@ class MemoryService:
|
||||
"""
|
||||
# Create new observations
|
||||
new_observations = ObservationsIn.model_validate(observations_in)
|
||||
print(f"\nAdding new observations to entity {new_observations.entity_id}")
|
||||
print(f"New observations to add: {len(new_observations.observations)}")
|
||||
|
||||
# Read entity from filesystem
|
||||
entity = await read_entity_file(self.entities_path, new_observations.entity_id)
|
||||
print(f"Entity {entity.id} from file has {len(entity.observations)} observations")
|
||||
|
||||
|
||||
# Create new observations for the entity
|
||||
for obs in new_observations.observations:
|
||||
entity.observations.append(obs)
|
||||
print(f"After appending, entity has {len(entity.observations)} observations")
|
||||
|
||||
# Write updated entity file
|
||||
await write_entity_file(self.entities_path, entity)
|
||||
|
||||
# Update database index
|
||||
added_observations = await self.observation_service.add_observations(entity, new_observations.observations)
|
||||
print(f"Added {len(added_observations)} observations to DB")
|
||||
|
||||
|
||||
db_entity = await self.entity_service.get_entity(entity.id)
|
||||
print(f"Entity {entity.id} in DB now has {len(db_entity.observations)} observations")
|
||||
return added_observations
|
||||
|
||||
async def delete_entities(self, entity_names: List[str]) -> None:
|
||||
|
||||
@@ -25,23 +25,31 @@ class ObservationService:
|
||||
Add multiple observations to an entity.
|
||||
Returns the created observations with IDs set.
|
||||
"""
|
||||
print(f"\nObservationService.add_observations called for entity {entity.id}")
|
||||
print(f"Adding {len(observations)} observations")
|
||||
|
||||
async def add_observation(observation: ObservationIn) -> Observation:
|
||||
try:
|
||||
return await self.observation_repo.create({
|
||||
obs = await self.observation_repo.create({
|
||||
'entity_id': entity.id,
|
||||
'content': observation.content,
|
||||
'context': observation.context,
|
||||
'created_at': datetime.now(UTC)
|
||||
})
|
||||
# Ensure each observation is flushed
|
||||
await self.observation_repo.session.flush()
|
||||
# Refresh to get latest state
|
||||
await self.observation_repo.session.refresh(obs)
|
||||
return obs
|
||||
except Exception as e:
|
||||
raise DatabaseSyncError(f"Failed to add observation to database: {str(e)}") from e
|
||||
|
||||
# Add each observation and collect the results
|
||||
created_observations = [await add_observation(obs) for obs in observations]
|
||||
print(f"Created {len(created_observations)} observations in DB")
|
||||
|
||||
# Make sure observations are in sync before returning
|
||||
# This helps ensure related entities see the new observations
|
||||
await self.observation_repo.session.flush()
|
||||
for obs in created_observations:
|
||||
await self.observation_repo.session.refresh(obs)
|
||||
|
||||
return created_observations
|
||||
|
||||
async def search_observations(self, query: str) -> List[Observation]:
|
||||
|
||||
@@ -25,8 +25,7 @@ class RelationService:
|
||||
try:
|
||||
db_data = relation.model_dump()
|
||||
db_data['created_at'] = datetime.now(UTC)
|
||||
await self.relation_repo.create(db_data)
|
||||
return relation
|
||||
return await self.relation_repo.create(db_data)
|
||||
except Exception as e:
|
||||
raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e
|
||||
|
||||
|
||||
Reference in New Issue
Block a user