From ac95036d0e9edfdbf0d707182a821caf97a3f552 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 7 Dec 2024 20:09:09 -0600 Subject: [PATCH] add_obserations wip --- src/basic_memory/schemas.py | 36 +++++++---- src/basic_memory/services/memory_service.py | 59 +++++++++++++------ .../services/observation_service.py | 43 +++++++++----- tests/test_memory_service.py | 51 +++++++++++++++- 4 files changed, 141 insertions(+), 48 deletions(-) diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 580a9967..ec7664f7 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -3,7 +3,6 @@ Core pydantic models for basic-memory entities, observations, and relations. These models define the schema for our core data types while remaining independent from storage/persistence concerns. """ - from datetime import datetime, UTC from typing import List, Optional, Dict, Any from uuid import uuid4 @@ -11,30 +10,43 @@ from uuid import uuid4 from pydantic import BaseModel, model_validator +class ObservationIn(BaseModel): + """Schema for creating a single observation.""" + content: str + context: Optional[str] = None + +class ObservationsIn(BaseModel): + """Schema for adding observations to an entity.""" + entity_id: str # Maps to Entity.id + observations: List[ObservationIn] + +class ObservationOut(ObservationIn): + """Schema for observation data returned from the service.""" + id: int + +class ObservationsOut(BaseModel): + """Schema for bulk observation operation results.""" + entity_id: str + observations: List[ObservationOut] + +# Original schemas kept for now until we migrate everything class Observation(BaseModel): """An atomic piece of information about an entity.""" id: Optional[int] = None # Let the database handle ID generation content: str context: Optional[str] = None - -class ObservationCreate(BaseModel): - """Schema for creating a new observation.""" - content: str - - class Relation(BaseModel): """ Represents a directed edge between entities in the knowledge graph. Relations are always stored in active voice (e.g. "created", "teaches", etc.) """ - id: Optional[int] = None # Let the database handle ID generation - from_id: str # Reference to Entity text ID - to_id: str # Reference to Entity text ID + id: Optional[int] = None + from_id: str + to_id: str relation_type: str context: Optional[str] = None - class Entity(BaseModel): """ Represents a node in our knowledge graph - could be a person, project, @@ -44,8 +56,6 @@ class Entity(BaseModel): id: str # Text ID for filesystem references name: str entity_type: str - description: str = "" # Match DB default - references: str = "" # Match DB default observations: List[Observation] = [] relations: List[Relation] = [] diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index 3acef6a0..c4c5256f 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -3,7 +3,10 @@ import asyncio from typing import List, Dict, Any, Optional from pathlib import Path -from basic_memory.schemas import Entity, Observation, Relation +from basic_memory.schemas import ( + Entity, Observation, Relation, + ObservationsIn, ObservationsOut, ObservationOut +) from basic_memory.fileio import write_entity_file, read_entity_file, delete_entity_file from basic_memory.services import EntityService, RelationService, ObservationService @@ -65,26 +68,46 @@ class MemoryService: return 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) + async def add_observations(self, observations_in: Dict[str, Any]) -> ObservationsOut: + """Add observations to an existing entity. + + Args: + observations_in: input containing entity_name and observations + + Returns: + ObservationsOut containing the created observations with IDs + """ + # Create new observations + new_observations = ObservationsIn.model_validate(observations_in) - # Write updated entities in parallel - async def write_file(entity: Entity): - await write_entity_file(self.entities_path, entity) + # Read entity from filesystem + entity = await read_entity_file(self.entities_path, new_observations.entity_id) - file_writes = [write_file(entity) for entity in entity_updates] - await asyncio.gather(*file_writes) + # Convert ObservationIn to Observation before adding to entity + entity_observations = [ + Observation(content=obs.content, context=obs.context) + for obs in new_observations.observations + ] + entity.observations.extend(entity_observations) - # Update database indexes sequentially - for entity in entity_updates: - await self.entity_service.rebuild_index(entity) + # 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) + + # Create and return output model + return ObservationsOut( + entity_id=entity.id, + observations=[ + ObservationOut( + id=obs.id, + content=obs.content, + context=obs.context + ) + for obs in added_observations + ] + ) async def delete_entities(self, entity_names: List[str]) -> None: """Delete multiple entities and their associated data.""" diff --git a/src/basic_memory/services/observation_service.py b/src/basic_memory/services/observation_service.py index 8c967691..8df8f949 100644 --- a/src/basic_memory/services/observation_service.py +++ b/src/basic_memory/services/observation_service.py @@ -7,7 +7,8 @@ from sqlalchemy import select, delete from basic_memory.models import Observation as DbObservation from basic_memory.repository import ObservationRepository -from basic_memory.schemas import Entity, Observation +from basic_memory.schemas import Entity, Observation, ObservationIn +from basic_memory.models import Observation as ObservationModel from . import ServiceError, DatabaseSyncError @@ -21,28 +22,38 @@ class ObservationService: self.project_path = project_path self.observation_repo = observation_repo - async def add_observations(self, entity: Entity, observations: List[Observation]) -> List[Observation]: + async def add_observations(self, entity: Entity, observations: List[ObservationIn]) -> List[Observation]: """ Add multiple observations to an entity. - Updates database indexes only - filesystem write handled by MemoryService. + Returns the created observations with IDs set. """ - # Update database index - for observation in observations: + created_observations = [] + + async def add_observation(observation: ObservationIn) -> Observation: try: - await self.observation_repo.create({ - 'id': f"{entity.id}-obs-{uuid4().hex[:8]}", + db_observation = await self.observation_repo.create({ 'entity_id': entity.id, 'content': observation.content, + 'context': observation.context, 'created_at': datetime.now(UTC) }) + # Convert db model to schema + return Observation( + id=db_observation.id, + content=db_observation.content, + context=db_observation.context + ) except Exception as e: - raise DatabaseSyncError(f"Failed to sync observation to database: {str(e)}") from e - - # Add to entity in memory - entity.observations.extend(observations) - return observations - - async def search_observations(self, query: str) -> List[Observation]: + 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] + + # Update entity in memory with the created observations that have IDs + entity.observations.extend(created_observations) + return created_observations + + async def search_observations(self, query: str) -> List[ObservationModel]: """ Search for observations across all entities. @@ -62,7 +73,7 @@ class ObservationService: for obs in result.scalars().all() ] - async def get_observations_by_context(self, context: str) -> List[Observation]: + async def get_observations_by_context(self, context: str) -> List[ObservationModel]: """Get all observations with a specific context.""" db_observations = await self.observation_repo.find_by_context(context) return [ @@ -83,8 +94,8 @@ class ObservationService: # Rebuild from entity's observations 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, + 'context': obs.context, 'created_at': datetime.now(UTC) }) \ No newline at end of file diff --git a/tests/test_memory_service.py b/tests/test_memory_service.py index 9f68ccc2..de694bc4 100644 --- a/tests/test_memory_service.py +++ b/tests/test_memory_service.py @@ -1,9 +1,9 @@ """Tests for the MemoryService class.""" import pytest - from basic_memory.services import MemoryService from basic_memory.fileio import read_entity_file from basic_memory.models import Entity as EntityModel +from basic_memory.schemas import ObservationsIn, ObservationIn test_entities_data = [ { @@ -47,6 +47,55 @@ async def test_create_entities(memory_service: MemoryService): assert entity1_path.exists() assert entity2_path.exists() +@pytest.mark.asyncio +async def test_add_observations(memory_service: MemoryService): + """Should add observations to an existing entity.""" + # First create an entity + entities = await memory_service.create_entities([test_entities_data[0]]) + entity = entities[0] + + # Create observations input + observations_data = { + "entity_id": entity.id, + "observations": [ + {"content": "New observation 1"}, + {"content": "New observation 2", "context": "test context"} + ] + } + + # Add observations + result = await memory_service.add_observations(observations_data) + + # Check the result + assert result.entity_id == entity.id + assert len(result.observations) == 2 + assert result.observations[0].content == "New observation 1" + assert result.observations[0].context is None + assert result.observations[1].content == "New observation 2" + assert result.observations[1].context == "test context" + + # Verify file was updated + updated_entity = await read_entity_file(memory_service.entities_path, entity.id) + assert len(updated_entity.observations) == 4 # 2 original + 2 new + assert updated_entity.observations[2].content == "New observation 1" + assert updated_entity.observations[3].content == "New observation 2" + assert updated_entity.observations[3].context == "test context" + + # Verify database was updated via observation service + db_entity = await memory_service.entity_service.get_entity(entity.id) + assert len(db_entity.observations) == 4 + +@pytest.mark.asyncio +async def test_add_observations_nonexistent_entity(memory_service: MemoryService): + """Should raise an appropriate error when adding observations to a non-existent entity.""" + observations_data = { + "entity_id": "nonexistent-id", + "observations": [{"content": "Test observation"}] + } + + with pytest.raises(Exception) as exc: # We might want to define a specific error type + await memory_service.add_observations(observations_data) + @pytest.mark.asyncio async def test_create_relations(memory_service: MemoryService): """Should create relations between entities and update both filesystem and database."""