cfc/docs/scripts/kb_check_host.py

# -*- coding: utf-8 -*-
"""宿主机会话内批量自检：逐句对知识库做 cosine 检索（复刻 verify 逻辑）。

- 数据：/tmp/cv4（生产 chroma 一致性副本，生产文件零接触）
- embedding：生产同款 BAAI/bge-large-zh-v1.5（siliconflow）
- 首次 query 需加载 HNSW 索引（约 80s），后续每句约 1s
用法：cd /app/cfc/cfc-langgraph && .venv/bin/python /tmp/kb_check_host.py <sentences.json> [out.json]
"""
import json
import math
import sys
import time

from langchain_chroma import Chroma
from app.rag.embeddings import get_embeddings

CHROMA_PATH = "/tmp/cv4"
SOURCE = "microbiome"
K = 3
THRESHOLD = 0.7


def _cosine(a, b) -> float:
    a = list(a)
    b = list(b)
    if not a or not b or len(a) != len(b):
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(float(x) * float(x) for x in a))
    nb = math.sqrt(sum(float(y) * float(y) for y in b))
    if na == 0.0 or nb == 0.0:
        return 0.0
    return float(dot) / (float(na) * float(nb))


def main():
    in_path = sys.argv[1]
    out_path = sys.argv[2] if len(sys.argv) > 2 else "/tmp/kb_check_result.json"
    with open(in_path, encoding="utf-8") as f:
        data = json.load(f)
    sentences = data["sentences"] if isinstance(data, dict) else data

    t0 = time.time()
    store = Chroma(collection_name="cfc_knowledge",
                   embedding_function=get_embeddings(),
                   persist_directory=CHROMA_PATH)
    embedder = get_embeddings()
    print("store ready", round(time.time() - t0, 1), "s", flush=True)

    results = []
    all_passed = True
    for i, s in enumerate(sentences, 1):
        text = (s or "").strip()
        best = {"max_similarity": 0.0, "matched_title": None, "matched_chunk": None}
        if text:
            try:
                cands = store.similarity_search_with_score(text, k=K,
                                                          filter={"source": SOURCE})
                query_vec = embedder.embed_query(text)
                for doc, _dist in cands:
                    meta = doc.metadata or {}
                    doc_id = meta.get("id")
                    if not doc_id:
                        continue
                    got = store._collection.get(ids=[doc_id], include=["embeddings"])
                    vecs = got.get("embeddings")
                    if vecs is None or len(vecs) == 0:
                        continue
                    sim = _cosine(query_vec, vecs[0])
                    if sim > best["max_similarity"]:
                        best = {"max_similarity": sim,
                                "matched_title": meta.get("title"),
                                "matched_chunk": (doc.page_content[:120]
                                                  if doc.page_content else None)}
            except Exception as e:
                print(f"  !! sentence {i} error: {e}", flush=True)
        best["id"] = f"s{i}"
        best["text"] = text
        best["passed"] = best["max_similarity"] >= THRESHOLD
        results.append(best)
        if not best["passed"]:
            all_passed = False
        flag = "OK  " if best["passed"] else "FAIL"
        print(f"[{i:2d}/{len(sentences)}] {flag} sim={best['max_similarity']:.4f} "
              f"| {best['matched_title']} | {text[:40]}", flush=True)

    print(f"\nall_passed={all_passed}  elapsed={round(time.time()-t0,1)}s")
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump({"results": results, "all_passed": all_passed}, f,
                  ensure_ascii=False, indent=2)
    return 0 if all_passed else 1


if __name__ == "__main__":
    sys.exit(main())