feat: add semantic search module

实现语义检索模块 Searcher 类,支持:
- search(): 向量语义检索,返回格式化结果(含 content、source_file、section_title、heading_level、chunk_index、score 字段)
- get_collection_info(): 获取 collection 统计信息
- list_sources(): 列出已入库源文件
- delete_by_source(): 按源文件删除文档

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
2026-07-05 01:06:08 +08:00
parent 757bd8b336
commit 7309b0a437
2 changed files with 164 additions and 0 deletions
+88
View File
@@ -0,0 +1,88 @@
"""语义检索模块."""
from src.core.db import VectorDB
from src.core.embedder import Embedder
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:
existing = self.collection.get(
where={"source_file": file_name}
)
if existing and existing["ids"]:
self.collection.delete(ids=existing["ids"])
return True
except Exception:
pass
return False