wehub-resource-sync
1334 lines
53 KiB
Python
1334 lines
53 KiB
Python
import asyncio
|
||
import os
|
||
import time
|
||
import traceback
|
||
import weakref
|
||
from dataclasses import MISSING, dataclass, field, fields
|
||
from functools import partial
|
||
from typing import Any
|
||
|
||
from pymilvus import (
|
||
AnnSearchRequest,
|
||
Collection,
|
||
CollectionSchema,
|
||
DataType,
|
||
FieldSchema,
|
||
Function,
|
||
FunctionType,
|
||
WeightedRanker,
|
||
connections,
|
||
db,
|
||
utility,
|
||
)
|
||
|
||
from yuxi.knowledge.base import FileStatus, KnowledgeBase
|
||
from yuxi.knowledge.chunking.ragflow_like.dispatcher import chunk_markdown
|
||
from yuxi.knowledge.chunking.ragflow_like.nlp import count_tokens
|
||
from yuxi.knowledge.parser.unified import Parser
|
||
from yuxi.knowledge.utils.kb_utils import resolve_processing_params
|
||
from yuxi.models.providers.cache import model_cache
|
||
from yuxi.repositories.knowledge_chunk_repository import KnowledgeChunkRepository
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
from yuxi.utils import hashstr, logger
|
||
|
||
MILVUS_AVAILABLE = True
|
||
CONTENT_SPARSE_FIELD = "content_sparse"
|
||
CONTENT_ANALYZER_PARAMS = {"type": "chinese"}
|
||
VECTOR_METRIC_TYPE = "COSINE"
|
||
MILVUS_CHUNK_EMBED_BATCH_SIZE = 200
|
||
MILVUS_QUERY_OFFLOAD_LIMIT = 8
|
||
_milvus_query_offload_semaphore_refs: dict[
|
||
int,
|
||
tuple[weakref.ReferenceType[asyncio.AbstractEventLoop], weakref.ReferenceType[asyncio.Semaphore]],
|
||
] = {}
|
||
|
||
|
||
def _get_milvus_query_offload_semaphore() -> asyncio.Semaphore:
|
||
loop = asyncio.get_running_loop()
|
||
loop_id = id(loop)
|
||
entry = _milvus_query_offload_semaphore_refs.get(loop_id)
|
||
if entry is not None:
|
||
loop_ref, semaphore_ref = entry
|
||
semaphore = semaphore_ref()
|
||
if loop_ref() is loop and semaphore is not None:
|
||
return semaphore
|
||
|
||
semaphore = asyncio.Semaphore(MILVUS_QUERY_OFFLOAD_LIMIT)
|
||
|
||
def cleanup(ref, stale_loop_id=loop_id):
|
||
current_entry = _milvus_query_offload_semaphore_refs.get(stale_loop_id)
|
||
if current_entry is not None and current_entry[1] is ref:
|
||
_milvus_query_offload_semaphore_refs.pop(stale_loop_id, None)
|
||
|
||
_milvus_query_offload_semaphore_refs[loop_id] = (weakref.ref(loop), weakref.ref(semaphore, cleanup))
|
||
return semaphore
|
||
|
||
|
||
async def _run_milvus_query_io(func, /, *args, **kwargs):
|
||
semaphore = _get_milvus_query_offload_semaphore()
|
||
await semaphore.acquire()
|
||
task = asyncio.create_task(asyncio.to_thread(func, *args, **kwargs))
|
||
|
||
def release_capacity(completed_task: asyncio.Task):
|
||
semaphore.release()
|
||
if completed_task.cancelled():
|
||
return
|
||
completed_task.exception()
|
||
|
||
task.add_done_callback(release_capacity)
|
||
return await asyncio.shield(task)
|
||
|
||
|
||
@dataclass(kw_only=True)
|
||
class MilvusRetrievalConfig:
|
||
search_mode: str = field(
|
||
default="vector",
|
||
metadata={
|
||
"label": "检索模式",
|
||
"type": "select",
|
||
"options": [
|
||
{"value": "vector", "label": "向量检索", "description": "仅使用向量相似度检索"},
|
||
{"value": "keyword", "label": "BM25 全文检索", "description": "仅使用 Milvus BM25 检索"},
|
||
{"value": "hybrid", "label": "混合检索", "description": "Milvus 向量检索与 BM25 融合检索"},
|
||
],
|
||
"description": "选择检索模式",
|
||
},
|
||
)
|
||
final_top_k: int = field(
|
||
default=10,
|
||
metadata={
|
||
"label": "最终返回 Chunk 数",
|
||
"type": "number",
|
||
"min": 1,
|
||
"max": 100,
|
||
"description": "重排序后返回给前端的文档数量",
|
||
},
|
||
)
|
||
similarity_threshold: float = field(
|
||
default=0.0,
|
||
metadata={
|
||
"label": "相似度阈值（0-1）",
|
||
"type": "number",
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.1,
|
||
"description": "过滤相似度低于此值的结果",
|
||
},
|
||
)
|
||
bm25_top_k: int = field(
|
||
default=50,
|
||
metadata={
|
||
"label": "BM25 召回数量",
|
||
"type": "number",
|
||
"min": 1,
|
||
"max": 200,
|
||
"description": "BM25 全文检索和混合检索中的 BM25 候选数量",
|
||
},
|
||
)
|
||
vector_weight: float = field(
|
||
default=0.7,
|
||
metadata={
|
||
"label": "向量检索权重",
|
||
"type": "number",
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.1,
|
||
"description": "混合检索中向量召回结果的融合权重",
|
||
},
|
||
)
|
||
bm25_weight: float = field(
|
||
default=0.3,
|
||
metadata={
|
||
"label": "BM25 权重",
|
||
"type": "number",
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.1,
|
||
"description": "混合检索中 BM25 召回结果的融合权重",
|
||
},
|
||
)
|
||
bm25_drop_ratio_search: float = field(
|
||
default=0.0,
|
||
metadata={
|
||
"label": "BM25 稀疏项丢弃比例",
|
||
"type": "number",
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.1,
|
||
"description": "BM25 检索时丢弃低分稀疏项的比例，数值越大检索越快但可能降低召回",
|
||
},
|
||
)
|
||
include_distances: bool = field(
|
||
default=True,
|
||
metadata={"label": "显示相似度", "type": "boolean", "description": "在结果中显示相似度分数"},
|
||
)
|
||
use_graph_retrieval: bool = field(
|
||
default=False,
|
||
metadata={"label": "启用图检索", "type": "boolean", "description": "是否启用实体和三元组扩散检索"},
|
||
)
|
||
graph_entity_top_k: int = field(
|
||
default=10,
|
||
metadata={
|
||
"label": "图实体召回数量",
|
||
"type": "number",
|
||
"min": 1,
|
||
"max": 100,
|
||
"depend_on": ("use_graph_retrieval", True),
|
||
"description": "通过 Query 召回的实体数量",
|
||
},
|
||
)
|
||
graph_triple_top_k: int = field(
|
||
default=10,
|
||
metadata={
|
||
"label": "图三元组召回数量",
|
||
"type": "number",
|
||
"min": 1,
|
||
"max": 100,
|
||
"depend_on": ("use_graph_retrieval", True),
|
||
"description": "通过 Query 召回的三元组数量",
|
||
},
|
||
)
|
||
graph_max_nodes: int = field(
|
||
default=10000,
|
||
metadata={
|
||
"label": "图检索最大节点数",
|
||
"type": "number",
|
||
"min": 100,
|
||
"max": 50000,
|
||
"depend_on": ("use_graph_retrieval", True),
|
||
"description": "2-hop 扩散子图最多读取的节点数量",
|
||
},
|
||
)
|
||
graph_top_k: int = field(
|
||
default=20,
|
||
metadata={
|
||
"label": "图召回 Chunk 数",
|
||
"type": "number",
|
||
"min": 1,
|
||
"max": 200,
|
||
"depend_on": ("use_graph_retrieval", True),
|
||
"description": "PPR 后从图谱路径召回的 Chunk 数量",
|
||
},
|
||
)
|
||
graph_weight: float = field(
|
||
default=1.0,
|
||
metadata={
|
||
"label": "图检索融合权重",
|
||
"type": "number",
|
||
"min": 0.0,
|
||
"max": 5.0,
|
||
"step": 0.1,
|
||
"depend_on": ("use_graph_retrieval", True),
|
||
"description": "排名融合时图检索结果的权重",
|
||
},
|
||
)
|
||
ppr_damping: float = field(
|
||
default=0.85,
|
||
metadata={
|
||
"label": "PPR 阻尼系数",
|
||
"type": "number",
|
||
"min": 0.1,
|
||
"max": 0.99,
|
||
"step": 0.01,
|
||
"depend_on": ("use_graph_retrieval", True),
|
||
"description": "Personalized PageRank 的阻尼系数",
|
||
},
|
||
)
|
||
use_reranker: bool = field(
|
||
default=False,
|
||
metadata={"label": "启用重排序", "type": "boolean", "description": "是否使用精排模型对检索结果进行重排序"},
|
||
)
|
||
reranker_model: str = field(
|
||
default="",
|
||
metadata={
|
||
"label": "重排序模型",
|
||
"type": "select",
|
||
"depend_on": ("use_reranker", True),
|
||
"description": "选择用于本次查询的重排序模型",
|
||
"options_provider": "rerank_models",
|
||
},
|
||
)
|
||
recall_top_k: int = field(
|
||
default=50,
|
||
metadata={
|
||
"label": "召回数量",
|
||
"type": "number",
|
||
"min": 10,
|
||
"max": 200,
|
||
"depend_on": ("use_reranker", True),
|
||
"description": "向量检索或混合检索保留的候选数量（启用重排序时有效）",
|
||
},
|
||
)
|
||
|
||
|
||
def _retrieval_config_options() -> list[dict[str, Any]]:
|
||
options = []
|
||
for config_field in fields(MilvusRetrievalConfig):
|
||
metadata = dict(config_field.metadata)
|
||
options_provider = metadata.pop("options_provider", None)
|
||
default = None if config_field.default is MISSING else config_field.default
|
||
option = {
|
||
"key": config_field.name,
|
||
"default": default,
|
||
**metadata,
|
||
}
|
||
if options_provider == "rerank_models":
|
||
option["options"] = [
|
||
{"label": info.display_name, "value": info.spec} for info in model_cache.get_all_specs("rerank")
|
||
]
|
||
options.append(option)
|
||
return options
|
||
|
||
|
||
class MilvusKB(KnowledgeBase):
|
||
"""基于 Milvus 的生产级向量库"""
|
||
|
||
kb_type = "milvus"
|
||
name = "Milvus"
|
||
description = "基于 Milvus 的生产级向量知识库，适合高性能部署"
|
||
|
||
def __init__(self, work_dir: str, **kwargs):
|
||
"""
|
||
初始化 Milvus 知识库
|
||
|
||
Args:
|
||
work_dir: 工作目录
|
||
**kwargs: 其他配置参数
|
||
"""
|
||
super().__init__(work_dir)
|
||
|
||
if not MILVUS_AVAILABLE:
|
||
raise ImportError("pymilvus is not installed. Please install it with: pip install pymilvus")
|
||
|
||
# Milvus 配置
|
||
# self.milvus_host = kwargs.get('milvus_host', os.getenv('MILVUS_HOST', 'localhost'))
|
||
# self.milvus_port = kwargs.get('milvus_port', int(os.getenv('MILVUS_PORT', '19530')))
|
||
self.milvus_token = kwargs.get("milvus_token", os.getenv("MILVUS_TOKEN") or "")
|
||
self.milvus_uri = kwargs.get("milvus_uri", os.getenv("MILVUS_URI") or "http://localhost:19530")
|
||
self.milvus_db = kwargs.get("milvus_db") or "yuxi"
|
||
|
||
# 连接名称
|
||
self.connection_alias = f"milvus_{hashstr(work_dir, 6)}"
|
||
|
||
# 存储集合映射 {kb_id: Collection}
|
||
self.collections: dict[str, Any] = {}
|
||
|
||
# 初始化连接
|
||
self._init_connection()
|
||
|
||
logger.info("MilvusKB initialized")
|
||
|
||
def _init_connection(self):
|
||
"""初始化 Milvus 连接"""
|
||
try:
|
||
# 连接到 Milvus
|
||
connections.connect(alias=self.connection_alias, uri=self.milvus_uri, token=self.milvus_token)
|
||
|
||
# 创建数据库（如果不存在）
|
||
try:
|
||
if self.milvus_db not in db.list_database():
|
||
db.create_database(self.milvus_db)
|
||
db.using_database(self.milvus_db)
|
||
except Exception as e:
|
||
logger.warning(f"Database operation failed, using default: {e}")
|
||
|
||
logger.info(f"Connected to Milvus at {self.milvus_uri}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to connect to Milvus: {e}")
|
||
raise
|
||
|
||
async def _create_kb_instance(self, kb_id: str, kb_config: dict) -> Any:
|
||
"""创建 Milvus 集合"""
|
||
logger.info(f"Creating Milvus collection for {kb_id}")
|
||
|
||
if not (metadata := self.databases_meta.get(kb_id)):
|
||
raise ValueError(f"Database {kb_id} not found")
|
||
|
||
embedding_model_spec = metadata.get("embedding_model_spec")
|
||
if not embedding_model_spec:
|
||
raise ValueError(f"Embedding model spec not found for database {kb_id}")
|
||
|
||
embedding_info = model_cache.get_model_info(embedding_model_spec)
|
||
if not embedding_info or embedding_info.model_type != "embedding":
|
||
raise ValueError(f"Unsupported embedding model: {embedding_model_spec}")
|
||
|
||
collection_name = kb_id
|
||
|
||
try:
|
||
# 检查集合是否存在
|
||
if utility.has_collection(collection_name, using=self.connection_alias):
|
||
collection = Collection(name=collection_name, using=self.connection_alias)
|
||
|
||
# 检查嵌入模型是否匹配
|
||
description = collection.description
|
||
expected_model = embedding_info.model_id
|
||
|
||
if expected_model not in description:
|
||
logger.warning(
|
||
f"Collection {collection_name} model mismatch: "
|
||
f"expected='{expected_model}', found_in_description='{description}'"
|
||
)
|
||
utility.drop_collection(collection_name, using=self.connection_alias)
|
||
return self._create_new_collection(collection_name, embedding_info, kb_id)
|
||
|
||
if not self._collection_supports_bm25(collection):
|
||
logger.warning(f"Collection {collection_name} schema does not support BM25, recreating")
|
||
utility.drop_collection(collection_name, using=self.connection_alias)
|
||
return self._create_new_collection(collection_name, embedding_info, kb_id)
|
||
|
||
logger.info(f"Retrieved existing collection: {collection_name}")
|
||
return collection
|
||
else:
|
||
logger.info(f"Collection {collection_name} not found, creating new one")
|
||
return self._create_new_collection(collection_name, embedding_info, kb_id)
|
||
|
||
except (connections.MilvusException, RuntimeError) as e:
|
||
logger.error(f"Error checking collection {collection_name}: {e}")
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Unexpected error while managing collection {collection_name}: {e}")
|
||
logger.debug(f"Traceback: {traceback.format_exc()}")
|
||
raise
|
||
|
||
def _create_new_collection(self, collection_name: str, embedding_info: Any, kb_id: str) -> Collection:
|
||
"""创建新的 Milvus 集合"""
|
||
embedding_dim = embedding_info.dimension or 1024
|
||
model_name = embedding_info.model_id
|
||
|
||
# 定义集合Schema
|
||
fields = [
|
||
FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True),
|
||
FieldSchema(
|
||
name="content",
|
||
dtype=DataType.VARCHAR,
|
||
max_length=65535,
|
||
enable_analyzer=True,
|
||
analyzer_params=CONTENT_ANALYZER_PARAMS,
|
||
),
|
||
FieldSchema(name="chunk_id", dtype=DataType.VARCHAR, max_length=100),
|
||
FieldSchema(name="file_id", dtype=DataType.VARCHAR, max_length=100),
|
||
FieldSchema(name="chunk_index", dtype=DataType.INT64),
|
||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim),
|
||
FieldSchema(name=CONTENT_SPARSE_FIELD, dtype=DataType.SPARSE_FLOAT_VECTOR),
|
||
]
|
||
bm25_function = Function(
|
||
name="content_bm25",
|
||
input_field_names=["content"],
|
||
output_field_names=[CONTENT_SPARSE_FIELD],
|
||
function_type=FunctionType.BM25,
|
||
)
|
||
|
||
schema = CollectionSchema(
|
||
fields=fields,
|
||
description=f"Knowledge base collection for {kb_id} using {model_name}",
|
||
functions=[bm25_function],
|
||
)
|
||
|
||
# 创建集合
|
||
collection = Collection(name=collection_name, schema=schema, using=self.connection_alias)
|
||
|
||
# 创建索引
|
||
index_params = {"metric_type": VECTOR_METRIC_TYPE, "index_type": "IVF_FLAT", "params": {"nlist": 1024}}
|
||
collection.create_index("embedding", index_params)
|
||
sparse_index_params = {
|
||
"metric_type": "BM25",
|
||
"index_type": "SPARSE_INVERTED_INDEX",
|
||
"params": {"inverted_index_algo": "DAAT_MAXSCORE"},
|
||
}
|
||
collection.create_index(CONTENT_SPARSE_FIELD, sparse_index_params)
|
||
|
||
logger.info(f"Created new Milvus collection: {collection_name} '{model_name=}', {embedding_dim=}")
|
||
|
||
return collection
|
||
|
||
def _collection_supports_bm25(self, collection: Collection) -> bool:
|
||
"""检查集合是否具备 Milvus 内置 BM25 所需的 schema。"""
|
||
fields = {field.name: field for field in collection.schema.fields}
|
||
content_field = fields.get("content")
|
||
sparse_field = fields.get(CONTENT_SPARSE_FIELD)
|
||
if not content_field or content_field.dtype != DataType.VARCHAR:
|
||
return False
|
||
if content_field.params.get("enable_analyzer") is not True:
|
||
return False
|
||
if not sparse_field or sparse_field.dtype != DataType.SPARSE_FLOAT_VECTOR:
|
||
return False
|
||
|
||
for function in collection.schema.functions:
|
||
if (
|
||
function.type == FunctionType.BM25
|
||
and function.input_field_names == ["content"]
|
||
and function.output_field_names == [CONTENT_SPARSE_FIELD]
|
||
):
|
||
return True
|
||
return False
|
||
|
||
async def _initialize_kb_instance(self, instance: Any) -> None:
|
||
"""初始化 Milvus 集合（加载到内存）"""
|
||
try:
|
||
instance.load()
|
||
logger.info("Milvus collection loaded into memory")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to load collection into memory: {e}")
|
||
|
||
def _get_embedding_function(self, embedding_model_spec: str, *, sync: bool = False):
|
||
"""获取 embedding 编码函数。sync=True 返回同步版本，否则返回异步版本。"""
|
||
from yuxi.models.embed import select_embedding_model
|
||
|
||
model = select_embedding_model(embedding_model_spec)
|
||
batch_size = int(getattr(model, "batch_size", 40) or 40)
|
||
method = model.batch_encode if sync else model.abatch_encode
|
||
return partial(method, batch_size=batch_size)
|
||
|
||
async def _get_milvus_collection(self, kb_id: str):
|
||
"""获取或创建 Milvus 集合"""
|
||
if kb_id in self.collections:
|
||
return self.collections[kb_id]
|
||
|
||
if kb_id not in self.databases_meta:
|
||
return None
|
||
|
||
try:
|
||
# 创建集合
|
||
collection = await self._create_kb_instance(kb_id, {})
|
||
await self._initialize_kb_instance(collection)
|
||
|
||
self.collections[kb_id] = collection
|
||
return collection
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to create Milvus collection for {kb_id}: {e}")
|
||
logger.error(f"Traceback: {traceback.format_exc()}")
|
||
return None
|
||
|
||
def _split_text_into_chunks(self, text: str, file_id: str, filename: str, params: dict) -> list[dict]:
|
||
"""将文本分割成块"""
|
||
return chunk_markdown(text, file_id, filename, params)
|
||
|
||
def _calculate_chunk_stats(self, chunks: list[dict]) -> dict[str, int]:
|
||
return {
|
||
"chunk_count": len(chunks),
|
||
"token_count": sum(count_tokens(chunk["content"]) for chunk in chunks),
|
||
}
|
||
|
||
def _build_chunk_pg_records(self, kb_id: str, chunks: list[dict]) -> list[dict[str, Any]]:
|
||
return [
|
||
{
|
||
"chunk_id": chunk["chunk_id"],
|
||
"file_id": chunk["file_id"],
|
||
"kb_id": kb_id,
|
||
"chunk_index": chunk["chunk_index"],
|
||
"content": chunk["content"],
|
||
"start_char_pos": chunk.get("start_char_pos"),
|
||
"end_char_pos": chunk.get("end_char_pos"),
|
||
"start_token_pos": chunk.get("start_token_pos"),
|
||
"end_token_pos": chunk.get("end_token_pos"),
|
||
"graph_indexed": bool(chunk.get("graph_indexed", False)),
|
||
"ent_ids": chunk.get("ent_ids"),
|
||
"tags": chunk.get("tags"),
|
||
"extraction_result": chunk.get("extraction_result"),
|
||
}
|
||
for chunk in chunks
|
||
]
|
||
|
||
async def _insert_chunks_to_stores(
|
||
self,
|
||
kb_id: str,
|
||
file_id: str,
|
||
collection: Collection,
|
||
chunks: list[dict],
|
||
embeddings: list,
|
||
) -> None:
|
||
if not chunks:
|
||
return
|
||
|
||
entities = [
|
||
[chunk["id"] for chunk in chunks],
|
||
[chunk["content"] for chunk in chunks],
|
||
[chunk["chunk_id"] for chunk in chunks],
|
||
[chunk["file_id"] for chunk in chunks],
|
||
[chunk["chunk_index"] for chunk in chunks],
|
||
embeddings,
|
||
]
|
||
chunk_repo = KnowledgeChunkRepository()
|
||
|
||
def _insert_milvus_records():
|
||
collection.insert(entities)
|
||
|
||
pg_task = chunk_repo.batch_upsert(self._build_chunk_pg_records(kb_id, chunks))
|
||
milvus_task = asyncio.to_thread(_insert_milvus_records)
|
||
results = await asyncio.gather(pg_task, milvus_task, return_exceptions=True)
|
||
errors = [result for result in results if isinstance(result, Exception)]
|
||
if not errors:
|
||
return
|
||
|
||
logger.error(f"Chunk double-write failed for file {file_id}, rolling back PostgreSQL and Milvus chunks")
|
||
try:
|
||
await chunk_repo.delete_by_file_id(file_id)
|
||
except Exception as cleanup_error:
|
||
logger.error(f"Failed to rollback PostgreSQL chunks for {file_id}: {cleanup_error}")
|
||
try:
|
||
await self._delete_file_chunks_from_milvus(collection, file_id)
|
||
except Exception as cleanup_error:
|
||
logger.error(f"Failed to rollback Milvus chunks for {file_id}: {cleanup_error}")
|
||
raise errors[0]
|
||
|
||
async def _embed_and_store_chunks(
|
||
self,
|
||
kb_id: str,
|
||
file_id: str,
|
||
collection: Collection,
|
||
chunks: list[dict],
|
||
embedding_function,
|
||
*,
|
||
chunk_batch_size: int = MILVUS_CHUNK_EMBED_BATCH_SIZE,
|
||
) -> None:
|
||
"""对 chunks 进行分批嵌入并存储到 Milvus 和 PostgreSQL"""
|
||
if not chunks:
|
||
return
|
||
|
||
chunk_batch_size = max(int(chunk_batch_size), 1)
|
||
for start in range(0, len(chunks), chunk_batch_size):
|
||
batch_chunks = chunks[start : start + chunk_batch_size]
|
||
texts = [chunk["content"] for chunk in batch_chunks]
|
||
embeddings = await embedding_function(texts)
|
||
await self._insert_chunks_to_stores(
|
||
kb_id,
|
||
file_id,
|
||
collection,
|
||
batch_chunks,
|
||
embeddings,
|
||
)
|
||
|
||
async def _delete_file_chunks_from_milvus(self, collection: Collection, file_id: str) -> None:
|
||
expr = f'file_id == "{file_id}"'
|
||
results = collection.query(expr=expr, output_fields=["id"], limit=1)
|
||
|
||
if not results:
|
||
logger.info(f"File {file_id} not found in Milvus, skipping delete operation")
|
||
return
|
||
|
||
def _delete_from_milvus():
|
||
collection.delete(expr)
|
||
logger.info(f"Deleted chunks for file {file_id} from Milvus")
|
||
|
||
await asyncio.to_thread(_delete_from_milvus)
|
||
|
||
async def _hydrate_chunk_sources(self, kb_id: str, chunks: list[dict]) -> None:
|
||
file_ids = sorted(
|
||
{str(file_id) for chunk in chunks if (file_id := (chunk.get("metadata") or {}).get("file_id"))}
|
||
)
|
||
if not file_ids:
|
||
return
|
||
|
||
filenames = await KnowledgeFileRepository().get_filenames_by_file_ids(kb_id=kb_id, file_ids=file_ids)
|
||
for chunk in chunks:
|
||
metadata = chunk.get("metadata")
|
||
if not isinstance(metadata, dict):
|
||
continue
|
||
metadata["source"] = filenames.get(str(metadata.get("file_id") or ""), "") or "未知来源"
|
||
|
||
async def _build_file_name_expr(self, kb_id: str, file_name: str | None) -> str | None:
|
||
if not file_name:
|
||
return None
|
||
|
||
matched_file_ids = await KnowledgeFileRepository().list_file_ids_by_filename_contains(
|
||
kb_id=kb_id,
|
||
filename_pattern=file_name,
|
||
)
|
||
if not matched_file_ids:
|
||
return 'file_id == "__no_matching_file__"'
|
||
escaped_ids = [file_id.replace('"', '\\"') for file_id in matched_file_ids]
|
||
if len(escaped_ids) == 1:
|
||
return f'file_id == "{escaped_ids[0]}"'
|
||
joined_ids = '", "'.join(escaped_ids)
|
||
return f'file_id in ["{joined_ids}"]'
|
||
|
||
async def index_file(
|
||
self, kb_id: str, file_id: str, operator_id: str | None = None, params: dict | None = None
|
||
) -> dict:
|
||
"""
|
||
Index parsed file (Status: INDEXING -> INDEXED/ERROR_INDEXING)
|
||
|
||
Args:
|
||
kb_id: Database ID
|
||
file_id: File ID
|
||
operator_id: ID of the user performing the operation
|
||
params: Override processing params to apply during indexing (merged on top of stored params)
|
||
|
||
Returns:
|
||
Updated file metadata
|
||
"""
|
||
if kb_id not in self.databases_meta:
|
||
raise ValueError(f"Database {kb_id} not found")
|
||
|
||
# Get/Create collection
|
||
collection = await self._get_milvus_collection(kb_id)
|
||
if not collection:
|
||
raise ValueError(f"Failed to get Milvus collection for {kb_id}")
|
||
|
||
embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
|
||
embedding_function = self._get_embedding_function(embedding_model_spec)
|
||
|
||
file_meta = await self._load_file_meta(kb_id, file_id)
|
||
allowed_statuses = {
|
||
FileStatus.PARSED,
|
||
FileStatus.ERROR_INDEXING,
|
||
FileStatus.INDEXED,
|
||
"done",
|
||
}
|
||
params = resolve_processing_params(
|
||
kb_additional_params=self.databases_meta.get(kb_id, {}).get("metadata"),
|
||
file_processing_params=file_meta.get("processing_params"),
|
||
request_params=params,
|
||
)
|
||
|
||
claim_data = {
|
||
"status": FileStatus.INDEXING,
|
||
"processing_params": params,
|
||
"error_message": None,
|
||
}
|
||
if operator_id:
|
||
claim_data["updated_by"] = operator_id
|
||
|
||
claimed_record = await KnowledgeFileRepository().update_fields_if_status(
|
||
kb_id=kb_id,
|
||
file_id=file_id,
|
||
allowed_statuses=allowed_statuses,
|
||
data=claim_data,
|
||
)
|
||
if claimed_record is None:
|
||
current_meta = await self._load_file_meta(kb_id, file_id)
|
||
current_status = current_meta.get("status")
|
||
raise ValueError(
|
||
f"Cannot index file with status '{current_status}'. "
|
||
f"File must be parsed first (status should be one of: {', '.join(allowed_statuses)})"
|
||
)
|
||
|
||
file_meta = self._file_record_to_meta(claimed_record)
|
||
if not file_meta.get("markdown_file"):
|
||
await self._mark_file_unparsed(kb_id, file_id, operator_id)
|
||
raise ValueError("File has not been parsed yet (no markdown_file)")
|
||
|
||
logger.debug(f"[index_file] file_id={file_id}, processing_params={params}")
|
||
|
||
try:
|
||
# Read markdown
|
||
markdown_content = await self._read_markdown_from_minio(file_meta["markdown_file"])
|
||
filename = file_meta.get("filename")
|
||
|
||
# Split
|
||
chunks = self._split_text_into_chunks(markdown_content, file_id, filename, params)
|
||
logger.info(
|
||
f"Split {filename} into {len(chunks)} chunks with params: "
|
||
f"chunk_preset_id={params.get('chunk_preset_id')}, "
|
||
f"chunk_parser_config={params.get('chunk_parser_config')}"
|
||
)
|
||
|
||
chunk_stats = self._calculate_chunk_stats(chunks)
|
||
|
||
# Clean up existing chunks if any (for re-indexing)
|
||
await self.delete_file_chunks_only(kb_id, file_id)
|
||
|
||
if chunks:
|
||
await self._embed_and_store_chunks(kb_id, file_id, collection, chunks, embedding_function)
|
||
|
||
logger.info(f"Indexed file {file_id} into Milvus")
|
||
|
||
# Update status
|
||
update_data = {"status": FileStatus.INDEXED, "error_message": None, **chunk_stats}
|
||
if operator_id:
|
||
update_data["updated_by"] = operator_id
|
||
updated_record = await KnowledgeFileRepository().update_fields(
|
||
file_id=file_id,
|
||
kb_id=kb_id,
|
||
data=update_data,
|
||
)
|
||
result = (
|
||
self._file_record_to_meta(updated_record)
|
||
if updated_record is not None
|
||
else {
|
||
**file_meta,
|
||
**chunk_stats,
|
||
"status": FileStatus.INDEXED,
|
||
"error": None,
|
||
}
|
||
)
|
||
|
||
await self.refresh_database_stats(kb_id)
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"Indexing failed for {file_id}: {e}")
|
||
update_data = {"status": FileStatus.ERROR_INDEXING, "error_message": str(e)}
|
||
if operator_id:
|
||
update_data["updated_by"] = operator_id
|
||
await KnowledgeFileRepository().update_fields(file_id=file_id, kb_id=kb_id, data=update_data)
|
||
raise
|
||
|
||
async def update_content(self, kb_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]:
|
||
"""更新内容 - 根据file_ids重新解析文件并更新向量库"""
|
||
if kb_id not in self.databases_meta:
|
||
raise ValueError(f"Database {kb_id} not found")
|
||
|
||
collection = await self._get_milvus_collection(kb_id)
|
||
if not collection:
|
||
raise ValueError(f"Failed to get Milvus collection for {kb_id}")
|
||
|
||
embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
|
||
embedding_function = self._get_embedding_function(embedding_model_spec)
|
||
|
||
# 处理默认参数
|
||
if params is None:
|
||
params = {}
|
||
processed_items_info = []
|
||
|
||
for file_id in file_ids:
|
||
try:
|
||
file_meta = await self._load_file_meta(kb_id, file_id)
|
||
except ValueError:
|
||
logger.warning(f"File {file_id} not found in metadata, skipping")
|
||
continue
|
||
|
||
file_path = file_meta.get("path")
|
||
filename = file_meta.get("filename")
|
||
|
||
if not file_path:
|
||
logger.warning(f"File path not found for {file_id}, skipping")
|
||
continue
|
||
|
||
try:
|
||
# 更新状态为处理中
|
||
resolved_params = resolve_processing_params(
|
||
kb_additional_params=self.databases_meta.get(kb_id, {}).get("metadata"),
|
||
file_processing_params=file_meta.get("processing_params"),
|
||
request_params=params,
|
||
)
|
||
file_meta["processing_params"] = resolved_params
|
||
file_meta["status"] = FileStatus.INDEXING
|
||
await KnowledgeFileRepository().update_fields(
|
||
file_id=file_id,
|
||
kb_id=kb_id,
|
||
data={"status": FileStatus.INDEXING, "processing_params": resolved_params},
|
||
)
|
||
|
||
# 重新解析文件为 markdown
|
||
parse_params = {**resolved_params, "image_bucket": "public", "image_prefix": f"{kb_id}/kb-images"}
|
||
markdown_content = await Parser.aparse(source=file_path, params=parse_params)
|
||
|
||
# 重新生成 chunks
|
||
chunks = self._split_text_into_chunks(markdown_content, file_id, filename, resolved_params)
|
||
logger.info(f"Split {filename} into {len(chunks)} chunks")
|
||
chunk_stats = self._calculate_chunk_stats(chunks)
|
||
|
||
# 先删除现有 chunks，保留文件元数据
|
||
await self.delete_file_chunks_only(kb_id, file_id)
|
||
|
||
if chunks:
|
||
await self._embed_and_store_chunks(kb_id, file_id, collection, chunks, embedding_function)
|
||
|
||
logger.info(f"Updated file {file_path} in Milvus. Done.")
|
||
|
||
# 更新元数据状态
|
||
file_meta["status"] = FileStatus.INDEXED
|
||
file_meta.update(chunk_stats)
|
||
await KnowledgeFileRepository().update_fields(
|
||
file_id=file_id,
|
||
kb_id=kb_id,
|
||
data={"status": FileStatus.INDEXED, "error_message": None, **chunk_stats},
|
||
)
|
||
await self.refresh_database_stats(kb_id)
|
||
|
||
# 返回更新后的文件信息
|
||
updated_file_meta = file_meta.copy()
|
||
updated_file_meta["status"] = FileStatus.INDEXED
|
||
updated_file_meta.update(chunk_stats)
|
||
updated_file_meta["file_id"] = file_id
|
||
processed_items_info.append(updated_file_meta)
|
||
|
||
except Exception as e:
|
||
logger.error(f"更新file {file_path} 失败: {e}, {traceback.format_exc()}")
|
||
await KnowledgeFileRepository().update_fields(
|
||
file_id=file_id,
|
||
kb_id=kb_id,
|
||
data={"status": FileStatus.ERROR_INDEXING, "error_message": str(e)},
|
||
)
|
||
|
||
# 返回失败的文件信息
|
||
failed_file_meta = file_meta.copy()
|
||
failed_file_meta["status"] = FileStatus.ERROR_INDEXING
|
||
failed_file_meta["error"] = str(e)
|
||
failed_file_meta["file_id"] = file_id
|
||
processed_items_info.append(failed_file_meta)
|
||
|
||
return processed_items_info
|
||
|
||
def _build_chunk_from_hit(
|
||
self,
|
||
hit: Any,
|
||
score: float,
|
||
include_distances: bool,
|
||
score_field: str | None = None,
|
||
) -> dict:
|
||
"""将 Milvus Hit 转成知识库统一返回的 Chunk 结构。"""
|
||
entity = hit.entity
|
||
file_id = entity.get("file_id")
|
||
metadata = {
|
||
"source": "未知来源",
|
||
"chunk_id": entity.get("chunk_id"),
|
||
"file_id": file_id,
|
||
"chunk_index": entity.get("chunk_index"),
|
||
}
|
||
chunk = {"content": entity.get("content", ""), "metadata": metadata, "score": float(score or 0.0)}
|
||
if score_field:
|
||
chunk[score_field] = float(score or 0.0)
|
||
if include_distances:
|
||
chunk["distance"] = hit.distance
|
||
return chunk
|
||
|
||
async def aquery(self, query_text: str, kb_id: str, agent_call: bool = False, **kwargs) -> list[dict]:
|
||
"""异步查询知识库"""
|
||
collection = await self._get_milvus_collection(kb_id)
|
||
if not collection:
|
||
raise ValueError(f"Database {kb_id} not found")
|
||
|
||
query_params = self._get_query_params(kb_id)
|
||
# 合并查询参数：kwargs（临时参数）优先级高于 query_params（持久化参数）
|
||
# 这样允许用户在单次查询中临时覆盖持久化配置
|
||
merged_kwargs = {**query_params, **kwargs}
|
||
|
||
try:
|
||
# 查询参数（从 merged_kwargs 读取）
|
||
logger.debug(f"Query params: {merged_kwargs}")
|
||
final_top_k = int(merged_kwargs.get("final_top_k", 10))
|
||
final_top_k = max(final_top_k, 1)
|
||
similarity_threshold = float(merged_kwargs.get("similarity_threshold", 0.2))
|
||
metric_type = VECTOR_METRIC_TYPE
|
||
include_distances = bool(merged_kwargs.get("include_distances", True))
|
||
search_mode = str(merged_kwargs.get("search_mode", "vector")).lower()
|
||
if search_mode not in {"vector", "keyword", "hybrid"}:
|
||
search_mode = "vector"
|
||
|
||
use_reranker = bool(merged_kwargs.get("use_reranker", False))
|
||
use_graph_retrieval = bool(merged_kwargs.get("use_graph_retrieval", False))
|
||
if use_reranker or use_graph_retrieval:
|
||
recall_top_k = int(merged_kwargs.get("recall_top_k", 50))
|
||
recall_top_k = max(recall_top_k, final_top_k)
|
||
else:
|
||
recall_top_k = final_top_k
|
||
|
||
file_expr = await self._build_file_name_expr(kb_id, merged_kwargs.get("file_name"))
|
||
if file_expr:
|
||
logger.debug(f"Using filter expression: {file_expr}")
|
||
|
||
output_fields = ["content", "chunk_id", "file_id", "chunk_index"]
|
||
retrieved_chunks: list[dict] = []
|
||
if search_mode == "vector":
|
||
embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
|
||
embedding_function = self._get_embedding_function(embedding_model_spec, sync=True)
|
||
query_embedding = await _run_milvus_query_io(embedding_function, [query_text])
|
||
|
||
search_params = {"metric_type": metric_type, "params": {"nprobe": 10}}
|
||
|
||
results = await _run_milvus_query_io(
|
||
collection.search,
|
||
data=query_embedding,
|
||
anns_field="embedding",
|
||
param=search_params,
|
||
limit=recall_top_k,
|
||
expr=file_expr,
|
||
output_fields=output_fields,
|
||
)
|
||
|
||
if results and len(results) > 0 and len(results[0]) > 0:
|
||
for hit in results[0]:
|
||
similarity = hit.distance if metric_type == VECTOR_METRIC_TYPE else 1 / (1 + hit.distance)
|
||
if similarity < similarity_threshold:
|
||
continue
|
||
|
||
retrieved_chunks.append(self._build_chunk_from_hit(hit, similarity, include_distances))
|
||
|
||
logger.debug(
|
||
f"Milvus vector query response: {len(retrieved_chunks)} chunks found (after similarity filtering)"
|
||
)
|
||
|
||
elif search_mode == "keyword":
|
||
bm25_top_k = int(merged_kwargs.get("bm25_top_k", recall_top_k))
|
||
bm25_top_k = max(bm25_top_k, 1)
|
||
bm25_drop_ratio_search = float(merged_kwargs.get("bm25_drop_ratio_search", 0.0))
|
||
bm25_search_params = {
|
||
"metric_type": "BM25",
|
||
"params": {"drop_ratio_search": bm25_drop_ratio_search},
|
||
}
|
||
|
||
results = await _run_milvus_query_io(
|
||
collection.search,
|
||
data=[query_text],
|
||
anns_field=CONTENT_SPARSE_FIELD,
|
||
param=bm25_search_params,
|
||
limit=bm25_top_k,
|
||
expr=file_expr,
|
||
output_fields=output_fields,
|
||
)
|
||
|
||
if results and len(results) > 0 and len(results[0]) > 0:
|
||
for hit in results[0]:
|
||
retrieved_chunks.append(
|
||
self._build_chunk_from_hit(hit, hit.distance, include_distances, score_field="bm25_score")
|
||
)
|
||
|
||
logger.debug(f"Milvus BM25 query response: {len(retrieved_chunks)} chunks found")
|
||
else:
|
||
embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
|
||
embedding_function = self._get_embedding_function(embedding_model_spec, sync=True)
|
||
query_embedding = await _run_milvus_query_io(embedding_function, [query_text])
|
||
bm25_top_k = int(merged_kwargs.get("bm25_top_k", recall_top_k))
|
||
bm25_top_k = max(bm25_top_k, 1)
|
||
bm25_drop_ratio_search = float(merged_kwargs.get("bm25_drop_ratio_search", 0.0))
|
||
vector_weight = float(merged_kwargs.get("vector_weight", 0.7))
|
||
bm25_weight = float(merged_kwargs.get("bm25_weight", 0.3))
|
||
|
||
vector_request = AnnSearchRequest(
|
||
data=query_embedding,
|
||
anns_field="embedding",
|
||
param={"metric_type": metric_type, "params": {"nprobe": 10}},
|
||
limit=recall_top_k,
|
||
expr=file_expr,
|
||
)
|
||
bm25_request = AnnSearchRequest(
|
||
data=[query_text],
|
||
anns_field=CONTENT_SPARSE_FIELD,
|
||
param={
|
||
"metric_type": "BM25",
|
||
"params": {"drop_ratio_search": bm25_drop_ratio_search},
|
||
},
|
||
limit=bm25_top_k,
|
||
expr=file_expr,
|
||
)
|
||
results = await _run_milvus_query_io(
|
||
collection.hybrid_search,
|
||
reqs=[vector_request, bm25_request],
|
||
rerank=WeightedRanker(vector_weight, bm25_weight),
|
||
limit=recall_top_k,
|
||
output_fields=output_fields,
|
||
)
|
||
if results and len(results) > 0 and len(results[0]) > 0:
|
||
for hit in results[0]:
|
||
score = float(hit.distance or 0.0)
|
||
if score < similarity_threshold:
|
||
continue
|
||
retrieved_chunks.append(
|
||
self._build_chunk_from_hit(hit, score, include_distances, score_field="hybrid_score")
|
||
)
|
||
|
||
logger.debug(f"Milvus hybrid query response: {len(retrieved_chunks)} chunks found")
|
||
|
||
if use_graph_retrieval:
|
||
graph_chunks = await self._retrieve_graph_chunks(query_text, kb_id, retrieved_chunks, merged_kwargs)
|
||
if graph_chunks:
|
||
graph_weight = float(merged_kwargs.get("graph_weight", 1.0))
|
||
retrieved_chunks = self._fuse_chunk_rankings(retrieved_chunks, graph_chunks, graph_weight)
|
||
|
||
if not retrieved_chunks:
|
||
return []
|
||
|
||
await self._hydrate_chunk_sources(kb_id, retrieved_chunks)
|
||
|
||
if not use_reranker:
|
||
return retrieved_chunks[:final_top_k]
|
||
|
||
# 使用重排序模型
|
||
reranker_model = merged_kwargs.get("reranker_model")
|
||
if not reranker_model:
|
||
raise ValueError(
|
||
"Reranker model must be specified when use_reranker=True. "
|
||
"Please provide reranker_model in query parameters."
|
||
)
|
||
|
||
try:
|
||
from yuxi.models.rerank import get_reranker
|
||
|
||
reranker = get_reranker(reranker_model)
|
||
try:
|
||
rerank_start = time.time()
|
||
documents_text = [chunk["content"] for chunk in retrieved_chunks]
|
||
rerank_scores = await reranker.acompute_score([query_text, documents_text], normalize=True)
|
||
|
||
for chunk, rerank_score in zip(retrieved_chunks, rerank_scores):
|
||
chunk["rerank_score"] = float(rerank_score)
|
||
|
||
retrieved_chunks.sort(
|
||
key=lambda item: item.get("rerank_score", item.get("score", 0.0)), reverse=True
|
||
)
|
||
elapsed = time.time() - rerank_start
|
||
logger.info(f"Reranking completed for {kb_id} in {elapsed:.3f}s with model {reranker_model}")
|
||
finally:
|
||
await reranker.aclose()
|
||
|
||
except Exception as exc: # noqa: BLE001
|
||
logger.error(f"Reranking failed: {exc}, falling back to vector scores")
|
||
|
||
# 统一返回结果
|
||
return retrieved_chunks[:final_top_k]
|
||
|
||
except Exception as e:
|
||
logger.error(f"Milvus query error: {e}, {traceback.format_exc()}")
|
||
return []
|
||
|
||
async def _retrieve_graph_chunks(
|
||
self,
|
||
query_text: str,
|
||
kb_id: str,
|
||
base_chunks: list[dict],
|
||
query_params: dict[str, Any],
|
||
) -> list[dict]:
|
||
try:
|
||
from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService
|
||
from yuxi.knowledge.graphs.milvus_graph_vector_store import MilvusGraphVectorStore
|
||
|
||
embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
|
||
if not embedding_model_spec:
|
||
return []
|
||
|
||
entity_top_k = max(int(query_params.get("graph_entity_top_k", 10)), 1)
|
||
triple_top_k = max(int(query_params.get("graph_triple_top_k", 10)), 1)
|
||
graph_top_k = max(int(query_params.get("graph_top_k", 20)), 1)
|
||
graph_max_nodes = max(int(query_params.get("graph_max_nodes", 10000)), 1)
|
||
|
||
vector_store = await _run_milvus_query_io(MilvusGraphVectorStore)
|
||
entity_hits, triple_hits = await asyncio.gather(
|
||
vector_store.search_entities(
|
||
kb_id=kb_id,
|
||
query_text=query_text,
|
||
embedding_model_spec=embedding_model_spec,
|
||
top_k=entity_top_k,
|
||
),
|
||
vector_store.search_triples(
|
||
kb_id=kb_id,
|
||
query_text=query_text,
|
||
embedding_model_spec=embedding_model_spec,
|
||
top_k=triple_top_k,
|
||
),
|
||
)
|
||
seed_weights = await self._build_graph_seed_weights(kb_id, base_chunks, entity_hits, triple_hits)
|
||
if not seed_weights:
|
||
return []
|
||
|
||
graph_service = MilvusGraphService()
|
||
graph_scores = await graph_service.query_and_rank_chunks_by_ppr(
|
||
kb_id,
|
||
seed_weights,
|
||
max_nodes=graph_max_nodes,
|
||
top_k=graph_top_k,
|
||
damping=float(query_params.get("ppr_damping", 0.85)),
|
||
)
|
||
if not graph_scores:
|
||
return []
|
||
|
||
chunks = await KnowledgeChunkRepository().list_by_chunk_ids([chunk_id for chunk_id, _ in graph_scores])
|
||
score_by_chunk_id = dict(graph_scores)
|
||
return [
|
||
self._build_chunk_from_record(chunk, score_by_chunk_id[chunk.chunk_id], score_field="graph_score")
|
||
for chunk in chunks
|
||
]
|
||
except Exception as exc: # noqa: BLE001
|
||
logger.error(f"Graph retrieval failed for {kb_id}: {exc}")
|
||
return []
|
||
|
||
async def _build_graph_seed_weights(
|
||
self,
|
||
kb_id: str,
|
||
base_chunks: list[dict],
|
||
entity_hits: list[dict[str, Any]],
|
||
triple_hits: list[dict[str, Any]],
|
||
) -> dict[str, float]:
|
||
seed_weights: dict[str, float] = {}
|
||
|
||
def add_seed(entity_id: str | None, score: float, weight: float) -> None:
|
||
if not entity_id:
|
||
return
|
||
seed_weights[entity_id] = seed_weights.get(entity_id, 0.0) + max(float(score or 0.0), 0.0) * weight
|
||
|
||
for hit in entity_hits:
|
||
add_seed(hit.get("id"), hit.get("score", 0.0), 1.0)
|
||
|
||
for hit in triple_hits:
|
||
score = float(hit.get("score") or 0.0)
|
||
add_seed(hit.get("source_id"), score, 0.8)
|
||
add_seed(hit.get("target_id"), score, 0.8)
|
||
|
||
chunk_scores = {
|
||
chunk.get("metadata", {}).get("chunk_id"): float(chunk.get("score") or 0.0)
|
||
for chunk in base_chunks
|
||
if chunk.get("metadata", {}).get("chunk_id")
|
||
}
|
||
if chunk_scores:
|
||
chunks = await KnowledgeChunkRepository().list_by_chunk_ids(list(chunk_scores))
|
||
for chunk in chunks:
|
||
for entity_id in chunk.ent_ids or []:
|
||
add_seed(entity_id, chunk_scores.get(chunk.chunk_id, 0.0), 0.3)
|
||
|
||
total = sum(seed_weights.values())
|
||
if total <= 0:
|
||
return {}
|
||
return {entity_id: weight / total for entity_id, weight in seed_weights.items()}
|
||
|
||
def _build_chunk_from_record(self, chunk: Any, score: float, score_field: str | None = None) -> dict:
|
||
metadata = {
|
||
"source": "未知来源",
|
||
"chunk_id": chunk.chunk_id,
|
||
"file_id": chunk.file_id,
|
||
"chunk_index": chunk.chunk_index,
|
||
}
|
||
result = {"content": chunk.content, "metadata": metadata, "score": float(score or 0.0)}
|
||
if score_field:
|
||
result[score_field] = float(score or 0.0)
|
||
return result
|
||
|
||
def _fuse_chunk_rankings(
|
||
self,
|
||
base_chunks: list[dict],
|
||
graph_chunks: list[dict],
|
||
graph_weight: float,
|
||
) -> list[dict]:
|
||
fused: dict[str, dict[str, Any]] = {}
|
||
rrf_k = 60.0
|
||
|
||
def merge_chunk(chunk: dict, rank: int, weight: float, source: str) -> None:
|
||
chunk_id = chunk.get("metadata", {}).get("chunk_id")
|
||
if not chunk_id:
|
||
return
|
||
score = weight / (rrf_k + rank)
|
||
existing = fused.get(chunk_id)
|
||
if existing is None:
|
||
existing = {**chunk, "fusion_score": 0.0, "fusion_sources": []}
|
||
fused[chunk_id] = existing
|
||
existing["fusion_score"] += score
|
||
existing["score"] = existing["fusion_score"]
|
||
existing["fusion_sources"].append(source)
|
||
if source == "graph" and "graph_score" in chunk:
|
||
existing["graph_score"] = chunk["graph_score"]
|
||
|
||
for rank, chunk in enumerate(base_chunks, start=1):
|
||
merge_chunk(chunk, rank, 1.0, "chunk")
|
||
for rank, chunk in enumerate(graph_chunks, start=1):
|
||
merge_chunk(chunk, rank, max(graph_weight, 0.0), "graph")
|
||
|
||
return sorted(fused.values(), key=lambda item: item.get("fusion_score", 0.0), reverse=True)
|
||
|
||
async def delete_file_chunks_only(self, kb_id: str, file_id: str) -> None:
|
||
"""仅删除文件的chunks数据，保留元数据（用于更新操作）"""
|
||
chunk_repo = KnowledgeChunkRepository()
|
||
if await chunk_repo.count_graph_indexed_by_file_id(file_id):
|
||
from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService
|
||
|
||
try:
|
||
await MilvusGraphService().delete_file_graph(kb_id, file_id)
|
||
except Exception as e:
|
||
logger.error(f"Failed to delete graph data for file {file_id}: {e}")
|
||
await chunk_repo.delete_by_file_id(file_id)
|
||
collection = await self._get_milvus_collection(kb_id)
|
||
|
||
if collection:
|
||
# 先查询文件是否存在，避免不必要的删除操作
|
||
try:
|
||
await self._delete_file_chunks_from_milvus(collection, file_id)
|
||
except Exception as e:
|
||
logger.error(f"Error checking file existence in Milvus: {e}")
|
||
await KnowledgeFileRepository().update_fields(
|
||
file_id=file_id,
|
||
kb_id=kb_id,
|
||
data={"chunk_count": 0, "token_count": 0},
|
||
)
|
||
await self.refresh_database_stats(kb_id)
|
||
|
||
async def delete_file(self, kb_id: str, file_id: str) -> None:
|
||
"""删除文件（包括元数据）"""
|
||
# 先删除 Milvus 中的 chunks 数据
|
||
await self.delete_file_chunks_only(kb_id, file_id)
|
||
|
||
await KnowledgeFileRepository().delete(file_id)
|
||
await self.refresh_database_stats(kb_id)
|
||
|
||
async def get_file_basic_info(self, kb_id: str, file_id: str) -> dict:
|
||
"""获取文件基本信息（仅元数据）"""
|
||
return {"meta": await self._load_file_meta(kb_id, file_id)}
|
||
|
||
async def _get_file_content_from_meta(self, file_id: str, file_meta: dict) -> dict:
|
||
content_info = {"lines": []}
|
||
try:
|
||
chunks = await KnowledgeChunkRepository().list_by_file_id(file_id)
|
||
content_info["lines"] = [
|
||
{
|
||
"id": chunk.chunk_id,
|
||
"content": chunk.content,
|
||
"chunk_order_index": chunk.chunk_index,
|
||
"start_char_pos": chunk.start_char_pos,
|
||
"end_char_pos": chunk.end_char_pos,
|
||
"start_token_pos": chunk.start_token_pos,
|
||
"end_token_pos": chunk.end_token_pos,
|
||
"graph_indexed": chunk.graph_indexed,
|
||
"ent_ids": chunk.ent_ids,
|
||
"tags": chunk.tags,
|
||
"extraction_result": chunk.extraction_result,
|
||
}
|
||
for chunk in chunks
|
||
]
|
||
except Exception as e:
|
||
logger.error(f"Failed to get file content from PostgreSQL: {e}")
|
||
|
||
if not content_info["lines"]:
|
||
logger.warning(f"No chunks found in PostgreSQL for file {file_id}, file may not have been indexed")
|
||
|
||
# Try to read markdown content if available
|
||
if file_meta.get("markdown_file"):
|
||
try:
|
||
content = await self._read_markdown_from_minio(file_meta["markdown_file"])
|
||
content_info["content"] = content
|
||
except Exception as e:
|
||
logger.error(f"Failed to read markdown file for {file_id}: {e}")
|
||
|
||
return content_info
|
||
|
||
async def get_file_content(self, kb_id: str, file_id: str) -> dict:
|
||
"""获取文件内容信息（chunks和lines）"""
|
||
file_meta = await self._load_file_meta(kb_id, file_id)
|
||
return await self._get_file_content_from_meta(file_id, file_meta)
|
||
|
||
async def get_file_info(self, kb_id: str, file_id: str) -> dict:
|
||
"""获取文件完整信息（基本信息+内容信息）"""
|
||
file_meta = await self._load_file_meta(kb_id, file_id)
|
||
content_info = await self._get_file_content_from_meta(file_id, file_meta)
|
||
return {"meta": file_meta, **content_info}
|
||
|
||
def delete_database(self, kb_id: str) -> dict:
|
||
"""删除数据库，同时清除Milvus中的集合"""
|
||
# Drop Milvus collection
|
||
try:
|
||
if utility.has_collection(kb_id, using=self.connection_alias):
|
||
utility.drop_collection(kb_id, using=self.connection_alias)
|
||
logger.info(f"Dropped Milvus collection for {kb_id}")
|
||
else:
|
||
logger.info(f"Milvus collection {kb_id} does not exist, skipping")
|
||
except Exception as e:
|
||
logger.error(f"Failed to drop Milvus collection {kb_id}: {e}")
|
||
|
||
from yuxi.knowledge.graphs.milvus_graph_vector_store import MilvusGraphVectorStore
|
||
|
||
MilvusGraphVectorStore().drop_graph_collections(kb_id)
|
||
|
||
# Call base method to delete local files and metadata
|
||
return super().delete_database(kb_id)
|
||
|
||
def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
|
||
"""获取 Milvus 知识库的查询参数配置"""
|
||
return {"type": "milvus", "options": _retrieval_config_options()}
|
||
|
||
def __del__(self):
|
||
"""清理连接"""
|
||
try:
|
||
if hasattr(self, "connection_alias"):
|
||
connections.disconnect(self.connection_alias)
|
||
except Exception: # noqa: S110
|
||
pass
|