wx-agent/tools/agent_registration.py

87 lines
4.1 KiB
Python

import json
import time
from typing import Any, Dict
from config import POSTGRES_CONNECTION_STRING
WIKI_ADMIN_AGENT_ID = "6f0b5c4f-7d4a-4db4-9c4c-8f89b2a4f7a1"
WIKI_ADMIN_ROLE_ID = "5"
WIKI_ADMIN_AGENT: Dict[str, Any] = {
"id": WIKI_ADMIN_AGENT_ID,
"name": "Wiki知识库管理",
"description": "管理 Wiki 指针层的同步、检查、重建、清理和删除,帮助维护 11000 知识库与 Wiki 数据的一致性。",
"prompt": "您好,我是 Wiki 知识库管理助手。您可以直接用知识库名称让我检查同步情况、同步指定文件或知识库、清理 11000 已删除但 Wiki 仍存在的数据。我会先生成待执行计划,列出知识库、文件和动作;您确认后才会执行。删除类操作只删除 Wiki 指针层数据,不会删除 11000 原始知识库文件。",
"background": "Wiki 管理助手用于维护 wx-agent 内置 Wiki 指针层,不负责回答业务知识问题。",
"access_control": {},
"default_query": [
{"title": "检查知识库同步情况 →", "content": "检查知识库 163舰 的 Wiki 同步情况"},
{"title": "同步指定知识库 →", "content": "同步知识库 163舰 的前 20 个文件"},
{"title": "清理脏数据预览 →", "content": "清理知识库 163舰 中 11000 已删除但 Wiki 还存在的数据"},
{"title": "确认执行计划 →", "content": "确认执行"},
],
}
async def register_wiki_admin_agent() -> Dict[str, Any]:
"""Register the Wiki admin agent in an_webui tables and bind it to role 5."""
try:
from psycopg_pool import AsyncConnectionPool
except ImportError:
print("[agent_registration] psycopg_pool is not installed; skip wiki admin agent registration")
return {"success": False, "error": "psycopg_pool is not installed"}
now = int(time.time())
async with AsyncConnectionPool(
POSTGRES_CONNECTION_STRING,
kwargs={"autocommit": True},
) as pool:
async with pool.connection() as conn:
async with conn.cursor() as cur:
await cur.execute(
"""
INSERT INTO agent (
id, name, description, prompt, background, access_control,
created_time, updated_time, default_query
)
VALUES (%s, %s, %s, %s, %s, %s::json, %s, %s, %s::json)
ON CONFLICT (id) DO UPDATE SET
name = EXCLUDED.name,
description = EXCLUDED.description,
prompt = EXCLUDED.prompt,
background = EXCLUDED.background,
access_control = EXCLUDED.access_control,
updated_time = EXCLUDED.updated_time,
default_query = EXCLUDED.default_query
""",
(
WIKI_ADMIN_AGENT["id"],
WIKI_ADMIN_AGENT["name"],
WIKI_ADMIN_AGENT["description"],
WIKI_ADMIN_AGENT["prompt"],
WIKI_ADMIN_AGENT["background"],
json.dumps(WIKI_ADMIN_AGENT["access_control"], ensure_ascii=False),
now,
now,
json.dumps(WIKI_ADMIN_AGENT["default_query"], ensure_ascii=False),
),
)
await cur.execute(
"""
INSERT INTO agent_role_association (agent_id, agent_role_id, is_active)
VALUES (%s, %s, TRUE)
ON CONFLICT (agent_id, agent_role_id) DO UPDATE SET
is_active = EXCLUDED.is_active
""",
(WIKI_ADMIN_AGENT_ID, WIKI_ADMIN_ROLE_ID),
)
print(f"[agent_registration] wiki admin agent registered: {WIKI_ADMIN_AGENT_ID} -> role {WIKI_ADMIN_ROLE_ID}")
return {
"success": True,
"agent_id": WIKI_ADMIN_AGENT_ID,
"agent_role_id": WIKI_ADMIN_ROLE_ID,
}