227 lines
7.8 KiB
Python
227 lines
7.8 KiB
Python
import json
|
|
import logging
|
|
import re
|
|
import threading
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
import chromadb
|
|
from chromadb.config import Settings
|
|
|
|
from app.chunker import chunk_text
|
|
from app.config import COLLECTION_NAME, DATA_DIR, EMBEDDING_MODEL
|
|
from app.embeddings import Embedder
|
|
from app.seed import SAMPLES
|
|
|
|
logger = logging.getLogger("learn_rag")
|
|
|
|
|
|
class DisabledEmbedding:
|
|
"""占位嵌入函数。向量全部由 fastembed 预先算好，避免 Chroma 再下载一套模型。"""
|
|
|
|
@staticmethod
|
|
def name() -> str:
|
|
return "precomputed"
|
|
|
|
def __call__(self, input):
|
|
raise RuntimeError("向量由 fastembed 预先计算")
|
|
|
|
def embed_documents(self, input):
|
|
raise RuntimeError("向量由 fastembed 预先计算")
|
|
|
|
def embed_query(self, input):
|
|
raise RuntimeError("向量由 fastembed 预先计算")
|
|
|
|
|
|
def to_score(distance: float) -> float:
|
|
if distance <= 0:
|
|
return 1.0
|
|
if distance <= 1:
|
|
return 1 - distance
|
|
return 1 / (1 + distance)
|
|
|
|
|
|
class KnowledgeBase:
|
|
def __init__(self, embedder: Embedder) -> None:
|
|
self.embedder = embedder
|
|
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
|
self.path = DATA_DIR / "documents.json"
|
|
self._lock = threading.Lock()
|
|
self.client = chromadb.PersistentClient(
|
|
path=str(DATA_DIR / "chroma"),
|
|
settings=Settings(anonymized_telemetry=False),
|
|
)
|
|
self.collection = self._open_collection()
|
|
self.docs: list[dict] = []
|
|
|
|
def _open_collection(self):
|
|
return self.client.get_or_create_collection(
|
|
name=COLLECTION_NAME,
|
|
embedding_function=DisabledEmbedding(),
|
|
metadata={"hnsw:space": "cosine"},
|
|
)
|
|
|
|
def _reset_collection(self) -> None:
|
|
try:
|
|
self.client.delete_collection(COLLECTION_NAME)
|
|
except Exception:
|
|
logger.info("向量集合尚不存在，直接创建")
|
|
self.collection = self._open_collection()
|
|
|
|
def startup(self) -> None:
|
|
needs_seed = not self.path.exists()
|
|
payload = {"model": EMBEDDING_MODEL, "documents": []}
|
|
if self.path.exists():
|
|
payload = json.loads(self.path.read_text(encoding="utf-8"))
|
|
self.docs = payload.get("documents", [])
|
|
if needs_seed:
|
|
logger.info("写入示例知识库：云栖书店")
|
|
self._reset_collection()
|
|
for sample in SAMPLES:
|
|
self._index(sample["title"], sample["content"])
|
|
self._save()
|
|
return
|
|
if payload.get("model") != EMBEDDING_MODEL:
|
|
logger.info("嵌入模型已更换，重新构建向量库")
|
|
previous = list(self.docs)
|
|
self.docs = []
|
|
self._reset_collection()
|
|
for doc in previous:
|
|
self._index(doc["title"], doc["content"], doc["id"], doc["created_at"])
|
|
self._save()
|
|
return
|
|
self._reconcile()
|
|
|
|
def _save(self) -> None:
|
|
payload = {"model": EMBEDDING_MODEL, "documents": self.docs}
|
|
temporary = self.path.with_suffix(".json.tmp")
|
|
temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
temporary.replace(self.path)
|
|
|
|
def _summary(self, doc: dict) -> dict:
|
|
preview = re.sub(r"\s+", " ", doc["content"]).strip()[:80]
|
|
return {
|
|
"id": doc["id"],
|
|
"title": doc["title"],
|
|
"chunk_count": doc["chunk_count"],
|
|
"created_at": doc["created_at"],
|
|
"preview": preview,
|
|
}
|
|
|
|
def _index(
|
|
self,
|
|
title: str,
|
|
content: str,
|
|
doc_id: str | None = None,
|
|
created_at: str | None = None,
|
|
) -> dict:
|
|
chunks = chunk_text(content)
|
|
if not chunks:
|
|
raise ValueError("正文无法切分")
|
|
doc_id = doc_id or uuid.uuid4().hex
|
|
created_at = created_at or datetime.now(timezone.utc).isoformat(timespec="seconds")
|
|
vectors = self.embedder.embed_documents(chunks)
|
|
ids = [f"{doc_id}:{index}" for index in range(len(chunks))]
|
|
metadatas = [
|
|
{
|
|
"doc_id": doc_id,
|
|
"title": title,
|
|
"chunk_index": index,
|
|
"chunk_total": len(chunks),
|
|
}
|
|
for index in range(len(chunks))
|
|
]
|
|
self.collection.add(
|
|
ids=ids,
|
|
documents=chunks,
|
|
embeddings=vectors,
|
|
metadatas=metadatas,
|
|
)
|
|
record = {
|
|
"id": doc_id,
|
|
"title": title,
|
|
"content": content,
|
|
"created_at": created_at,
|
|
"chunk_count": len(chunks),
|
|
}
|
|
self.docs.append(record)
|
|
return record
|
|
|
|
def _reconcile(self) -> None:
|
|
stored = self.collection.get(include=["metadatas"])
|
|
present = {
|
|
meta["doc_id"]
|
|
for meta in (stored.get("metadatas") or [])
|
|
if meta and meta.get("doc_id")
|
|
}
|
|
wanted = {doc["id"] for doc in self.docs}
|
|
for doc_id in present - wanted:
|
|
self.collection.delete(where={"doc_id": doc_id})
|
|
for doc in self.docs:
|
|
if doc["id"] not in present:
|
|
logger.info("重新索引《%s》", doc["title"])
|
|
self.docs = [item for item in self.docs if item["id"] != doc["id"]]
|
|
self._index(doc["title"], doc["content"], doc["id"], doc["created_at"])
|
|
self._save()
|
|
|
|
def list_documents(self) -> list[dict]:
|
|
with self._lock:
|
|
return [self._summary(doc) for doc in self.docs]
|
|
|
|
def counts(self) -> tuple[int, int]:
|
|
with self._lock:
|
|
return len(self.docs), int(self.collection.count())
|
|
|
|
def add(self, title: str, content: str) -> dict:
|
|
with self._lock:
|
|
record = self._index(title, content)
|
|
try:
|
|
self._save()
|
|
except Exception:
|
|
self.docs.pop()
|
|
self.collection.delete(where={"doc_id": record["id"]})
|
|
raise
|
|
logger.info("新增《%s》，%s 个切片", title, record["chunk_count"])
|
|
return self._summary(record)
|
|
|
|
def delete(self, doc_id: str) -> bool:
|
|
with self._lock:
|
|
if not any(doc["id"] == doc_id for doc in self.docs):
|
|
return False
|
|
self.collection.delete(where={"doc_id": doc_id})
|
|
self.docs = [doc for doc in self.docs if doc["id"] != doc_id]
|
|
self._save()
|
|
return True
|
|
|
|
def search(self, question: str, top_k: int) -> list[dict]:
|
|
with self._lock:
|
|
count = int(self.collection.count())
|
|
if count == 0:
|
|
return []
|
|
vector = self.embedder.embed_query(question)
|
|
found = self.collection.query(
|
|
query_embeddings=[vector],
|
|
n_results=min(top_k, count),
|
|
include=["documents", "metadatas", "distances"],
|
|
)
|
|
documents = (found.get("documents") or [[]])[0] or []
|
|
metadatas = (found.get("metadatas") or [[]])[0] or []
|
|
distances = (found.get("distances") or [[]])[0] or []
|
|
chunks = []
|
|
for text, meta, distance in zip(documents, metadatas, distances):
|
|
meta = meta or {}
|
|
chunks.append(
|
|
{
|
|
"id": f"{meta.get('doc_id', '')}:{meta.get('chunk_index', 0)}",
|
|
"title": meta.get("title") or "未命名资料",
|
|
"text": text or "",
|
|
"score": round(to_score(float(distance)), 4),
|
|
"chunk_index": int(meta.get("chunk_index") or 0) + 1,
|
|
"chunk_total": int(meta.get("chunk_total") or 1),
|
|
}
|
|
)
|
|
chunks.sort(key=lambda item: item["score"], reverse=True)
|
|
for number, chunk in enumerate(chunks, start=1):
|
|
chunk["n"] = number
|
|
return chunks
|