From 0bbcf636f69a2f0b53459db38ca011102d2c76fb Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 24 Dec 2024 18:15:59 -0600 Subject: [PATCH] change /knowledge endpoints to use ids --- src/basic_memory/api/routers/knowledge.py | 17 +- src/basic_memory/models/knowledge.py | 2 +- .../repository/entity_repository.py | 36 +- .../repository/relation_repository.py | 2 +- src/basic_memory/schemas/__init__.py | 1 - src/basic_memory/schemas/base.py | 47 ++- src/basic_memory/schemas/delete.py | 6 +- src/basic_memory/schemas/request.py | 4 +- src/basic_memory/schemas/response.py | 65 ++-- src/basic_memory/services/entity_service.py | 46 ++- src/basic_memory/services/file_service.py | 2 +- .../services/knowledge/entities.py | 33 +- src/basic_memory/services/knowledge/files.py | 9 +- .../services/knowledge/observations.py | 28 +- .../services/knowledge/relations.py | 65 ++-- .../services/observation_service.py | 2 +- src/basic_memory/services/relation_service.py | 20 +- tests/api/test_knowledge_router.py | 312 ++++++++---------- tests/repository/test_entity_repository.py | 127 +++++-- tests/schemas/test_schemas.py | 17 +- tests/services/test_entity_service.py | 200 ++++++++--- 21 files changed, 620 insertions(+), 421 deletions(-) diff --git a/src/basic_memory/api/routers/knowledge.py b/src/basic_memory/api/routers/knowledge.py index e94a5492..39cbf05c 100644 --- a/src/basic_memory/api/routers/knowledge.py +++ b/src/basic_memory/api/routers/knowledge.py @@ -22,6 +22,7 @@ from basic_memory.schemas import ( DeleteRelationsRequest, DeleteEntitiesRequest, ) +from basic_memory.schemas.base import PathId from basic_memory.services.exceptions import EntityNotFoundError router = APIRouter(prefix="/knowledge", tags=["knowledge"]) @@ -66,14 +67,14 @@ async def add_observations( ## Read endpoints -@router.get("/entities/{entity_id}", response_model=EntityResponse) -async def get_entity(entity_id: int, entity_service: EntityServiceDep) -> EntityResponse: +@router.get("/entities/{path_id:path}", response_model=EntityResponse) +async def get_entity(path_id: PathId, entity_service: EntityServiceDep) -> EntityResponse: """Get a specific entity by ID.""" try: - entity = await entity_service.get_entity(entity_id) + entity = await entity_service.get_by_path_id(path_id) return EntityResponse.model_validate(entity) except EntityNotFoundError: - raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found") + raise HTTPException(status_code=404, detail=f"Entity with {path_id} not found") @router.post("/search", response_model=SearchNodesResponse) @@ -106,7 +107,7 @@ async def open_nodes(data: OpenNodesRequest, entity_service: EntityServiceDep) - async def delete_entity( data: DeleteEntitiesRequest, knowledge_service: KnowledgeServiceDep ) -> DeleteEntitiesResponse: - """Delete a specific entity by ID.""" + """Delete a specific entity by PathId.""" deleted = await knowledge_service.delete_entities(data.entity_ids) return DeleteEntitiesResponse(deleted=deleted) @@ -116,8 +117,8 @@ async def delete_observations( data: DeleteObservationsRequest, knowledge_service: KnowledgeServiceDep ) -> EntityResponse: """Delete observations from an entity.""" - entity_id = data.entity_id - updated_entity = await knowledge_service.delete_observations(entity_id, data.deletions) + path_id = data.entity_id + updated_entity = await knowledge_service.delete_observations(path_id, data.deletions) return EntityResponse.model_validate(updated_entity) @@ -137,4 +138,4 @@ async def delete_relations( updated_entities = await knowledge_service.delete_relations(to_delete) return CreateEntityResponse( entities=[EntityResponse.model_validate(entity) for entity in updated_entities] - ) \ No newline at end of file + ) diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py index e9c8b136..67edc5be 100644 --- a/src/basic_memory/models/knowledge.py +++ b/src/basic_memory/models/knowledge.py @@ -32,10 +32,10 @@ class Entity(Base): id: Mapped[int] = mapped_column(Integer, primary_key=True) name: Mapped[str] = mapped_column(String) entity_type: Mapped[str] = mapped_column(String) + path_id: Mapped[str] = mapped_column(String, index=True) # Content and validation description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) - path: Mapped[Optional[str]] = mapped_column(String, nullable=True) checksum: Mapped[Optional[str]] = mapped_column(String, nullable=True) # Metadata and tracking diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index fa4f9340..f4988e80 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -18,13 +18,9 @@ class EntityRepository(Repository[Entity]): """Initialize with session maker.""" super().__init__(session_maker, Entity) - async def get_entity_by_type_and_name(self, entity_type: str, name: str) -> Optional[Entity]: + async def get_by_path_id(self, path_id: str) -> Optional[Entity]: """Get entity by type and name.""" - query = ( - self.select() - .options(*self.get_load_options()) - .where(Entity.entity_type == entity_type, Entity.name == name) - ) + query = self.select().where(Entity.path_id == path_id).options(*self.get_load_options()) return await self.find_one(query) async def list_entities( @@ -86,3 +82,31 @@ class EntityRepository(Repository[Entity]): selectinload(Entity.from_relations).selectinload(Relation.to_entity), selectinload(Entity.to_relations).selectinload(Relation.from_entity), ] + + async def find_by_path_ids(self, path_ids: List[str]) -> Sequence[Entity]: + """Find multiple entities by their entity_type and name pairs.""" + + # Handle empty input explicitly + if not path_ids: + return [] + + # Use existing select pattern + query = self.select().options(*self.get_load_options()).where(Entity.path_id.in_(path_ids)) + + result = await self.execute_query(query) + return list(result.scalars().all()) + + async def delete_by_path_ids(self, path_ids: List[str]) -> int: + """Delete multiple entities by entity_type and name pairs.""" + + # Handle empty input explicitly + if not path_ids: + return 0 + + # Find matching entities + entities = await self.find_by_path_ids(path_ids) + if not entities: + return 0 + + # Use existing delete_by_ids + return await self.delete_by_ids([entity.id for entity in entities]) diff --git a/src/basic_memory/repository/relation_repository.py b/src/basic_memory/repository/relation_repository.py index 041b9d27..c1d4ef49 100644 --- a/src/basic_memory/repository/relation_repository.py +++ b/src/basic_memory/repository/relation_repository.py @@ -7,7 +7,7 @@ from sqlalchemy.ext.asyncio import async_sessionmaker from sqlalchemy.orm import selectinload from sqlalchemy.orm.interfaces import LoaderOption -from basic_memory.models import Relation +from basic_memory.models import Relation, Entity from basic_memory.repository.repository import Repository diff --git a/src/basic_memory/schemas/__init__.py b/src/basic_memory/schemas/__init__.py index 74742572..9399b4bd 100644 --- a/src/basic_memory/schemas/__init__.py +++ b/src/basic_memory/schemas/__init__.py @@ -34,7 +34,6 @@ from basic_memory.schemas.request import ( from basic_memory.schemas.response import ( SQLAlchemyModel, ObservationResponse, - ObservationsResponse, RelationResponse, EntityResponse, CreateEntityResponse, diff --git a/src/basic_memory/schemas/base.py b/src/basic_memory/schemas/base.py index a8364351..a95b3fcb 100644 --- a/src/basic_memory/schemas/base.py +++ b/src/basic_memory/schemas/base.py @@ -32,25 +32,36 @@ Common Relation Types: - 'tested_by': Test coverage """ +import re from typing import List, Optional, Annotated from annotated_types import MinLen, MaxLen from pydantic import BaseModel, BeforeValidator -# Strip whitespace -def strip_whitespace(obs: str) -> str: - return obs.strip() +def to_snake_case(name: str) -> str: + """Convert a string to snake_case. + + Examples: + BasicMemory -> basic_memory + Memory Service -> memory_service + memory-service -> memory_service + Memory_Service -> memory_service + """ + # Replace spaces and hyphens with underscores + s1 = re.sub(r"[\s\-]", "_", name) + + # Insert underscore between camelCase + s2 = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", s1) + + # Convert to lowercase + return s2.lower() -def lower_strip_whitespace(val: str) -> str: - return strip_whitespace(val.lower()) - - -PathId = Annotated[str, BeforeValidator(lower_strip_whitespace)] +PathId = Annotated[str, BeforeValidator(to_snake_case)] """Unique identifier in format '{path}/{normalized_name}'.""" -Observation = Annotated[str, BeforeValidator(strip_whitespace), MinLen(1), MaxLen(1000)] +Observation = Annotated[str, MinLen(1), MaxLen(1000)] """A single piece of information about an entity. Must be non-empty and under 1000 characters. Best Practices: @@ -65,7 +76,7 @@ Examples: - "Depends on SQLAlchemy for database operations" """ -EntityType = Annotated[str, BeforeValidator(strip_whitespace), MinLen(1), MaxLen(200)] +EntityType = Annotated[str, BeforeValidator(to_snake_case), MinLen(1), MaxLen(200)] """Classification of entity (e.g., 'person', 'project', 'concept'). The type serves multiple purposes: @@ -77,7 +88,7 @@ The type serves multiple purposes: Common types are listed in the module docstring. """ -RelationType = Annotated[str, BeforeValidator(strip_whitespace), MinLen(1), MaxLen(200)] +RelationType = Annotated[str, BeforeValidator(to_snake_case), MinLen(1), MaxLen(200)] """Type of relationship between entities. Always use active voice present tense. Guidelines: @@ -127,8 +138,8 @@ class Relation(BaseModel): } """ - from_id: int - to_id: int + from_id: PathId + to_id: PathId relation_type: RelationType context: Optional[str] = None @@ -147,7 +158,7 @@ class Entity(BaseModel): 1. Project Entity: { - "name": "Basic_Memory", + "name": "BasicMemory", "entity_type": "project", "description": "Knowledge graph system for AI-human collaboration", "observations": [ @@ -193,13 +204,13 @@ class Entity(BaseModel): } """ - id: Optional[int] = None name: str entity_type: EntityType description: Optional[str] = None observations: List[Observation] = [] @property - def file_path(self) -> str: - """The relative file path for this entity.""" - return f"{id}.md" + def path_id(self) -> PathId: + """Get the path ID in format {type}/{snake_case_name}.""" + normalized_name = to_snake_case(self.name) + return f"{self.entity_type}/{normalized_name}" diff --git a/src/basic_memory/schemas/delete.py b/src/basic_memory/schemas/delete.py index 5d2bf97a..60e580bd 100644 --- a/src/basic_memory/schemas/delete.py +++ b/src/basic_memory/schemas/delete.py @@ -21,7 +21,7 @@ from typing import List, Annotated from annotated_types import MinLen from pydantic import BaseModel -from basic_memory.schemas.base import Relation, Observation +from basic_memory.schemas.base import Relation, Observation, PathId class DeleteEntitiesRequest(BaseModel): @@ -56,7 +56,7 @@ class DeleteEntitiesRequest(BaseModel): 5. Create relations to replacement entities if applicable """ - entity_ids: Annotated[List[int], MinLen(1)] + entity_ids: Annotated[List[PathId], MinLen(1)] class DeleteRelationsRequest(BaseModel): @@ -132,5 +132,5 @@ class DeleteObservationsRequest(BaseModel): 5. Updating implementation details """ - entity_id: int + entity_id: PathId deletions: Annotated[List[Observation], MinLen(1)] diff --git a/src/basic_memory/schemas/request.py b/src/basic_memory/schemas/request.py index 24474426..092bf2ab 100644 --- a/src/basic_memory/schemas/request.py +++ b/src/basic_memory/schemas/request.py @@ -45,7 +45,7 @@ class AddObservationsRequest(BaseModel): 4. Add observations in logical groups for better history tracking """ - entity_id: int + entity_id: PathId context: Optional[str] = None observations: List[Observation] @@ -152,7 +152,7 @@ class OpenNodesRequest(BaseModel): relations between entities that interest you. """ - entity_ids: Annotated[List[int], MinLen(1)] + entity_ids: Annotated[List[PathId], MinLen(1)] class CreateRelationsRequest(BaseModel): diff --git a/src/basic_memory/schemas/response.py b/src/basic_memory/schemas/response.py index 36918e6d..5d523435 100644 --- a/src/basic_memory/schemas/response.py +++ b/src/basic_memory/schemas/response.py @@ -14,9 +14,9 @@ Key Features: import datetime from typing import List, Optional, Dict, Any -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field, AliasPath, AliasChoices -from basic_memory.schemas.base import Observation, Relation +from basic_memory.schemas.base import Observation, Relation, PathId class SQLAlchemyModel(BaseModel): @@ -43,36 +43,9 @@ class ObservationResponse(SQLAlchemyModel): } """ - id: int content: Observation -class ObservationsResponse(SQLAlchemyModel): - """Response schema for bulk observation operations. - - Returns all added/affected observations with their IDs and - the entity they were added to. - - Example Response: - { - "entity_id": "component/memory_service", - "observations": [ - { - "id": 123, - "content": "Added async support" - }, - { - "id": 124, - "content": "Improved error handling" - } - ] - } - """ - - entity_id: int - observations: List[ObservationResponse] - - class RelationResponse(Relation, SQLAlchemyModel): """Response schema for relation operations. @@ -81,15 +54,28 @@ class RelationResponse(Relation, SQLAlchemyModel): Example Response: { - "id": 45, "from_id": "test/memory_test", "to_id": "component/memory_service", "relation_type": "validates", "context": "Comprehensive test suite" } """ - - id: int + from_id: PathId = Field( + # use the path_id from the associated Entity + # or the from_id value + validation_alias=AliasChoices( + AliasPath('from_entity', 'path_id'), + 'from_id', + ) + ) + to_id: PathId = Field( + # use the path_id from the associated Entity + # or the to_id value + validation_alias=AliasChoices( + AliasPath('to_entity', 'path_id'), + 'to_id', + ) + ) class EntityResponse(SQLAlchemyModel): @@ -103,23 +89,20 @@ class EntityResponse(SQLAlchemyModel): Example Response: { - "id": "component/memory_service", + "path_id": "component/memory_service", "name": "MemoryService", "entity_type": "component", "description": "Core persistence service", "observations": [ { - "id": 123, "content": "Uses SQLite storage" }, { - "id": 124, "content": "Implements async operations" } ], "relations": [ { - "id": 45, "from_id": "test/memory_test", "to_id": "component/memory_service", "relation_type": "validates", @@ -129,7 +112,7 @@ class EntityResponse(SQLAlchemyModel): } """ - id: int + path_id: PathId name: str entity_type: str description: Optional[str] = None @@ -148,7 +131,7 @@ class CreateEntityResponse(SQLAlchemyModel): { "entities": [ { - "id": "component/search_service", + "path_id": "component/search_service", "name": "SearchService", "entity_type": "component", "description": "Knowledge graph search", @@ -161,7 +144,7 @@ class CreateEntityResponse(SQLAlchemyModel): "relations": [] }, { - "id": "document/api_docs", + "path_id": "document/api_docs", "name": "API_Documentation", "entity_type": "document", "description": "API Reference", @@ -190,7 +173,7 @@ class SearchNodesResponse(SQLAlchemyModel): { "matches": [ { - "id": "component/memory_service", + "path_id": "component/memory_service", "name": "MemoryService", "entity_type": "component", "description": "Core service", @@ -219,7 +202,7 @@ class OpenNodesResponse(SQLAlchemyModel): { "entities": [ { - "id": "component/memory_service", + "path_id": "component/memory_service", "name": "MemoryService", "entity_type": "component", "description": "Core service", diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index c0713519..00bc7c7d 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -15,6 +15,7 @@ def entity_model(entity): model = EntityModel( name=entity.name, entity_type=entity.entity_type, + path_id=entity.path_id, description=entity.description, observations=[Observation(content=observation) for observation in entity.observations], ) @@ -44,46 +45,41 @@ class EntityService(BaseService[EntityRepository]): created = await self.repository.add_all([entity_model(entity) for entity in entities_in]) return created - async def update_entity(self, entity_id: int, update_data: Dict[str, Any]) -> EntityModel: + async def update_entity(self, path_id: str, update_data: Dict[str, Any]) -> EntityModel: """Update an entity's fields.""" - logger.debug(f"Updating entity {entity_id} with data: {update_data}") - updated = await self.repository.update(entity_id, update_data) + logger.debug(f"Updating entity path_id: {path_id} with data: {update_data}") + entity = await self.get_by_path_id(path_id) + + updated = await self.repository.update(entity.id, update_data) if not updated: - raise EntityNotFoundError(f"Entity not found: {entity_id}") + raise EntityNotFoundError(f"Entity not found: {path_id}") return updated - async def get_entity(self, entity_id: int) -> EntityModel: - """Get entity by ID.""" - logger.debug(f"Getting entity by ID: {entity_id}") - db_entity = await self.repository.find_by_id(entity_id) - if not db_entity: - raise EntityNotFoundError(f"Entity not found: {entity_id}") - return db_entity - - async def get_by_type_and_name(self, entity_type: str, name: str) -> EntityModel: + async def get_by_path_id(self, path_id: str) -> EntityModel: """Get entity by type and name combination.""" - logger.debug(f"Getting entity by type/name: {entity_type}/{name}") - db_entity = await self.repository.get_entity_by_type_and_name(entity_type, name) + logger.debug(f"Getting entity by path_id: {path_id}") + db_entity = await self.repository.get_by_path_id(path_id) if not db_entity: - raise EntityNotFoundError(f"Entity not found: {entity_type}/{name}") + raise EntityNotFoundError(f"Entity not found: {path_id}") return db_entity async def get_all(self) -> Sequence[EntityModel]: """Get all entities.""" return await self.repository.find_all() - async def delete_entity(self, entity_id: int) -> bool: + async def delete_entity(self, path_id: str) -> bool: """Delete entity from database.""" - logger.debug(f"Deleting entity: {entity_id}") - return await self.repository.delete(entity_id) + logger.debug(f"Deleting entity path_id: {path_id}") + entity = await self.get_by_path_id(path_id) + return await self.repository.delete(entity.id) - async def open_nodes(self, entity_ids: List[int]) -> Sequence[EntityModel]: + async def open_nodes(self, path_ids: List[str]) -> Sequence[EntityModel]: """Get specific nodes and their relationships.""" - logger.debug(f"Opening nodes entity_ids: {entity_ids}") - return await self.repository.find_by_ids(entity_ids) + logger.debug(f"Opening nodes path_ids: {path_ids}") + return await self.repository.find_by_path_ids(path_ids) - async def delete_entities(self, entity_ids: List[int]) -> bool: + async def delete_entities(self, path_ids: List[str]) -> bool: """Delete entities and their files.""" - logger.debug(f"Deleting entities: {entity_ids}") - deleted_count = await self.repository.delete_by_ids(entity_ids) + logger.debug(f"Deleting entities: {path_ids}") + deleted_count = await self.repository.delete_by_path_ids(path_ids) return deleted_count > 0 diff --git a/src/basic_memory/services/file_service.py b/src/basic_memory/services/file_service.py index df6342a3..93212c7e 100644 --- a/src/basic_memory/services/file_service.py +++ b/src/basic_memory/services/file_service.py @@ -109,7 +109,7 @@ class FileService: async def add_frontmatter( self, *, - id: int, + id: str, content: str, created: datetime | None = None, updated: datetime | None = None, diff --git a/src/basic_memory/services/knowledge/entities.py b/src/basic_memory/services/knowledge/entities.py index 3754e0d3..22e18d7f 100644 --- a/src/basic_memory/services/knowledge/entities.py +++ b/src/basic_memory/services/knowledge/entities.py @@ -15,24 +15,29 @@ class EntityOperations(FileOperations): async def create_entity(self, entity: EntitySchema) -> EntityModel: """Create a new entity and write to filesystem.""" logger.debug(f"Creating entity: {entity}") + + db_entity = None + file_path = None try: # 1. Create entity in DB db_entity = await self.entity_service.create_entity(entity) # 2. Write file and get checksum - checksum = await self.write_entity_file(db_entity) + file_path, checksum = await self.write_entity_file(db_entity) # 3. Update DB with checksum - updated = await self.entity_service.update_entity(db_entity.id, {"checksum": checksum}) + updated = await self.entity_service.update_entity( + db_entity.path_id, {"checksum": checksum} + ) return updated except Exception as e: # Clean up on any failure - if "db_entity" in locals(): - await self.entity_service.delete_entity(db_entity.id) # pyright: ignore [reportPossiblyUnboundVariable] - if "path" in locals(): - await self.file_service.delete_file(path) # pyright: ignore [reportUndefinedVariable] # noqa: F821 + if db_entity: + await self.entity_service.delete_entity(db_entity.path_id) + if file_path: + await self.file_service.delete_file(file_path) logger.error(f"Failed to create entity: {e}") raise @@ -47,13 +52,13 @@ class EntityOperations(FileOperations): return created - async def delete_entity(self, entity_id: int) -> bool: + async def delete_entity(self, path_id: str) -> bool: """Delete entity and its file.""" - logger.debug(f"Deleting entity: {entity_id}") + logger.debug(f"Deleting entity: {path_id}") try: # Get entity first for file deletion - entity = await self.entity_service.get_entity(entity_id) + entity = await self.entity_service.get_by_path_id(path_id) if not entity: return True # Already deleted @@ -62,20 +67,20 @@ class EntityOperations(FileOperations): await self.file_service.delete_file(path) # Delete from DB (this will cascade to observations/relations) - return await self.entity_service.delete_entity(entity_id) + return await self.entity_service.delete_entity(path_id) except Exception as e: logger.error(f"Failed to delete entity: {e}") raise - async def delete_entities(self, entity_ids: List[int]) -> bool: + async def delete_entities(self, path_ids: List[str]) -> bool: """Delete multiple entities and their files.""" - logger.debug(f"Deleting entities: {entity_ids}") + logger.debug(f"Deleting entities: {path_ids}") success = True # Let errors bubble up - for entity_id in entity_ids: - await self.delete_entity(entity_id) + for path_id in path_ids: + await self.delete_entity(path_id) success = True return success diff --git a/src/basic_memory/services/knowledge/files.py b/src/basic_memory/services/knowledge/files.py index 790bf32d..2b257c36 100644 --- a/src/basic_memory/services/knowledge/files.py +++ b/src/basic_memory/services/knowledge/files.py @@ -1,6 +1,7 @@ """File operations for knowledge service.""" from pathlib import Path +from typing import Tuple from loguru import logger @@ -30,24 +31,24 @@ class FileOperations: """Generate filesystem path for entity.""" return self.base_path / entity.entity_type / f"{entity.name}.md" - async def write_entity_file(self, entity: EntityModel) -> str: + async def write_entity_file(self, entity: EntityModel) -> Tuple[Path, str]: """Write entity to filesystem and return checksum.""" try: # Ensure we have a fresh entity with all relations loaded - entity = await self.entity_service.get_entity(entity.id) + entity = await self.entity_service.get_by_path_id(entity.path_id) # Format content path = self.get_entity_path(entity) entity_content = await self.knowledge_writer.format_content(entity) file_content = await self.file_service.add_frontmatter( - id=entity.id, + id=entity.path_id, content=entity_content, created=entity.created_at, updated=entity.updated_at, ) # Write and get checksum - return await self.file_service.write_file(path, file_content) + return path, await self.file_service.write_file(path, file_content) except Exception as e: logger.error(f"Failed to write entity file: {e}") diff --git a/src/basic_memory/services/knowledge/observations.py b/src/basic_memory/services/knowledge/observations.py index 1f49a9f7..6d3a8046 100644 --- a/src/basic_memory/services/knowledge/observations.py +++ b/src/basic_memory/services/knowledge/observations.py @@ -18,26 +18,26 @@ class ObservationOperations(RelationOperations): self.observation_service = observation_service async def add_observations( - self, entity_id: int, observations: List[str], context: str | None = None + self, path_id: str, observations: List[str], context: str | None = None ) -> EntityModel: """Add observations to entity and update its file.""" - logger.debug(f"Adding observations to entity {entity_id}") + logger.debug(f"Adding observations to entity {path_id}") try: # Get entity to update - entity = await self.entity_service.get_entity(entity_id) + entity = await self.entity_service.get_by_path_id(path_id) if not entity: - raise EntityNotFoundError(f"Entity not found: {entity_id}") + raise EntityNotFoundError(f"Entity not found: {path_id}") # Add observations to DB - await self.observation_service.add_observations(entity_id, observations, context) + await self.observation_service.add_observations(entity.id, observations, context) # Get updated entity - updated_entity = await self.entity_service.get_entity(entity_id) + updated_entity = await self.entity_service.get_by_path_id(path_id) # Write updated file and checksum checksum = await self.write_entity_file(entity) - await self.entity_service.update_entity(entity_id, {"checksum": checksum}) + await self.entity_service.update_entity(path_id, {"checksum": checksum}) return updated_entity @@ -45,25 +45,25 @@ class ObservationOperations(RelationOperations): logger.error(f"Failed to add observations: {e}") raise - async def delete_observations(self, entity_id: int, observations: List[str]) -> EntityModel: + async def delete_observations(self, path_id: str, observations: List[str]) -> EntityModel: """Delete observations from entity and update its file.""" - logger.debug(f"Deleting observations from entity {entity_id}") + logger.debug(f"Deleting observations from entity {path_id}") try: # Get updated entity - entity = await self.entity_service.get_entity(entity_id) + entity = await self.entity_service.get_by_path_id(path_id) if not entity: - raise EntityNotFoundError(f"Entity not found: {entity_id}") + raise EntityNotFoundError(f"Entity not found: {path_id}") # Delete observations from DB - await self.observation_service.delete_observations(entity_id, observations) + await self.observation_service.delete_observations(entity.id, observations) # Write updated file checksum = await self.write_entity_file(entity) - await self.entity_service.update_entity(entity_id, {"checksum": checksum}) + await self.entity_service.update_entity(path_id, {"checksum": checksum}) # Get final entity with all updates - return await self.entity_service.get_entity(entity_id) + return await self.entity_service.get_by_path_id(path_id) except Exception as e: logger.error(f"Failed to delete observations: {e}") diff --git a/src/basic_memory/services/knowledge/relations.py b/src/basic_memory/services/knowledge/relations.py index a4eb90a8..1156c1df 100644 --- a/src/basic_memory/services/knowledge/relations.py +++ b/src/basic_memory/services/knowledge/relations.py @@ -5,6 +5,7 @@ from typing import Sequence, List, Dict, Any from loguru import logger from basic_memory.models import Entity as EntityModel +from basic_memory.models import Relation as RelationModel from basic_memory.schemas import Relation as RelationSchema from basic_memory.services.exceptions import EntityNotFoundError from basic_memory.services.relation_service import RelationService @@ -21,78 +22,88 @@ class RelationOperations(EntityOperations): async def create_relations(self, relations: List[RelationSchema]) -> Sequence[EntityModel]: """Create relations and return updated entities.""" logger.debug(f"Creating {len(relations)} relations") - created_entities = [] - updated_entity_ids = set() + updated_entities = [] + entities_to_update = set() - for relation in relations: + for rs in relations: try: - # Create relation in DB + from_entity = await self.entity_service.get_by_path_id(rs.from_id) + to_entity = await self.entity_service.get_by_path_id(rs.to_id) + + relation = RelationModel( + from_id=from_entity.id, + to_id=to_entity.id, + relation_type=rs.relation_type, + context=rs.context, + ) + # Create rs in DB await self.relation_service.create_relation(relation) # Keep track of entities we need to update - updated_entity_ids.add(relation.from_id) - updated_entity_ids.add(relation.to_id) + entities_to_update.add(rs.from_id) + entities_to_update.add(rs.to_id) except Exception as e: - logger.error(f"Failed to create relation: {e}") + logger.error(f"Failed to create rs: {e}") continue # Get fresh copies of all updated entities - for entity_id in updated_entity_ids: + for path_id in entities_to_update: try: # Get fresh entity - entity = await self.entity_service.get_entity(entity_id) - if not entity: - raise EntityNotFoundError(f"Entity not found: {entity_id}") + entity = await self.entity_service.get_by_path_id(path_id) # Write updated file - checksum = await self.write_entity_file(entity) - updated = await self.entity_service.update_entity(entity_id, {"checksum": checksum}) + _, checksum = await self.write_entity_file(entity) + updated = await self.entity_service.update_entity(path_id, {"checksum": checksum}) - created_entities.append(updated) + updated_entities.append(updated) except Exception as e: - logger.error(f"Failed to update entity {entity_id}: {e}") + logger.error(f"Failed to update entity {path_id}: {e}") continue - - return created_entities + + # select again to eagerly load all relations + return await self.entity_service.open_nodes([e.path_id for e in updated_entities]) async def delete_relations(self, to_delete: List[Dict[str, Any]]) -> Sequence[EntityModel]: """Delete relations and return all updated entities.""" logger.debug(f"Deleting {len(to_delete)} relations") - updated_entity_ids = set() + updated_entities = [] + entities_to_update = set() try: # Delete relations from DB for relation in to_delete: - updated_entity_ids.add(relation["from_id"]) - updated_entity_ids.add(relation["to_id"]) + entities_to_update.add(relation["from_id"]) + entities_to_update.add(relation["to_id"]) deleted = await self.relation_service.delete_relations(to_delete) if not deleted: logger.warning("No relations were deleted") # Get fresh copies of all updated entities - updated_entities = [] - for entity_id in updated_entity_ids: + for path_id in entities_to_update: try: # Get fresh entity - entity = await self.entity_service.get_entity(entity_id) + entity = await self.entity_service.get_by_path_id(path_id) if not entity: - raise EntityNotFoundError(f"Entity not found: {entity_id}") + raise EntityNotFoundError(f"Entity not found: {path_id}") # Write updated file checksum = await self.write_entity_file(entity) - updated = await self.entity_service.update_entity(entity_id, {"checksum": checksum}) + updated = await self.entity_service.update_entity( + path_id, {"checksum": checksum} + ) updated_entities.append(updated) except Exception as e: - logger.error(f"Failed to update entity {entity_id}: {e}") + logger.error(f"Failed to update entity {path_id}: {e}") continue return updated_entities except Exception as e: logger.error(f"Failed to delete relations: {e}") - raise \ No newline at end of file + raise diff --git a/src/basic_memory/services/observation_service.py b/src/basic_memory/services/observation_service.py index 689d10e1..f9900e16 100644 --- a/src/basic_memory/services/observation_service.py +++ b/src/basic_memory/services/observation_service.py @@ -21,7 +21,7 @@ class ObservationService(BaseService[ObservationRepository]): async def add_observations( self, entity_id: int, observations: List[str], context: str | None = None - ) -> List[ObservationModel]: + ) -> Sequence[ObservationModel]: """Add multiple observations to an entity.""" logger.debug(f"Adding {len(observations)} observations to entity: {entity_id}") return await self.repository.create_all( diff --git a/src/basic_memory/services/relation_service.py b/src/basic_memory/services/relation_service.py index 2d2e646b..55605566 100644 --- a/src/basic_memory/services/relation_service.py +++ b/src/basic_memory/services/relation_service.py @@ -1,12 +1,11 @@ """Service for managing relations in the database.""" -from typing import List, Dict, Any +from typing import List, Dict, Any, Sequence from loguru import logger -from basic_memory.models import Relation as RelationModel +from basic_memory.models import Entity, Relation from basic_memory.repository.relation_repository import RelationRepository -from basic_memory.schemas import Entity as EntitySchema, Relation from .service import BaseService @@ -19,13 +18,13 @@ class RelationService(BaseService[RelationRepository]): def __init__(self, relation_repository: RelationRepository): super().__init__(relation_repository) - async def create_relation(self, relation: Relation) -> RelationModel: + async def create_relation(self, relation: Relation) -> Relation: """Create a new relation in the database.""" logger.debug(f"Creating relation: {relation}") - return await self.repository.create(relation.model_dump()) + return await self.repository.add(relation) async def delete_relation( - self, from_entity: EntitySchema, to_entity: EntitySchema, relation_type: str + self, from_entity: Entity, to_entity: Entity, relation_type: str ) -> bool: """Delete a specific relation between entities.""" logger.debug(f"Deleting relation between {from_entity.id} and {to_entity.id}") @@ -53,9 +52,8 @@ class RelationService(BaseService[RelationRepository]): return deleted - async def create_relations(self, relations_data: List[Relation]) -> List[RelationModel]: + async def create_relations(self, relations: List[Relation]) -> Sequence[Relation]: """Create multiple relations between entities.""" - logger.debug(f"Creating {len(relations_data)} relations") - return await self.repository.create_all( - [Relation.model_dump(relation) for relation in relations_data] - ) + logger.debug(f"Creating {len(relations)} relations") + return await self.repository.add_all(relations) + diff --git a/tests/api/test_knowledge_router.py b/tests/api/test_knowledge_router.py index 90b454fa..be818b8c 100644 --- a/tests/api/test_knowledge_router.py +++ b/tests/api/test_knowledge_router.py @@ -1,6 +1,7 @@ """Tests for knowledge graph API routes.""" from typing import List +from urllib.parse import quote import pytest from httpx import AsyncClient @@ -15,7 +16,7 @@ from basic_memory.schemas import ( async def create_entity(client) -> EntityResponse: data = { - "name": "Test Entity", + "name": "TestEntity", "entity_type": "test", "observations": ["First observation", "Second observation"], } @@ -37,10 +38,10 @@ async def create_entity(client) -> EntityResponse: return create_response.entities[0] -async def add_observations(client, entity_id) -> List[ObservationResponse]: +async def add_observations(client, path_id: str) -> List[ObservationResponse]: response = await client.post( "/knowledge/observations", - json={"entity_id": entity_id, "observations": ["First observation", "Second observation"]}, + json={"entity_id": path_id, "observations": ["First observation", "Second observation"]}, ) # Verify observations were added assert response.status_code == 200 @@ -53,19 +54,21 @@ async def add_observations(client, entity_id) -> List[ObservationResponse]: async def create_related_entities(client) -> List[RelationResponse]: # pyright: ignore [reportReturnType] # Create two entities to relate entities = [ - {"name": "Source Entity", "entity_type": "test"}, - {"name": "Target Entity", "entity_type": "test"}, + {"name": "SourceEntity", "entity_type": "test"}, + {"name": "TargetEntity", "entity_type": "test"}, ] create_response = await client.post("/knowledge/entities", json={"entities": entities}) created = create_response.json()["entities"] - source_id = created[0]["id"] - target_id = created[1]["id"] + source_path_id = "test/source_entity" + target_path_id = "test/target_entity" # Create relation between them response = await client.post( "/knowledge/relations", json={ - "relations": [{"from_id": source_id, "to_id": target_id, "relation_type": "related_to"}] + "relations": [ + {"from_id": source_path_id, "to_id": target_path_id, "relation_type": "related_to"} + ] }, ) @@ -74,7 +77,6 @@ async def create_related_entities(client) -> List[RelationResponse]: # pyright: data = response.json() relation_response = CreateEntityResponse.model_validate(data) - assert len(relation_response.entities) == 2 source_entity = relation_response.entities[0] @@ -82,13 +84,13 @@ async def create_related_entities(client) -> List[RelationResponse]: # pyright: assert len(source_entity.relations) == 1 source_relation = source_entity.relations[0] - assert source_relation.from_id == source_id - assert source_relation.to_id == target_id + assert source_relation.from_id == source_path_id + assert source_relation.to_id == target_path_id assert len(target_entity.relations) == 1 target_relation = target_entity.relations[0] - assert target_relation.from_id == source_id - assert target_relation.to_id == target_id + assert target_relation.from_id == source_path_id + assert target_relation.to_id == target_path_id return source_entity.relations + target_entity.relations @@ -101,32 +103,23 @@ async def test_create_entities(client: AsyncClient): @pytest.mark.asyncio async def test_get_entity(client: AsyncClient): - """Should retrieve an entity by ID.""" + """Should retrieve an entity by path ID.""" # First create an entity - create_response = await client.post( - "/knowledge/entities", - json={ - "entities": [ - { - "name": "Test Entity", - "entity_type": "test", - } - ] - }, - ) - entity_id = create_response.json()["entities"][0]["id"] + data = {"name": "TestEntity", "entity_type": "test"} + response = await client.post("/knowledge/entities", json={"entities": [data]}) + assert response.status_code == 200 + data = response.json() - # Now get it by ID - response = await client.get(f"/knowledge/entities/{entity_id}") + # Now get it by path + path_id = data["entities"][0]["path_id"] + response = await client.get(f"/knowledge/entities/{path_id}") # Verify retrieval assert response.status_code == 200 entity = response.json() - assert entity["id"] == entity_id - assert entity["name"] == "Test Entity" - - entity_type = entity.get("entity_type") or entity.get("entityType") - assert entity_type == "test" + assert entity["name"] == "TestEntity" + assert entity["entity_type"] == "test" + assert entity["path_id"] == "test/test_entity" @pytest.mark.asyncio @@ -139,16 +132,15 @@ async def test_create_relations(client: AsyncClient): async def test_add_observations(client: AsyncClient): """Should add observations to an entity.""" # Create an entity first - create_response = await client.post( - "/knowledge/entities", json={"entities": [{"name": "Test Entity", "entity_type": "test"}]} - ) - entity_id = create_response.json()["entities"][0]["id"] + data = {"name": "TestEntity", "entity_type": "test"} + await client.post("/knowledge/entities", json={"entities": [data]}) + path_id = "test/test_entity" # Add observations - await add_observations(client, entity_id) + await add_observations(client, path_id) # Verify observations appear in entity - entity_response = await client.get(f"/knowledge/entities/{entity_id}") + entity_response = await client.get(f"/knowledge/entities/{path_id}") entity = entity_response.json() assert len(entity["observations"]) == 2 @@ -158,10 +150,10 @@ async def test_search_nodes(client: AsyncClient): """Should search for entities in the knowledge graph.""" # Create a few entities with different names entities = [ - {"name": "Not found", "entity_type": "negative"}, - {"name": "Alpha Test", "entity_type": "test"}, - {"name": "Beta Test", "entity_type": "test"}, - {"name": "Gamma Production", "entity_type": "test"}, # match entity_type + {"name": "NotFound", "entity_type": "negative"}, + {"name": "AlphaTest", "entity_type": "test"}, + {"name": "BetaTest", "entity_type": "test"}, + {"name": "GammaProduction", "entity_type": "test"}, # match entity_type ] await client.post("/knowledge/entities", json={"entities": entities}) @@ -173,117 +165,110 @@ async def test_search_nodes(client: AsyncClient): data = response.json() assert data["query"] == "Test" assert len(data["matches"]) == 3 - names = [entity["name"] for entity in data["matches"]] - assert "Alpha Test" in names - assert "Beta Test" in names - assert "Gamma Production" in names + names = {entity["name"] for entity in data["matches"]} + assert names == {"AlphaTest", "BetaTest", "GammaProduction"} @pytest.mark.asyncio async def test_open_nodes(client: AsyncClient): - """Should search for entities in the knowledge graph.""" + """Should open multiple nodes by path IDs.""" # Create a few entities with different names entities = [ - {"name": "Alpha Test", "entity_type": "test"}, - {"name": "Beta Test", "entity_type": "test"}, + {"name": "AlphaTest", "entity_type": "test"}, + {"name": "BetaTest", "entity_type": "test"}, ] - create_response = await client.post("/knowledge/entities", json={"entities": entities}) - created_entities = create_response.json() - assert len(created_entities["entities"]) == 2 + await client.post("/knowledge/entities", json={"entities": entities}) - # open nodes + # Open nodes by path IDs response = await client.post( "/knowledge/nodes", - json={ - "entity_ids": [ - created_entities["entities"][0]["id"], - ] - }, + json={"entity_ids": ["test/alpha_test"]}, ) - # Verify search results + # Verify results assert response.status_code == 200 data = response.json() assert len(data["entities"]) == 1 entity = data["entities"][0] - assert entity["name"] == "Alpha Test" + assert entity["name"] == "AlphaTest" assert entity["entity_type"] == "test" + assert entity["path_id"] == "test/alpha_test" @pytest.mark.asyncio async def test_delete_entity(client: AsyncClient): - """Test DELETE /knowledge/entities/{entity_id}.""" + """Test DELETE /knowledge/entities with path ID.""" # Create test entity - entity = await create_entity(client) + entity_data = {"name": "TestEntity", "entity_type": "test"} + await client.post("/knowledge/entities", json={"entities": [entity_data]}) # Test deletion - response = await client.post("/knowledge/entities/delete", json={"entity_ids": [entity.id]}) + response = await client.post( + "/knowledge/entities/delete", json={"entity_ids": ["test/TestEntity"]} + ) assert response.status_code == 200 assert response.json() == {"deleted": True} # Verify entity is gone - response = await client.get(f"/knowledge/entities/{entity.id}") + path_id = quote("test/TestEntity") + response = await client.get(f"/knowledge/entities/{path_id}") assert response.status_code == 404 @pytest.mark.asyncio async def test_delete_entity_bulk(client: AsyncClient): - """Test DELETE /knowledge/entities/{entity_id}.""" - # Create test entity - entity1 = await create_entity(client) - - e2_response = await client.post( - "/knowledge/entities", json={"entities": [{"name": "Test Entity2", "entity_type": "test"}]} - ) - create_response = CreateEntityResponse.model_validate(e2_response.json()) - entity2 = create_response.entities[0] + """Test bulk entity deletion using path IDs.""" + # Create test entities + entities = [ + {"name": "Entity1", "entity_type": "test"}, + {"name": "Entity2", "entity_type": "test"}, + ] + await client.post("/knowledge/entities", json={"entities": entities}) # Test deletion response = await client.post( - "/knowledge/entities/delete", json={"entity_ids": [entity1.id, entity2.id]} + "/knowledge/entities/delete", json={"entity_ids": ["test/Entity1", "test/Entity2"]} ) assert response.status_code == 200 assert response.json() == {"deleted": True} # Verify entities are gone - response = await client.get(f"/knowledge/entities/{entity1.id}") - assert response.status_code == 404 - - response = await client.get(f"/knowledge/entities/{entity2.id}") - assert response.status_code == 404 + for name in ["Entity1", "Entity2"]: + path_id = quote(f"test/{name}") + response = await client.get(f"/knowledge/entities/{path_id}") + assert response.status_code == 404 @pytest.mark.asyncio async def test_delete_entity_with_observations(client, observation_repository): """Test cascading delete with observations.""" - # Create test data - entity = await create_entity(client) - await add_observations(client, entity.id) + # Create test entity and add observations + entity_data = {"name": "TestEntity", "entity_type": "test"} + await client.post("/knowledge/entities", json={"entities": [entity_data]}) + await add_observations(client, "test/TestEntity") # Delete entity - response = await client.post("/knowledge/entities/delete", json={"entity_ids": [entity.id]}) + response = await client.post( + "/knowledge/entities/delete", json={"entity_ids": ["test/TestEntity"]} + ) assert response.status_code == 200 assert response.json() == {"deleted": True} # Verify observations are gone - remaining = await observation_repository.find_by_entity(entity.id) - assert len(remaining) == 0 + observations = await observation_repository.find_all() + assert len(observations) == 0 @pytest.mark.asyncio async def test_delete_observations(client, observation_repository): - """Test DELETE /knowledge/entities/{entity_id}/observations.""" - - # Create an entity first - create_response = await client.post( - "/knowledge/entities", json={"entities": [{"name": "Test Entity", "entity_type": "test"}]} - ) - entity_id = create_response.json()["entities"][0]["id"] - - observations = await add_observations(client, entity_id) # adds 2 + """Test deleting specific observations.""" + # Create entity and add observations + entity_data = {"name": "TestEntity", "entity_type": "test"} + await client.post("/knowledge/entities", json={"entities": [entity_data]}) + observations = await add_observations(client, "test/TestEntity") # adds 2 # Delete specific observations - request_data = {"entity_id": entity_id, "deletions": [observations[0].content]} + request_data = {"entity_id": "test/TestEntity", "deletions": [observations[0].content]} response = await client.post("/knowledge/observations/delete", json=request_data) assert response.status_code == 200 data = response.json() @@ -295,11 +280,10 @@ async def test_delete_observations(client, observation_repository): @pytest.mark.asyncio async def test_delete_relations(client, relation_repository): - """Test DELETE /knowledge/relations.""" - # Create test data - relatations = await create_related_entities(client) - assert len(relatations) == 2 - relation = relatations[0] + """Test deleting relations between entities.""" + relations = await create_related_entities(client) + assert len(relations) == 2 + relation = relations[0] # Delete relation request_data = { @@ -317,43 +301,43 @@ async def test_delete_relations(client, relation_repository): del_response = CreateEntityResponse.model_validate(data) assert len(del_response.entities) == 2 - assert len(del_response.entities[0].relations) == 0 - assert len(del_response.entities[1].relations) == 0 + assert all(len(e.relations) == 0 for e in del_response.entities) @pytest.mark.asyncio -async def test_delete_nonexistent_entity(client): - """Test deleting an entity that doesn't exist.""" - response = await client.post("/knowledge/entities/delete", json={"entity_ids": [0]}) - +async def test_delete_nonexistent_entity(client: AsyncClient): + """Test deleting a nonexistent entity by path ID.""" + response = await client.post( + "/knowledge/entities/delete", json={"entity_ids": ["test/NonExistent"]} + ) assert response.status_code == 200 assert response.json() == {"deleted": False} @pytest.mark.asyncio -async def test_delete_nonexistent_observations(client): - """Test deleting observations that don't exist.""" +async def test_delete_nonexistent_observations(client: AsyncClient): + """Test deleting nonexistent observations.""" # Create test entity - entity = await create_entity(client) + entity_data = {"name": "TestEntity", "entity_type": "test"} + await client.post("/knowledge/entities", json={"entities": [entity_data]}) - request_data = {"entity_id": entity.id, "deletions": ["Nonexistent observation"]} + request_data = {"entity_id": "test/TestEntity", "deletions": ["Nonexistent observation"]} response = await client.post("/knowledge/observations/delete", json=request_data) assert response.status_code == 200 data = response.json() - # no op - del_response = EntityResponse.model_validate(data) - assert del_response == entity + entity = EntityResponse.model_validate(data) + assert len(entity.observations) == 0 @pytest.mark.asyncio -async def test_delete_nonexistent_relations(client): - """Test deleting relations that don't exist.""" +async def test_delete_nonexistent_relations(client: AsyncClient): + """Test deleting nonexistent relations.""" request_data = { "relations": [ { - "from_id": 0, - "to_id": 0, + "from_id": "test/NonExistent1", + "to_id": "test/NonExistent2", "relation_type": "nonexistent", } ] @@ -366,47 +350,38 @@ async def test_delete_nonexistent_relations(client): assert del_response.entities == [] +@pytest.mark.asyncio +async def test_invalid_path_id_format(client: AsyncClient): + """Test handling of invalid path ID formats.""" + invalid_path_ids = [ + "no_type_separator", + "/missing_type/name", + "type//extra_separator", + "/", + "", + ] + for invalid_id in invalid_path_ids: + path_id = quote(invalid_id) + response = await client.get(f"/knowledge/entities/{path_id}") + assert response.status_code == 404 + + @pytest.mark.asyncio async def test_full_knowledge_flow(client: AsyncClient): - """Test a complete knowledge graph flow with multiple operations.""" - # 1. Create main entity - main_response = await client.post( - "/knowledge/entities", - json={ - "entities": [ - {"name": "Main Entity", "entity_type": "test"}, - {"name": "Non Entity", "entity_type": "n_a"}, - ] - }, - ) - assert main_response.status_code == 200 - - main_entity_id = None - non_entity_id = None - for entity in main_response.json()["entities"]: - if entity["name"] == "Main Entity": - main_entity_id = entity["id"] - - if entity["name"] == "Non Entity": - non_entity_id = entity["id"] - - assert main_entity_id is not None - assert non_entity_id is not None + """Test complete knowledge graph flow with path IDs.""" + # 1. Create main entities + main_entities = [ + {"name": "MainEntity", "entity_type": "test"}, + {"name": "NonEntity", "entity_type": "n_a"}, + ] + await client.post("/knowledge/entities", json={"entities": main_entities}) # 2. Create related entities - related_response = await client.post( - "/knowledge/entities", - json={ - "entities": [ - {"name": "Related One", "entity_type": "test"}, - {"name": "Related Two", "entity_type": "test"}, - ] - }, - ) - related = related_response.json()["entities"] - related_ids = [e["id"] for e in related] - assert related_response.status_code == 200 - assert len(related_ids) == 2 + related_entities = [ + {"name": "RelatedOne", "entity_type": "test"}, + {"name": "RelatedTwo", "entity_type": "test"}, + ] + await client.post("/knowledge/entities", json={"entities": related_entities}) # 3. Add relations relations_response = await client.post( @@ -414,13 +389,13 @@ async def test_full_knowledge_flow(client: AsyncClient): json={ "relations": [ { - "from_id": main_entity_id, - "to_id": related_ids[0], + "from_id": "test/MainEntity", + "to_id": "test/RelatedOne", "relation_type": "connects_to", }, { - "from_id": main_entity_id, - "to_id": related_ids[1], + "from_id": "test/MainEntity", + "to_id": "test/RelatedTwo", "relation_type": "connects_to", }, ] @@ -432,7 +407,7 @@ async def test_full_knowledge_flow(client: AsyncClient): await client.post( "/knowledge/observations", json={ - "entity_id": main_entity_id, + "entity_id": "test/MainEntity", "observations": [ "Connected to first related entity", "Connected to second related entity", @@ -441,23 +416,28 @@ async def test_full_knowledge_flow(client: AsyncClient): ) # 5. Verify full graph structure - main_get = await client.get(f"/knowledge/entities/{main_entity_id}") + path_id = quote("test/MainEntity") + main_get = await client.get(f"/knowledge/entities/{path_id}") main_entity = main_get.json() # Check entity structure - assert main_entity["name"] == "Main Entity" + assert main_entity["name"] == "MainEntity" assert len(main_entity["observations"]) == 2 assert len(main_entity["relations"]) == 2 # 6. Search should find all related entities search = await client.post("/knowledge/search", json={"query": "Related"}) matches = search.json()["matches"] - assert len(matches) == 3 # Should find both related entities + assert len(matches) == 2 # Should find both related entities - # 7. delete entities + # 7. Delete main entity response = await client.post( - "/knowledge/entities/delete", json={"entity_ids": [main_entity_id, non_entity_id]} + "/knowledge/entities/delete", json={"entity_ids": ["test/MainEntity", "test/NonEntity"]} ) - assert response.status_code == 200 assert response.json() == {"deleted": True} + + # Verify deletion + path_id = quote("test/MainEntity") + response = await client.get(f"/knowledge/entities/{path_id}") + assert response.status_code == 404 diff --git a/tests/repository/test_entity_repository.py b/tests/repository/test_entity_repository.py index 8f20ce3a..aab55e82 100644 --- a/tests/repository/test_entity_repository.py +++ b/tests/repository/test_entity_repository.py @@ -180,26 +180,6 @@ async def test_find_by_id(entity_repository: EntityRepository, sample_entity: En assert db_entity.description == found.description -@pytest.mark.asyncio -async def test_find_by_type_and_name(entity_repository: EntityRepository, sample_entity: Entity): - """Test finding an entity by name""" - found = await entity_repository.get_entity_by_type_and_name( - sample_entity.entity_type, sample_entity.name - ) - assert found is not None - assert found.id == sample_entity.id - assert found.name == sample_entity.name - - # Verify against direct database query - async with db.scoped_session(entity_repository.session_maker) as session: - stmt = select(Entity).where(Entity.name == sample_entity.name) - result = await session.execute(stmt) - db_entity = result.scalar_one() - assert db_entity.id == found.id - assert db_entity.name == found.name - assert db_entity.description == found.description - - @pytest.mark.asyncio async def test_update_entity(entity_repository: EntityRepository, sample_entity: Entity): """Test updating an entity""" @@ -360,3 +340,110 @@ async def test_search(session_maker, entity_repository: EntityRepository): results = await entity_repository.search("searchable") assert len(results) == 1 assert results[0].id == entity1.id + + +@pytest_asyncio.fixture +async def test_entities(session_maker): + """Create multiple test entities.""" + async with db.scoped_session(session_maker) as session: + entities = [ + Entity( + name="entity1", + entity_type="type1", + description="First test entity", + path_id="type1/entity1", + ), + Entity( + name="entity2", + entity_type="type1", + description="Second test entity", + path_id="type1/entity2", + ), + Entity( + name="entity3", + entity_type="type2", + description="Third test entity", + path_id="type2/entity3", + ), + ] + session.add_all(entities) + return entities + + +@pytest.mark.asyncio +async def test_find_by_path_ids(entity_repository: EntityRepository, test_entities): + """Test finding multiple entities by their type/name pairs.""" + # Test finding multiple entities + path_ids = [e.path_id for e in test_entities] + found = await entity_repository.find_by_path_ids(path_ids) + assert len(found) == 3 + names = {e.name for e in found} + assert names == {"entity1", "entity2", "entity3"} + + # Test finding subset of entities + path_ids = [e.path_id for e in test_entities if e.name != "entity2"] + found = await entity_repository.find_by_path_ids(path_ids) + assert len(found) == 2 + names = {e.name for e in found} + assert names == {"entity1", "entity3"} + + # Test with non-existent entities + path_ids = ["type1/entity1", "type3/nonexistent"] + found = await entity_repository.find_by_path_ids(path_ids) + assert len(found) == 1 + assert found[0].name == "entity1" + + # Test empty input + found = await entity_repository.find_by_path_ids([]) + assert len(found) == 0 + + +@pytest.mark.asyncio +async def test_delete_by_path_ids(entity_repository: EntityRepository, test_entities): + """Test deleting entities by type/name pairs.""" + # Test deleting multiple entities + path_ids = [e.path_id for e in test_entities if e.name != "entity3"] + deleted_count = await entity_repository.delete_by_path_ids(path_ids) + assert deleted_count == 2 + + # Verify deletions + remaining = await entity_repository.find_all() + assert len(remaining) == 1 + assert remaining[0].name == "entity3" + + # Test deleting non-existent entities + path__ids = ["type3/nonexistent"] + deleted_count = await entity_repository.delete_by_path_ids(path__ids) + assert deleted_count == 0 + + # Test empty input + deleted_count = await entity_repository.delete_by_path_ids([]) + assert deleted_count == 0 + + +@pytest.mark.asyncio +async def test_delete_by_path_ids_with_observations( + entity_repository: EntityRepository, test_entities, session_maker +): + """Test deleting entities with observations by type/name pairs.""" + # Add observations + async with db.scoped_session(session_maker) as session: + observations = [ + Observation(entity_id=test_entities[0].id, content="First observation"), + Observation(entity_id=test_entities[1].id, content="Second observation"), + ] + session.add_all(observations) + + # Delete entities + path_ids = [e.path_id for e in test_entities] + deleted_count = await entity_repository.delete_by_path_ids(path_ids) + assert deleted_count == 3 + + # Verify observations were cascaded + async with db.scoped_session(session_maker) as session: + query = select(Observation).filter( + Observation.entity_id.in_([e.id for e in test_entities[:2]]) + ) + result = await session.execute(query) + remaining_observations = result.scalars().all() + assert len(remaining_observations) == 0 diff --git a/tests/schemas/test_schemas.py b/tests/schemas/test_schemas.py index ea966b62..0ee1a0c3 100644 --- a/tests/schemas/test_schemas.py +++ b/tests/schemas/test_schemas.py @@ -9,7 +9,7 @@ from basic_memory.schemas import ( Relation, CreateEntityRequest, SearchNodesRequest, - OpenNodesRequest, + OpenNodesRequest, RelationResponse, ) @@ -53,10 +53,10 @@ def test_entity_in_validation(): def test_relation_in_validation(): """Test RelationIn validation.""" - data = {"from_id": 123, "to_id": 456, "relation_type": "test"} + data = {"from_id": "123", "to_id": "456", "relation_type": "test"} relation = Relation.model_validate(data) - assert relation.from_id == 123 - assert relation.to_id == 456 + assert relation.from_id == "123" + assert relation.to_id == "456" assert relation.relation_type == "test" assert relation.context is None @@ -69,6 +69,15 @@ def test_relation_in_validation(): with pytest.raises(ValidationError): Relation.model_validate({"from_id": "123", "to_id": "456"}) # Missing relationType +def test_relation_response(): + """Test RelationResponse validation.""" + data = {"from_id": 123, "to_id": 456, "relation_type": "test", "from_entity":{"path_id": "123"}, "to_entity":{"path_id": "456"}} + relation = RelationResponse.model_validate(data) + assert relation.from_id == "123" + assert relation.to_id == "456" + assert relation.relation_type == "test" + assert relation.context is None + def test_create_entities_input(): """Test CreateEntitiesInput validation.""" diff --git a/tests/services/test_entity_service.py b/tests/services/test_entity_service.py index 7658fb9e..97620495 100644 --- a/tests/services/test_entity_service.py +++ b/tests/services/test_entity_service.py @@ -28,7 +28,7 @@ async def entity_service(entity_repository: EntityRepository) -> EntityService: async def test_create_entity(entity_service: EntityService): """Test successful entity creation.""" entity_data = Entity( - name="Test Entity", + name="TestEntity", entity_type="test", description="A test entity description", observations=["this is a test observation"], @@ -39,17 +39,17 @@ async def test_create_entity(entity_service: EntityService): # Assert Entity assert isinstance(entity, EntityModel) - assert entity.name == "Test Entity" + assert entity.name == "TestEntity" assert entity.entity_type == "test" assert entity.description == "A test entity description" assert entity.created_at is not None assert entity.observations[0].content == "this is a test observation" assert len(entity.relations) == 0 - # Verify we can retrieve it - retrieved = await entity_service.get_entity(entity.id) + # Verify we can retrieve it using path_id + retrieved = await entity_service.get_by_path_id(entity_data.path_id) assert retrieved.description == "A test entity description" - assert retrieved.name == "Test Entity" + assert retrieved.name == "TestEntity" assert retrieved.entity_type == "test" assert retrieved.description == "A test entity description" assert retrieved.created_at is not None @@ -60,13 +60,13 @@ async def test_create_entities(entity_service: EntityService): """Test successful entity creation.""" entity_data = [ Entity( - name="Test Entity_1", + name="TestEntity1", entity_type="test", description="A test entity description", observations=["this is a test observation"], ), Entity( - name="Test Entity_2", + name="TestEntity2", entity_type="test", description="A test entity description", observations=["this is a test observation"], @@ -80,7 +80,7 @@ async def test_create_entities(entity_service: EntityService): assert len(entities) == 2 entity1 = entities[0] assert isinstance(entity1, EntityModel) - assert entity1.name == "Test Entity_1" + assert entity1.name == "TestEntity1" assert entity1.entity_type == "test" assert entity1.description == "A test entity description" assert entity1.created_at is not None @@ -89,25 +89,25 @@ async def test_create_entities(entity_service: EntityService): entity2 = entities[1] assert isinstance(entity1, EntityModel) - assert entity2.name == "Test Entity_2" + assert entity2.name == "TestEntity2" assert entity2.entity_type == "test" assert entity2.description == "A test entity description" assert entity2.created_at is not None assert entity2.observations[0].content == "this is a test observation" - # Verify we can retrieve them - retrieved1 = await entity_service.get_entity(entity1.id) + # Verify we can retrieve them using path_ids + retrieved1 = await entity_service.get_by_path_id(entity_data[0].path_id) assert retrieved1.description == "A test entity description" - retrieved2 = await entity_service.get_entity(entity2.id) + retrieved2 = await entity_service.get_by_path_id(entity_data[1].path_id) assert retrieved2.description == "A test entity description" -async def test_get_by_type_and_name(entity_service: EntityService): +async def test_get_by_path_id(entity_service: EntityService): """Test finding entity by type and name combination.""" # Create two entities with same name but different types entity1_data = Entity( - name="Test Entity", + name="TestEntity", entity_type="type1", description="First test entity", observations=[], @@ -115,7 +115,7 @@ async def test_get_by_type_and_name(entity_service: EntityService): entity1 = await entity_service.create_entity(entity1_data) entity2_data = Entity( - name="Test Entity", # Same name + name="TestEntity", # Same name entity_type="type2", # Different type description="Second test entity", observations=[], @@ -123,14 +123,14 @@ async def test_get_by_type_and_name(entity_service: EntityService): entity2 = await entity_service.create_entity(entity2_data) # Find by type1 and name - found = await entity_service.get_by_type_and_name("type1", "Test Entity") + found = await entity_service.get_by_path_id(entity1_data.path_id) assert found is not None assert found.id == entity1.id assert found.entity_type == "type1" assert found.description == "First test entity" # Find by type2 and name - found = await entity_service.get_by_type_and_name("type2", "Test Entity") + found = await entity_service.get_by_path_id(entity2_data.path_id) assert found is not None assert found.id == entity2.id assert found.entity_type == "type2" @@ -138,113 +138,113 @@ async def test_get_by_type_and_name(entity_service: EntityService): # Test not found case with pytest.raises(EntityNotFoundError): - await entity_service.get_by_type_and_name("nonexistent", "Test Entity") + await entity_service.get_by_path_id("nonexistent/test_entity") async def test_create_entity_no_description(entity_service: EntityService): """Test creating entity without description (should be None).""" - entity_data = Entity(name="Test Entity", entity_type="test", observations=[]) + entity_data = Entity(name="TestEntity", entity_type="test", observations=[]) entity = await entity_service.create_entity(entity_data) assert entity.description is None # Verify after retrieval - retrieved = await entity_service.get_entity(entity.id) + retrieved = await entity_service.get_by_path_id(entity_data.path_id) assert retrieved.description is None async def test_get_entity_success(entity_service: EntityService): """Test successful entity retrieval.""" entity_data = Entity( - name="Test Entity", + name="TestEntity", entity_type="test", description="Test description", observations=[], ) - created = await entity_service.create_entity(entity_data) + await entity_service.create_entity(entity_data) - # Act - retrieved = await entity_service.get_entity(created.id) + # Get by path ID + retrieved = await entity_service.get_by_path_id(entity_data.path_id) - # Assert assert isinstance(retrieved, EntityModel) - assert retrieved.id == created.id - assert retrieved.name == created.name - assert retrieved.entity_type == created.entity_type + assert retrieved.name == "TestEntity" + assert retrieved.entity_type == "test" assert retrieved.description == "Test description" async def test_update_entity_description(entity_service: EntityService): """Test updating an entity's description.""" entity_data = Entity( - name="Test Entity", + name="TestEntity", entity_type="test", description="Initial description", observations=[], ) - entity = await entity_service.create_entity(entity_data) + await entity_service.create_entity(entity_data) - # Update description - updated = await entity_service.update_entity(entity.id, {"description": "Updated description"}) + # Update description using path_id + updated = await entity_service.update_entity( + entity_data.path_id, {"description": "Updated description"} + ) assert updated.description == "Updated description" # Verify after retrieval - retrieved = await entity_service.get_entity(entity.id) + retrieved = await entity_service.get_by_path_id(entity_data.path_id) assert retrieved.description == "Updated description" async def test_update_entity_description_to_none(entity_service: EntityService): """Test updating an entity's description to None.""" entity_data = Entity( - name="Test Entity", + name="TestEntity", entity_type="test", description="Initial description", observations=[], ) - entity = await entity_service.create_entity(entity_data) + await entity_service.create_entity(entity_data) - # Update description to None - updated = await entity_service.update_entity(entity.id, {"description": None}) + # Update description to None using path_id + updated = await entity_service.update_entity(entity_data.path_id, {"description": None}) assert updated.description is None # Verify after retrieval - retrieved = await entity_service.get_entity(entity.id) + retrieved = await entity_service.get_by_path_id(entity_data.path_id) assert retrieved.description is None async def test_delete_entity_success(entity_service: EntityService): """Test successful entity deletion.""" entity_data = Entity( - name="Test Entity", + name="TestEntity", entity_type="test", observations=[], ) - entity = await entity_service.create_entity(entity_data) + await entity_service.create_entity(entity_data) - # Act - result = await entity_service.delete_entity(entity.id) + # Act using path_id + result = await entity_service.delete_entity(entity_data.path_id) # Assert assert result is True with pytest.raises(EntityNotFoundError): - await entity_service.get_entity(entity.id) + await entity_service.get_by_path_id(entity_data.path_id) -async def test_get_entity_not_found(entity_service: EntityService): +async def test_get_entity_by_path_id_not_found(entity_service: EntityService): """Test handling of non-existent entity retrieval.""" with pytest.raises(EntityNotFoundError): - await entity_service.get_entity(0) + await entity_service.get_by_path_id("test/non_existent") async def test_delete_nonexistent_entity(entity_service: EntityService): """Test deleting an entity that doesn't exist.""" - result = await entity_service.delete_entity(0) - assert result is False + with pytest.raises(EntityNotFoundError): + await entity_service.delete_entity("test/non_existent") async def test_create_entity_with_special_chars(entity_service: EntityService): """Test entity creation with special characters in name and description.""" - name = "Test & Entity! With @ Special #Chars" + name = "TestEntity_Special" # Note: Using valid path characters description = "Description with $pecial chars & symbols!" entity_data = Entity( name=name, @@ -256,8 +256,8 @@ async def test_create_entity_with_special_chars(entity_service: EntityService): assert entity.name == name assert entity.description == description - # Verify after retrieval - retrieved = await entity_service.get_entity(entity.id) + # Verify after retrieval using path_id + retrieved = await entity_service.get_by_path_id(entity_data.path_id) assert retrieved.description == description @@ -265,7 +265,7 @@ async def test_create_entity_long_description(entity_service: EntityService): """Test creating entity with a long description.""" long_description = "A" * 1000 # 1000 character description entity_data = Entity( - name="Test Entity", + name="TestEntity", entity_type="test", description=long_description, observations=[], @@ -274,6 +274,100 @@ async def test_create_entity_long_description(entity_service: EntityService): entity = await entity_service.create_entity(entity_data) assert entity.description == long_description - # Verify after retrieval - retrieved = await entity_service.get_entity(entity.id) + # Verify after retrieval using path_id + retrieved = await entity_service.get_by_path_id(entity_data.path_id) assert retrieved.description == long_description + + +async def test_open_nodes_by_path_ids(entity_service: EntityService): + """Test opening multiple nodes by path IDs.""" + # Create test entities + entity1_data = Entity( + name="Entity1", + entity_type="type1", + description="First entity", + observations=[], + ) + entity2_data = Entity( + name="Entity2", + entity_type="type2", + description="Second entity", + observations=[], + ) + await entity_service.create_entity(entity1_data) + await entity_service.create_entity(entity2_data) + + # Open nodes by path IDs + path_ids = [entity1_data.path_id, entity2_data.path_id] + found = await entity_service.open_nodes(path_ids) + + assert len(found) == 2 + names = {e.name for e in found} + assert names == {"Entity1", "Entity2"} + + +async def test_open_nodes_empty_input(entity_service: EntityService): + """Test opening nodes with empty path ID list.""" + found = await entity_service.open_nodes([]) + assert len(found) == 0 + + +async def test_open_nodes_some_not_found(entity_service: EntityService): + """Test opening nodes with mix of existing and non-existent path IDs.""" + # Create one test entity + entity_data = Entity( + name="Entity1", + entity_type="type1", + description="Test entity", + observations=[], + ) + await entity_service.create_entity(entity_data) + + # Try to open two nodes, one exists, one doesn't + path_ids = [entity_data.path_id, "type1/non_existent"] + found = await entity_service.open_nodes(path_ids) + + assert len(found) == 1 + assert found[0].name == "Entity1" + + +async def test_delete_entities_by_path_ids(entity_service: EntityService): + """Test deleting multiple entities by path IDs.""" + # Create test entities + entity1_data = Entity( + name="Entity1", + entity_type="type1", + description="First entity", + observations=[], + ) + entity2_data = Entity( + name="Entity2", + entity_type="type2", + description="Second entity", + observations=[], + ) + await entity_service.create_entity(entity1_data) + await entity_service.create_entity(entity2_data) + + # Delete by path IDs + path_ids = [entity1_data.path_id, entity2_data.path_id] + result = await entity_service.delete_entities(path_ids) + assert result is True + + # Verify both are deleted + for path_id in path_ids: + with pytest.raises(EntityNotFoundError): + await entity_service.get_by_path_id(path_id) + + +async def test_delete_entities_empty_input(entity_service: EntityService): + """Test deleting entities with empty path ID list.""" + result = await entity_service.delete_entities([]) + assert result is False + + +async def test_delete_entities_none_found(entity_service: EntityService): + """Test deleting non-existent entities by path IDs.""" + path_ids = ["type1/NonExistent1", "type2/NonExistent2"] + result = await entity_service.delete_entities(path_ids) + assert result is False