add_obserations wip

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] = []
 

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."""

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)
             })

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."""