leocrm/app/plugins/builtins/graph_rag/models.py

"""EntityRelationship model for GraphRAG — stores relationships between entities."""

from __future__ import annotations

import uuid
from typing import Any

from sqlalchemy import Index, String
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.dialects.postgresql import UUID as PGUUID
from sqlalchemy.orm import Mapped, mapped_column

from app.core.db import Base, TenantMixin
from app.models.owned_mixin import OwnedMixin


class EntityRelationship(Base, TenantMixin, OwnedMixin):
    """Directed relationship between two entities in the knowledge graph.

    Stores typed relationships (e.g., 'works_for', 'has_email', 'related_to')
    between any two entities in the system. Used for GraphRAG traversal.
    """

    __tablename__ = "entity_relationships"
    __table_args__ = (
        Index("ix_entity_relationships_source", "tenant_id", "source_type", "source_id"),
        Index("ix_entity_relationships_target", "tenant_id", "target_type", "target_id"),
        Index("ix_entity_relationships_type", "tenant_id", "relationship_type"),
        Index("ix_entity_relationships_source_target", "tenant_id", "source_type", "source_id", "target_type", "target_id"),
    )

    id: Mapped[uuid.UUID] = mapped_column(
        PGUUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    source_type: Mapped[str] = mapped_column(
        String(50), nullable=False, comment="Entity type of the source (e.g. 'contact', 'company', 'file')"
    )
    source_id: Mapped[uuid.UUID] = mapped_column(
        PGUUID(as_uuid=True), nullable=False, comment="UUID of the source entity"
    )
    target_type: Mapped[str] = mapped_column(
        String(50), nullable=False, comment="Entity type of the target (e.g. 'contact', 'email', 'task')"
    )
    target_id: Mapped[uuid.UUID] = mapped_column(
        PGUUID(as_uuid=True), nullable=False, comment="UUID of the target entity"
    )
    relationship_type: Mapped[str] = mapped_column(
        String(50), nullable=False, comment="Type of relationship (e.g. 'works_for', 'has_email', 'related_to')"
    )
    meta: Mapped[dict[str, Any] | None] = mapped_column(
        "metadata", JSONB, nullable=True, default=dict, comment="Arbitrary metadata about the relationship"
    )