feat: Searcher 集成 HybridRetriever + SearchConfig 配置
This commit is contained in:
@@ -21,3 +21,9 @@ server:
|
|||||||
port: 8000
|
port: 8000
|
||||||
# ssl_keyfile: "" # HTTPS 私钥路径(设置后启用 HTTPS)
|
# ssl_keyfile: "" # HTTPS 私钥路径(设置后启用 HTTPS)
|
||||||
# ssl_certfile: "" # 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 重排序
|
||||||
|
|||||||
@@ -65,6 +65,16 @@ class ServerConfig:
|
|||||||
ssl_certfile: str = "" # HTTPS 证书文件路径 (空则使用 HTTP)
|
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
|
@dataclass
|
||||||
class AppConfig:
|
class AppConfig:
|
||||||
"""应用总配置."""
|
"""应用总配置."""
|
||||||
@@ -73,6 +83,7 @@ class AppConfig:
|
|||||||
embed: EmbedConfig = field(default_factory=EmbedConfig)
|
embed: EmbedConfig = field(default_factory=EmbedConfig)
|
||||||
chunk: ChunkConfig = field(default_factory=ChunkConfig)
|
chunk: ChunkConfig = field(default_factory=ChunkConfig)
|
||||||
server: ServerConfig = field(default_factory=ServerConfig)
|
server: ServerConfig = field(default_factory=ServerConfig)
|
||||||
|
search: SearchConfig = field(default_factory=SearchConfig)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_dict(cls, data: dict) -> "AppConfig":
|
def from_dict(cls, data: dict) -> "AppConfig":
|
||||||
@@ -82,6 +93,7 @@ class AppConfig:
|
|||||||
embed=EmbedConfig(**data.get("embed", {})),
|
embed=EmbedConfig(**data.get("embed", {})),
|
||||||
chunk=ChunkConfig(**data.get("chunk", {})),
|
chunk=ChunkConfig(**data.get("chunk", {})),
|
||||||
server=ServerConfig(**data.get("server", {})),
|
server=ServerConfig(**data.get("server", {})),
|
||||||
|
search=SearchConfig(**data.get("search", {})),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
+93
-3
@@ -1,30 +1,62 @@
|
|||||||
"""语义检索模块."""
|
"""语义检索模块."""
|
||||||
|
import csv
|
||||||
|
import io
|
||||||
|
import json as json_lib
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
|
from src.core.config import SearchConfig
|
||||||
from src.core.db import VectorDB
|
from src.core.db import VectorDB
|
||||||
from src.core.embedder import Embedder
|
from src.core.embedder import Embedder
|
||||||
|
from src.core.retriever import HybridRetriever
|
||||||
|
|
||||||
logger = logging.getLogger("md-vector-db")
|
logger = logging.getLogger("md-vector-db")
|
||||||
|
|
||||||
|
|
||||||
class Searcher:
|
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.db = db
|
||||||
self.embedder = embedder
|
self.embedder = embedder
|
||||||
self.collection_name = collection_name
|
self.collection_name = collection_name
|
||||||
|
self._search_config = search_config or SearchConfig()
|
||||||
|
self._hybrid: HybridRetriever | None = None
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def collection(self):
|
def collection(self):
|
||||||
return self.db.get_or_create_collection(self.collection_name)
|
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(
|
def search(
|
||||||
self,
|
self,
|
||||||
query: str,
|
query: str,
|
||||||
top_k: int = 10,
|
top_k: int = 10,
|
||||||
source_file: str | None = None,
|
source_file: str | None = None,
|
||||||
) -> list[dict]:
|
) -> 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])
|
embeddings = self.embedder.embed([query])
|
||||||
if not embeddings:
|
if not embeddings:
|
||||||
raise RuntimeError("嵌入器返回空结果, 无法进行检索")
|
raise RuntimeError("嵌入器返回空结果, 无法进行检索")
|
||||||
@@ -83,3 +115,61 @@ class Searcher:
|
|||||||
def delete_by_source(self, file_name: str) -> bool:
|
def delete_by_source(self, file_name: str) -> bool:
|
||||||
"""按文件名删除文档 (委托 VectorDB)."""
|
"""按文件名删除文档 (委托 VectorDB)."""
|
||||||
return self.db.delete_by_source(self.collection_name, file_name)
|
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
|
||||||
|
|||||||
+4
-1
@@ -38,7 +38,10 @@ class AppState:
|
|||||||
name = collection or self.default_collection
|
name = collection or self.default_collection
|
||||||
with self._cache_lock:
|
with self._cache_lock:
|
||||||
if name not in self._searchers:
|
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]
|
return self._searchers[name]
|
||||||
|
|
||||||
def get_ingestor(self, collection: str | None = None) -> DocumentIngestor:
|
def get_ingestor(self, collection: str | None = None) -> DocumentIngestor:
|
||||||
|
|||||||
Reference in New Issue
Block a user