admin/basicmachines-co-basic-memory
1
0
Fork 0
mirror of https://github.com/basicmachines-co/basic-memory synced 2026-06-21 13:47:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

change /knowledge endpoints to use ids

Browse Source
This commit is contained in:
phernandez
2024-12-24 18:15:59 -06:00
parent bde381c9e2
commit 0bbcf636f6
21 changed files with 620 additions and 421 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_memory/api/routers/knowledge.py
+9 -8
View File
@@ -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]
)
)
src/basic_memory/models/knowledge.py
+1 -1
View File
@@ -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
src/basic_memory/repository/entity_repository.py
+30 -6
View File
@@ -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])
src/basic_memory/repository/relation_repository.py
+1 -1
View File
@@ -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
src/basic_memory/schemas/__init__.py
-1
View File
@@ -34,7 +34,6 @@ from basic_memory.schemas.request import (
from basic_memory.schemas.response import (
SQLAlchemyModel,
ObservationResponse,
ObservationsResponse,
RelationResponse,
EntityResponse,
CreateEntityResponse,
src/basic_memory/schemas/base.py
+29 -18
View File
@@ -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}"
src/basic_memory/schemas/delete.py
+3 -3
View File
@@ -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)]
src/basic_memory/schemas/request.py
+2 -2
View File
@@ -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):
src/basic_memory/schemas/response.py
+24 -41
View File
@@ -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",
src/basic_memory/services/entity_service.py
+21 -25
View File
@@ -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
src/basic_memory/services/file_service.py
+1 -1
View File
@@ -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,
src/basic_memory/services/knowledge/entities.py
+19 -14
View File
@@ -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
src/basic_memory/services/knowledge/files.py
+5 -4
View File
@@ -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}")
src/basic_memory/services/knowledge/observations.py
+14 -14
View File
@@ -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}")
src/basic_memory/services/knowledge/relations.py
+38 -27
View File
@@ -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
raise
src/basic_memory/services/observation_service.py
+1 -1
View File
@@ -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(
src/basic_memory/services/relation_service.py
+9 -11
View File
@@ -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)
tests/api/test_knowledge_router.py
+146 -166
View File
@@ -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
tests/repository/test_entity_repository.py
+107 -20
View File
@@ -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
tests/schemas/test_schemas.py
+13 -4
View File
@@ -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."""
tests/services/test_entity_service.py
+147 -53
View File
@@ -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