PlantAI-rag-new/tooling/aqua_kb/qdrant_loader.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""tooling/aqua_kb/qdrant_loader.py —— Qdrant 批量灌库器

读 phase2b_converter.py 产出的 qdrant_points.jsonl（一行一个 point：id/text/payload），
用与现网 direct_flow_retriever.embed() 同款端点（POST {embed_url} {"model":..., "input":...}）
嵌入 text 字段，分批 upsert 到 Qdrant collection。

特性：
  - 批量嵌入：input 传数组（--embed-batch 条/批），一次请求嵌入一批文本，
    再按批 upsert（实测单条约 2.7 texts/s → 批量约 37 texts/s，约 14x 吞吐）
  - checkpoint JSONL：每成功 upsert 一个 assertion_id 记一行 → 断点续跑（kill 后重跑跳过已传）
  - 同 id 重放覆盖（Qdrant upsert 语义 + UUID 确定性 id → 幂等）
  - 确定性处理顺序（文件路径 + 文件内行序排序）

用法：
  python qdrant_loader.py --points <dir|file> --qdrant http://127.0.0.1:27435 \
      --collection aqua_srt_v1_20260930_normalized --batch 200 \
      --checkpoint <file.jsonl> \
      [--embed-url http://127.0.0.1:11434/api/embed] [--embed-model bge-m3:567m] \
      [--embed-batch 32]
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path
from typing import Any, Dict, Iterator, List

import requests

DEFAULT_EMBED_URL = "http://127.0.0.1:11434/api/embed"
DEFAULT_EMBED_MODEL = "bge-m3:567m"
EMBED_INPUT_LIMIT = 2400  # 与 direct_flow_retriever.embed() 一致


def iter_points(paths: List[Path]) -> Iterator[Dict[str, Any]]:
    files = sorted(
        p for p in paths if p.is_file() and p.name == "qdrant_points.jsonl"
    )
    for file_path in files:
        with open(file_path, "r", encoding="utf-8") as handle:
            for line_number, line in enumerate(handle, 1):
                line = line.strip()
                if not line:
                    continue
                try:
                    record = json.loads(line)
                except json.JSONDecodeError as exc:
                    print(f"[loader] WARN bad json {file_path}:{line_number}: {exc}",
                          file=sys.stderr)
                    continue
                yield record


class EmbedClient:
    def __init__(self, url: str, model: str) -> None:
        self.url = url.rstrip("/")
        self.model = model
        self.dim: int | None = None

    @staticmethod
    def _extract_vector(item: Any) -> List[float] | None:
        """从 embedding 数组元素中提取向量（兼容 embedding 对象与裸数组两种形态）。"""
        if isinstance(item, dict):
            item = item.get("embedding")
        if not isinstance(item, (list, tuple)):
            return None
        try:
            return [float(value) for value in item]
        except (TypeError, ValueError):
            return None

    def embed(self, text: str) -> List[float]:
        """单条嵌入（预热 probe / 失败回退路径）。"""
        vector = self.embed_batch([text])
        if not vector:
            raise RuntimeError("embedding endpoint returned no vector")
        return vector[0]

    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """批量嵌入：input 传数组，一次请求嵌入多条（约 60 texts/s）。"""
        if not texts:
            return []
        payload = {
            "model": self.model,
            "input": [t[:EMBED_INPUT_LIMIT] for t in texts],
        }
        for attempt in range(4):
            try:
                response = requests.post(self.url, json=payload, timeout=120)
                response.raise_for_status()
                data = response.json()
                raw = data.get("embeddings")
                if not isinstance(raw, list) or not raw:
                    raise RuntimeError("embedding endpoint returned no vectors")
                vectors = [self._extract_vector(item) for item in raw]
                if any(v is None for v in vectors):
                    raise RuntimeError(
                        f"embedding endpoint returned {len(vectors)} vectors "
                        f"for {len(texts)} inputs (bad shapes)"
                    )
                if len(vectors) != len(texts):
                    raise RuntimeError(
                        f"embedding endpoint returned {len(vectors)} vectors "
                        f"for {len(texts)} inputs"
                    )
                self.dim = len(vectors[0])
                return vectors
            except Exception as exc:
                if attempt == 3:
                    raise
                time.sleep(2 * (attempt + 1))
        raise RuntimeError("unreachable")


def ensure_collection(qdrant_url: str, collection: str, dim: int) -> None:
    info = requests.get(
        f"{qdrant_url.rstrip('/')}/collections/{collection}", timeout=15
    )
    if info.status_code == 200:
        return
    if info.status_code != 404:
        info.raise_for_status()
    create = requests.put(
        f"{qdrant_url.rstrip('/')}/collections/{collection}",
        json={
            "vectors": {"size": dim, "distance": "Cosine"},
            "shard_number": 1,
            "replication_factor": 1,
            "write_consistency_factor": 1,
            "on_disk_payload": True,
        },
        timeout=30,
    )
    create.raise_for_status()
    print(f"[loader] created collection {collection} (dim={dim}, COSINE)")


def load_checkpoint(checkpoint: Path) -> set:
    done: set = set()
    if checkpoint.exists():
        with open(checkpoint, "r", encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if not line:
                    continue
                try:
                    done.add(json.loads(line))  # 兼容 json.dumps 写入的带引号行
                except json.JSONDecodeError:
                    done.add(line)
    return done


def embed_group(
    embed: "EmbedClient", records: List[Dict[str, Any]], embed_batch: int
) -> List[List[float] | None]:
    """按 embed_batch 条/批批量嵌入 records；单条失败时按单条回退重试，仍失败则该条置 None。"""
    vectors: List[List[float] | None] = [None] * len(records)
    for start in range(0, len(records), embed_batch):
        group = records[start:start + embed_batch]
        try:
            vectors[start:start + embed_batch] = embed.embed_batch(
                [r.get("text") or "" for r in group]
            )
        except Exception as exc:
            print(f"[loader] WARN embed batch failed ({exc}); retrying singly",
                  file=sys.stderr)
            for i, record in enumerate(group):
                idx = start + i
                try:
                    vectors[idx] = embed.embed(record.get("text") or "")
                except Exception:
                    vectors[idx] = None
    return vectors


def upsert_file(
    points_path: Path,
    qdrant_url: str,
    collection: str,
    batch_size: int,
    *,
    checkpoint: Path,
    embed: "EmbedClient",
    embed_batch: int = 32,
) -> Dict[str, int]:
    """上传单个 qdrant_points.jsonl（断点续跑，checkpoint 追加）。

    嵌入按 embed_batch 条/批批量调用 /api/embed，再攒 batch_size 条 upsert；
    批内单条失败时逐条回退重试。checkpoint 记录口径不变（assertion_id）。
    """
    checkpoint.parent.mkdir(parents=True, exist_ok=True)
    done = load_checkpoint(checkpoint)
    qdrant = qdrant_url.rstrip("/")
    upsert_url = f"{qdrant}/collections/{collection}/points?wait=true"

    total = skipped = uploaded = 0
    batch: List[Dict[str, Any]] = []
    pending: List[Dict[str, Any]] = []  # 待嵌入的未 checkpoint 记录
    started = time.time()
    with open(checkpoint, "a", encoding="utf-8") as checkpoint_handle:
        def flush_batch() -> None:
            nonlocal batch
            if not batch:
                return
            response = requests.put(upsert_url, json={"points": batch}, timeout=180)
            response.raise_for_status()
            for record in batch:
                checkpoint_handle.write(
                    json.dumps(record["payload"]["assertion_id"]) + "\n"
                )
                checkpoint_handle.flush()
            batch = []

        def flush_pending() -> None:
            nonlocal pending, uploaded
            if not pending:
                return
            vectors = embed_group(embed, pending, embed_batch)
            for record, vector in zip(pending, vectors):
                if vector is None:
                    raise RuntimeError(
                        f"embed failed for assertion "
                        f"{record['payload']['assertion_id']} "
                        f"(batch and single retries exhausted)"
                    )
                record["vector"] = vector
                batch.append(record)
                uploaded += 1
                if len(batch) >= batch_size:
                    flush_batch()
            pending = []

        for record in iter_points([points_path]):
            total += 1
            assertion_id = record["payload"]["assertion_id"]
            if assertion_id in done:
                skipped += 1
                continue
            pending.append(record)
            if len(pending) >= embed_batch:
                flush_pending()
        flush_pending()
        flush_batch()
    return {
        "uploaded": uploaded,
        "skipped": skipped,
        "total": total,
        "elapsed": round(time.time() - started, 1),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Qdrant bulk loader for aqua KB")
    parser.add_argument("--points", required=True, help="dir or qdrant_points.jsonl file")
    parser.add_argument("--qdrant", required=True, help="e.g. http://127.0.0.1:27435")
    parser.add_argument("--collection", default="aqua_srt_v1_20260930_normalized")
    parser.add_argument("--batch", type=int, default=200)
    parser.add_argument("--checkpoint", required=True, help="JSONL of uploaded assertion_ids")
    parser.add_argument("--embed-url", default=DEFAULT_EMBED_URL)
    parser.add_argument("--embed-model", default=DEFAULT_EMBED_MODEL)
    parser.add_argument("--embed-batch", type=int, default=32,
                        help="texts per /api/embed request (batch embedding)")
    args = parser.parse_args()

    points_arg = Path(args.points)
    if points_arg.is_dir():
        files = sorted(points_arg.rglob("qdrant_points.jsonl"))
    else:
        files = [points_arg]
    checkpoint = Path(args.checkpoint)

    embed = EmbedClient(args.embed_url, args.embed_model)
    qdrant = args.qdrant.rstrip("/")

    # 先预热一发拿到维度（与现网 bge-m3 一致；collection 不存在则按此维度创建）
    probe = embed.embed("aquaculture warm-up probe")
    print(f"[loader] embed dim={len(probe)} (model={args.embed_model}, url={args.embed_url})")
    ensure_collection(qdrant, args.collection, len(probe))

    grand = {"uploaded": 0, "skipped": 0, "total": 0}
    started = time.time()
    for points_path in files:
        result = upsert_file(
            points_path,
            qdrant,
            args.collection,
            args.batch,
            checkpoint=checkpoint,
            embed=embed,
            embed_batch=args.embed_batch,
        )
        for key in grand:
            grand[key] += result[key]
        print(
            f"[loader] {points_path.name}: {result['uploaded']} uploaded, "
            f"{result['skipped']} skipped, {result['total']} seen, "
            f"{result['elapsed']}s"
        )

    print(
        f"[loader] DONE: {grand['uploaded']} uploaded, {grand['skipped']} skipped, "
        f"{grand['total']} seen in {time.time() - started:.1f}s"
    )


if __name__ == "__main__":
    main()