src/personalization_module/preference_acquisition/preference_memory/models.py

from __future__ import annotations

from datetime import datetime

from sqlalchemy import Column, Float, Integer, String, ForeignKey, DateTime
from sqlalchemy.orm import relationship

from .db import Base


class Node(Base):
    __tablename__ = "nodes"

    id = Column(Integer, primary_key=True)
    name = Column(String, unique=True, nullable=False)
    node_type = Column(String, nullable=False)

    outgoing_edges = relationship(
        "Edge",
        foreign_keys="Edge.src_id",
        back_populates="src",
    )

    incoming_edges = relationship(
        "Edge",
        foreign_keys="Edge.dst_id",
        back_populates="dst",
    )


class Edge(Base):
    __tablename__ = "edges"

    id = Column(Integer, primary_key=True)

    src_id = Column(Integer, ForeignKey("nodes.id"), nullable=False)
    relation = Column(String, nullable=False)
    dst_id = Column(Integer, ForeignKey("nodes.id"), nullable=False)

    created_at = Column(DateTime, default=datetime.utcnow)
    weight = Column(Float, nullable=True)

    src = relationship("Node", foreign_keys=[src_id], back_populates="outgoing_edges")
    dst = relationship("Node", foreign_keys=[dst_id], back_populates="incoming_edges")