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

fix entity service tests

Browse Source
This commit is contained in:
phernandez
2024-12-22 12:00:02 -06:00
parent ecb46ddceb
commit c2d16acb59
6 changed files with 119 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_memory/models/documents.py
+1 -1
View File
@@ -18,7 +18,7 @@ class Document(Base):
is the real source of truth.
"""
__tablename__ = "documents"
__tablename__ = "document"
id: Mapped[int] = mapped_column(primary_key=True)
path: Mapped[str] = mapped_column(String, unique=True, nullable=False)
src/basic_memory/models/knowledge.py
+5 -3
View File
@@ -7,6 +7,7 @@ from sqlalchemy import Integer, String, Text, ForeignKey, UniqueConstraint, text
from sqlalchemy.orm import Mapped, mapped_column, relationship
from basic_memory.models.base import Base
from basic_memory.models.documents import Document
class Entity(Base):
@@ -45,7 +46,7 @@ class Entity(Base):
# Relations
doc_id: Mapped[Optional[int]] = mapped_column(
Integer, ForeignKey("documents.id", ondelete="SET NULL"), nullable=True
Integer, ForeignKey("document.id", ondelete="SET NULL"), nullable=True
)
# Relationships
@@ -64,6 +65,7 @@ class Entity(Base):
foreign_keys="[Relation.to_id]",
cascade="all, delete-orphan",
)
document: Mapped[Optional[Document]] = relationship(Document, back_populates="entities")
@property
def relations(self):
@@ -80,7 +82,7 @@ class Observation(Base):
Observations are atomic facts or notes about an entity.
"""
__tablename__ = "observations"
__tablename__ = "observation"
id: Mapped[int] = mapped_column(Integer, primary_key=True)
entity_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id"))
@@ -99,7 +101,7 @@ class Relation(Base):
A directed relation between two entities.
"""
__tablename__ = "relations"
__tablename__ = "relation"
__table_args__ = (
UniqueConstraint("from_id", "to_id", "relation_type", name="uix_relation"),
Index("ix_relation_type", "relation_type"),
src/basic_memory/repository/entity_repository.py
+90 -45
View File
@@ -1,14 +1,15 @@
"""Repository for managing entities in the knowledge graph."""
from typing import List, Optional, Dict, Any, Sequence
from datetime import datetime
from loguru import logger
from sqlalchemy import select, delete
from sqlalchemy import select, or_
from sqlalchemy.exc import NoResultFound
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from sqlalchemy.sql import text
from sqlalchemy.orm import selectinload
from basic_memory import db
from basic_memory.models.knowledge import Entity
from basic_memory.models.knowledge import Entity, Observation
from basic_memory.repository.repository import Repository
@@ -19,49 +20,85 @@ class EntityRepository(Repository[Entity]):
"""Initialize with session maker."""
super().__init__(session_maker, Entity)
async def create_entity(
self,
name: str,
entity_type: str,
description: Optional[str] = None,
path: Optional[str] = None,
checksum: Optional[str] = None,
doc_id: Optional[int] = None,
) -> Entity:
"""Create a new entity."""
data = {
"name": name,
"entity_type": entity_type,
"description": description,
"path": path,
"checksum": checksum,
"doc_id": doc_id,
}
return await self.create(data)
async def create(self, data: dict) -> Entity: # pyright: ignore [reportIncompatibleMethodOverride]
"""Create a new entity in the database from the provided data."""
created = await super().create(data)
async def get_entity_by_type_and_name(
self, entity_type: str, name: str
) -> Optional[Entity]:
# we have to find to get relations
found = await self.find_by_id(created.id)
assert found is not None, f"Created entity {created} should not be None"
return found
async def create_all(self, data_list: List[dict]) -> Sequence[Entity]: # pyright: ignore [reportIncompatibleMethodOverride]
"""Create a new entity in the database from the provided data."""
created = await super().create_all(data_list)
# we have to find to get relations
return await self.find_by_ids([e.id for e in created])
async def find_by_id(self, entity_id: int) -> Optional[Entity]:
"""Find entity by ID with all relationships eagerly loaded."""
logger.debug(f"Finding entity by ID: {entity_id}")
async with db.scoped_session(self.session_maker) as session:
try:
result = await session.execute(
select(Entity)
.filter(Entity.id == entity_id)
.options(
selectinload(Entity.observations),
selectinload(Entity.from_relations),
selectinload(Entity.to_relations),
)
)
entity = result.scalars().one()
logger.debug(f"Found entity: {entity.id}")
return entity
except NoResultFound:
logger.debug(f"No entity found with ID: {entity_id}")
return None
async def find_by_ids(self, ids: List[int]) -> Sequence[Entity]:
"""Search for entities of a specific type."""
logger.debug(f"Find entities by ids: {ids}")
async with db.scoped_session(self.session_maker) as session:
result = await session.execute(
select(Entity)
.where(self.primary_key.in_(ids))
.options(
selectinload(Entity.observations),
selectinload(Entity.from_relations),
selectinload(Entity.to_relations),
)
)
entities = result.scalars().all()
logger.debug(f"Found {len(entities)}")
return entities
async def get_entity_by_type_and_name(self, entity_type: str, name: str) -> Optional[Entity]:
"""Get entity by type and name."""
query = self.select().where(
Entity.entity_type == entity_type,
Entity.name == name
query = (
self.select()
.options(
selectinload(Entity.observations),
selectinload(Entity.from_relations),
selectinload(Entity.to_relations),
)
.where(Entity.entity_type == entity_type, Entity.name == name)
)
return await self.find_one(query)
async def list_entities(
self,
self,
entity_type: Optional[str] = None,
doc_id: Optional[int] = None,
) -> Sequence[Entity]:
"""List all entities, optionally filtered by type."""
query = self.select()
if entity_type:
query = query.where(Entity.entity_type == entity_type)
if doc_id:
query = query.where(Entity.doc_id == doc_id)
async with db.scoped_session(self.session_maker) as session:
result = await session.execute(query)
return list(result.scalars().all())
@@ -73,32 +110,40 @@ class EntityRepository(Repository[Entity]):
result = await session.execute(query)
return [r[0] for r in result.all()]
async def search_entities(self, query_str: str) -> List[Entity]:
async def search(self, query_str: str) -> List[Entity]:
"""
Search for entities.
Searches across:
- Entity names
- Entity types
- Entity types
- Entity descriptions
- Associated Observations content
"""
search_term = f"%{query_str}%"
query = self.select().where(
(Entity.name.ilike(search_term)) |
(Entity.entity_type.ilike(search_term)) |
(Entity.description.ilike(search_term))
query = (
self.select()
.where(
or_(
Entity.name.ilike(search_term),
Entity.entity_type.ilike(search_term),
Entity.description.ilike(search_term),
Entity.observations.any(Observation.content.ilike(search_term)),
)
)
.options(
selectinload(Entity.observations),
selectinload(Entity.from_relations),
selectinload(Entity.to_relations),
)
)
result = await self.execute_query(query)
return list(result.scalars().all())
async def update_entity(
self,
entity_id: int,
updates: Dict[str, Any]
) -> Optional[Entity]:
async def update_entity(self, entity_id: int, updates: Dict[str, Any]) -> Optional[Entity]:
"""Update an entity with the given fields."""
return await self.update(entity_id, updates)
async def delete_entities_by_doc_id(self, doc_id: int) -> bool:
"""Delete all entities associated with a document."""
return await self.delete_by_fields(doc_id=doc_id)
return await self.delete_by_fields(doc_id=doc_id)
src/basic_memory/services/entity_service.py
+6 -7
View File
@@ -13,7 +13,6 @@ from .service import BaseService
def entity_model(entity):
model = EntityModel(
id=EntityModel.generate_id(entity.entity_type, entity.name),
name=entity.name,
entity_type=entity.entity_type,
description=entity.description,
@@ -45,7 +44,7 @@ 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: str, update_data: Dict[str, Any]) -> EntityModel:
async def update_entity(self, entity_id: int, 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)
@@ -53,7 +52,7 @@ class EntityService(BaseService[EntityRepository]):
raise EntityNotFoundError(f"Entity not found: {entity_id}")
return updated
async def get_entity(self, entity_id: str) -> EntityModel:
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)
@@ -64,7 +63,7 @@ class EntityService(BaseService[EntityRepository]):
async def get_by_type_and_name(self, entity_type: str, name: 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.find_by_type_and_name(entity_type, name)
db_entity = await self.repository.get_entity_by_type_and_name(entity_type, name)
if not db_entity:
raise EntityNotFoundError(f"Entity not found: {entity_type}/{name}")
return db_entity
@@ -73,17 +72,17 @@ class EntityService(BaseService[EntityRepository]):
"""Get all entities."""
return await self.repository.find_all()
async def delete_entity(self, entity_id: str) -> bool:
async def delete_entity(self, entity_id: int) -> bool:
"""Delete entity from database."""
logger.debug(f"Deleting entity: {entity_id}")
return await self.repository.delete(entity_id)
async def open_nodes(self, entity_ids: List[str]) -> Sequence[EntityModel]:
async def open_nodes(self, entity_ids: List[int]) -> 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)
async def delete_entities(self, entity_ids: List[str]) -> bool:
async def delete_entities(self, entity_ids: List[int]) -> bool:
"""Delete entities and their files."""
logger.debug(f"Deleting entities: {entity_ids}")
deleted_count = await self.repository.delete_by_ids(entity_ids)
tests/repository/test_entity_repository.py
+6 -4
View File
@@ -47,7 +47,9 @@ async def related_entities(session_maker):
entity_type="target",
description="Target entity",
)
session.add_all([source, target])
session.add(source)
session.add(target)
await session.flush()
relation = Relation(from_id=source.id, to_id=target.id, relation_type="connects_to")
session.add(relation)
@@ -343,18 +345,18 @@ async def test_search(session_maker, entity_repository: EntityRepository):
)
# Test search by name
results = await entity_repository.search_entities("Search Test")
results = await entity_repository.search("Search Test")
assert len(results) == 2
names = {e.name for e in results}
assert "Search Test 1" in names
assert "Search Test 2" in names
# Test search by type
results = await entity_repository.search_entities("other")
results = await entity_repository.search("other")
assert len(results) == 1
assert results[0].entity_type == "other"
# Test search by observation content
results = await entity_repository.search_entities("searchable")
results = await entity_repository.search("searchable")
assert len(results) == 1
assert results[0].id == entity1.id
tests/services/test_entity_service.py
+11 -20
View File
@@ -143,7 +143,7 @@ async def test_get_by_type_and_name(entity_service: EntityService):
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=[], relations=[])
entity_data = Entity(name="Test Entity", entity_type="test", observations=[])
entity = await entity_service.create_entity(entity_data)
assert entity.description is None
@@ -214,7 +214,11 @@ async def test_update_entity_description_to_none(entity_service: EntityService):
async def test_delete_entity_success(entity_service: EntityService):
"""Test successful entity deletion."""
entity_data = Entity(name="Test Entity", entity_type="test", observations=[], relations=[])
entity_data = Entity(
name="Test Entity",
entity_type="test",
observations=[],
)
entity = await entity_service.create_entity(entity_data)
# Act
@@ -229,12 +233,12 @@ async def test_delete_entity_success(entity_service: EntityService):
async def test_get_entity_not_found(entity_service: EntityService):
"""Test handling of non-existent entity retrieval."""
with pytest.raises(EntityNotFoundError):
await entity_service.get_entity("nonexistent-id")
await entity_service.get_entity(0)
async def test_delete_nonexistent_entity(entity_service: EntityService):
"""Test deleting an entity that doesn't exist."""
result = await entity_service.delete_entity("nonexistent-id")
result = await entity_service.delete_entity(0)
assert result is False
@@ -243,7 +247,9 @@ async def test_create_entity_with_special_chars(entity_service: EntityService):
name = "Test & Entity! With @ Special #Chars"
description = "Description with $pecial chars & symbols!"
entity_data = Entity(
name=name, entity_type="test", description=description, observations=[], relations=[]
name=name,
entity_type="test",
description=description,
)
entity = await entity_service.create_entity(entity_data)
@@ -255,21 +261,6 @@ async def test_create_entity_with_special_chars(entity_service: EntityService):
assert retrieved.description == description
async def test_entity_id_generation(entity_service: EntityService):
"""Test that entities get unique IDs generated correctly."""
entity_data = Entity(
name="Test Entity",
entity_type="test",
description="Test description",
observations=[],
)
entity = await entity_service.create_entity(entity_data)
assert entity.id # ID should be generated
assert "test/test_entity" == entity.id # Should contain normalized name
async def test_create_entity_long_description(entity_service: EntityService):
"""Test creating entity with a long description."""
long_description = "A" * 1000 # 1000 character description