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:
@@ -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
|
||||||
@@ -0,0 +1,76 @@
|
|||||||
|
"""检索模块测试."""
|
||||||
|
import tempfile
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from src.core.config import EmbedConfig
|
||||||
|
from src.core.db import VectorDB
|
||||||
|
from src.core.embedder import create_embedder
|
||||||
|
from src.core.ingest import DocumentIngestor
|
||||||
|
from src.core.search import Searcher
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def searcher():
|
||||||
|
"""创建带测试数据的 Searcher."""
|
||||||
|
tmpdir = tempfile.mkdtemp()
|
||||||
|
db = VectorDB(persist_dir=tmpdir)
|
||||||
|
embedder = create_embedder(EmbedConfig(mode="local"))
|
||||||
|
ingestor = DocumentIngestor(db, embedder, "test_search")
|
||||||
|
|
||||||
|
# 入库一些测试文档
|
||||||
|
content = """# Python 入门
|
||||||
|
Python 是一种解释型编程语言。
|
||||||
|
|
||||||
|
## 安装 Python
|
||||||
|
从 python.org 下载安装包。
|
||||||
|
|
||||||
|
# 向量数据库
|
||||||
|
ChromaDB 是一个轻量级向量数据库。
|
||||||
|
|
||||||
|
## ChromaDB 安装
|
||||||
|
使用 pip install chromadb 安装。"""
|
||||||
|
ingestor.ingest_content(content, "guide.md")
|
||||||
|
|
||||||
|
return Searcher(db, embedder, "test_search")
|
||||||
|
|
||||||
|
|
||||||
|
class TestSearcher:
|
||||||
|
"""检索器测试."""
|
||||||
|
|
||||||
|
def test_search_returns_results(self, searcher):
|
||||||
|
"""搜索返回至少一条结果."""
|
||||||
|
results = searcher.search("Python 编程", top_k=3)
|
||||||
|
assert len(results) > 0
|
||||||
|
for r in results:
|
||||||
|
assert r["content"]
|
||||||
|
assert r["source_file"]
|
||||||
|
assert "score" in r
|
||||||
|
|
||||||
|
def test_search_scores_are_descending(self, searcher):
|
||||||
|
"""搜索结果按相似度降序排列."""
|
||||||
|
results = searcher.search("向量数据库", top_k=5)
|
||||||
|
scores = [r["score"] for r in results]
|
||||||
|
assert scores == sorted(scores, reverse=True)
|
||||||
|
|
||||||
|
def test_search_respects_top_k(self, searcher):
|
||||||
|
"""top_k 参数限制返回数量."""
|
||||||
|
results = searcher.search("安装", top_k=2)
|
||||||
|
assert len(results) <= 2
|
||||||
|
|
||||||
|
def test_search_returns_all_fields(self, searcher):
|
||||||
|
"""搜索结果包含完整字段."""
|
||||||
|
results = searcher.search("ChromaDB", top_k=1)
|
||||||
|
if results:
|
||||||
|
r = results[0]
|
||||||
|
assert "content" in r
|
||||||
|
assert "source_file" in r
|
||||||
|
assert "section_title" in r
|
||||||
|
assert "heading_level" in r
|
||||||
|
assert "chunk_index" in r
|
||||||
|
assert "score" in r
|
||||||
|
|
||||||
|
def test_search_no_results(self, searcher):
|
||||||
|
"""无语义匹配时不崩溃."""
|
||||||
|
results = searcher.search("xyzxyz不存在的内容abcabc", top_k=3)
|
||||||
|
assert isinstance(results, list)
|
||||||
Reference in New Issue
Block a user