Serendipity
832201186d
H1: AppConfig 自定义 __init__ 改用 from_dict() 类方法 H2: list_collections_with_stats 改为从 ChromaDB 直接查询 H3: RateLimiter 过期 key 自动清理, 防止内存泄漏 M1: delete_by_source 区分 ValueError 与真实异常, 记日志 M2: VectorDB 新增 list_collections() 封装方法 M3: _remove_by_source 异常记日志, 不再静默吞掉 M4: CLI 集合回退支持 MD_VECTOR_DB_COLLECTION 环境变量 L1: DEFAULT_CONFIG_PATH 自动从项目根目录解析 L2: ingest 命令内重复 import 移至模块顶部 L3: 新增死循环回归测试 + 密集分隔符分块测试 L4: 统一 logger 名称为 md-vector-db 删除旧版审计文档 测试: 48 passed
95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
"""语义检索模块."""
|
|
import logging
|
|
from src.core.db import VectorDB
|
|
from src.core.embedder import Embedder
|
|
|
|
logger = logging.getLogger("md-vector-db")
|
|
|
|
|
|
class Searcher:
|
|
"""向量检索器."""
|
|
|
|
def __init__(self, db: VectorDB, embedder: Embedder, collection_name: str):
|
|
self.db = db
|
|
self.embedder = embedder
|
|
self.collection_name = collection_name
|
|
|
|
@property
|
|
def collection(self):
|
|
return self.db.get_or_create_collection(self.collection_name)
|
|
|
|
def search(
|
|
self,
|
|
query: str,
|
|
top_k: int = 10,
|
|
source_file: str | None = None,
|
|
) -> list[dict]:
|
|
"""语义检索, 返回格式化结果列表."""
|
|
query_embedding = self.embedder.embed([query])[0]
|
|
|
|
where_filter = None
|
|
if source_file:
|
|
where_filter = {"source_file": source_file}
|
|
|
|
results = self.collection.query(
|
|
query_embeddings=[query_embedding],
|
|
n_results=top_k,
|
|
where=where_filter,
|
|
include=["documents", "metadatas", "distances"],
|
|
)
|
|
|
|
formatted = []
|
|
if results["ids"] and results["ids"][0]:
|
|
for i, doc_id in enumerate(results["ids"][0]):
|
|
metadata = results["metadatas"][0][i] if results["metadatas"] else {}
|
|
distance = results["distances"][0][i] if results["distances"] else 0.0
|
|
# ChromaDB 默认用余弦距离, 转为相似度分数 (0~1)
|
|
score = round(1.0 - distance, 4)
|
|
|
|
formatted.append({
|
|
"id": doc_id,
|
|
"content": results["documents"][0][i] if results["documents"] else "",
|
|
"source_file": metadata.get("source_file", ""),
|
|
"section_title": metadata.get("section_title", ""),
|
|
"heading_level": metadata.get("heading_level", 0),
|
|
"chunk_index": metadata.get("chunk_index", 0),
|
|
"score": max(0.0, score),
|
|
})
|
|
|
|
return formatted
|
|
|
|
def get_collection_info(self) -> dict:
|
|
"""获取 collection 信息."""
|
|
return {
|
|
"name": self.collection_name,
|
|
"count": self.collection.count(),
|
|
}
|
|
|
|
def list_sources(self) -> list[str]:
|
|
"""列出所有已入库的源文件."""
|
|
if self.collection.count() == 0:
|
|
return []
|
|
result = self.collection.get(include=["metadatas"])
|
|
sources = set()
|
|
if result and result["metadatas"]:
|
|
for m in result["metadatas"]:
|
|
if m and "source_file" in m:
|
|
sources.add(m["source_file"])
|
|
return sorted(sources)
|
|
|
|
def delete_by_source(self, file_name: str) -> bool:
|
|
"""按文件名删除文档 (线程安全)."""
|
|
try:
|
|
with self.db.write_lock:
|
|
existing = self.collection.get(
|
|
where={"source_file": file_name}
|
|
)
|
|
if existing and existing["ids"]:
|
|
self.collection.delete(ids=existing["ids"])
|
|
return True
|
|
except ValueError:
|
|
pass # collection 为空时 ChromaDB 抛 ValueError
|
|
except Exception:
|
|
logger.exception("删除文档失败: %s", file_name)
|
|
return False
|