Tool_Agent/scripts/merge_unconnected_tools.py

#!/usr/bin/env python3
import asyncio
import json
import logging
import argparse
from psycopg2.extras import RealDictCursor

# 导入应用内的模块
from tool_agent.tool.tool_store import PostgreSQLToolStore
from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient, AssistantMessage, TextBlock

logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
logger = logging.getLogger(__name__)

SYSTEM_PROMPT = """你是一个专业的 API 工具库去重专家。
你的任务是将一份【未接入工具列表】与一份核心的【已接入工具列表】进行比对归并。
如果某一个【未接入工具】（哪怕换了名字叫别的）在【已接入工具】中已经有同样的本质功能，请找出这种映射关系。

例如：
未接入："通过 Midjourney 生成图片" 或 "Midjourney 查进度"
已接入：由于我们有 `midjourney_submit_job` 和 `midjourney_query_job_status`，请寻找最合理的对应。

【重要领域等价字典】请你在判定时，严格遵从以下我们架构中的内置别名等价逻辑：
1. "Nano Banana" 或者 "nanobanana" 等同于 "Google Gemini 生图" 或 "Imagen 3"
2. "BFL" 等同于 "Flux 官方生图 API"
3. "LiblibAI" 或 "哩布哩布" 等同于任何有关 Controlnet (不论是 OpenPose/Canny 等变体) 的支持
4. "RunComfy" 等同于 "ComfyUI 远程实例或工作流执行"
5. "即梦" 或 "ji_meng" 同等对应其系列任务

如果未能找到合理的已接入工具应对，则不输出映射。

请输出一个合法的 JSON 数组，必须是一个 list of objects：
[
  {
    "unconnected_id": "<未接入的工具id>",
    "connected_id": "<对应最合理的已接入的工具id>",
    "reason": "<非常简短的一句话匹配理由>"
  }
]

请绝对不要包含任何 Markdown 语法、```json 标签，直接输出原生 JSON 字符串。
"""

async def match_tools_with_claude(connected_tools, unconnected_batch):
    conn_str = json.dumps([{"id": t["id"], "name": t["name"], "desc": t["introduction"]} for t in connected_tools], ensure_ascii=False)
    unconn_str = json.dumps([{"id": t["id"], "name": t["name"], "desc": t["introduction"]} for t in unconnected_batch], ensure_ascii=False)
    
    prompt = f"【已接入工具库（作为对标参考）】:\n{conn_str}\n\n【本次需判定匹配的未接入工具】:\n{unconn_str}\n\n请输出JSON结果："
    
    options = ClaudeAgentOptions(model="claude-sonnet-4-5", system_prompt=SYSTEM_PROMPT)
    
    # 建立客户端调用
    result_text = ""
    try:
        async with ClaudeSDKClient(options=options) as client:
            await client.query(prompt)
            async for msg in client.receive_response():
                if isinstance(msg, AssistantMessage):
                    for block in msg.content:
                        if isinstance(block, TextBlock):
                            result_text += block.text
    except Exception as e:
        logger.error(f"Claude API failed: {e}")
        return []

    # 毫无保留地保存最原始的字符串落盘（防崩备用方案）
    with open("raw_claude_responses.log", "a", encoding="utf-8") as f:
        f.write(f"\n--- Batch Output ---\n{result_text}\n")

    # 尝试解析 JSON
    try:
        # 去掉可能的 Markdown 包装
        clean_json = result_text.strip()
        if clean_json.startswith("```json"):
            clean_json = clean_json[7:]
        elif clean_json.startswith("```"):
            clean_json = clean_json[3:]
        if clean_json.endswith("```"):
            clean_json = clean_json[:-3]
        
        data = json.loads(clean_json.strip())
        if isinstance(data, dict):
            if "unconnected_id" in data:
                return [data]
            return []
        elif isinstance(data, list):
            return [item for item in data if isinstance(item, dict)]
        return []
    except Exception as e:
        logger.error(f"Failed to parse JSON response: {e}\nResponse was: {result_text[:200]}")
        return []

