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

entity repository tests

Browse Source
This commit is contained in:
phernandez
2024-12-22 11:10:35 -06:00
parent 239753cf77
commit ecb46ddceb
5 changed files with 53 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_memory/models/knowledge.py
+26 -27
View File
@@ -1,21 +1,25 @@
"""Knowledge graph models."""
from datetime import datetime
from typing import Optional, List
from typing import Optional
from sqlalchemy import Integer, String, Text, ForeignKey, UniqueConstraint, text, DateTime, Index
from sqlalchemy.orm import Mapped, mapped_column, relationship
from basic_memory.models.base import Base
class Entity(Base):
"""
Core entity in the knowledge graph.
Entities represent semantic nodes maintained by the AI layer. Each entity:
- Has a unique numeric ID (database-generated)
- Maps to a document file on disk (optional)
- Maps to a document file on disk (optional)
- Maintains a checksum for change detection
- Tracks both source document and semantic properties
"""
__tablename__ = "entity"
__table_args__ = (
UniqueConstraint("entity_type", "name", name="uix_entity_type_name"),
@@ -27,32 +31,27 @@ 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)
# 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
created_at: Mapped[datetime] = mapped_column(
DateTime,
server_default=text("CURRENT_TIMESTAMP")
)
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"))
updated_at: Mapped[datetime] = mapped_column(
DateTime,
server_default=text("CURRENT_TIMESTAMP"),
onupdate=text("CURRENT_TIMESTAMP")
DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP")
)
# Relations
doc_id: Mapped[Optional[int]] = mapped_column(
Integer,
ForeignKey("documents.id", ondelete="SET NULL"),
nullable=True
Integer, ForeignKey("documents.id", ondelete="SET NULL"), nullable=True
)
# Relationships
observations = relationship("Observation", back_populates="entity", cascade="all, delete-orphan")
observations = relationship(
"Observation", back_populates="entity", cascade="all, delete-orphan"
)
from_relations = relationship(
"Relation",
back_populates="from_entity",
@@ -66,6 +65,10 @@ class Entity(Base):
cascade="all, delete-orphan",
)
@property
def relations(self):
return self.to_relations + self.from_relations
def __repr__(self) -> str:
return f"Entity(id={self.id}, name='{self.name}', type='{self.entity_type}')"
@@ -73,18 +76,16 @@ class Entity(Base):
class Observation(Base):
"""
An observation about an entity.
Observations are atomic facts or notes about an entity.
"""
__tablename__ = "observations"
id: Mapped[int] = mapped_column(Integer, primary_key=True)
entity_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id"))
content: Mapped[str] = mapped_column(Text)
created_at: Mapped[datetime] = mapped_column(
DateTime,
server_default=text("CURRENT_TIMESTAMP")
)
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"))
# Relationships
entity = relationship("Entity", back_populates="observations")
@@ -97,6 +98,7 @@ class Relation(Base):
"""
A directed relation between two entities.
"""
__tablename__ = "relations"
__table_args__ = (
UniqueConstraint("from_id", "to_id", "relation_type", name="uix_relation"),
@@ -107,14 +109,11 @@ class Relation(Base):
from_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id"))
to_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id"))
relation_type: Mapped[str] = mapped_column(String)
created_at: Mapped[datetime] = mapped_column(
DateTime,
server_default=text("CURRENT_TIMESTAMP")
)
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"))
# Relationships
from_entity = relationship("Entity", foreign_keys=[from_id], back_populates="from_relations")
to_entity = relationship("Entity", foreign_keys=[to_id], back_populates="to_relations")
def __repr__(self) -> str:
return f"Relation(from_id={self.from_id}, to_id={self.to_id}, type='{self.relation_type}')"
return f"Relation(from_id={self.from_id}, to_id={self.to_id}, type='{self.relation_type}')"
src/basic_memory/repository/entity_repository.py
+6 -2
View File
@@ -15,6 +15,10 @@ from basic_memory.repository.repository import Repository
class EntityRepository(Repository[Entity]):
"""Repository for Entity model."""
def __init__(self, session_maker: async_sessionmaker[AsyncSession]):
"""Initialize with session maker."""
super().__init__(session_maker, Entity)
async def create_entity(
self,
name: str,
@@ -93,8 +97,8 @@ class EntityRepository(Repository[Entity]):
updates: Dict[str, Any]
) -> Optional[Entity]:
"""Update an entity with the given fields."""
return await self.update(str(entity_id), updates)
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/repository/repository.py
+7 -7
View File
@@ -91,7 +91,7 @@ class Repository[T: Base]:
logger.debug(f"Found {len(items)} {self.Model.__name__} records")
return items
async def find_by_id(self, entity_id: str) -> Optional[T]:
async def find_by_id(self, entity_id: int) -> Optional[T]:
"""Fetch an entity by its unique identifier."""
logger.debug(f"Finding {self.Model.__name__} by ID: {entity_id}")
async with db.scoped_session(self.session_maker) as session:
@@ -117,7 +117,7 @@ class Repository[T: Base]:
logger.debug(f"No {self.Model.__name__} found")
return entity
async def find_by_ids(self, ids: List[str]) -> Sequence[T]:
async def find_by_ids(self, ids: List[int]) -> Sequence[T]:
"""Fetch multiple entities by their identifiers in a single query."""
logger.debug(f"Finding {self.Model.__name__} by IDs: {ids}")
async with db.scoped_session(self.session_maker) as session:
@@ -146,7 +146,7 @@ class Repository[T: Base]:
session.add_all(model_list)
return model_list
async def update(self, entity_id: str, entity_data: dict) -> Optional[T]:
async def update(self, entity_id: int, entity_data: dict) -> Optional[T]:
"""Update an entity with the given data."""
logger.debug(f"Updating {self.Model.__name__} {entity_id} with data: {entity_data}")
async with db.scoped_session(self.session_maker) as session:
@@ -169,7 +169,7 @@ class Repository[T: Base]:
logger.debug(f"No {self.Model.__name__} found to update: {entity_id}")
return None
async def delete(self, entity_id: str) -> bool:
async def delete(self, entity_id: int) -> bool:
"""Delete an entity from the database."""
logger.debug(f"Deleting {self.Model.__name__}: {entity_id}")
async with db.scoped_session(self.session_maker) as session:
@@ -186,8 +186,8 @@ class Repository[T: Base]:
logger.debug(f"No {self.Model.__name__} found to delete: {entity_id}")
return False
async def delete_by_ids(self, ids: List[str]) -> int:
"""Delete records matching given field values."""
async def delete_by_ids(self, ids: List[int]) -> int:
"""Delete records matching given IDs."""
logger.debug(f"Deleting {self.Model.__name__} by ids: {ids}")
async with db.scoped_session(self.session_maker) as session:
query = delete(self.Model).where(self.primary_key.in_(ids))
@@ -223,4 +223,4 @@ class Repository[T: Base]:
async with db.scoped_session(self.session_maker) as session:
result = await session.execute(query)
logger.debug("Query executed successfully")
return result
return result
tests/conftest.py
+3 -11
View File
@@ -15,12 +15,12 @@ from sqlalchemy.ext.asyncio import (
from basic_memory import db
from basic_memory.config import ProjectConfig
from basic_memory.db import DatabaseType
from basic_memory.models import Base, Entity as EntityModel
from basic_memory.models import Base
from basic_memory.models.knowledge import Entity
from basic_memory.repository.document_repository import DocumentRepository
from basic_memory.repository.entity_repository import EntityRepository
from basic_memory.repository.observation_repository import ObservationRepository
from basic_memory.repository.relation_repository import RelationRepository
from basic_memory.schemas import Entity
from basic_memory.services import (
EntityService,
ObservationService,
@@ -136,19 +136,11 @@ async def observation_service(observation_repository: ObservationRepository) ->
@pytest_asyncio.fixture(scope="function")
async def sample_entity(entity_repository: EntityRepository) -> EntityModel:
async def sample_entity(entity_repository: EntityRepository) -> Entity:
"""Create a sample entity for testing."""
entity_data = {
"id": "test/test_entity",
"name": "Test Entity",
"entity_type": "test",
"description": "A test entity",
}
return await entity_repository.create(entity_data)
@pytest_asyncio.fixture
async def test_entity(entity_service: EntityService) -> EntityModel:
"""Create a test entity for reuse in tests."""
entity_data = Entity(name="Test Entity", entity_type="test", observations=[])
return await entity_service.create_entity(entity_data)
tests/repository/test_entity_repository.py
+11 -17
View File
@@ -16,9 +16,7 @@ from basic_memory.repository.entity_repository import EntityRepository
async def test_entity(session_maker):
"""Create a test entity."""
async with db.scoped_session(session_maker) as session:
entity = Entity(
id="test/test_entity", name="test_entity", entity_type="test", description="Test entity"
)
entity = Entity(name="test_entity", entity_type="test", description="Test entity")
session.add(entity)
return entity
@@ -40,13 +38,11 @@ async def related_entities(session_maker):
"""Create entities with relations between them."""
async with db.scoped_session(session_maker) as session:
source = Entity(
id="source/test_entity",
name="source",
entity_type="source",
description="Source entity",
)
target = Entity(
id="target/test_entity",
name="target",
entity_type="target",
description="Target entity",
@@ -70,10 +66,11 @@ async def test_create_entity(entity_repository: EntityRepository):
entity = await entity_repository.create(entity_data)
# Verify returned object
assert entity.id == "test/test"
assert entity.id is not None
assert entity.name == "Test"
assert entity.description == "Test description"
assert isinstance(entity.created_at, datetime)
assert isinstance(entity.updated_at, datetime)
# Verify in database
found = await entity_repository.find_by_id(entity.id)
@@ -126,7 +123,6 @@ async def test_entity_type_name_unique_constraint(entity_repository: EntityRepos
"""Test the unique constraint on entity_type + name combination."""
# Create first entity
entity1_data = {
"id": "20240102-test1",
"name": "Test Entity",
"entity_type": "type1",
"description": "First entity",
@@ -135,7 +131,6 @@ async def test_entity_type_name_unique_constraint(entity_repository: EntityRepos
# Try to create another entity with same type and name
entity2_data = {
"id": "20240102-test2",
"name": "Test Entity", # Same name
"entity_type": "type1", # Same type
"description": "Second entity",
@@ -151,7 +146,6 @@ async def test_entity_type_name_unique_constraint(entity_repository: EntityRepos
async def test_create_entity_null_description(session_maker, entity_repository: EntityRepository):
"""Test creating an entity with null description"""
entity_data = {
"id": "20240102-test",
"name": "Test",
"entity_type": "test",
"description": None,
@@ -185,9 +179,11 @@ async def test_find_by_id(entity_repository: EntityRepository, sample_entity: En
@pytest.mark.asyncio
async def test_find_by_name(entity_repository: EntityRepository, sample_entity: Entity):
async def test_find_by_type_and_name(entity_repository: EntityRepository, sample_entity: Entity):
"""Test finding an entity by name"""
found = await entity_repository.find_by_name(sample_entity.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
@@ -307,7 +303,7 @@ async def test_delete_entity_with_relations(entity_repository: EntityRepository,
@pytest.mark.asyncio
async def test_delete_nonexistent_entity(entity_repository: EntityRepository):
"""Test deleting an entity that doesn't exist."""
result = await entity_repository.delete("nonexistent/id")
result = await entity_repository.delete(0)
assert result is False
@@ -317,13 +313,11 @@ async def test_search(session_maker, entity_repository: EntityRepository):
# First create and commit the entities
async with db.scoped_session(session_maker) as session:
entity1 = Entity(
id="20240102-test1",
name="Search Test 1",
entity_type="test",
description="First test entity",
)
entity2 = Entity(
id="20240102-test2",
name="Search Test 2",
entity_type="other",
description="Second test entity",
@@ -349,18 +343,18 @@ async def test_search(session_maker, entity_repository: EntityRepository):
)
# Test search by name
results = await entity_repository.search("Search Test")
results = await entity_repository.search_entities("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("other")
results = await entity_repository.search_entities("other")
assert len(results) == 1
assert results[0].entity_type == "other"
# Test search by observation content
results = await entity_repository.search("searchable")
results = await entity_repository.search_entities("searchable")
assert len(results) == 1
assert results[0].id == entity1.id