Agent/knowhub/knowhub_db/create_tables.py

#!/usr/bin/env python3
"""
创建 tool_table, skill_table, requirement_table 三张新表
"""

import os
import psycopg2
from psycopg2.extras import RealDictCursor
from dotenv import load_dotenv

load_dotenv()

CREATE_TOOL_TABLE = """
CREATE TABLE IF NOT EXISTS tool_table (
    id              VARCHAR(64) PRIMARY KEY,
    name            VARCHAR(255) NOT NULL,
    version         VARCHAR(64),
    introduction    TEXT,
    tutorial        TEXT,
    input           JSONB,
    output          JSONB,
    updated_time    BIGINT,
    status          VARCHAR(32) DEFAULT '未接入',  -- 未接入/可用/异常
    knowledge       JSONB DEFAULT '[]',            -- 关联知识条目
    case_knowledge  JSONB DEFAULT '[]',            -- 用例知识
    process_knowledge JSONB DEFAULT '[]'           -- 工序知识
) WITH (appendoptimized=false);
"""

CREATE_SKILL_TABLE = """
CREATE TABLE IF NOT EXISTS skill_table (
    id              VARCHAR(64) PRIMARY KEY,
    name            VARCHAR(255) NOT NULL,
    version         VARCHAR(64),
    introduction    TEXT,
    input           JSONB,
    output          JSONB,
    updated_time    BIGINT,
    resource_id     VARCHAR(64),                   -- 关联的工具id (resource.id)
    knowledge       JSONB DEFAULT '[]',            -- 关联知识条目
    case_knowledge  JSONB DEFAULT '[]',            -- 用例知识
    process_knowledge JSONB DEFAULT '[]',          -- 工序知识
    status          VARCHAR(32) DEFAULT '未验证'   -- 未验证/可用/异常
) WITH (appendoptimized=false);
"""

CREATE_REQUIREMENT_TABLE = """
CREATE TABLE IF NOT EXISTS requirement_table (
    id                VARCHAR(64) PRIMARY KEY,
    task              TEXT,                          -- 任务需求
    type              VARCHAR(64),                   -- 制作/...
    source_type       VARCHAR(64),                   -- itemset
    source_itemset_id VARCHAR(64),                   -- 关联pattern的id
    source_items      JSONB DEFAULT '[]',            -- 包含的特征/分类
    tools             JSONB DEFAULT '[]',            -- 关联的工具id与简介
    knowledge         JSONB DEFAULT '[]',            -- 关联知识条目
    case_knowledge    JSONB DEFAULT '[]',            -- 用例知识
    process_knowledge JSONB DEFAULT '[]',            -- 工序知识
    trace             JSONB DEFAULT '[]',            -- 关联的任务执行log
    body              TEXT,                          -- str/json
    embedding         float4[]                       -- 向量（1536维）
) WITH (appendoptimized=false);
"""

CREATE_REQUIREMENT_EMBEDDING_INDEX = """
CREATE INDEX IF NOT EXISTS idx_requirement_embedding
ON requirement_table USING ann(embedding)
WITH (dim=1536, hnsw_m=16, pq_enable=0);
"""

def create_tables():
    conn = psycopg2.connect(
        host=os.getenv('KNOWHUB_DB'),
        port=int(os.getenv('KNOWHUB_PORT', 5432)),
        user=os.getenv('KNOWHUB_USER'),
        password=os.getenv('KNOWHUB_PASSWORD'),
        database=os.getenv('KNOWHUB_DB_NAME')
    )
    conn.autocommit = True
    cursor = conn.cursor(cursor_factory=RealDictCursor)

    tables = [
        ("tool_table", CREATE_TOOL_TABLE),
        ("skill_table", CREATE_SKILL_TABLE),
        ("requirement_table", CREATE_REQUIREMENT_TABLE),
    ]

    for name, sql in tables:
        try:
            cursor.execute(sql)
            print(f"✅ {name} 创建成功")
        except Exception as e:
            print(f"❌ {name} 创建失败: {e}")

    try:
        cursor.execute(CREATE_REQUIREMENT_EMBEDDING_INDEX)
        print("✅ requirement_table fastann 索引创建成功")
    except Exception as e:
        print(f"❌ requirement_table 索引创建失败: {e}")

    # 验证
    cursor.execute("""
        SELECT tablename FROM pg_tables
        WHERE schemaname = 'public'
        ORDER BY tablename;
    """)
    tables_now = cursor.fetchall()
    print(f"\n当前所有表 ({len(tables_now)}个):")
    for t in tables_now:
        print(f"  - {t['tablename']}")

    cursor.close()
    conn.close()

if __name__ == "__main__":
    create_tables()