def apply_database_merge(store: PostgreSQLToolStore, match_plan: list):
    """根据生成的合法匹配图，执行数据库替换"""
    if not match_plan:
        logger.info("没有任何匹配数据可执行落地！")
        return
        
    c = store.conn.cursor(cursor_factory=RealDictCursor)
    
    merged_count = 0
    for match in match_plan:
        uid = match.get("unconnected_id")
        cid = match.get("connected_id")
        
        if not uid or not cid:
            continue
            
        logger.info(f"==> 执行合并映射 [未接入: {uid} -> 已接入: {cid}]")
        
        try:
            # 1. 查询出原未接入工具的特征介绍准备拼接
            c.execute("SELECT introduction, tutorial, version FROM tool WHERE id = %s", (uid,))
            u_tool = c.fetchone()
            
            c.execute("SELECT introduction, tutorial, version FROM tool WHERE id = %s", (cid,))
            c_tool = c.fetchone()
            
            if u_tool and c_tool:
                u_ver = u_tool.get('version') or ''
                c_ver = c_tool.get('version') or ''
                # Extract best version, favor c_ver if exists and not placeholder
                latest_ver = c_ver if c_ver and c_ver != "1.0.0" else (u_ver if u_ver else "1.0.0")
                
                version_prefix = f"已支持的最新版本: {latest_ver}\n\n"
                
                new_intro = version_prefix + (c_tool['introduction'] or '') + "\n\n[融合补充描述]:\n" + (u_tool['introduction'] or '')
                new_tutorial = (c_tool['tutorial'] or '') + "\n\n" + (u_tool['tutorial'] or '')
                c.execute("UPDATE tool SET introduction=%s, tutorial=%s WHERE id=%s", (new_intro[:4000], new_tutorial[:4000], cid))
                
            # 2. 跟原表里外键绑定的迁移过去（只搬迁关系，利用 ON CONFLICT 防止主键重复）
            tables_to_migrate = [ 
                ("capability_tool", "capability_id", "tool_id"), 
                ("tool_knowledge", "knowledge_id", "tool_id"),
                ("tool_provider", "provider_id", "tool_id")
            ]
            
            for table, _, tool_col in tables_to_migrate:
                # 找出旧工具的所有关联
                c.execute(f"SELECT * FROM {table} WHERE {tool_col} = %s", (uid,))
                rows = c.fetchall()
                for row in rows:
                    if table == "capability_tool":
                        c.execute("INSERT INTO capability_tool (capability_id, tool_id) VALUES (%s, %s) ON CONFLICT DO NOTHING", (row['capability_id'], cid))
                    elif table == "tool_knowledge":
                        c.execute("INSERT INTO tool_knowledge (knowledge_id, tool_id) VALUES (%s, %s) ON CONFLICT DO NOTHING", (row['knowledge_id'], cid))
                    elif table == "tool_provider":
                        c.execute("INSERT INTO tool_provider (provider_id, tool_id) VALUES (%s, %s) ON CONFLICT DO NOTHING", (row['provider_id'], cid))
                        
            # 3. 将原未接入工具删除 (cascade delete if properly constructed by DB, else manually delete relations first)
            # 既然我们已经读了并插入到了新的关联里，我们需要手动删除旧的外键绑定好让 CASCADE 清晰或者直接删
            c.execute("DELETE FROM capability_tool WHERE tool_id = %s", (uid,))
            c.execute("DELETE FROM tool_knowledge WHERE tool_id = %s", (uid,))
            c.execute("DELETE FROM tool_provider WHERE tool_id = %s", (uid,))
            c.execute("DELETE FROM tool WHERE id = %s", (uid,))
            
            merged_count += 1
        except Exception as e:
            logger.error(f"处理 {uid} 的融合时出错: {e}")
            store.conn.rollback() # 出错则回滚该条
            continue
            
        # 单条合并成功可提交
        store.conn.commit()
    
    c.close()
    logger.info(f"所有指定合并执行完毕，已彻底抹除融合了 {merged_count} 个废弃节点工具！")

async def run_workflow(apply=False):
    store = PostgreSQLToolStore()
    tools = store.list_all(limit=2000)
    
    connected = [t for t in tools if t.get("status") == "已接入"]
    unconnected = [t for t in tools if t.get("status") != "已接入"]
    
    logger.info(f"提取出 {len(connected)} 个基础已接入工具，以及 {len(unconnected)} 个未接入待判决工具。")
    if not unconnected:
        logger.info("没有任何未接入工具需要处理。")
        store.close()
        return

    # 尝试从本地缓存文件读取，避免重复消耗大模型调用
    import os
    match_file = "matches.json"
    
    if os.path.exists(match_file):
        logger.info(f"发现本地存在的缓存文件 {match_file}，将直接读取现有分析结果！")
        with open(match_file, "r", encoding="utf-8") as f:
            all_matches = json.load(f)
    else:
        # 切割 batch，降低大模型的单次信息处理压力，提升专注度
        batch_size = 10
        all_matches = []
        
        total_batches = (len(unconnected) // batch_size) + (1 if len(unconnected) % batch_size != 0 else 0)
        for i in range(0, len(unconnected), batch_size):
            batch = unconnected[i:i + batch_size]
            logger.info(f"正在交给 Claude 引擎评估第 {i//batch_size + 1}/{total_batches} 批 ({len(batch)} items) ...")
            
            matches = await match_tools_with_claude(connected, batch)
            if matches:
                all_matches.extend(matches)
                
        # 缓存映射结果
        with open(match_file, "w", encoding="utf-8") as f:
            json.dump(all_matches, f, ensure_ascii=False, indent=2)
        
    logger.info(f"\n====================== 预览报告 ======================")
    if not all_matches:
        logger.info("Claude 没有找到任何可以安全结合的工具！")
    else:
        logger.info(f"Claude 匹配上了 {len(all_matches)} 个工具，建议的更换路线如下：")
        for m in all_matches:
            logger.info(f" - [冗余废除] {m.get('unconnected_id')}  ==> [主工具] {m.get('connected_id')}")
            logger.info(f"   理由: {m.get('reason','')}")
            
    logger.info(f"====================================================\n")
    
    if apply:
        logger.info("启动真实的数据库落地置换 (Apply mode is ON)...")
        apply_database_merge(store, all_matches)
    else:
        logger.info("目前为 Dry-Run 模式，如果审查匹配无误觉得干的漂亮，请携带 --apply 开关挂载实际融合执行！")

    store.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--apply", action="store_true", help="Apply DB migrations physically")
    args = parser.parse_args()
    
    asyncio.run(run_workflow(apply=args.apply))