Init: 第一个版本
This commit is contained in:
+226
@@ -0,0 +1,226 @@
|
||||
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
|
||||
Reference in New Issue
Block a user