Agent/knowhub/knowhub_db/clean_tool_research_data.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
清理之前导入的 tool_research 数据
"""
import sys
from pathlib import Path

# 设置 Windows 控制台编码
if sys.platform == 'win32':
    import codecs
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, 'strict')
    sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, 'strict')

# 添加父目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

# 加载环境变量
from dotenv import load_dotenv
project_root = Path(__file__).parent.parent.parent
env_path = project_root / '.env'
load_dotenv(env_path)

from knowhub.knowhub_db.pg_store import PostgreSQLStore
from knowhub.knowhub_db.pg_requirement_store import PostgreSQLRequirementStore


def clean_requirements(req_store: PostgreSQLRequirementStore):
    """清理需求表中的 tool_research 数据"""
    print("\n=== 清理需求表 ===")

    cursor = req_store._get_cursor()
    try:
        # 删除 midjourney 和 seedream 的需求
        cursor.execute("""
            DELETE FROM requirement_table
            WHERE id LIKE 'req_midjourney_%' OR id LIKE 'req_seedream_%'
        """)
        deleted_count = cursor.rowcount
        req_store.conn.commit()
        print(f"删除了 {deleted_count} 条需求记录")
    except Exception as e:
        print(f"清理需求失败: {e}")
        req_store.conn.rollback()
    finally:
        cursor.close()


def clean_cases(knowledge_store: PostgreSQLStore):
    """清理知识表中的 case 数据"""
    print("\n=== 清理知识表中的 cases ===")

    cursor = knowledge_store._get_cursor()
    try:
        # 删除 tool_research_agent 创建的 case 类型知识
        cursor.execute("""
            DELETE FROM knowledge
            WHERE owner = 'tool_research_agent'
            AND 'case' = ANY(types)
        """)
        deleted_count = cursor.rowcount
        knowledge_store.conn.commit()
        print(f"删除了 {deleted_count} 条 case 知识记录")
    except Exception as e:
        print(f"清理 cases 失败: {e}")
        knowledge_store.conn.rollback()
    finally:
        cursor.close()


def main():
    print("开始清理之前导入的 tool_research 数据...")

    knowledge_store = PostgreSQLStore()
    req_store = PostgreSQLRequirementStore()

    clean_requirements(req_store)
    clean_cases(knowledge_store)

    print("\n清理完成!")


if __name__ == '__main__':
    main()