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,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