diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py index 6c8a3f9d..561b0b96 100644 --- a/src/basic_memory/models/knowledge.py +++ b/src/basic_memory/models/knowledge.py @@ -1,56 +1,58 @@ """Knowledge graph models.""" - from datetime import datetime -from typing import Optional - +from typing import Optional, List 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"), - Index("ix_entity_type", "entity_type"), # index on entity_type + Index("ix_entity_type", "entity_type"), + Index("ix_entity_doc_id", "doc_id"), ) # Core identity 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")) - updated_at: Mapped[datetime] = mapped_column( - DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=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") + ) + # 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", @@ -71,16 +73,18 @@ 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") @@ -93,22 +97,24 @@ class Relation(Base): """ A directed relation between two entities. """ - __tablename__ = "relations" __table_args__ = ( UniqueConstraint("from_id", "to_id", "relation_type", name="uix_relation"), - Index("ix_relation_type", "relation_type"), # index on relation_type + Index("ix_relation_type", "relation_type"), ) id: Mapped[int] = mapped_column(Integer, primary_key=True) 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}')" \ No newline at end of file diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index c4e6d3e4..e164a576 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -1,165 +1,100 @@ -"""Repository for managing Entity objects.""" - -from typing import Optional, Sequence, List +"""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, or_ -from sqlalchemy.exc import NoResultFound -from sqlalchemy.ext.asyncio import async_sessionmaker -from sqlalchemy.orm import selectinload +from sqlalchemy import select, delete +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from sqlalchemy.sql import text from basic_memory import db -from basic_memory.models import Entity, Observation +from basic_memory.models.knowledge import Entity from basic_memory.repository.repository import Repository class EntityRepository(Repository[Entity]): - """Repository for Entity model with memory-specific operations.""" + """Repository for Entity model.""" - def __init__(self, session_maker: async_sessionmaker): - super().__init__(session_maker, Entity) - logger.debug("Initialized EntityRepository") + 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.""" - entity_id = Entity.generate_id(data["entity_type"], data["name"]) - await super().create({**data, "id": entity_id}) + 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 + ) + return await self.find_one(query) - # we have to find to get relations - created = await self.find_by_id(entity_id) - assert created is not None, f"Created entity {entity_id} should not be None" - return created - - 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.""" - for data in data_list: - entity_id = Entity.generate_id(data["entity_type"], data["name"]) - data["id"] = entity_id - 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: str) -> 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.outgoing_relations), - selectinload(Entity.incoming_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[str]) -> 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.outgoing_relations), - selectinload(Entity.incoming_relations), - ) - ) - entities = result.scalars().all() - logger.debug(f"Found {len(entities)}") - return entities - - async def find_by_name(self, name: str) -> Optional[Entity]: - """Find an entity by its unique name.""" - logger.debug(f"Finding entity by name: {name}") - async with db.scoped_session(self.session_maker) as session: - result = await session.execute( - select(Entity) - .filter(Entity.name == name) - .options( - selectinload(Entity.observations), - selectinload(Entity.outgoing_relations), - selectinload(Entity.incoming_relations), - ) - ) - entity = result.scalars().one_or_none() - if entity: - logger.debug(f"Found entity: {entity.id}") - else: - logger.debug(f"No entity found with name: {name}") - return entity - - async def find_by_type_and_name(self, entity_type: str, name: str) -> Optional[Entity]: - """Find an entity by its type and name combination.""" - logger.debug(f"Finding entity by type and name: {entity_type}/{name}") - async with db.scoped_session(self.session_maker) as session: - result = await session.execute( - select(Entity) - .filter(Entity.entity_type == entity_type) - .filter(Entity.name == name) - .options( - selectinload(Entity.observations), - selectinload(Entity.outgoing_relations), - selectinload(Entity.incoming_relations), - ) - ) - entity = result.scalars().one_or_none() - if entity: - logger.debug(f"Found entity: {entity.id}") - else: - logger.debug(f"No entity found with type/name: {entity_type}/{name}") - return entity - - async def search_by_type( - self, entity_type: str, skip: int = 0, limit: int = 100 + async def list_entities( + self, + entity_type: Optional[str] = None, + doc_id: Optional[int] = None, ) -> Sequence[Entity]: - """Search for entities of a specific type.""" - logger.debug(f"Searching entities by type: {entity_type} (skip={skip}, limit={limit})") + """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( - select(Entity) - .filter(Entity.entity_type == entity_type) - .options( - selectinload(Entity.observations), - selectinload(Entity.outgoing_relations), - selectinload(Entity.incoming_relations), - ) - .offset(skip) - .limit(limit) - ) - entities = result.scalars().all() - logger.debug(f"Found {len(entities)} entities of type {entity_type}") - return entities + result = await session.execute(query) + return list(result.scalars().all()) - async def search(self, query: str) -> Sequence[Entity]: - """Search entities using LIKE pattern matching.""" - logger.debug(f"Searching entities with query: {query}") + async def get_entity_types(self) -> List[str]: + """Get list of distinct entity types.""" + query = select(Entity.entity_type).distinct() async with db.scoped_session(self.session_maker) as session: - stmt = ( - select(Entity) - .distinct() - .where( - or_( - Entity.name.ilike(f"%{query}%"), - Entity.entity_type.ilike(f"%{query}%"), - Entity.observations.any(Observation.content.ilike(f"%{query}%")), - ) - ) - .options( - selectinload(Entity.observations), - selectinload(Entity.outgoing_relations), - selectinload(Entity.incoming_relations), - ) - ) - result = await session.execute(stmt) - entities = list(result.scalars()) - logger.debug(f"Found {len(entities)} matching entities") - return entities + result = await session.execute(query) + return [r[0] for r in result.all()] + + async def search_entities(self, query_str: str) -> List[Entity]: + """ + Search for entities. + + Searches across: + - Entity names + - Entity types + - Entity descriptions + """ + 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)) + ) + 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]: + """Update an entity with the given fields.""" + return await self.update(str(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)