geo/backend/app/models/knowledge_graph.py

"""知识图谱数据模型"""
import uuid
from datetime import datetime
import enum

from sqlalchemy import String, Text, DateTime, ForeignKey, Index, Enum, func
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import Mapped, mapped_column, relationship

from app.database import Base, JSONType


class EntityType(str, enum.Enum):
    """实体类型"""
    ORGANIZATION = "ORGANIZATION"      # 组织/公司
    PRODUCT = "PRODUCT"                # 产品
    PERSON = "PERSON"                  # 人物
    LOCATION = "LOCATION"              # 地点
    TECHNOLOGY = "TECHNOLOGY"          # 技术
    BRAND = "BRAND"                    # 品牌
    EVENT = "EVENT"                    # 事件
    CONCEPT = "CONCEPT"               # 概念
    OTHER = "OTHER"                   # 其他


class RelationType(str, enum.Enum):
    """关系类型"""
    # 组织关系
    COMPETES_WITH = "COMPETES_WITH"          # 竞争对手
    PARTNERS_WITH = "PARTNERS_WITH"          # 合作伙伴
    ACQUIRES = "ACQUIRES"                    # 收购
    SUBSIDIARY_OF = "SUBSIDIARY_OF"          # 子公司

    # 产品关系
    PRODUCES = "PRODUCES"                    # 生产
    USES_TECHNOLOGY = "USES_TECHNOLOGY"      # 使用技术
    PART_OF = "PART_OF"                      # 属于（产品线）

    # 地点关系
    LOCATED_IN = "LOCATED_IN"                # 位于
    FOUNDED_IN = "FOUNDED_IN"                # 成立于

    # 人物关系
    CEO_OF = "CEO_OF"                        # CEO
    FOUNDER_OF = "FOUNDER_OF"                # 创始人

    # 通用关系
    RELATED_TO = "RELATED_TO"                 # 相关
    MENTIONED_IN = "MENTIONED_IN"            # 提及于
    ALSO_KNOWN_AS = "ALSO_KNOWN_AS"         # 又名


class KnowledgeEntity(Base):
    """知识实体"""
    __tablename__ = "knowledge_entities"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        primary_key=True,
        default=uuid.uuid4,
    )
    knowledge_base_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("knowledge_bases.id", ondelete="CASCADE"),
        nullable=False,
    )

    # 实体信息
    name: Mapped[str] = mapped_column(String(500), nullable=False, index=True)
    entity_type: Mapped[EntityType] = mapped_column(Enum(EntityType), nullable=False, index=True)
    description: Mapped[str | None] = mapped_column(Text, nullable=True)

    # 扩展属性（JSON）
    properties: Mapped[dict | None] = mapped_column(JSONType, default=dict)

    # 来源信息
    source_chunk_id: Mapped[uuid.UUID | None] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("knowledge_chunks.id", ondelete="SET NULL"),
        nullable=True,
    )
    confidence: Mapped[str | None] = mapped_column(String(20), nullable=True)  # 置信度描述：high/medium/low

    # 元数据
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        nullable=False,
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        onupdate=func.now(),
        nullable=False,
    )

    # 索引
    __table_args__ = (
        Index("ix_entities_kb_name", "knowledge_base_id", "name"),
        Index("ix_entities_kb_type", "knowledge_base_id", "entity_type"),
    )

    # 关系
    outgoing_relations: Mapped[list["KnowledgeRelation"]] = relationship(
        "KnowledgeRelation",
        foreign_keys="KnowledgeRelation.source_entity_id",
        back_populates="source_entity",
        cascade="all, delete-orphan",
    )
    incoming_relations: Mapped[list["KnowledgeRelation"]] = relationship(
        "KnowledgeRelation",
        foreign_keys="KnowledgeRelation.target_entity_id",
        back_populates="target_entity",
        cascade="all, delete-orphan",
    )


class KnowledgeRelation(Base):
    """知识关系"""
    __tablename__ = "knowledge_relations"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        primary_key=True,
        default=uuid.uuid4,
    )

    # 关系两端
    source_entity_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("knowledge_entities.id", ondelete="CASCADE"),
        nullable=False,
    )
    target_entity_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("knowledge_entities.id", ondelete="CASCADE"),
        nullable=False,
    )

    # 关系信息
    relation_type: Mapped[RelationType] = mapped_column(Enum(RelationType), nullable=False, index=True)

    # 扩展属性
    properties: Mapped[dict | None] = mapped_column(JSONType, default=dict)

    # 来源信息
    source_chunk_id: Mapped[uuid.UUID | None] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("knowledge_chunks.id", ondelete="SET NULL"),
        nullable=True,
    )
    confidence: Mapped[str | None] = mapped_column(String(20), nullable=True)

    # 元数据
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        nullable=False,
    )

    # 关系
    source_entity: Mapped["KnowledgeEntity"] = relationship(
        "KnowledgeEntity",
        foreign_keys=[source_entity_id],
        back_populates="outgoing_relations",
    )
    target_entity: Mapped["KnowledgeEntity"] = relationship(
        "KnowledgeEntity",
        foreign_keys=[target_entity_id],
        back_populates="incoming_relations",
    )

    # 索引
    __table_args__ = (
        Index("ix_relations_source", "source_entity_id"),
        Index("ix_relations_target", "target_entity_id"),
        Index("ix_relations_type", "relation_type"),
    )