143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
"""把 kb_build/pages/*.txt 的站点文档写入 nonebot-plugin-wtfllm 知识库。
|
|
|
|
用法:
|
|
python kb_build/import_kb.py --agent-id <机器人QQ号>
|
|
|
|
首次运行会自动下载 Qdrant 与 embedding 模型,需要几分钟。
|
|
"""
|
|
import argparse
|
|
import asyncio
|
|
import re
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
# NoneBot 必须先初始化,wtfllm 的配置才可用
|
|
import nonebot
|
|
|
|
nonebot.init()
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
PAGES_DIR = Path(__file__).parent / "pages"
|
|
VISION_DIR = Path(__file__).parent / "vision_pages"
|
|
|
|
# 每条知识的目标长度(字符),按标题切块后超过再按段落切
|
|
CHUNK_SIZE = 600
|
|
|
|
CATEGORY_BY_DIR = {
|
|
"功能帮助": "功能帮助",
|
|
"脚本进阶": "脚本进阶",
|
|
"简单使用": "简单使用",
|
|
"其它": "其它",
|
|
}
|
|
|
|
|
|
def parse_page(path: Path) -> list[dict]:
|
|
"""把单页文本按标题切成知识条目。"""
|
|
raw = path.read_text(encoding="utf-8")
|
|
lines = raw.splitlines()
|
|
page_title = lines[0].lstrip("# ").replace("\u200b", "") if lines else path.stem
|
|
source = lines[1].lstrip("> 来源: ").strip() if len(lines) > 1 else ""
|
|
|
|
parts: list[dict] = []
|
|
cur_head = page_title
|
|
cur_body: list[str] = []
|
|
|
|
def flush() -> None:
|
|
text = "\n".join(cur_body).strip()
|
|
if text:
|
|
parts.append({"title": cur_head, "content": text})
|
|
|
|
for line in lines[2:]:
|
|
m = re.match(r"^(#{1,4})\s+(.+)$", line)
|
|
if m:
|
|
flush()
|
|
cur_body = []
|
|
h = m.group(2).replace("\u200b", "").strip()
|
|
cur_head = h if h == page_title else f"{page_title} · {h}"
|
|
else:
|
|
cur_body.append(line)
|
|
flush()
|
|
|
|
# 超长条目按段落二次切分
|
|
chunks: list[dict] = []
|
|
for p in parts:
|
|
if len(p["content"]) <= CHUNK_SIZE * 1.5:
|
|
chunks.append(p)
|
|
continue
|
|
paras = p["content"].split("\n\n")
|
|
buf: list[str] = []
|
|
for para in paras:
|
|
if buf and sum(len(b) for b in buf) + len(para) > CHUNK_SIZE:
|
|
chunks.append({"title": p["title"], "content": "\n\n".join(buf)})
|
|
buf = []
|
|
buf.append(para)
|
|
if buf:
|
|
chunks.append({"title": p["title"], "content": "\n\n".join(buf)})
|
|
|
|
# 附上页面来源与分类
|
|
first_dir = path.stem.split("_")[0]
|
|
category = CATEGORY_BY_DIR.get(first_dir, "general")
|
|
for c in chunks:
|
|
c["category"] = category
|
|
c["source"] = source
|
|
c["page_title"] = page_title
|
|
return chunks
|
|
|
|
|
|
async def main(agent_id: str, dry_run: bool, include_vision: bool = False) -> None:
|
|
entries_spec: list[dict] = []
|
|
dirs = [PAGES_DIR] + ([VISION_DIR] if include_vision else [])
|
|
for d in dirs:
|
|
for f in sorted(d.glob("*.txt")):
|
|
entries_spec.extend(parse_page(f))
|
|
print(f"共解析出 {len(entries_spec)} 条知识条目")
|
|
if dry_run:
|
|
for e in entries_spec[:10]:
|
|
print(f" [{e['category']}] {e['title']} ({len(e['content'])} chars)")
|
|
return
|
|
|
|
# 必须先以插件形式加载,让 localstore 能定位数据目录
|
|
nonebot.load_plugin("nonebot_plugin_wtfllm")
|
|
|
|
from nonebot_plugin_wtfllm.memory.items.knowledge_base import KnowledgeEntry
|
|
from nonebot_plugin_wtfllm.v_db.lifecycle import on_startup as vdb_startup
|
|
from nonebot_plugin_wtfllm.v_db.repositories import KnowledgeBaseRepository
|
|
|
|
print("初始化 Qdrant(首次会自动下载)...")
|
|
await vdb_startup()
|
|
|
|
now = int(time.time())
|
|
entries = [
|
|
KnowledgeEntry.create(
|
|
content=e["content"],
|
|
title=e["title"],
|
|
agent_id=agent_id,
|
|
category=e["category"],
|
|
tags=[e["page_title"], "danding.vip"],
|
|
token_count=len(e["content"]),
|
|
)
|
|
for e in entries_spec
|
|
]
|
|
repo = KnowledgeBaseRepository()
|
|
ids = await repo.save_many_knowledge(entries)
|
|
print(f"已写入 {len(ids)} 条知识")
|
|
|
|
# 验证检索
|
|
for q in ["怎么设置定时调度", "积分怎么获得", "下载安装"]:
|
|
results = await repo.search_relevant(agent_id=agent_id, query=q, limit=3)
|
|
print(f"\n检索「{q}」:")
|
|
for r in results:
|
|
print(f" {r.score:.3f} {r.item.title}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--agent-id", required=True, help="机器人 QQ 号(wtfllm 的 agent_id)")
|
|
parser.add_argument("--dry-run", action="store_true", help="只解析不写入")
|
|
parser.add_argument("--include-vision", action="store_true", help="包含截图提取的知识")
|
|
args = parser.parse_args()
|
|
asyncio.run(main(args.agent_id, args.dry_run, args.include_vision))
|