Return entities when modifiying observations/relations in /knowledge routes

src/basic_memory/api/routers/knowledge.py

@@ -13,19 +13,13 @@ from basic_memory.schemas import (
     SearchNodesRequest,
     SearchNodesResponse,
     CreateRelationsRequest,
-    CreateRelationsResponse,
     EntityResponse,
     AddObservationsRequest,
-    ObservationResponse,
     OpenNodesRequest,
     OpenNodesResponse,
     DeleteEntitiesResponse,
     DeleteObservationsRequest,
-    DeleteObservationsResponse,
     DeleteRelationsRequest,
-    DeleteRelationsResponse,
-    AddObservationsResponse,
-    RelationResponse,
     DeleteEntitiesRequest,
 )
 from basic_memory.services.exceptions import EntityNotFoundError
@@ -46,32 +40,29 @@ async def create_entities(
     )
 
 
-@router.post("/relations", response_model=CreateRelationsResponse)
+@router.post("/relations", response_model=CreateEntityResponse)
 async def create_relations(
     data: CreateRelationsRequest, knowledge_service: KnowledgeServiceDep
-) -> CreateRelationsResponse:
+) -> CreateEntityResponse:
     """Create relations between entities."""
-    relations = await knowledge_service.create_relations(data.relations)
-    return CreateRelationsResponse(
-        relations=[RelationResponse.model_validate(relation) for relation in relations]
+
+    # TODO knowledge_service.create_relations should return updated Entities
+    updated_entities = await knowledge_service.create_relations(data.relations)
+    return CreateEntityResponse(
+        entities=[EntityResponse.model_validate(entity) for entity in updated_entities]
     )
 
 
-@router.post("/observations", response_model=AddObservationsResponse)
+@router.post("/observations", response_model=EntityResponse)
 async def add_observations(
     data: AddObservationsRequest, knowledge_service: KnowledgeServiceDep
-) -> AddObservationsResponse:
+) -> EntityResponse:
     """Add observations to an entity."""
     logger.debug(f"Adding observations to entity: {data.entity_id}")
-    observations = await knowledge_service.add_observations(
+    updated_entity = await knowledge_service.add_observations(
         data.entity_id, data.observations, data.context
     )
-    return AddObservationsResponse(
-        entity_id=data.entity_id,
-        observations=[
-            ObservationResponse.model_validate(observation) for observation in observations
-        ],
-    )
+    return EntityResponse.model_validate(updated_entity)
 
 
 ## Read endpoints
@@ -122,20 +113,22 @@ async def delete_entity(
     return DeleteEntitiesResponse(deleted=deleted)
 
 
-@router.post("/observations/delete", response_model=DeleteObservationsResponse)
+@router.post("/observations/delete", response_model=EntityResponse)
 async def delete_observations(
     data: DeleteObservationsRequest, knowledge_service: KnowledgeServiceDep
-) -> DeleteObservationsResponse:
+) -> EntityResponse:
     """Delete observations from an entity."""
     entity_id = data.entity_id
-    deleted = await knowledge_service.delete_observations(entity_id, data.deletions)
-    return DeleteObservationsResponse(deleted=deleted)
+
+    # TODO add knowledge_service.delete_observations
+    updated_entity = await knowledge_service.delete_observations(entity_id, data.deletions)
+    return EntityResponse.model_validate(updated_entity)
 
 
-@router.post("/relations/delete", response_model=DeleteRelationsResponse)
+@router.post("/relations/delete", response_model=EntityResponse)
 async def delete_relations(
     data: DeleteRelationsRequest, knowledge_service: KnowledgeServiceDep
-) -> DeleteRelationsResponse:
+) -> EntityResponse:
     """Delete relations between entities."""
     to_delete = [
         {
@@ -145,5 +138,6 @@ async def delete_relations(
         }
         for relation in data.relations
     ]
-    deleted = await knowledge_service.delete_relations(to_delete)
-    return DeleteRelationsResponse(deleted=deleted)
+    # TODO add knowledge_service.delete_relations
+    updated_entity = await knowledge_service.delete_relations(to_delete)
+    return EntityResponse.model_validate(updated_entity)