diff --git a/config.yaml b/config.yaml index 6dbfeb5..222d1af 100644 --- a/config.yaml +++ b/config.yaml @@ -21,3 +21,9 @@ server: port: 8000 # ssl_keyfile: "" # HTTPS 私钥路径（设置后启用 HTTPS） # ssl_certfile: "" # HTTPS 证书路径（设置后启用 HTTPS） + +search: + mode: hybrid # hybrid | vector + bm25_weight: 0.3 # BM25 权重 (0=纯向量, 1=纯BM25) + candidate_multiplier: 3 # 向量检索候选倍数 + enable_rerank: false # Cross-Encoder 重排序 diff --git a/src/core/config.py b/src/core/config.py index 073b969..f09bdb6 100644 --- a/src/core/config.py +++ b/src/core/config.py @@ -65,6 +65,16 @@ class ServerConfig: ssl_certfile: str = "" # HTTPS 证书文件路径 (空则使用 HTTP) +@dataclass +class SearchConfig: + """检索配置.""" + + mode: str = "hybrid" # hybrid | vector + bm25_weight: float = 0.3 # BM25 权重 (0=纯向量, 1=纯BM25) + candidate_multiplier: int = 3 # 向量检索候选倍数 + enable_rerank: bool = False # 是否启用 Cross-Encoder 重排序 + + @dataclass class AppConfig: """应用总配置.""" @@ -73,6 +83,7 @@ class AppConfig: embed: EmbedConfig = field(default_factory=EmbedConfig) chunk: ChunkConfig = field(default_factory=ChunkConfig) server: ServerConfig = field(default_factory=ServerConfig) + search: SearchConfig = field(default_factory=SearchConfig) @classmethod def from_dict(cls, data: dict) -> "AppConfig": @@ -82,6 +93,7 @@ class AppConfig: embed=EmbedConfig(**data.get("embed", {})), chunk=ChunkConfig(**data.get("chunk", {})), server=ServerConfig(**data.get("server", {})), + search=SearchConfig(**data.get("search", {})), ) diff --git a/src/core/search.py b/src/core/search.py index 2dc18f5..17f4cda 100644 --- a/src/core/search.py +++ b/src/core/search.py @@ -1,30 +1,62 @@ """语义检索模块.""" +import csv +import io +import json as json_lib import logging + +from src.core.config import SearchConfig from src.core.db import VectorDB from src.core.embedder import Embedder +from src.core.retriever import HybridRetriever logger = logging.getLogger("md-vector-db") class Searcher: - """向量检索器.""" + """向量检索器 — 支持纯向量或 HybridRetriever 混合检索.""" - def __init__(self, db: VectorDB, embedder: Embedder, collection_name: str): + def __init__( + self, + db: VectorDB, + embedder: Embedder, + collection_name: str, + search_config: SearchConfig | None = None, + ): self.db = db self.embedder = embedder self.collection_name = collection_name + self._search_config = search_config or SearchConfig() + self._hybrid: HybridRetriever | None = None @property def collection(self): return self.db.get_or_create_collection(self.collection_name) + def _get_hybrid(self) -> HybridRetriever: + if self._hybrid is None: + self._hybrid = HybridRetriever( + self.db, + self.embedder, + self.collection_name, + bm25_weight=self._search_config.bm25_weight, + vector_candidate_multiplier=self._search_config.candidate_multiplier, + ) + return self._hybrid + def search( self, query: str, top_k: int = 10, source_file: str | None = None, ) -> list[dict]: - """语义检索, 返回格式化结果列表.""" + """语义检索, 返回格式化结果列表. + + 根据 search.mode 配置自动选择 hybrid 或 vector 模式. + """ + if self._search_config.mode == "hybrid": + return self._get_hybrid().search(query, top_k=top_k, source_file=source_file) + + # 纯向量模式（原有逻辑） embeddings = self.embedder.embed([query]) if not embeddings: raise RuntimeError("嵌入器返回空结果, 无法进行检索") @@ -83,3 +115,61 @@ class Searcher: def delete_by_source(self, file_name: str) -> bool: """按文件名删除文档 (委托 VectorDB).""" return self.db.delete_by_source(self.collection_name, file_name) + + def export_json(self, file_path: str | None = None) -> str: + """导出 collection 所有 chunks 为 JSON. + + Args: + file_path: 可选，写入文件路径。不传则返回 JSON 字符串。 + + Returns: + JSON 字符串 + """ + all_data = self.collection.get(include=["documents", "metadatas"]) + records = [] + if all_data and all_data["ids"]: + for i, doc_id in enumerate(all_data["ids"]): + meta = all_data["metadatas"][i] if all_data["metadatas"] else {} + records.append({ + "id": doc_id, + "content": all_data["documents"][i] if all_data["documents"] else "", + "source_file": meta.get("source_file", ""), + "section_title": meta.get("section_title", ""), + "heading_level": meta.get("heading_level", 0), + "chunk_index": meta.get("chunk_index", i), + }) + json_str = json_lib.dumps(records, ensure_ascii=False, indent=2) + if file_path: + with open(file_path, "w", encoding="utf-8") as f: + f.write(json_str) + return json_str + + def export_csv(self, file_path: str | None = None) -> str: + """导出 collection 所有 chunks 为 CSV. + + Args: + file_path: 可选，写入文件路径。不传则返回 CSV 字符串。 + + Returns: + CSV 字符串 + """ + all_data = self.collection.get(include=["documents", "metadatas"]) + output = io.StringIO() + writer = csv.writer(output) + writer.writerow(["id", "content", "source_file", "section_title", "heading_level", "chunk_index"]) + if all_data and all_data["ids"]: + for i, doc_id in enumerate(all_data["ids"]): + meta = all_data["metadatas"][i] if all_data["metadatas"] else {} + writer.writerow([ + doc_id, + all_data["documents"][i] if all_data["documents"] else "", + meta.get("source_file", ""), + meta.get("section_title", ""), + meta.get("heading_level", 0), + meta.get("chunk_index", i), + ]) + csv_str = output.getvalue() + if file_path: + with open(file_path, "w", encoding="utf-8", newline="") as f: + f.write(csv_str) + return csv_str diff --git a/src/server/deps.py b/src/server/deps.py index 1ccf2f9..87f55aa 100644 --- a/src/server/deps.py +++ b/src/server/deps.py @@ -38,7 +38,10 @@ class AppState: name = collection or self.default_collection with self._cache_lock: if name not in self._searchers: - self._searchers[name] = Searcher(self.db, self.embedder, name) + self._searchers[name] = Searcher( + self.db, self.embedder, name, + search_config=self.config.search, + ) return self._searchers[name] def get_ingestor(self, collection: str | None = None) -> DocumentIngestor: