lambdagent/agentexample/qaagent67lite/scripts/setup.py

#!/usr/bin/env python3
"""
qaagent67lite 一键部署脚本

用法:
  # 1. 创建 instance.yml 指定数据目录
  # 2. 将原始文件放入 {baseDir}/raw/
  # 3. 运行:
  python scripts/setup.py

自动执行:
  Step 1: 提取文本 (PDF/Word/PPT → txt)
  Step 2: 建 BM25 索引 (关键词检索)
  Step 3: Wiki 编译 (LLM 知识编译)
  Step 4: 验证
"""
import json
import os
import sys
import time
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
AGENT_DIR = SCRIPT_DIR.parent
PROJECT_ROOT = AGENT_DIR.parent.parent

sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(AGENT_DIR.parent))


def load_config():
    """加载配置 (agent-config.yml + instance.yml)"""
    import yaml

    config = {}
    agent_config = AGENT_DIR / "agent-config.yml"
    if agent_config.exists():
        with open(agent_config) as f:
            config = yaml.safe_load(f) or {}

    # instance.yml 覆盖
    for candidate in [
        os.environ.get("INSTANCE_CONFIG", ""),
        str(SCRIPT_DIR / "instance.yml"),
        str(AGENT_DIR / "instance.yml"),
    ]:
        if candidate and Path(candidate).is_file():
            with open(candidate) as f:
                instance = yaml.safe_load(f) or {}
            # 深度合并
            for k, v in instance.items():
                if k in config and isinstance(config[k], dict) and isinstance(v, dict):
                    config[k].update(v)
                else:
                    config[k] = v
            break

    return config


def setup():
    cfg = load_config()
    knowledge = cfg.get("knowledge", {})
    wiki_cfg = cfg.get("wiki", {})

    base_dir = Path(knowledge.get("baseDir", "./knowledge"))
    raw_dir = Path(knowledge.get("rawDir", str(base_dir / "raw")))
    processed_dir = Path(knowledge.get("processedDir", str(base_dir / "processed")))
    wiki_dir = Path(wiki_cfg.get("dir", "./wiki"))

    print(f"{'='*60}")
    print(f"  qaagent67lite 一键部署")
    print(f"{'='*60}")
    print(f"  知识库: {base_dir}")
    print(f"  原始文件: {raw_dir}")
    print(f"  Wiki: {wiki_dir}")
    print()

    # 检查原始文件
    if not raw_dir.exists():
        print(f"  ❌ 原始文件目录不存在: {raw_dir}")
        print(f"  请先将文件放入 {raw_dir}/")
        return

    raw_files = list(raw_dir.rglob("*"))
    raw_files = [f for f in raw_files if f.is_file() and not f.name.startswith(".")]
    print(f"  📁 发现 {len(raw_files)} 个原始文件")

    if len(raw_files) == 0:
        print(f"  ❌ 无文件可处理")
        return

    if len(raw_files) > 500:
        print(f"  ⚠️  文件数 > 500，建议使用 qaagent67lambda (带向量检索)")

    # Step 1: 提取文本
    print(f"\n  📄 Step 1/3: 提取文本...")
    processed_dir.mkdir(parents=True, exist_ok=True)
    extracted = 0
    for f in raw_files:
        out_path = processed_dir / f"{f.stem}.txt"
        if out_path.exists():
            extracted += 1
            continue
        try:
            if f.suffix.lower() == ".pdf":
                _extract_pdf(f, out_path)
            elif f.suffix.lower() in (".docx", ".doc"):
                _extract_docx(f, out_path)
            elif f.suffix.lower() in (".txt", ".md", ".csv", ".json", ".yaml", ".yml"):
                import shutil
                shutil.copy2(f, out_path)
            else:
                continue
            extracted += 1
        except Exception as e:
            print(f"    ⚠️ {f.name}: {e}")
    print(f"  ✅ 提取完成: {extracted} 个文件")

    # Step 2: BM25 索引
    print(f"\n  🔍 Step 2/3: 建 BM25 索引...")
    t0 = time.time()
    try:
        _build_bm25_index(processed_dir, base_dir, cfg)
        print(f"  ✅ BM25 索引完成 ({time.time()-t0:.1f}s)")
    except Exception as e:
        print(f"  ❌ BM25 索引失败: {e}")

    # Step 3: Wiki 编译
    print(f"\n  📚 Step 3/3: Wiki 编译...")
    t0 = time.time()
    try:
        _wiki_compile(processed_dir, wiki_dir)
        print(f"  ✅ Wiki 编译完成 ({time.time()-t0:.1f}s)")
    except Exception as e:
        print(f"  ⚠️ Wiki 编译失败 (可后续补): {e}")

    # 验证
    print(f"\n{'='*60}")
    print(f"  ✅ 部署完成!")
    print(f"  BM25 索引: {base_dir / 'rag_index.json'}")
    wiki_pages = len(list(wiki_dir.rglob("*.md"))) if wiki_dir.exists() else 0
    print(f"  Wiki 页面: {wiki_pages}")
    print(f"\n  启动: python3 -m agentpaas agent create --name lite-qa \\")
    print(f"          --config agentexample/qaagent67lite/agent-config.yml")
    print(f"        python3 -m agentpaas chat lite-qa")
    print(f"{'='*60}")


