diff --git a/src/basic_memory/models/documents.py b/src/basic_memory/models/documents.py index 9c2e1276..b049cc7a 100644 --- a/src/basic_memory/models/documents.py +++ b/src/basic_memory/models/documents.py @@ -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) diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py index cd01dc21..e33ca7ae 100644 --- a/src/basic_memory/models/knowledge.py +++ b/src/basic_memory/models/knowledge.py @@ -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"), diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index bb488753..7390ac4b 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -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) \ No newline at end of file + return await self.delete_by_fields(doc_id=doc_id) diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index bba9b73f..c0713519 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -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) diff --git a/tests/repository/test_entity_repository.py b/tests/repository/test_entity_repository.py index 2a7d797d..8f20ce3a 100644 --- a/tests/repository/test_entity_repository.py +++ b/tests/repository/test_entity_repository.py @@ -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 diff --git a/tests/services/test_entity_service.py b/tests/services/test_entity_service.py index 5a33d129..7658fb9e 100644 --- a/tests/services/test_entity_service.py +++ b/tests/services/test_entity_service.py @@ -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