From 3168024e4813b9a023ef4e570dc07021709fcc03 Mon Sep 17 00:00:00 2001
From: phernandez <phernandez@basicmachines.co>
Date: Fri, 20 Dec 2024 20:19:15 -0600
Subject: [PATCH] fix FK from entity to doc

---
 src/basic_memory/models/documents.py | 11 +++++++++--
 src/basic_memory/models/knowledge.py |  6 +++++-
 2 files changed, 14 insertions(+), 3 deletions(-)

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