|
|
@@ -0,0 +1,95 @@
|
|
|
+# -*- 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())
|