add metadata and errors to /documents api

This commit is contained in:
phernandez
2024-12-21 19:56:54 -06:00
parent c170aed855
commit ddf2d2fd0e
2 changed files with 44 additions and 70 deletions
+22 -50
View File
@@ -3,10 +3,7 @@
from datetime import datetime
from typing import List, Optional
from sqlalchemy import (
String, DateTime, ForeignKey, Text, Integer,
text, UniqueConstraint
)
from sqlalchemy import String, DateTime, ForeignKey, Text, Integer, text, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column, relationship
from basic_memory.models.base import Base
@@ -24,53 +21,42 @@ class Entity(Base):
- A description (optional)
- A list of observations
"""
__tablename__ = "entity"
__table_args__ = (
UniqueConstraint("entity_type", "name", name="uix_entity_type_name"),
)
__table_args__ = (UniqueConstraint("entity_type", "name", name="uix_entity_type_name"),)
id: Mapped[str] = mapped_column(String, primary_key=True)
name: Mapped[str] = mapped_column(String)
entity_type: Mapped[str] = mapped_column(String)
description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
created_at: Mapped[datetime] = mapped_column(
DateTime,
server_default=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")
DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP")
)
# Link to source document
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
document: Mapped["Document"] = relationship(
"Document",
back_populates="entities"
document: Mapped["Document"] = relationship( # pyright: ignore [reportUndefinedVariable] # noqa: F821
"Document", back_populates="entities"
)
observations: Mapped[List["Observation"]] = relationship(
"Observation",
back_populates="entity",
cascade="all, delete-orphan"
"Observation", back_populates="entity", cascade="all, delete-orphan"
)
outgoing_relations: Mapped[List["Relation"]] = relationship(
"Relation",
foreign_keys="[Relation.from_id]",
back_populates="from_entity",
cascade="all, delete-orphan"
cascade="all, delete-orphan",
)
incoming_relations: Mapped[List["Relation"]] = relationship(
"Relation",
foreign_keys="[Relation.to_id]",
back_populates="to_entity",
cascade="all, delete-orphan"
cascade="all, delete-orphan",
)
@property
@@ -100,19 +86,15 @@ class Observation(Base):
- Can be added or removed independently
- Should be atomic (one fact per observation)
"""
__tablename__ = "observation"
id: Mapped[int] = mapped_column(Integer, primary_key=True)
entity_id: Mapped[str] = mapped_column(
String,
ForeignKey("entity.id", ondelete="CASCADE"),
index=True
String, ForeignKey("entity.id", ondelete="CASCADE"), index=True
)
content: 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"))
context: Mapped[Optional[str]] = mapped_column(String, nullable=True)
# Relationships
@@ -126,40 +108,30 @@ class Observation(Base):
class Relation(Base):
"""
Relations define directed connections between entities.
They are always stored in active voice and describe how entities
They are always stored in active voice and describe how entities
interact or relate to each other.
"""
__tablename__ = "relation"
id: Mapped[int] = mapped_column(Integer, primary_key=True)
from_id: Mapped[str] = mapped_column(
String,
ForeignKey("entity.id", ondelete="CASCADE"),
index=True
String, ForeignKey("entity.id", ondelete="CASCADE"), index=True
)
to_id: Mapped[str] = mapped_column(
String,
ForeignKey("entity.id", ondelete="CASCADE"),
index=True
String, ForeignKey("entity.id", ondelete="CASCADE"), index=True
)
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"))
context: Mapped[Optional[str]] = mapped_column(String, nullable=True)
# Relationships
from_entity: Mapped[Entity] = relationship(
"Entity",
foreign_keys=[from_id],
back_populates="outgoing_relations"
"Entity", foreign_keys=[from_id], back_populates="outgoing_relations"
)
to_entity: Mapped[Entity] = relationship(
"Entity",
foreign_keys=[to_id],
back_populates="incoming_relations"
"Entity", foreign_keys=[to_id], back_populates="incoming_relations"
)
def __repr__(self) -> str:
return f"Relation(id={self.id}, from='{self.from_id}', type='{self.relation_type}', to='{self.to_id}')"
return f"Relation(id={self.id}, from='{self.from_id}', type='{self.relation_type}', to='{self.to_id}')"