From 7309b0a4378972e264e558a7e611012deb66d74f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E8=88=AA=E5=AE=87?= <3364451258@qq.com> Date: Sun, 5 Jul 2026 01:06:08 +0800 Subject: [PATCH] feat: add semantic search module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 实现语义检索模块 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 --- src/core/search.py | 88 ++++++++++++++++++++++++++++++++++++++++++++ tests/test_search.py | 76 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 164 insertions(+) create mode 100644 src/core/search.py create mode 100644 tests/test_search.py diff --git a/src/core/search.py b/src/core/search.py new file mode 100644 index 0000000..48e883d --- /dev/null +++ b/src/core/search.py @@ -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 diff --git a/tests/test_search.py b/tests/test_search.py new file mode 100644 index 0000000..42611f4 --- /dev/null +++ b/tests/test_search.py @@ -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)