144 lines
5.1 KiB
Python
144 lines
5.1 KiB
Python
"""Add knowledge graph tables
|
|
|
|
Revision ID: f7a8b9c0de56
|
|
Revises: e5f7g9h1cd45
|
|
Create Date: 2026-05-24 12:00:00.000000
|
|
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
import sqlalchemy as sa
|
|
from alembic import op
|
|
from sqlalchemy.dialects import postgresql
|
|
import enum
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = "f7a8b9c0de56"
|
|
down_revision: Union[str, None] = "810a29804f5a"
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# ------------------------------------------------------------------ #
|
|
# 1. 创建实体类型枚举
|
|
# ------------------------------------------------------------------ #
|
|
entity_type_enum = enum.Enum(
|
|
'EntityType',
|
|
[
|
|
'ORGANIZATION', 'PRODUCT', 'PERSON', 'LOCATION',
|
|
'TECHNOLOGY', 'BRAND', 'EVENT', 'CONCEPT', 'OTHER'
|
|
]
|
|
)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# 2. 创建关系类型枚举
|
|
# ------------------------------------------------------------------ #
|
|
relation_type_enum = enum.Enum(
|
|
'RelationType',
|
|
[
|
|
'COMPETES_WITH', 'PARTNERS_WITH', 'ACQUIRES', 'SUBSIDIARY_OF',
|
|
'PRODUCES', 'USES_TECHNOLOGY', 'PART_OF',
|
|
'LOCATED_IN', 'FOUNDED_IN',
|
|
'CEO_OF', 'FOUNDER_OF',
|
|
'RELATED_TO', 'MENTIONED_IN', 'ALSO_KNOWN_AS'
|
|
]
|
|
)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# 3. knowledge_entities 表
|
|
# ------------------------------------------------------------------ #
|
|
op.create_table(
|
|
"knowledge_entities",
|
|
sa.Column(
|
|
"id",
|
|
postgresql.UUID(as_uuid=True),
|
|
primary_key=True,
|
|
nullable=False,
|
|
),
|
|
sa.Column(
|
|
"knowledge_base_id",
|
|
postgresql.UUID(as_uuid=True),
|
|
sa.ForeignKey("knowledge_bases.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
),
|
|
sa.Column("name", sa.String(500), nullable=False),
|
|
sa.Column("entity_type", sa.Enum(entity_type_enum, name="entitytype"), nullable=False),
|
|
sa.Column("description", sa.Text, nullable=True),
|
|
sa.Column("properties", postgresql.JSONB(astext_type=sa.Text()), nullable=True, server_default="{}"),
|
|
sa.Column(
|
|
"source_chunk_id",
|
|
postgresql.UUID(as_uuid=True),
|
|
sa.ForeignKey("knowledge_chunks.id", ondelete="SET NULL"),
|
|
nullable=True,
|
|
),
|
|
sa.Column("confidence", sa.String(20), nullable=True),
|
|
sa.Column(
|
|
"created_at",
|
|
sa.DateTime(timezone=True),
|
|
server_default=sa.text("now()"),
|
|
nullable=False,
|
|
),
|
|
sa.Column(
|
|
"updated_at",
|
|
sa.DateTime(timezone=True),
|
|
server_default=sa.text("now()"),
|
|
nullable=False,
|
|
),
|
|
)
|
|
op.create_index("ix_entities_name", "knowledge_entities", ["name"])
|
|
op.create_index("ix_entities_kb_name", "knowledge_entities", ["knowledge_base_id", "name"])
|
|
op.create_index("ix_entities_kb_type", "knowledge_entities", ["knowledge_base_id", "entity_type"])
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# 4. knowledge_relations 表
|
|
# ------------------------------------------------------------------ #
|
|
op.create_table(
|
|
"knowledge_relations",
|
|
sa.Column(
|
|
"id",
|
|
postgresql.UUID(as_uuid=True),
|
|
primary_key=True,
|
|
nullable=False,
|
|
),
|
|
sa.Column(
|
|
"source_entity_id",
|
|
postgresql.UUID(as_uuid=True),
|
|
sa.ForeignKey("knowledge_entities.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
),
|
|
sa.Column(
|
|
"target_entity_id",
|
|
postgresql.UUID(as_uuid=True),
|
|
sa.ForeignKey("knowledge_entities.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
),
|
|
sa.Column("relation_type", sa.Enum(relation_type_enum, name="relationtype"), nullable=False),
|
|
sa.Column("properties", postgresql.JSONB(astext_type=sa.Text()), nullable=True, server_default="{}"),
|
|
sa.Column(
|
|
"source_chunk_id",
|
|
postgresql.UUID(as_uuid=True),
|
|
sa.ForeignKey("knowledge_chunks.id", ondelete="SET NULL"),
|
|
nullable=True,
|
|
),
|
|
sa.Column("confidence", sa.String(20), nullable=True),
|
|
sa.Column(
|
|
"created_at",
|
|
sa.DateTime(timezone=True),
|
|
server_default=sa.text("now()"),
|
|
nullable=False,
|
|
),
|
|
)
|
|
op.create_index("ix_relations_source", "knowledge_relations", ["source_entity_id"])
|
|
op.create_index("ix_relations_target", "knowledge_relations", ["target_entity_id"])
|
|
op.create_index("ix_relations_type", "knowledge_relations", ["relation_type"])
|
|
|
|
|
|
def downgrade() -> None:
|
|
# 删除表（注意外键约束会自动处理）
|
|
op.drop_table("knowledge_relations")
|
|
op.drop_table("knowledge_entities")
|
|
# 删除枚举类型
|
|
op.execute("DROP TYPE IF EXISTS relationtype")
|
|
op.execute("DROP TYPE IF EXISTS entitytype")
|