diff --git a/src/basic_memory/models/documents.py b/src/basic_memory/models/documents.py index dc5282d7..b5bb250b 100644 --- a/src/basic_memory/models/documents.py +++ b/src/basic_memory/models/documents.py @@ -1,10 +1,10 @@ """Document model for tracking files in the knowledge base.""" from datetime import datetime -from typing import Optional +from typing import Optional, List from sqlalchemy import String, DateTime, text, JSON -from sqlalchemy.orm import Mapped, mapped_column +from sqlalchemy.orm import Mapped, mapped_column, relationship from basic_memory.models.base import Base @@ -33,5 +33,12 @@ class Document(Base): onupdate=text("CURRENT_TIMESTAMP") ) + # Relationships + entities: Mapped[List["Entity"]] = relationship( + "Entity", + back_populates="document", + cascade="all, delete-orphan" + ) + def __repr__(self) -> str: return f"Document(id={self.id}, path='{self.path}')" \ No newline at end of file diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py index 43d8a328..4452ddb7 100644 --- a/src/basic_memory/models/knowledge.py +++ b/src/basic_memory/models/knowledge.py @@ -5,7 +5,7 @@ from typing import List, Optional from sqlalchemy import ( String, DateTime, ForeignKey, Text, Integer, - text, UniqueConstraint, ForeignKeyConstraint + text, UniqueConstraint ) from sqlalchemy.orm import Mapped, mapped_column, relationship @@ -51,6 +51,10 @@ class Entity(Base): ) # Relationships + document: Mapped["Document"] = relationship( + "Document", + back_populates="entities" + ) observations: Mapped[List["Observation"]] = relationship( "Observation", back_populates="entity",