|
@@ -0,0 +1,362 @@
|
|
|
|
|
+"""Milvus 访问层:无状态 Store + 短生命周期 gRPC 连接(避免长期持有失效 channel)。"""
|
|
|
|
|
+from __future__ import annotations
|
|
|
|
|
+
|
|
|
|
|
+import os
|
|
|
|
|
+from contextlib import contextmanager
|
|
|
|
|
+from dataclasses import dataclass
|
|
|
|
|
+from typing import Callable, Iterator, TypeVar
|
|
|
|
|
+
|
|
|
|
|
+from pymilvus import AnnSearchRequest, DataType, MilvusClient, RRFRanker, Function, FunctionType
|
|
|
|
|
+
|
|
|
|
|
+QUERY_MAX_LIMIT = 16384
|
|
|
|
|
+T = TypeVar("T")
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@dataclass(frozen=True)
|
|
|
|
|
+class MilvusSettings:
|
|
|
|
|
+ host: str
|
|
|
|
|
+ port: str
|
|
|
|
|
+ collection_name: str
|
|
|
|
|
+ uri: str
|
|
|
|
|
+ timeout: float
|
|
|
|
|
+
|
|
|
|
|
+ @classmethod
|
|
|
|
|
+ def from_env(cls) -> "MilvusSettings":
|
|
|
|
|
+ host = os.getenv("MILVUS_HOST", "localhost")
|
|
|
|
|
+ port = os.getenv("MILVUS_PORT", "19530")
|
|
|
|
|
+ collection = os.getenv("MILVUS_COLLECTION", "kb_documents")
|
|
|
|
|
+ timeout = float(os.getenv("MILVUS_TIMEOUT", "30"))
|
|
|
|
|
+ return cls(
|
|
|
|
|
+ host=host,
|
|
|
|
|
+ port=port,
|
|
|
|
|
+ collection_name=collection,
|
|
|
|
|
+ uri=f"http://{host}:{port}",
|
|
|
|
|
+ timeout=timeout,
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@contextmanager
|
|
|
|
|
+def milvus_client_session(settings: MilvusSettings | None = None) -> Iterator[MilvusClient]:
|
|
|
|
|
+ """一次 RPC 会话:创建连接,用完后关闭,不缓存 gRPC channel。"""
|
|
|
|
|
+ cfg = settings or MilvusSettings.from_env()
|
|
|
|
|
+ client = MilvusClient(uri=cfg.uri, timeout=cfg.timeout)
|
|
|
|
|
+ try:
|
|
|
|
|
+ yield client
|
|
|
|
|
+ finally:
|
|
|
|
|
+ client.close()
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def _normalize_filter(filter_expr: str) -> str:
|
|
|
|
|
+ return filter_expr.strip() if filter_expr.strip() else "id >= 0"
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+class MilvusStore:
|
|
|
|
|
+ """Milvus 集合读写;本身不持有连接,所有 IO 经 milvus_client_session。"""
|
|
|
|
|
+
|
|
|
|
|
+ def __init__(self, settings: MilvusSettings | None = None):
|
|
|
|
|
+ self._settings = settings or MilvusSettings.from_env()
|
|
|
|
|
+
|
|
|
|
|
+ @property
|
|
|
|
|
+ def collection_name(self) -> str:
|
|
|
|
|
+ return self._settings.collection_name
|
|
|
|
|
+
|
|
|
|
|
+ def _run(self, operation: Callable[[MilvusClient], T]) -> T:
|
|
|
|
|
+ with milvus_client_session(self._settings) as client:
|
|
|
|
|
+ return operation(client)
|
|
|
|
|
+
|
|
|
|
|
+ @contextmanager
|
|
|
|
|
+ def session(self) -> Iterator[MilvusClient]:
|
|
|
|
|
+ """同一业务流(如整次上传)内复用一条连接,用毕即关。"""
|
|
|
|
|
+ with milvus_client_session(self._settings) as client:
|
|
|
|
|
+ yield client
|
|
|
|
|
+
|
|
|
|
|
+ @staticmethod
|
|
|
|
|
+ def ensure_collection(client: MilvusClient, collection_name: str, dense_dim: int) -> None:
|
|
|
|
|
+ if client.has_collection(collection_name):
|
|
|
|
|
+ return
|
|
|
|
|
+
|
|
|
|
|
+ schema = client.create_schema(auto_id=True, enable_dynamic_field=True)
|
|
|
|
|
+ schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
|
|
|
|
+ schema.add_field("dense_embedding", DataType.FLOAT_VECTOR, dim=dense_dim)
|
|
|
|
|
+ schema.add_field("sparse_embedding", DataType.SPARSE_FLOAT_VECTOR)
|
|
|
|
|
+ schema.add_field(
|
|
|
|
|
+ "text",
|
|
|
|
|
+ DataType.VARCHAR,
|
|
|
|
|
+ max_length=65535,
|
|
|
|
|
+ enable_analyzer=True,
|
|
|
|
|
+ analyzer_params={"type": "standard"},
|
|
|
|
|
+ enable_match=True,
|
|
|
|
|
+ )
|
|
|
|
|
+ schema.add_field("document_id", DataType.INT64)
|
|
|
|
|
+ schema.add_field("filename", DataType.VARCHAR, max_length=255)
|
|
|
|
|
+ schema.add_field("file_type", DataType.VARCHAR, max_length=50)
|
|
|
|
|
+ schema.add_field("file_path", DataType.VARCHAR, max_length=1024)
|
|
|
|
|
+ schema.add_field("page_number", DataType.INT64)
|
|
|
|
|
+ schema.add_field("chunk_idx", DataType.INT64)
|
|
|
|
|
+ schema.add_field("chunk_id", DataType.VARCHAR, max_length=512)
|
|
|
|
|
+ schema.add_field("parent_chunk_id", DataType.VARCHAR, max_length=512)
|
|
|
|
|
+ schema.add_field("root_chunk_id", DataType.VARCHAR, max_length=512)
|
|
|
|
|
+ schema.add_field("chunk_level", DataType.INT64)
|
|
|
|
|
+
|
|
|
|
|
+ bm25_function = Function(
|
|
|
|
|
+ name="text_bm25_emb",
|
|
|
|
|
+ function_type=FunctionType.BM25,
|
|
|
|
|
+ input_field_names=["text"],
|
|
|
|
|
+ output_field_names=["sparse_embedding"],
|
|
|
|
|
+ )
|
|
|
|
|
+ schema.add_function(bm25_function)
|
|
|
|
|
+
|
|
|
|
|
+ index_params = client.prepare_index_params()
|
|
|
|
|
+ index_params.add_index(
|
|
|
|
|
+ field_name="dense_embedding",
|
|
|
|
|
+ index_type="HNSW",
|
|
|
|
|
+ metric_type="IP",
|
|
|
|
|
+ params={"M": 16, "efConstruction": 256},
|
|
|
|
|
+ )
|
|
|
|
|
+ index_params.add_index(
|
|
|
|
|
+ field_name="sparse_embedding",
|
|
|
|
|
+ index_type="SPARSE_INVERTED_INDEX",
|
|
|
|
|
+ metric_type="BM25",
|
|
|
|
|
+ params={"drop_ratio_build": 0.2},
|
|
|
|
|
+ )
|
|
|
|
|
+ try:
|
|
|
|
|
+ client.create_collection(
|
|
|
|
|
+ collection_name=collection_name,
|
|
|
|
|
+ schema=schema,
|
|
|
|
|
+ index_params=index_params,
|
|
|
|
|
+ )
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ # Milvus Lite on Windows 在 create_index 后重命名 manifest.json 时偶发
|
|
|
|
|
+ # WinError 183,但集合与索引实际已创建成功,因此若集合已存在则忽略。
|
|
|
|
|
+ if client.has_collection(collection_name):
|
|
|
|
|
+ return
|
|
|
|
|
+ raise
|
|
|
|
|
+
|
|
|
|
|
+ def init_collection(self, dense_dim: int | None = None) -> None:
|
|
|
|
|
+ if dense_dim is None:
|
|
|
|
|
+ dense_dim = int(os.getenv("DENSE_EMBEDDING_DIM", "1024"))
|
|
|
|
|
+
|
|
|
|
|
+ def _init(client: MilvusClient) -> None:
|
|
|
|
|
+ self.ensure_collection(client, self.collection_name, dense_dim)
|
|
|
|
|
+
|
|
|
|
|
+ self._run(_init)
|
|
|
|
|
+
|
|
|
|
|
+ def insert(self, data: list[dict]):
|
|
|
|
|
+ return self._run(lambda client: client.insert(self.collection_name, data))
|
|
|
|
|
+
|
|
|
|
|
+ def delete_by_document(self, document_id: int) -> int:
|
|
|
|
|
+ def _delete(client: MilvusClient) -> int:
|
|
|
|
|
+ result = client.delete(
|
|
|
|
|
+ collection_name=self.collection_name,
|
|
|
|
|
+ filter=f"document_id == {document_id}",
|
|
|
|
|
+ )
|
|
|
|
|
+ return getattr(result, "delete_count", len(result)) if result else 0
|
|
|
|
|
+
|
|
|
|
|
+ return self._run(_delete)
|
|
|
|
|
+
|
|
|
|
|
+ def delete_by_ids(self, ids: list[int | str]) -> int:
|
|
|
|
|
+ if not ids:
|
|
|
|
|
+ return 0
|
|
|
|
|
+
|
|
|
|
|
+ def _delete(client: MilvusClient) -> int:
|
|
|
|
|
+ result = client.delete(collection_name=self.collection_name, ids=ids)
|
|
|
|
|
+ return getattr(result, "delete_count", len(result)) if result else 0
|
|
|
|
|
+
|
|
|
|
|
+ return self._run(_delete)
|
|
|
|
|
+
|
|
|
|
|
+ def query(
|
|
|
|
|
+ self,
|
|
|
|
|
+ filter_expr: str = "",
|
|
|
|
|
+ output_fields: list[str] | None = None,
|
|
|
|
|
+ limit: int = 10000,
|
|
|
|
|
+ offset: int = 0,
|
|
|
|
|
+ ):
|
|
|
|
|
+ expr = _normalize_filter(filter_expr)
|
|
|
|
|
+ fields = output_fields or ["filename", "file_type"]
|
|
|
|
|
+
|
|
|
|
|
+ def _query(client: MilvusClient):
|
|
|
|
|
+ return client.query(
|
|
|
|
|
+ collection_name=self.collection_name,
|
|
|
|
|
+ filter=expr,
|
|
|
|
|
+ output_fields=fields,
|
|
|
|
|
+ limit=min(limit, QUERY_MAX_LIMIT),
|
|
|
|
|
+ offset=offset,
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ return self._run(_query)
|
|
|
|
|
+
|
|
|
|
|
+ def query_all(self, filter_expr: str = "", output_fields: list[str] | None = None) -> list:
|
|
|
|
|
+ """分页拉取;单次 session 内完成,避免每页新建连接。"""
|
|
|
|
|
+ fields = output_fields or ["filename", "file_type"]
|
|
|
|
|
+ expr = _normalize_filter(filter_expr)
|
|
|
|
|
+
|
|
|
|
|
+ def _query_all(client: MilvusClient) -> list:
|
|
|
|
|
+ out: list = []
|
|
|
|
|
+ offset = 0
|
|
|
|
|
+ while True:
|
|
|
|
|
+ batch = client.query(
|
|
|
|
|
+ collection_name=self.collection_name,
|
|
|
|
|
+ filter=expr,
|
|
|
|
|
+ output_fields=fields,
|
|
|
|
|
+ limit=QUERY_MAX_LIMIT,
|
|
|
|
|
+ offset=offset,
|
|
|
|
|
+ )
|
|
|
|
|
+ if not batch:
|
|
|
|
|
+ break
|
|
|
|
|
+ out.extend(batch)
|
|
|
|
|
+ if len(batch) < QUERY_MAX_LIMIT:
|
|
|
|
|
+ break
|
|
|
|
|
+ offset += len(batch)
|
|
|
|
|
+ return out
|
|
|
|
|
+
|
|
|
|
|
+ return self._run(_query_all)
|
|
|
|
|
+
|
|
|
|
|
+ def get_chunks_by_ids(self, chunk_ids: list[str]) -> list[dict]:
|
|
|
|
|
+ ids = [item for item in chunk_ids if item]
|
|
|
|
|
+ if not ids:
|
|
|
|
|
+ return []
|
|
|
|
|
+ quoted_ids = ", ".join(f'"{item}"' for item in ids)
|
|
|
|
|
+ return self.query(
|
|
|
|
|
+ filter_expr=f"chunk_id in [{quoted_ids}]",
|
|
|
|
|
+ output_fields=[
|
|
|
|
|
+ "text",
|
|
|
|
|
+ "filename",
|
|
|
|
|
+ "file_type",
|
|
|
|
|
+ "page_number",
|
|
|
|
|
+ "chunk_id",
|
|
|
|
|
+ "parent_chunk_id",
|
|
|
|
|
+ "root_chunk_id",
|
|
|
|
|
+ "chunk_level",
|
|
|
|
|
+ "chunk_idx",
|
|
|
|
|
+ ],
|
|
|
|
|
+ limit=len(ids),
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ def hybrid_retrieve(
|
|
|
|
|
+ self,
|
|
|
|
|
+ dense_embedding: list[float],
|
|
|
|
|
+ query: str,
|
|
|
|
|
+ top_k: int = 5,
|
|
|
|
|
+ rrf_k: int = 60,
|
|
|
|
|
+ filter_expr: str = "",
|
|
|
|
|
+ ) -> list[dict]:
|
|
|
|
|
+ output_fields = [
|
|
|
|
|
+ "text",
|
|
|
|
|
+ "filename",
|
|
|
|
|
+ "file_type",
|
|
|
|
|
+ "page_number",
|
|
|
|
|
+ "chunk_id",
|
|
|
|
|
+ "parent_chunk_id",
|
|
|
|
|
+ "root_chunk_id",
|
|
|
|
|
+ "chunk_level",
|
|
|
|
|
+ "chunk_idx",
|
|
|
|
|
+ ]
|
|
|
|
|
+ dense_search = AnnSearchRequest(
|
|
|
|
|
+ data=[dense_embedding],
|
|
|
|
|
+ anns_field="dense_embedding",
|
|
|
|
|
+ param={"metric_type": "IP", "params": {"ef": 64}},
|
|
|
|
|
+ limit=top_k * 2,
|
|
|
|
|
+ expr=filter_expr,
|
|
|
|
|
+ )
|
|
|
|
|
+ sparse_search = AnnSearchRequest(
|
|
|
|
|
+ data=[query],
|
|
|
|
|
+ anns_field="sparse_embedding",
|
|
|
|
|
+ param={"metric_type": "BM25", "params": {"drop_ratio_search": 0.2}},
|
|
|
|
|
+ limit=top_k * 2,
|
|
|
|
|
+ expr=filter_expr,
|
|
|
|
|
+ )
|
|
|
|
|
+ reranker = RRFRanker(k=rrf_k)
|
|
|
|
|
+
|
|
|
|
|
+ def _search(client: MilvusClient):
|
|
|
|
|
+ return client.hybrid_search(
|
|
|
|
|
+ collection_name=self.collection_name,
|
|
|
|
|
+ reqs=[dense_search, sparse_search],
|
|
|
|
|
+ ranker=reranker,
|
|
|
|
|
+ limit=top_k,
|
|
|
|
|
+ output_fields=output_fields,
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ results = self._run(_search)
|
|
|
|
|
+ formatted_results = []
|
|
|
|
|
+ for hits in results:
|
|
|
|
|
+ for hit in hits:
|
|
|
|
|
+ formatted_results.append({
|
|
|
|
|
+ "id": hit.get("id"),
|
|
|
|
|
+ "text": hit.get("text", ""),
|
|
|
|
|
+ "filename": hit.get("filename", ""),
|
|
|
|
|
+ "file_type": hit.get("file_type", ""),
|
|
|
|
|
+ "page_number": hit.get("page_number", 0),
|
|
|
|
|
+ "chunk_id": hit.get("chunk_id", ""),
|
|
|
|
|
+ "parent_chunk_id": hit.get("parent_chunk_id", ""),
|
|
|
|
|
+ "root_chunk_id": hit.get("root_chunk_id", ""),
|
|
|
|
|
+ "chunk_level": hit.get("chunk_level", 0),
|
|
|
|
|
+ "chunk_idx": hit.get("chunk_idx", 0),
|
|
|
|
|
+ "score": hit.get("distance", 0.0),
|
|
|
|
|
+ })
|
|
|
|
|
+ return formatted_results
|
|
|
|
|
+
|
|
|
|
|
+ def dense_retrieve(
|
|
|
|
|
+ self,
|
|
|
|
|
+ dense_embedding: list[float],
|
|
|
|
|
+ top_k: int = 5,
|
|
|
|
|
+ filter_expr: str = "",
|
|
|
|
|
+ ) -> list[dict]:
|
|
|
|
|
+ def _search(client: MilvusClient):
|
|
|
|
|
+ return client.search(
|
|
|
|
|
+ collection_name=self.collection_name,
|
|
|
|
|
+ data=[dense_embedding],
|
|
|
|
|
+ anns_field="dense_embedding",
|
|
|
|
|
+ search_params={"metric_type": "IP", "params": {"ef": 64}},
|
|
|
|
|
+ limit=top_k,
|
|
|
|
|
+ output_fields=[
|
|
|
|
|
+ "text",
|
|
|
|
|
+ "filename",
|
|
|
|
|
+ "file_type",
|
|
|
|
|
+ "page_number",
|
|
|
|
|
+ "chunk_id",
|
|
|
|
|
+ "parent_chunk_id",
|
|
|
|
|
+ "root_chunk_id",
|
|
|
|
|
+ "chunk_level",
|
|
|
|
|
+ "chunk_idx",
|
|
|
|
|
+ ],
|
|
|
|
|
+ filter=filter_expr,
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ results = self._run(_search)
|
|
|
|
|
+ formatted_results = []
|
|
|
|
|
+ for hits in results:
|
|
|
|
|
+ for hit in hits:
|
|
|
|
|
+ formatted_results.append({
|
|
|
|
|
+ "id": hit.get("id"),
|
|
|
|
|
+ "text": hit.get("entity", {}).get("text", ""),
|
|
|
|
|
+ "filename": hit.get("entity", {}).get("filename", ""),
|
|
|
|
|
+ "file_type": hit.get("entity", {}).get("file_type", ""),
|
|
|
|
|
+ "page_number": hit.get("entity", {}).get("page_number", 0),
|
|
|
|
|
+ "chunk_id": hit.get("entity", {}).get("chunk_id", ""),
|
|
|
|
|
+ "parent_chunk_id": hit.get("entity", {}).get("parent_chunk_id", ""),
|
|
|
|
|
+ "root_chunk_id": hit.get("entity", {}).get("root_chunk_id", ""),
|
|
|
|
|
+ "chunk_level": hit.get("entity", {}).get("chunk_level", 0),
|
|
|
|
|
+ "chunk_idx": hit.get("entity", {}).get("chunk_idx", 0),
|
|
|
|
|
+ "score": hit.get("distance", 0.0),
|
|
|
|
|
+ })
|
|
|
|
|
+ return formatted_results
|
|
|
|
|
+
|
|
|
|
|
+ def drop_collection(self) -> None:
|
|
|
|
|
+ def _drop(client: MilvusClient) -> None:
|
|
|
|
|
+ if client.has_collection(self.collection_name):
|
|
|
|
|
+ client.drop_collection(self.collection_name)
|
|
|
|
|
+
|
|
|
|
|
+ self._run(_drop)
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+# 兼容旧名;全项目共用同一无状态 Store 实例即可(不缓存连接)
|
|
|
|
|
+MilvusManager = MilvusStore
|
|
|
|
|
+
|
|
|
|
|
+_store: MilvusStore | None = None
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def get_milvus_store() -> MilvusStore:
|
|
|
|
|
+ global _store
|
|
|
|
|
+ if _store is None:
|
|
|
|
|
+ _store = MilvusStore()
|
|
|
|
|
+ return _store
|