def _extract_pdf(pdf_path, out_path):
    """PDF → txt"""
    try:
        import PyPDF2
        with open(pdf_path, "rb") as f:
            reader = PyPDF2.PdfReader(f)
            text = "\n\n".join(
                page.extract_text() or "" for page in reader.pages
            )
        with open(out_path, "w", encoding="utf-8") as f:
            f.write(text)
    except ImportError:
        # 降级到 pdftotext CLI
        import subprocess
        subprocess.run(["pdftotext", str(pdf_path), str(out_path)],
                      capture_output=True, timeout=30)


def _extract_docx(docx_path, out_path):
    """Word → txt"""
    try:
        import docx
        doc = docx.Document(str(docx_path))
        text = "\n\n".join(p.text for p in doc.paragraphs if p.text.strip())
        with open(out_path, "w", encoding="utf-8") as f:
            f.write(text)
    except ImportError:
        raise RuntimeError("需要 python-docx: pip install python-docx")


def _build_bm25_index(processed_dir, base_dir, cfg):
    """构建 BM25 索引 (简化版，不需要 embedding)"""
    import re
    from collections import Counter

    chunk_size = cfg.get("rag", {}).get("chunkSize", 512)
    chunk_overlap = cfg.get("rag", {}).get("chunkOverlap", 64)

    chunks = []
    keywords_index = {}

    for txt_file in sorted(processed_dir.glob("*.txt")):
        with open(txt_file, "r", encoding="utf-8", errors="replace") as f:
            content = f.read()

        if not content.strip():
            continue

        # 分块
        file_chunks = _chunk_text(content, chunk_size, chunk_overlap)
        for i, chunk in enumerate(file_chunks):
            chunk_id = f"{txt_file.stem}_chunk_{i}"
            chunks.append({
                "source": txt_file.stem,
                "chunk_id": chunk_id,
                "text": chunk,
            })

            # 关键词索引
            words = re.findall(r'[\u4e00-\u9fff]+|[a-zA-Z]+', chunk.lower())
            word_counts = Counter(words)
            for word, count in word_counts.items():
                if len(word) >= 2:
                    if word not in keywords_index:
                        keywords_index[word] = []
                    keywords_index[word].append({
                        "chunk_id": chunk_id,
                        "count": count,
                    })

    # 保存
    index_path = base_dir / "rag_index.json"
    with open(index_path, "w", encoding="utf-8") as f:
        json.dump({"chunks": chunks, "total": len(chunks)}, f, ensure_ascii=False)

    kw_path = base_dir / "rag_index.keywords.json"
    with open(kw_path, "w", encoding="utf-8") as f:
        json.dump(keywords_index, f, ensure_ascii=False)

    print(f"    {len(chunks)} 个文本块, {len(keywords_index)} 个关键词")


def _chunk_text(text, size=512, overlap=64):
    """按字符数分块"""
    chunks = []
    start = 0
    while start < len(text):
        end = start + size
        chunk = text[start:end]
        if chunk.strip():
            chunks.append(chunk.strip())
        start = end - overlap
    return chunks


def _wiki_compile(processed_dir, wiki_dir):
    """调用 WikiIngest 编译 wiki"""
    from lambdagent.builtin_tools.wiki_tools import wiki_ingest

    wiki_dir.mkdir(parents=True, exist_ok=True)

    txt_files = sorted(processed_dir.glob("*.txt"))
    total = len(txt_files)
    done = 0

    for txt_file in txt_files:
        try:
            wiki_ingest(json.dumps({
                "path": str(txt_file),
                "wiki_root": str(wiki_dir),
            }))
            done += 1
            if done % 10 == 0:
                print(f"    Wiki: {done}/{total} ({done*100//total}%)")
        except Exception as e:
            print(f"    ⚠️ {txt_file.name}: {e}")

    print(f"    Wiki 编译: {done}/{total} 完成")


if __name__ == "__main__":
    setup()