diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py index 561b0b96..cd01dc21 100644 --- a/src/basic_memory/models/knowledge.py +++ b/src/basic_memory/models/knowledge.py @@ -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}')" \ No newline at end of file + return f"Relation(from_id={self.from_id}, to_id={self.to_id}, type='{self.relation_type}')" diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index e164a576..bb488753 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -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) \ No newline at end of file diff --git a/src/basic_memory/repository/repository.py b/src/basic_memory/repository/repository.py index f7c69881..f91ef81b 100644 --- a/src/basic_memory/repository/repository.py +++ b/src/basic_memory/repository/repository.py @@ -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 \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index 8b040a32..2d03e0d1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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) diff --git a/tests/repository/test_entity_repository.py b/tests/repository/test_entity_repository.py index 02bcce1f..2a7d797d 100644 --- a/tests/repository/test_entity_repository.py +++ b/tests/repository/test_entity_repository.py @@ -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