xianjing-ai/scripts/bootstrap_ragflow.py

#!/usr/bin/env python3
"""注册 Ragflow 账号、建库、灌入 knowledge/corpus，并写回仓库 .env。"""

from __future__ import annotations

import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from engine.app.ragflow_client import RagflowClient, RagflowError  # noqa: E402
from engine.app.settings import settings  # noqa: E402

CORPUS = {
    "products": ["products-a.md", "products-b.md"],
    "solutions": ["solutions.md"],
    "faqs": ["faq.md", "policies.md"],
}
DESCRIPTIONS = {
    "products": "鲜境示例产品资料：A 类传感、B 类控制",
    "solutions": "鲜境示例方案判定说明",
    "faqs": "鲜境示例询盘与对外边界",
}


def wait_healthy(client: RagflowClient, timeout: int = 180) -> None:
    deadline = time.time() + timeout
    while time.time() < deadline:
        health = client.health()
        if health.get("ok"):
            print("ragflow health ok", health)
            return
        print("waiting ragflow", health)
        time.sleep(5)
    raise SystemExit("Ragflow 未在时限内就绪")


def auth_or_register(client: RagflowClient) -> str:
    try:
        return client.register(settings.ragflow_email, settings.ragflow_password, "鲜境管理员")
    except RagflowError as exc:
        text = str(exc)
        if "already" in text.lower() or "已经" in text or "registered" in text.lower():
            return client.login(settings.ragflow_email, settings.ragflow_password)
        raise


def upsert_datasets(client: RagflowClient) -> dict[str, str]:
    existing = {item.get("name"): item for item in client.list_datasets()}
    ids: dict[str, str] = {}
    for name, description in DESCRIPTIONS.items():
        found = existing.get(name)
        if found:
            ids[name] = found["id"]
            print("reuse dataset", name, found["id"])
            continue
        created = client.create_dataset(name, description)
        dataset_id = created.get("id")
        if not dataset_id:
            raise SystemExit(f"创建数据集失败: {created}")
        ids[name] = dataset_id
        print("created dataset", name, dataset_id)
    embed = {
        "embedding_model": "BAAI/bge-small-zh-v1.5@Builtin",
        "parser_config": {
            "graphrag": {"use_graphrag": False},
            "raptor": {"use_raptor": False},
        },
    }
    for dataset_id in ids.values():
        client.update_dataset(dataset_id, embed)
    return ids


def ingest(client: RagflowClient, dataset_ids: dict[str, str]) -> None:
    corpus_dir = ROOT / "knowledge" / "corpus"
    for name, files in CORPUS.items():
        dataset_id = dataset_ids[name]
        docs = client.list_documents(dataset_id)
        present = {item.get("name") or item.get("location") for item in docs}
        new_ids = []
        for filename in files:
            if filename in present:
                print("skip existing", filename)
                continue
            uploaded = client.upload_document(dataset_id, corpus_dir / filename)
            for item in uploaded:
                doc_id = item.get("id")
                if doc_id:
                    new_ids.append(doc_id)
            print("uploaded", filename, new_ids[-1] if new_ids else uploaded)
        if new_ids:
            client.parse_documents(dataset_id, new_ids)
            print("parse queued", name, new_ids)


def wait_chunks(client: RagflowClient, dataset_ids: dict[str, str], timeout: int = 300) -> None:
    deadline = time.time() + timeout
    while time.time() < deadline:
        pending = 0
        for dataset_id in dataset_ids.values():
            for doc in client.list_documents(dataset_id):
                run = str(doc.get("run") or "")
                chunks = int(doc.get("chunk_count") or 0)
                if chunks > 0 or run.upper() in {"DONE", "3", "SUCCESS"}:
                    continue
                if run.upper() in {"FAIL", "4", "FAILED"}:
                    raise SystemExit(f"解析失败 {doc.get('name')}: {(doc.get('progress_msg') or '')[-300:]}")
                pending += 1
                print("parsing", doc.get("name"), "run=", run, "chunks=", chunks)
        if pending == 0:
            print("all documents parsed")
            return
        time.sleep(6)
    print("warn: parse still running, continue anyway")


def write_env(dataset_ids: dict[str, str], api_key: str) -> None:
    env_path = ROOT / ".env"
    lines = [
        "RETRIEVAL_BACKEND=ragflow",
        "RAGFLOW_BASE_URL=http://127.0.0.1:9380",
        f"RAGFLOW_API_KEY={api_key}",
        "RAGFLOW_DATASET_IDS=" + ",".join(dataset_ids[name] for name in ("products", "solutions", "faqs")),
        "RAGFLOW_WEB_URL=http://127.0.0.1:18080",
        f"RAGFLOW_EMAIL={settings.ragflow_email}",
        f"RAGFLOW_PASSWORD={settings.ragflow_password}",
        "",
    ]
    env_path.write_text("\n".join(lines), encoding="utf-8")
    print("wrote", env_path)


def main() -> None:
    client = RagflowClient(api_key="")
    wait_healthy(client)
    session_auth = auth_or_register(client)
    token = client.ensure_api_token(session_auth)
    client.api_key = token
    dataset_ids = upsert_datasets(client)
    ingest(client, dataset_ids)
    wait_chunks(client, dataset_ids)
    write_env(dataset_ids, token)
    print("bootstrap done")


if __name__ == "__main__":
    main()