fix entity service tests

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)

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"),

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)
+        return await self.delete_by_fields(doc_id=doc_id)

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)