diff --git a/src/basic_memory/api/routers/knowledge.py b/src/basic_memory/api/routers/knowledge.py index d6d0b82d..5cf6997d 100644 --- a/src/basic_memory/api/routers/knowledge.py +++ b/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) diff --git a/src/basic_memory/schemas/response.py b/src/basic_memory/schemas/response.py index f4c02b8b..46126ad8 100644 --- a/src/basic_memory/schemas/response.py +++ b/src/basic_memory/schemas/response.py @@ -259,29 +259,6 @@ class AddObservationsResponse(SQLAlchemyModel): observations: List[ObservationResponse] -class CreateRelationsResponse(SQLAlchemyModel): - """Response for creating new relations. - - Returns complete information about all created relations, - including their generated IDs. - - Example Response: - { - "relations": [ - { - "id": 46, - "from_id": "component/memory_service", - "to_id": "component/database", - "relation_type": "depends_on", - "context": "Storage dependency" - } - ] - } - """ - - relations: List[Relation] - - class DeleteEntitiesResponse(SQLAlchemyModel): """Response indicating successful entity deletion.