mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
refactor entity/repo
This commit is contained in:
phernandez
@@ -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}')"
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user