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)