118 lines
3.8 KiB
Python
118 lines
3.8 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
import chromadb
|
|
from chromadb.utils import embedding_functions
|
|
|
|
basedir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
|
def _md5(text: str) -> str:
|
|
return hashlib.md5(text.encode("utf-8")).hexdigest()
|
|
|
|
|
|
class RAGStore:
|
|
"""RAG vector store for content moderation examples with MD5 dedup."""
|
|
|
|
def __init__(self, persist_dir=None):
|
|
persist_dir = persist_dir or os.path.join(basedir, "chroma_db")
|
|
self.client = chromadb.PersistentClient(path=persist_dir)
|
|
self.ef = embedding_functions.DefaultEmbeddingFunction()
|
|
|
|
existing = [c.name for c in self.client.list_collections()]
|
|
if "moderation_examples" in existing:
|
|
self.collection = self.client.get_collection(
|
|
name="moderation_examples",
|
|
embedding_function=self.ef,
|
|
)
|
|
else:
|
|
self.collection = self.client.create_collection(
|
|
name="moderation_examples",
|
|
embedding_function=self.ef,
|
|
)
|
|
# 无论首次还是后续启动，都增量同步 knowledge_base.json
|
|
self._sync_knowledge_base()
|
|
|
|
def _md5_exists(self, md5_hash: str) -> bool:
|
|
result = self.collection.get(where={"md5": md5_hash})
|
|
return len(result["ids"]) > 0
|
|
|
|
def _sync_knowledge_base(self):
|
|
"""增量同步：只加入 knowledge_base.json 中尚未入库的条目。"""
|
|
kb_path = os.path.join(basedir, "knowledge_base.json")
|
|
if not os.path.exists(kb_path):
|
|
return
|
|
|
|
with open(kb_path, "r", encoding="utf-8") as f:
|
|
entries = json.load(f)
|
|
|
|
new_entries = []
|
|
for e in entries:
|
|
h = _md5(e["text"])
|
|
if not self._md5_exists(h):
|
|
new_entries.append(e)
|
|
|
|
if new_entries:
|
|
self.collection.add(
|
|
ids=[e["id"] for e in new_entries],
|
|
documents=[e["text"] for e in new_entries],
|
|
metadatas=[
|
|
{
|
|
"type": e["type"],
|
|
"category": e["category"],
|
|
"md5": _md5(e["text"]),
|
|
}
|
|
for e in new_entries
|
|
],
|
|
)
|
|
|
|
def query(self, text: str, n_results: int = 3) -> list[dict]:
|
|
results = self.collection.query(
|
|
query_texts=[text],
|
|
n_results=n_results,
|
|
)
|
|
output = []
|
|
ids = results["ids"][0]
|
|
docs = results["documents"][0]
|
|
metas = results["metadatas"][0]
|
|
distances = results["distances"][0]
|
|
|
|
for i, doc, meta, dist in zip(ids, docs, metas, distances):
|
|
output.append({
|
|
"id": i,
|
|
"text": doc,
|
|
"type": meta.get("type", ""),
|
|
"category": meta.get("category", ""),
|
|
"distance": dist,
|
|
})
|
|
return output
|
|
|
|
def add_example(self, entry_id: str, text: str, entry_type: str, category: str):
|
|
"""运行时动态添加，MD5 去重。"""
|
|
h = _md5(text)
|
|
if self._md5_exists(h):
|
|
return False
|
|
|
|
self.collection.add(
|
|
ids=[entry_id],
|
|
documents=[text],
|
|
metadatas=[{"type": entry_type, "category": category, "md5": h}],
|
|
)
|
|
return True
|
|
|
|
def list_all(self) -> list[dict]:
|
|
"""返回 RAG 库中所有条目。"""
|
|
if self.collection.count() == 0:
|
|
return []
|
|
|
|
result = self.collection.get()
|
|
output = []
|
|
for i, doc, meta in zip(result["ids"], result["documents"], result["metadatas"]):
|
|
output.append({
|
|
"id": i,
|
|
"text": doc,
|
|
"type": meta.get("type", ""),
|
|
"category": meta.get("category", ""),
|
|
})
|
|
return output
|