kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

refactor: 知识库系统重构 - 统一命名、优化图检索与向量存储逻辑

Browse Source 
- base/manager: db_id → kb_id 命名统一，重构元数据加载与文件记录逻辑
- implementations: dify/milvus/notion/read_only 适配新命名与接口
- graphs: milvus_graph_service/vector_store/graph_utils 优化图构建与检索逻辑
- eval: benchmark_generation/evaluator/service 适配新命名并优化评估流程
- utils/kb_utils: 对应更新工具函数
This commit is contained in:
Wenjie Zhang 2026-05-21 19:32:25 +08:00
parent c772e5ca3a
commit 99c590d003
13 changed files with 584 additions and 580 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

316
backend/package/yuxi/knowledge/base.py
View File

@ -2,6 +2,8 @@ import asyncio
import mimetypes import mimetypes
import os import os
import re import re
import secrets
import string
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Any from typing import Any
@ -83,10 +85,10 @@ class KnowledgeBase(ABC):
"""由 KnowledgeBaseManager 调用，同步加载元数据""" """由 KnowledgeBaseManager 调用，同步加载元数据"""
# 过滤出当前 kb_type 的知识库 # 过滤出当前 kb_type 的知识库
self.databases_meta = {} self.databases_meta = {}
for db_id, meta in global_databases_meta.items(): for kb_id, meta in global_databases_meta.items():
if meta.get("kb_type") == self.kb_type: if meta.get("kb_type") == self.kb_type:
normalized_additional_params = self.normalize_additional_params(meta.get("additional_params")) normalized_additional_params = self.normalize_additional_params(meta.get("additional_params"))
self.databases_meta[db_id] = { self.databases_meta[kb_id] = {
"name": meta.get("name"), "name": meta.get("name"),
"description": meta.get("description"), "description": meta.get("description"),
"kb_type": meta.get("kb_type"), "kb_type": meta.get("kb_type"),
@ -100,9 +102,9 @@ class KnowledgeBase(ABC):
# 过滤文件 # 过滤文件
self.files_meta = {} self.files_meta = {}
for file_id, meta in files_meta.items(): for file_id, meta in files_meta.items():
if meta.get("database_id") in self.databases_meta: if meta.get("kb_id") in self.databases_meta:
db_id = meta.get("database_id") kb_id = meta.get("kb_id")
kb_additional_params = self.databases_meta.get(db_id, {}).get("metadata") or {} kb_additional_params = self.databases_meta.get(kb_id, {}).get("metadata") or {}
normalized_meta = dict(meta) normalized_meta = dict(meta)
normalized_meta["processing_params"] = resolve_processing_params( normalized_meta["processing_params"] = resolve_processing_params(
kb_additional_params=kb_additional_params, kb_additional_params=kb_additional_params,
@ -182,12 +184,12 @@ class KnowledgeBase(ABC):
return params return params
@abstractmethod @abstractmethod
async def _create_kb_instance(self, db_id: str, config: dict) -> Any: async def _create_kb_instance(self, kb_id: str, config: dict) -> Any:
""" """
创建底层知识库实例 创建底层知识库实例
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
config: 配置信息 config: 配置信息
Returns: Returns:
@ -206,13 +208,13 @@ class KnowledgeBase(ABC):
pass pass
async def add_file_record( async def add_file_record(
self, db_id: str, item: str, params: dict | None = None, operator_id: str | None = None self, kb_id: str, item: str, params: dict | None = None, operator_id: str | None = None
) -> dict: ) -> dict:
""" """
Add a file record to metadata (Status: UPLOADED) Add a file record to metadata (Status: UPLOADED)
Args: Args:
db_id: Database ID kb_id: Database ID
item: File path or URL item: File path or URL
params: Parameters params: Parameters
operator_id: Operator ID who created the file operator_id: Operator ID who created the file
@ -226,9 +228,9 @@ class KnowledgeBase(ABC):
content_type = params.get("content_type", "file") content_type = params.get("content_type", "file")
# Prepare metadata # Prepare metadata
metadata = await prepare_item_metadata(item, content_type, db_id, params=params) metadata = await prepare_item_metadata(item, content_type, kb_id, params=params)
file_id = metadata["file_id"] file_id = metadata["file_id"]
kb_additional_params = self.databases_meta.get(db_id, {}).get("metadata") or {} kb_additional_params = self.databases_meta.get(kb_id, {}).get("metadata") or {}
metadata["processing_params"] = resolve_processing_params( metadata["processing_params"] = resolve_processing_params(
kb_additional_params=kb_additional_params, kb_additional_params=kb_additional_params,
file_processing_params=metadata.get("processing_params"), file_processing_params=metadata.get("processing_params"),
@ -262,12 +264,12 @@ class KnowledgeBase(ABC):
return metadata return metadata
async def parse_file(self, db_id: str, file_id: str, operator_id: str | None = None) -> dict: async def parse_file(self, kb_id: str, file_id: str, operator_id: str | None = None) -> dict:
""" """
Parse file to Markdown and save to MinIO (Status: PARSING -> PARSED/ERROR_PARSING) Parse file to Markdown and save to MinIO (Status: PARSING -> PARSED/ERROR_PARSING)
Args: Args:
db_id: Database ID kb_id: Database ID
file_id: File ID file_id: File ID
operator_id: ID of the user performing the operation operator_id: ID of the user performing the operation
@ -317,7 +319,7 @@ class KnowledgeBase(ABC):
# Prepare params # Prepare params
params = file_meta.get("processing_params", {}) or {} params = file_meta.get("processing_params", {}) or {}
params["image_bucket"] = "public" params["image_bucket"] = "public"
params["image_prefix"] = f"{db_id}/kb-images" params["image_prefix"] = f"{kb_id}/kb-images"
markdown_content = await Parser.aparse( markdown_content = await Parser.aparse(
source=file_path, source=file_path,
@ -325,7 +327,7 @@ class KnowledgeBase(ABC):
) )
# Save Markdown to MinIO # Save Markdown to MinIO
markdown_file_path = await self._save_markdown_to_minio(db_id, file_id, markdown_content) markdown_file_path = await self._save_markdown_to_minio(kb_id, file_id, markdown_content)
# Update metadata # Update metadata
self.files_meta[file_id]["status"] = FileStatus.PARSED self.files_meta[file_id]["status"] = FileStatus.PARSED
@ -354,7 +356,7 @@ class KnowledgeBase(ABC):
# Remove from processing queue # Remove from processing queue
self._remove_from_processing_queue(file_id) self._remove_from_processing_queue(file_id)
async def update_file_params(self, db_id: str, file_id: str, params: dict, operator_id: str | None = None) -> None: async def update_file_params(self, kb_id: str, file_id: str, params: dict, operator_id: str | None = None) -> None:
"""Update file processing params""" """Update file processing params"""
if file_id not in self.files_meta: if file_id not in self.files_meta:
raise ValueError(f"File {file_id} not found") raise ValueError(f"File {file_id} not found")
@ -364,7 +366,7 @@ class KnowledgeBase(ABC):
return return
current_params = self.files_meta[file_id].get("processing_params", {}) or {} current_params = self.files_meta[file_id].get("processing_params", {}) or {}
kb_additional_params = self.databases_meta.get(db_id, {}).get("metadata") or {} kb_additional_params = self.databases_meta.get(kb_id, {}).get("metadata") or {}
logger.debug(f"[update_file_params] file_id={file_id}, current_params={current_params}, new_params={params}") logger.debug(f"[update_file_params] file_id={file_id}, current_params={current_params}, new_params={params}")
@ -395,7 +397,7 @@ class KnowledgeBase(ABC):
self.files_meta[file_id]["updated_by"] = operator_id self.files_meta[file_id]["updated_by"] = operator_id
await self._persist_file(file_id) await self._persist_file(file_id)
async def _save_markdown_to_minio(self, db_id: str, file_id: str, content: str) -> str: async def _save_markdown_to_minio(self, kb_id: str, file_id: str, content: str) -> str:
"""Save markdown content to MinIO and return HTTP URL""" """Save markdown content to MinIO and return HTTP URL"""
from yuxi.storage.minio import get_minio_client from yuxi.storage.minio import get_minio_client
@ -403,7 +405,7 @@ class KnowledgeBase(ABC):
bucket_name = minio_client.KB_BUCKETS["parsed"] bucket_name = minio_client.KB_BUCKETS["parsed"]
await asyncio.to_thread(minio_client.ensure_bucket_exists, bucket_name) await asyncio.to_thread(minio_client.ensure_bucket_exists, bucket_name)
object_name = f"{db_id}/parsed/{file_id}.md" object_name = f"{kb_id}/parsed/{file_id}.md"
data = content.encode("utf-8") data = content.encode("utf-8")
# Return standard HTTP URL from UploadResult # Return standard HTTP URL from UploadResult
@ -431,9 +433,9 @@ class KnowledgeBase(ABC):
content_bytes = await self._read_minio_bytes(file_path) content_bytes = await self._read_minio_bytes(file_path)
return content_bytes.decode("utf-8") return content_bytes.decode("utf-8")
def _get_file_meta(self, db_id: str, file_id: str) -> dict: def _get_file_meta(self, kb_id: str, file_id: str) -> dict:
file_meta = self.files_meta.get(file_id) file_meta = self.files_meta.get(file_id)
if not file_meta or file_meta.get("database_id") != db_id: if not file_meta or file_meta.get("kb_id") != kb_id:
raise ValueError(f"File {file_id} not found") raise ValueError(f"File {file_id} not found")
return file_meta return file_meta
@ -450,7 +452,7 @@ class KnowledgeBase(ABC):
variants.append({"key": "parsed", "label": "MD", "supported": True}) variants.append({"key": "parsed", "label": "MD", "supported": True})
return variants return variants
def _knowledge_file_entry(self, db_id: str, file_id: str, file_meta: dict) -> dict: def _knowledge_file_entry(self, kb_id: str, file_id: str, file_meta: dict) -> dict:
is_dir = bool(file_meta.get("is_folder")) is_dir = bool(file_meta.get("is_folder"))
variants = [] if is_dir else self._knowledge_preview_variants(file_meta) variants = [] if is_dir else self._knowledge_preview_variants(file_meta)
preview_modes = [item["key"] for item in variants] preview_modes = [item["key"] for item in variants]
@ -464,11 +466,11 @@ class KnowledgeBase(ABC):
path = f"{path}/" path = f"{path}/"
return { return {
"source": "knowledge", "source": "knowledge",
"db_id": db_id, "kb_id": kb_id,
"file_id": file_id, "file_id": file_id,
"parent_id": file_meta.get("parent_id"), "parent_id": file_meta.get("parent_id"),
"path": path, "path": path,
"virtual_path": f"/knowledge/{db_id}/{file_id}", "virtual_path": f"/knowledge/{kb_id}/{file_id}",
"name": file_meta.get("filename") or file_meta.get("original_filename") or file_id, "name": file_meta.get("filename") or file_meta.get("original_filename") or file_id,
"is_dir": is_dir, "is_dir": is_dir,
"size": 0 if is_dir else file_meta.get("size") or 0, "size": 0 if is_dir else file_meta.get("size") or 0,
@ -487,7 +489,7 @@ class KnowledgeBase(ABC):
def _list_knowledge_children( def _list_knowledge_children(
self, self,
db_id: str, kb_id: str,
parent_id: str | None, parent_id: str | None,
*, *,
recursive: bool, recursive: bool,
@ -496,16 +498,16 @@ class KnowledgeBase(ABC):
children = [ children = [
(file_id, meta) (file_id, meta)
for file_id, meta in self.files_meta.items() for file_id, meta in self.files_meta.items()
if meta.get("database_id") == db_id and meta.get("parent_id") == parent_id if meta.get("kb_id") == kb_id and meta.get("parent_id") == parent_id
] ]
entries = [] entries = []
for file_id, meta in children: for file_id, meta in children:
if not files_only or not meta.get("is_folder"): if not files_only or not meta.get("is_folder"):
entries.append(self._knowledge_file_entry(db_id, file_id, meta)) entries.append(self._knowledge_file_entry(kb_id, file_id, meta))
if recursive and meta.get("is_folder"): if recursive and meta.get("is_folder"):
entries.extend( entries.extend(
self._list_knowledge_children( self._list_knowledge_children(
db_id, kb_id,
file_id, file_id,
recursive=True, recursive=True,
files_only=files_only, files_only=files_only,
@ -515,20 +517,20 @@ class KnowledgeBase(ABC):
async def list_file_tree( async def list_file_tree(
self, self,
db_id: str, kb_id: str,
parent_id: str | None = None, parent_id: str | None = None,
recursive: bool = False, recursive: bool = False,
files_only: bool = False, files_only: bool = False,
) -> dict: ) -> dict:
if db_id not in self.databases_meta: if kb_id not in self.databases_meta:
raise ValueError(f"Database {db_id} not found") raise ValueError(f"Database {kb_id} not found")
if parent_id: if parent_id:
parent_meta = self._get_file_meta(db_id, parent_id) parent_meta = self._get_file_meta(kb_id, parent_id)
if not parent_meta.get("is_folder"): if not parent_meta.get("is_folder"):
raise ValueError("Parent is not a folder") raise ValueError("Parent is not a folder")
return { return {
"entries": self._list_knowledge_children( "entries": self._list_knowledge_children(
db_id, kb_id,
parent_id, parent_id,
recursive=recursive, recursive=recursive,
files_only=files_only, files_only=files_only,
@ -536,10 +538,10 @@ class KnowledgeBase(ABC):
"readonly": True, "readonly": True,
} }
async def read_file_preview(self, db_id: str, file_id: str, variant: str = "parsed") -> dict: async def read_file_preview(self, kb_id: str, file_id: str, variant: str = "parsed") -> dict:
from yuxi.services.viewer_filesystem_service import _detect_preview_type from yuxi.services.viewer_filesystem_service import _detect_preview_type
file_meta = self._get_file_meta(db_id, file_id) file_meta = self._get_file_meta(kb_id, file_id)
if file_meta.get("is_folder"): if file_meta.get("is_folder"):
raise ValueError("Cannot preview a folder") raise ValueError("Cannot preview a folder")
@ -551,7 +553,7 @@ class KnowledgeBase(ABC):
filename = file_meta.get("filename") or file_meta.get("original_filename") or file_id filename = file_meta.get("filename") or file_meta.get("original_filename") or file_id
response = { response = {
"source": "knowledge", "source": "knowledge",
"db_id": db_id, "kb_id": kb_id,
"file_id": file_id, "file_id": file_id,
"filename": filename, "filename": filename,
"variant": variant, "variant": variant,
@ -626,8 +628,8 @@ class KnowledgeBase(ABC):
"message": message, "message": message,
} }
async def get_file_download(self, db_id: str, file_id: str, variant: str = "original") -> dict: async def get_file_download(self, kb_id: str, file_id: str, variant: str = "original") -> dict:
file_meta = self._get_file_meta(db_id, file_id) file_meta = self._get_file_meta(kb_id, file_id)
if file_meta.get("is_folder"): if file_meta.get("is_folder"):
raise ValueError("Cannot download a folder") raise ValueError("Cannot download a folder")
if variant not in {"original", "parsed"}: if variant not in {"original", "parsed"}:
@ -675,7 +677,7 @@ class KnowledgeBase(ABC):
} }
@staticmethod @staticmethod
def build_search_output(resource_id: str, retrieval_results: Any) -> dict[str, Any] | Any: def build_search_output(kb_id: str, retrieval_results: Any) -> dict[str, Any] | Any:
if not isinstance(retrieval_results, list): if not isinstance(retrieval_results, list):
return retrieval_results return retrieval_results
@ -705,14 +707,14 @@ class KnowledgeBase(ABC):
results.append( results.append(
SearchResultSchema( SearchResultSchema(
id=str(chunk_id or f"{file_id}:{index + 1}"), id=str(chunk_id or f"{file_id}:{index + 1}"),
resource_id=str(resource_id), kb_id=str(kb_id),
file_id=str(file_id or ""), file_id=str(file_id or ""),
content=str(chunk.get("content") or ""), content=str(chunk.get("content") or ""),
metadata=metadata, metadata=metadata,
) )
) )
return SearchOutputSchema(resource_id=str(resource_id), results=results).model_dump() return SearchOutputSchema(kb_id=str(kb_id), results=results).model_dump()
@staticmethod @staticmethod
def _build_find_file_windows( def _build_find_file_windows(
@ -770,21 +772,21 @@ class KnowledgeBase(ABC):
break break
return FindOutputSchema( return FindOutputSchema(
resource_id="", kb_id="",
file_id="", file_id="",
semantic=False, semantic=False,
match_mode="regex" if use_regex else "keyword", match_mode="regex" if use_regex else "keyword",
total_matches=len(matched_indexes), total_matches=len(matched_indexes),
windows=windows, windows=windows,
).model_dump(exclude={"resource_id", "file_id"}) ).model_dump(exclude={"kb_id", "file_id"})
async def open_file_content(self, db_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict: async def open_file_content(self, kb_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict:
"""按行窗口打开文件解析后的 Markdown 内容""" """按行窗口打开文件解析后的 Markdown 内容"""
file_meta = self.files_meta.get(file_id) file_meta = self.files_meta.get(file_id)
if file_meta is None: if file_meta is None:
raise Exception(f"文件不存在: {file_id}") raise Exception(f"文件不存在: {file_id}")
if file_meta.get("database_id") != db_id: if file_meta.get("kb_id") != kb_id:
raise Exception(f"文件 {file_id} 不属于知识库 {db_id}") raise Exception(f"文件 {file_id} 不属于知识库 {kb_id}")
if file_meta.get("is_folder"): if file_meta.get("is_folder"):
raise Exception(f"文件 {file_id} 是文件夹") raise Exception(f"文件 {file_id} 是文件夹")
@ -797,7 +799,7 @@ class KnowledgeBase(ABC):
async def find_file_content( async def find_file_content(
self, self,
db_id: str, kb_id: str,
file_id: str, file_id: str,
patterns: list[str], patterns: list[str],
*, *,
@ -809,8 +811,8 @@ class KnowledgeBase(ABC):
file_meta = self.files_meta.get(file_id) file_meta = self.files_meta.get(file_id)
if file_meta is None: if file_meta is None:
raise Exception(f"文件不存在: {file_id}") raise Exception(f"文件不存在: {file_id}")
if file_meta.get("database_id") != db_id: if file_meta.get("kb_id") != kb_id:
raise Exception(f"文件 {file_id} 不属于知识库 {db_id}") raise Exception(f"文件 {file_id} 不属于知识库 {kb_id}")
if file_meta.get("is_folder"): if file_meta.get("is_folder"):
raise Exception(f"文件 {file_id} 是文件夹") raise Exception(f"文件 {file_id} 是文件夹")
@ -829,12 +831,12 @@ class KnowledgeBase(ABC):
) )
@abstractmethod @abstractmethod
async def index_file(self, db_id: str, file_id: str, operator_id: str | None = None) -> dict: async def index_file(self, kb_id: str, file_id: str, operator_id: str | None = None) -> dict:
""" """
Index parsed file (Status: INDEXING -> INDEXED/ERROR_INDEXING) Index parsed file (Status: INDEXING -> INDEXED/ERROR_INDEXING)
Args: Args:
db_id: Database ID kb_id: Database ID
file_id: File ID file_id: File ID
operator_id: ID of the user performing the operation operator_id: ID of the user performing the operation
@ -864,13 +866,15 @@ class KnowledgeBase(ABC):
Returns: Returns:
数据库信息字典 数据库信息字典
""" """
from yuxi.utils import hashstr
kwargs = self.normalize_additional_params(kwargs) kwargs = self.normalize_additional_params(kwargs)
db_id = f"kb_{hashstr(database_name, with_salt=True, length=32)}" alphabet = string.ascii_lowercase + string.digits
while True:
kb_id = "kb_" + "".join(secrets.choice(alphabet) for _ in range(10))
if kb_id not in self.databases_meta:
break
self.databases_meta[db_id] = { self.databases_meta[kb_id] = {
"name": database_name, "name": database_name,
"description": description, "description": description,
"kb_type": self.kb_type, "kb_type": self.kb_type,
@ -878,32 +882,32 @@ class KnowledgeBase(ABC):
"llm_model_spec": llm_model_spec, "llm_model_spec": llm_model_spec,
"metadata": kwargs, "metadata": kwargs,
"created_at": utc_isoformat(), "created_at": utc_isoformat(),
"query_params": self._get_default_query_params(db_id), "query_params": self._get_default_query_params(kb_id),
} }
await self._persist_kb(db_id) await self._persist_kb(kb_id)
# 创建工作目录 # 创建工作目录
working_dir = os.path.join(self.work_dir, db_id) working_dir = os.path.join(self.work_dir, kb_id)
os.makedirs(working_dir, exist_ok=True) os.makedirs(working_dir, exist_ok=True)
# 返回数据库信息 # 返回数据库信息
db_dict = self.databases_meta[db_id].copy() db_dict = self.databases_meta[kb_id].copy()
db_dict["db_id"] = db_id db_dict["kb_id"] = kb_id
db_dict["files"] = {} db_dict["files"] = {}
return db_dict return db_dict
async def delete_database(self, db_id: str) -> dict: async def delete_database(self, kb_id: str) -> dict:
""" """
删除数据库 删除数据库
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
Returns: Returns:
操作结果 操作结果
""" """
if db_id in self.databases_meta: if kb_id in self.databases_meta:
from yuxi.knowledge.utils.kb_utils import is_minio_url, parse_minio_url from yuxi.knowledge.utils.kb_utils import is_minio_url, parse_minio_url
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
from yuxi.storage.minio import get_minio_client from yuxi.storage.minio import get_minio_client
@ -911,7 +915,7 @@ class KnowledgeBase(ABC):
minio_client = get_minio_client() minio_client = get_minio_client()
# 1. 删除文件元数据中记录的 MinIO 文件 # 1. 删除文件元数据中记录的 MinIO 文件
files_to_delete = [fid for fid, finfo in self.files_meta.items() if finfo.get("database_id") == db_id] files_to_delete = [fid for fid, finfo in self.files_meta.items() if finfo.get("kb_id") == kb_id]
for file_id in files_to_delete: for file_id in files_to_delete:
file_path = self.files_meta[file_id].get("path") file_path = self.files_meta[file_id].get("path")
if file_path and is_minio_url(file_path): if file_path and is_minio_url(file_path):
@ -922,13 +926,13 @@ class KnowledgeBase(ABC):
logger.warning(f"Failed to delete MinIO file {file_path}: {e}") logger.warning(f"Failed to delete MinIO file {file_path}: {e}")
# 删除解析后的 markdown 文件 # 删除解析后的 markdown 文件
parsed_object = f"{db_id}/parsed/{file_id}.md" parsed_object = f"{kb_id}/parsed/{file_id}.md"
await minio_client.adelete_file(minio_client.KB_BUCKETS["parsed"], parsed_object) await minio_client.adelete_file(minio_client.KB_BUCKETS["parsed"], parsed_object)
del self.files_meta[file_id] del self.files_meta[file_id]
# 2. 并行删除所有知识库 bucket 中该 db_id 下的文件 # 2. 并行删除所有知识库 bucket 中该 kb_id 下的文件
prefix = f"{db_id}/" prefix = f"{kb_id}/"
cleanup_buckets = { cleanup_buckets = {
minio_client.KB_BUCKETS["parsed"], minio_client.KB_BUCKETS["parsed"],
minio_client.KB_BUCKETS["documents"], minio_client.KB_BUCKETS["documents"],
@ -940,13 +944,13 @@ class KnowledgeBase(ABC):
await asyncio.gather(*cleanup_tasks) await asyncio.gather(*cleanup_tasks)
# 3. 删除数据库记录 # 3. 删除数据库记录
del self.databases_meta[db_id] del self.databases_meta[kb_id]
kb_repo = KnowledgeBaseRepository() kb_repo = KnowledgeBaseRepository()
await kb_repo.delete(db_id) await kb_repo.delete(kb_id)
await self._save_metadata() await self._save_metadata()
# 删除工作目录 # 删除工作目录
working_dir = os.path.join(self.work_dir, db_id) working_dir = os.path.join(self.work_dir, kb_id)
if os.path.exists(working_dir): if os.path.exists(working_dir):
import shutil import shutil
@ -957,7 +961,7 @@ class KnowledgeBase(ABC):
return {"message": "删除成功"} return {"message": "删除成功"}
async def create_folder(self, db_id: str, folder_name: str, parent_id: str | None = None) -> dict: async def create_folder(self, kb_id: str, folder_name: str, parent_id: str | None = None) -> dict:
"""Create a folder in the database.""" """Create a folder in the database."""
import uuid import uuid
@ -968,7 +972,7 @@ class KnowledgeBase(ABC):
"filename": folder_name, "filename": folder_name,
"is_folder": True, "is_folder": True,
"parent_id": parent_id, "parent_id": parent_id,
"database_id": db_id, "kb_id": kb_id,
"created_at": utc_isoformat(), "created_at": utc_isoformat(),
"status": "done", "status": "done",
"path": folder_name, "path": folder_name,
@ -978,12 +982,12 @@ class KnowledgeBase(ABC):
return self.files_meta[folder_id] return self.files_meta[folder_id]
@abstractmethod @abstractmethod
async def update_content(self, db_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]: async def update_content(self, kb_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]:
""" """
更新内容 - 根据file_ids重新解析文件并更新向量库 更新内容 - 根据file_ids重新解析文件并更新向量库
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
file_ids: 文件ID列表 file_ids: 文件ID列表
params: 处理参数 params: 处理参数
@ -993,13 +997,13 @@ class KnowledgeBase(ABC):
pass pass
@abstractmethod @abstractmethod
async def aquery(self, query_text: str, db_id: str, **kwargs) -> list[dict]: async def aquery(self, query_text: str, kb_id: str, **kwargs) -> list[dict]:
""" """
异步查询知识库 异步查询知识库
Args: Args:
query_text: 查询文本 query_text: 查询文本
db_id: 数据库ID kb_id: 数据库ID
**kwargs: 查询参数 **kwargs: 查询参数
Returns: Returns:
@ -1008,12 +1012,12 @@ class KnowledgeBase(ABC):
pass pass
@abstractmethod @abstractmethod
def get_query_params_config(self, db_id: str, **kwargs) -> dict: def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
""" """
获取知识库类型的查询参数配置 获取知识库类型的查询参数配置
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
**kwargs: 额外参数 **kwargs: 额外参数
Returns: Returns:
@ -1037,54 +1041,54 @@ class KnowledgeBase(ABC):
""" """
pass pass
async def export_data(self, db_id: str, format: str = "zip", **kwargs) -> str: async def export_data(self, kb_id: str, format: str = "zip", **kwargs) -> str:
pass pass
def _get_query_params(self, db_id: str) -> dict: def _get_query_params(self, kb_id: str) -> dict:
"""从实例元数据中加载查询参数""" """从实例元数据中加载查询参数"""
if db_id in self.databases_meta: if kb_id in self.databases_meta:
query_params_meta = self.databases_meta[db_id].get("query_params") or {} query_params_meta = self.databases_meta[kb_id].get("query_params") or {}
return query_params_meta.get("options", {}) return query_params_meta.get("options", {})
return {} return {}
def _get_default_query_params(self, db_id: str) -> dict[str, Any]: def _get_default_query_params(self, kb_id: str) -> dict[str, Any]:
"""从 get_query_params_config 中提取所有参数的默认值，返回 {"options": {...}}""" """从 get_query_params_config 中提取所有参数的默认值，返回 {"options": {...}}"""
config = self.get_query_params_config(db_id) config = self.get_query_params_config(kb_id)
defaults = {} defaults = {}
for opt in config.get("options", []): for opt in config.get("options", []):
if "default" in opt: if "default" in opt:
defaults[opt["key"]] = opt["default"] defaults[opt["key"]] = opt["default"]
return {"options": defaults} return {"options": defaults}
def get_database_info(self, db_id: str, include_files: bool = True) -> dict | None: def get_database_info(self, kb_id: str, include_files: bool = True) -> dict | None:
""" """
获取数据库详细信息 获取数据库详细信息
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
include_files: 是否包含文件信息，默认为True include_files: 是否包含文件信息，默认为True
Returns: Returns:
数据库信息或None 数据库信息或None
""" """
if db_id not in self.databases_meta: if kb_id not in self.databases_meta:
return None return None
meta = self.databases_meta[db_id].copy() meta = self.databases_meta[kb_id].copy()
meta["db_id"] = db_id meta["kb_id"] = kb_id
# 检查并修复异常的processing状态 # 检查并修复异常的processing状态
self._check_and_fix_processing_status(db_id) self._check_and_fix_processing_status(kb_id)
# 统计文件数量（始终计算，即使不加载文件详情） # 统计文件数量（始终计算，即使不加载文件详情）
db_file_count = sum(1 for file_info in self.files_meta.values() if file_info.get("database_id") == db_id) db_file_count = sum(1 for file_info in self.files_meta.values() if file_info.get("kb_id") == kb_id)
meta["row_count"] = db_file_count meta["row_count"] = db_file_count
# 仅在需要时加载文件详情 # 仅在需要时加载文件详情
if include_files: if include_files:
db_files = {} db_files = {}
for file_id, file_info in self.files_meta.items(): for file_id, file_info in self.files_meta.items():
if file_info.get("database_id") == db_id: if file_info.get("kb_id") == kb_id:
created_at = self._normalize_timestamp(file_info.get("created_at")) created_at = self._normalize_timestamp(file_info.get("created_at"))
db_files[file_id] = { db_files[file_id] = {
"file_id": file_id, "file_id": file_id,
@ -1126,22 +1130,22 @@ class KnowledgeBase(ABC):
self._ensure_metadata_loaded() self._ensure_metadata_loaded()
databases = [] databases = []
for db_id, meta in self.databases_meta.items(): for kb_id, meta in self.databases_meta.items():
# 检查并修复异常的processing状态 # 检查并修复异常的processing状态
self._check_and_fix_processing_status(db_id) self._check_and_fix_processing_status(kb_id)
db_dict = meta.copy() db_dict = meta.copy()
db_dict["db_id"] = db_id db_dict["kb_id"] = kb_id
# 统计文件数量（始终计算，即使不加载文件详情） # 统计文件数量（始终计算，即使不加载文件详情）
db_file_count = sum(1 for file_info in self.files_meta.values() if file_info.get("database_id") == db_id) db_file_count = sum(1 for file_info in self.files_meta.values() if file_info.get("kb_id") == kb_id)
db_dict["row_count"] = db_file_count db_dict["row_count"] = db_file_count
# 仅在需要时加载文件详情 # 仅在需要时加载文件详情
if include_files: if include_files:
db_files = {} db_files = {}
for file_id, file_info in self.files_meta.items(): for file_id, file_info in self.files_meta.items():
if file_info.get("database_id") == db_id: if file_info.get("kb_id") == kb_id:
created_at = self._normalize_timestamp(file_info.get("created_at")) created_at = self._normalize_timestamp(file_info.get("created_at"))
db_files[file_id] = { db_files[file_id] = {
"file_id": file_id, "file_id": file_id,
@ -1209,13 +1213,13 @@ class KnowledgeBase(ABC):
with cls._processing_lock: with cls._processing_lock:
return file_id in cls._processing_files return file_id in cls._processing_files
def _check_and_fix_processing_status(self, db_id: str) -> None: def _check_and_fix_processing_status(self, kb_id: str) -> None:
""" """
检查并修复异常的处理中状态 检查并修复异常的处理中状态
如果文件状态为处理中但实际不在处理队列中，则修改为相应的错误状态 如果文件状态为处理中但实际不在处理队列中，则修改为相应的错误状态
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
""" """
try: try:
status_changed = False status_changed = False
@ -1228,7 +1232,7 @@ class KnowledgeBase(ABC):
# 检查该数据库下所有中间状态的文件 # 检查该数据库下所有中间状态的文件
for file_id, file_info in self.files_meta.items(): for file_id, file_info in self.files_meta.items():
if file_info.get("database_id") == db_id: if file_info.get("kb_id") == kb_id:
current_status = file_info.get("status") current_status = file_info.get("status")
if current_status in intermediate_states: if current_status in intermediate_states:
@ -1248,44 +1252,44 @@ class KnowledgeBase(ABC):
# 如果有状态变更，保存元数据 # 如果有状态变更，保存元数据
if status_changed: if status_changed:
logger.info(f"Fixed interrupted processing status for database {db_id}") logger.info(f"Fixed interrupted processing status for database {kb_id}")
except Exception as e: except Exception as e:
logger.error(f"Error checking processing status for database {db_id}: {e}") logger.error(f"Error checking processing status for database {kb_id}: {e}")
async def delete_folder(self, db_id: str, folder_id: str) -> None: async def delete_folder(self, kb_id: str, folder_id: str) -> None:
""" """
Recursively delete a folder and its content. Recursively delete a folder and its content.
Args: Args:
db_id: Database ID kb_id: Database ID
folder_id: Folder ID to delete folder_id: Folder ID to delete
""" """
# Find all children # Find all children
children = [ children = [
fid fid
for fid, meta in self.files_meta.items() for fid, meta in self.files_meta.items()
if meta.get("database_id") == db_id and meta.get("parent_id") == folder_id if meta.get("kb_id") == kb_id and meta.get("parent_id") == folder_id
] ]
for child_id in children: for child_id in children:
child_meta = self.files_meta.get(child_id) child_meta = self.files_meta.get(child_id)
if child_meta and child_meta.get("is_folder"): if child_meta and child_meta.get("is_folder"):
await self.delete_folder(db_id, child_id) await self.delete_folder(kb_id, child_id)
else: else:
await self.delete_file(db_id, child_id) await self.delete_file(kb_id, child_id)
# Delete the folder itself # Delete the folder itself
# We call delete_file which should handle the actual removal. # We call delete_file which should handle the actual removal.
# Implementations should ensure they handle folder deletion gracefully (e.g. skip vector deletion) # Implementations should ensure they handle folder deletion gracefully (e.g. skip vector deletion)
await self.delete_file(db_id, folder_id) await self.delete_file(kb_id, folder_id)
async def move_file(self, db_id: str, file_id: str, new_parent_id: str | None) -> dict: async def move_file(self, kb_id: str, file_id: str, new_parent_id: str | None) -> dict:
""" """
Move a file or folder to a new parent folder. Move a file or folder to a new parent folder.
Args: Args:
db_id: Database ID kb_id: Database ID
file_id: File/Folder ID to move file_id: File/Folder ID to move
new_parent_id: New parent folder ID (None for root) new_parent_id: New parent folder ID (None for root)
@ -1296,8 +1300,8 @@ class KnowledgeBase(ABC):
raise ValueError(f"File {file_id} not found") raise ValueError(f"File {file_id} not found")
meta = self.files_meta[file_id] meta = self.files_meta[file_id]
if meta.get("database_id") != db_id: if meta.get("kb_id") != kb_id:
raise ValueError(f"File {file_id} does not belong to database {db_id}") raise ValueError(f"File {file_id} does not belong to database {kb_id}")
# Basic cycle detection for folders # Basic cycle detection for folders
if meta.get("is_folder") and new_parent_id: if meta.get("is_folder") and new_parent_id:
@ -1320,23 +1324,23 @@ class KnowledgeBase(ABC):
return meta return meta
@abstractmethod @abstractmethod
async def delete_file(self, db_id: str, file_id: str) -> None: async def delete_file(self, kb_id: str, file_id: str) -> None:
""" """
删除文件 删除文件
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
file_id: 文件ID file_id: 文件ID
""" """
pass pass
@abstractmethod @abstractmethod
async def get_file_basic_info(self, db_id: str, file_id: str) -> dict: async def get_file_basic_info(self, kb_id: str, file_id: str) -> dict:
""" """
获取文件基本信息（仅元数据） 获取文件基本信息（仅元数据）
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
file_id: 文件ID file_id: 文件ID
Returns: Returns:
@ -1345,12 +1349,12 @@ class KnowledgeBase(ABC):
pass pass
@abstractmethod @abstractmethod
async def get_file_content(self, db_id: str, file_id: str) -> dict: async def get_file_content(self, kb_id: str, file_id: str) -> dict:
""" """
获取文件内容信息（chunks和lines） 获取文件内容信息（chunks和lines）
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
file_id: 文件ID file_id: 文件ID
Returns: Returns:
@ -1359,12 +1363,12 @@ class KnowledgeBase(ABC):
pass pass
@abstractmethod @abstractmethod
async def get_file_info(self, db_id: str, file_id: str) -> dict: async def get_file_info(self, kb_id: str, file_id: str) -> dict:
""" """
获取文件完整信息（基本信息+内容信息） 获取文件完整信息（基本信息+内容信息）
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
file_id: 文件ID file_id: 文件ID
Returns: Returns:
@ -1374,7 +1378,7 @@ class KnowledgeBase(ABC):
def update_database( def update_database(
self, self,
db_id: str, kb_id: str,
name: str, name: str,
description: str, description: str,
llm_model_spec: str | None = None, llm_model_spec: str | None = None,
@ -1384,7 +1388,7 @@ class KnowledgeBase(ABC):
更新数据库 更新数据库
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
name: 新名称 name: 新名称
description: 新描述 description: 新描述
llm_model_spec: LLM 模型 spec（可选） llm_model_spec: LLM 模型 spec（可选）
@ -1392,15 +1396,15 @@ class KnowledgeBase(ABC):
Returns: Returns:
更新后的数据库信息 更新后的数据库信息
""" """
if db_id not in self.databases_meta: if kb_id not in self.databases_meta:
raise ValueError(f"数据库 {db_id} 不存在") raise ValueError(f"数据库 {kb_id} 不存在")
self.databases_meta[db_id]["name"] = name self.databases_meta[kb_id]["name"] = name
self.databases_meta[db_id]["description"] = description self.databases_meta[kb_id]["description"] = description
if update_llm_model_spec: if update_llm_model_spec:
self.databases_meta[db_id]["llm_model_spec"] = llm_model_spec self.databases_meta[kb_id]["llm_model_spec"] = llm_model_spec
return self.get_database_info(db_id) return self.get_database_info(kb_id)
def get_retrievers(self) -> dict[str, dict]: def get_retrievers(self) -> dict[str, dict]:
""" """
@ -1410,19 +1414,19 @@ class KnowledgeBase(ABC):
检索器字典 检索器字典
""" """
retrievers = {} retrievers = {}
for db_id, meta in self.databases_meta.items(): for kb_id, meta in self.databases_meta.items():
def make_retriever(db_id): def make_retriever(kb_id):
async def retriever(query_text, **kwargs): async def retriever(query_text, **kwargs):
results = await self.aquery(query_text, db_id, agent_call=True, **kwargs) results = await self.aquery(query_text, kb_id, agent_call=True, **kwargs)
return self.build_search_output(db_id, results) return self.build_search_output(kb_id, results)
return retriever return retriever
retrievers[db_id] = { retrievers[kb_id] = {
"name": meta["name"], "name": meta["name"],
"description": meta["description"], "description": meta["description"],
"retriever": make_retriever(db_id), "retriever": make_retriever(kb_id),
"metadata": meta, "metadata": meta,
} }
return retrievers return retrievers
@ -1436,13 +1440,13 @@ class KnowledgeBase(ABC):
databases = [kb for kb in await kb_repo.get_all() if kb.kb_type == self.kb_type] databases = [kb for kb in await kb_repo.get_all() if kb.kb_type == self.kb_type]
self.databases_meta = { self.databases_meta = {
kb.db_id: { kb.kb_id: {
"name": kb.name, "name": kb.name,
"description": kb.description, "description": kb.description,
"kb_type": kb.kb_type, "kb_type": kb.kb_type,
"embedding_model_spec": kb.embedding_model_spec, "embedding_model_spec": kb.embedding_model_spec,
"llm_model_spec": kb.llm_model_spec, "llm_model_spec": kb.llm_model_spec,
"query_params": kb.query_params or self._get_default_query_params(kb.db_id), "query_params": kb.query_params or self._get_default_query_params(kb.kb_id),
"metadata": self.normalize_additional_params(kb.additional_params), "metadata": self.normalize_additional_params(kb.additional_params),
"created_at": utc_isoformat(kb.created_at) if kb.created_at else utc_isoformat(), "created_at": utc_isoformat(kb.created_at) if kb.created_at else utc_isoformat(),
} }
@ -1451,11 +1455,11 @@ class KnowledgeBase(ABC):
self.files_meta = {} self.files_meta = {}
for kb in databases: for kb in databases:
kb_additional_params = self.databases_meta.get(kb.db_id, {}).get("metadata") or {} kb_additional_params = self.databases_meta.get(kb.kb_id, {}).get("metadata") or {}
for record in await file_repo.list_by_db_id(kb.db_id): for record in await file_repo.list_by_kb_id(kb.kb_id):
self.files_meta[record.file_id] = { self.files_meta[record.file_id] = {
"file_id": record.file_id, "file_id": record.file_id,
"database_id": record.db_id, "kb_id": record.kb_id,
"parent_id": record.parent_id, "parent_id": record.parent_id,
"filename": record.filename, "filename": record.filename,
"file_type": record.file_type, "file_type": record.file_type,
@ -1542,11 +1546,11 @@ class KnowledgeBase(ABC):
self._normalize_metadata_state() self._normalize_metadata_state()
for db_id, meta in self.databases_meta.items(): for kb_id, meta in self.databases_meta.items():
existing = await kb_repo.get_by_id(db_id) existing = await kb_repo.get_by_kb_id(kb_id)
payload = { payload = {
"db_id": db_id, "kb_id": kb_id,
"name": meta.get("name") or db_id, "name": meta.get("name") or kb_id,
"description": meta.get("description"), "description": meta.get("description"),
"kb_type": meta.get("kb_type") or self.kb_type, "kb_type": meta.get("kb_type") or self.kb_type,
"embedding_model_spec": meta.get("embedding_model_spec"), "embedding_model_spec": meta.get("embedding_model_spec"),
@ -1558,13 +1562,13 @@ class KnowledgeBase(ABC):
await kb_repo.create(payload) await kb_repo.create(payload)
for file_id, meta in self.files_meta.items(): for file_id, meta in self.files_meta.items():
db_id = meta.get("database_id") kb_id = meta.get("kb_id")
if not db_id: if not kb_id:
continue continue
await file_repo.upsert( await file_repo.upsert(
file_id=file_id, file_id=file_id,
data={ data={
"db_id": db_id, "kb_id": kb_id,
"parent_id": meta.get("parent_id"), "parent_id": meta.get("parent_id"),
"filename": meta.get("filename") or "", "filename": meta.get("filename") or "",
"original_filename": meta.get("original_filename"), "original_filename": meta.get("original_filename"),
@ -1594,14 +1598,14 @@ class KnowledgeBase(ABC):
return return
meta = self.files_meta[file_id] meta = self.files_meta[file_id]
db_id = meta.get("database_id") kb_id = meta.get("kb_id")
if not db_id: if not kb_id:
return return
await file_repo.upsert( await file_repo.upsert(
file_id=file_id, file_id=file_id,
data={ data={
"db_id": db_id, "kb_id": kb_id,
"parent_id": meta.get("parent_id"), "parent_id": meta.get("parent_id"),
"filename": meta.get("filename") or "", "filename": meta.get("filename") or "",
"original_filename": meta.get("original_filename"), "original_filename": meta.get("original_filename"),
@ -1621,20 +1625,20 @@ class KnowledgeBase(ABC):
}, },
) )
async def _persist_kb(self, db_id: str) -> None: async def _persist_kb(self, kb_id: str) -> None:
"""只保存单个知识库到数据库，避免全量遍历""" """只保存单个知识库到数据库，避免全量遍历"""
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
kb_repo = KnowledgeBaseRepository() kb_repo = KnowledgeBaseRepository()
if db_id not in self.databases_meta: if kb_id not in self.databases_meta:
return return
meta = self.databases_meta[db_id] meta = self.databases_meta[kb_id]
existing = await kb_repo.get_by_id(db_id) existing = await kb_repo.get_by_kb_id(kb_id)
payload = { payload = {
"db_id": db_id, "kb_id": kb_id,
"name": meta.get("name") or db_id, "name": meta.get("name") or kb_id,
"description": meta.get("description"), "description": meta.get("description"),
"kb_type": meta.get("kb_type") or self.kb_type, "kb_type": meta.get("kb_type") or self.kb_type,
"embedding_model_spec": meta.get("embedding_model_spec"), "embedding_model_spec": meta.get("embedding_model_spec"),
@ -1647,7 +1651,7 @@ class KnowledgeBase(ABC):
await kb_repo.create(payload) await kb_repo.create(payload)
else: else:
await kb_repo.update( await kb_repo.update(
db_id, kb_id,
{ {
"name": payload["name"], "name": payload["name"],
"description": payload["description"], "description": payload["description"],

26
backend/package/yuxi/knowledge/eval/benchmark_generation.py
View File

@ -18,13 +18,13 @@ GRAPH_PPR_DAMPING = 0.85
GRAPH_PPR_MAX_NODES = 10000 GRAPH_PPR_MAX_NODES = 10000
async def collect_kb_chunks(kb_instance: Any, db_id: str) -> list[dict[str, Any]]: async def collect_kb_chunks(kb_instance: Any, kb_id: str) -> list[dict[str, Any]]:
chunks = [] chunks = []
for fid, finfo in kb_instance.files_meta.items(): for fid, finfo in kb_instance.files_meta.items():
if finfo.get("database_id") != db_id: if finfo.get("kb_id") != kb_id:
continue continue
try: try:
content_info = await kb_instance.get_file_content(db_id, fid) content_info = await kb_instance.get_file_content(kb_id, fid)
for line in content_info.get("lines", []): for line in content_info.get("lines", []):
chunks.append( chunks.append(
{ {
@ -75,7 +75,7 @@ def _is_anchor_chunk(candidate: dict[str, Any], anchor_chunk: dict[str, Any]) ->
async def select_neighbor_chunks_by_kb_query( async def select_neighbor_chunks_by_kb_query(
*, kb_instance: Any, db_id: str, anchor_chunk: dict[str, Any], neighbors_count: int *, kb_instance: Any, kb_id: str, anchor_chunk: dict[str, Any], neighbors_count: int
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
if neighbors_count <= 0: if neighbors_count <= 0:
return [] return []
@ -86,7 +86,7 @@ async def select_neighbor_chunks_by_kb_query(
candidates = await kb_instance.aquery( candidates = await kb_instance.aquery(
anchor_content, anchor_content,
db_id, kb_id,
search_mode="vector", search_mode="vector",
final_top_k=neighbors_count + 3, final_top_k=neighbors_count + 3,
use_reranker=False, use_reranker=False,
@ -120,7 +120,7 @@ async def select_neighbor_chunks_by_kb_query(
async def select_graph_enhanced_chunks( async def select_graph_enhanced_chunks(
*, *,
db_id: str, kb_id: str,
anchor_chunk: dict[str, Any], anchor_chunk: dict[str, Any],
chunks_by_id: dict[str, dict[str, Any]], chunks_by_id: dict[str, dict[str, Any]],
context_count: int, context_count: int,
@ -146,7 +146,7 @@ async def select_graph_enhanced_chunks(
seed_weights[entity_id] = 1.0 seed_weights[entity_id] = 1.0
ranked_chunks = await graph_service.query_and_rank_chunks_by_ppr( ranked_chunks = await graph_service.query_and_rank_chunks_by_ppr(
db_id, kb_id,
seed_weights, seed_weights,
max_nodes=GRAPH_PPR_MAX_NODES, max_nodes=GRAPH_PPR_MAX_NODES,
top_k=max(context_count * 5, 20), top_k=max(context_count * 5, 20),
@ -194,7 +194,7 @@ def build_benchmark_generation_prompt(ctx_items: list[tuple[str, str]]) -> str:
async def _generate_benchmark_item_once( async def _generate_benchmark_item_once(
*, *,
kb_instance: Any, kb_instance: Any,
db_id: str, kb_id: str,
all_chunks: list[dict[str, Any]], all_chunks: list[dict[str, Any]],
llm: Any, llm: Any,
context_count: int, context_count: int,
@ -210,7 +210,7 @@ async def _generate_benchmark_item_once(
raise ValueError("No graph indexed chunks with entities found in knowledge base") raise ValueError("No graph indexed chunks with entities found in knowledge base")
anchor_chunk = graph_anchor_chunks[random.randrange(len(graph_anchor_chunks))] anchor_chunk = graph_anchor_chunks[random.randrange(len(graph_anchor_chunks))]
ctx_chunks = await select_graph_enhanced_chunks( ctx_chunks = await select_graph_enhanced_chunks(
db_id=db_id, kb_id=kb_id,
anchor_chunk=anchor_chunk, anchor_chunk=anchor_chunk,
chunks_by_id=chunks_by_id, chunks_by_id=chunks_by_id,
context_count=context_count, context_count=context_count,
@ -222,7 +222,7 @@ async def _generate_benchmark_item_once(
anchor_chunk = all_chunks[random.randrange(len(all_chunks))] anchor_chunk = all_chunks[random.randrange(len(all_chunks))]
neighbor_chunks = await select_neighbor_chunks_by_kb_query( neighbor_chunks = await select_neighbor_chunks_by_kb_query(
kb_instance=kb_instance, kb_instance=kb_instance,
db_id=db_id, kb_id=kb_id,
anchor_chunk=anchor_chunk, anchor_chunk=anchor_chunk,
neighbors_count=context_count - 1, neighbors_count=context_count - 1,
) )
@ -254,7 +254,7 @@ async def _generate_benchmark_item_once(
async def iter_generated_benchmark_items( async def iter_generated_benchmark_items(
*, *,
kb_instance: Any, kb_instance: Any,
db_id: str, kb_id: str,
count: int, count: int,
neighbors_count: int, neighbors_count: int,
llm_model_spec: str | None, llm_model_spec: str | None,
@ -267,7 +267,7 @@ async def iter_generated_benchmark_items(
if progress_cb: if progress_cb:
await progress_cb(5, "加载chunks") await progress_cb(5, "加载chunks")
all_chunks = await collect_kb_chunks(kb_instance, db_id) all_chunks = await collect_kb_chunks(kb_instance, kb_id)
if not all_chunks: if not all_chunks:
raise ValueError("No chunks found in knowledge base") raise ValueError("No chunks found in knowledge base")
chunks_by_id = {str(chunk["id"]): chunk for chunk in all_chunks if chunk.get("id") is not None} chunks_by_id = {str(chunk["id"]): chunk for chunk in all_chunks if chunk.get("id") is not None}
@ -310,7 +310,7 @@ async def iter_generated_benchmark_items(
try: try:
item = await _generate_benchmark_item_once( item = await _generate_benchmark_item_once(
kb_instance=kb_instance, kb_instance=kb_instance,
db_id=db_id, kb_id=kb_id,
all_chunks=all_chunks, all_chunks=all_chunks,
llm=llm, llm=llm,
context_count=context_count, context_count=context_count,

4
backend/package/yuxi/knowledge/eval/evaluator.py
View File

@ -58,7 +58,7 @@ async def generate_answer_if_needed(
async def evaluate_question( async def evaluate_question(
*, *,
kb_instance: Any, kb_instance: Any,
db_id: str, kb_id: str,
question_data: dict[str, Any], question_data: dict[str, Any],
retrieval_config: dict[str, Any], retrieval_config: dict[str, Any],
has_gold_chunks: bool, has_gold_chunks: bool,
@ -67,7 +67,7 @@ async def evaluate_question(
select_model_fn: Callable[..., Any], select_model_fn: Callable[..., Any],
) -> dict[str, Any]: ) -> dict[str, Any]:
query = question_data["query"] query = question_data["query"]
query_result = await kb_instance.aquery(query, db_id, **retrieval_config) query_result = await kb_instance.aquery(query, kb_id, **retrieval_config)
generated_answer, retrieved_chunks = normalize_query_result(query_result) generated_answer, retrieved_chunks = normalize_query_result(query_result)
generated_answer = await generate_answer_if_needed( generated_answer = await generate_answer_if_needed(
query=query, query=query,

74
backend/package/yuxi/knowledge/eval/service.py
View File

@ -35,7 +35,7 @@ class EvaluationService:
"dataset_id": row.dataset_id, "dataset_id": row.dataset_id,
"name": row.name, "name": row.name,
"description": row.description, "description": row.description,
"db_id": row.db_id, "kb_id": row.kb_id,
"item_count": row.item_count, "item_count": row.item_count,
"has_gold_chunks": row.has_gold_chunks, "has_gold_chunks": row.has_gold_chunks,
"has_gold_answers": row.has_gold_answers, "has_gold_answers": row.has_gold_answers,
@ -93,13 +93,13 @@ class EvaluationService:
row.build_metadata = metadata row.build_metadata = metadata
def _build_dataset_items( def _build_dataset_items(
self, dataset_id: str, db_id: str, questions: list[dict[str, Any]] self, dataset_id: str, kb_id: str, questions: list[dict[str, Any]]
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
return [ return [
{ {
"item_id": f"dataset_item_{uuid.uuid4().hex[:12]}", "item_id": f"dataset_item_{uuid.uuid4().hex[:12]}",
"dataset_id": dataset_id, "dataset_id": dataset_id,
"db_id": db_id, "kb_id": kb_id,
"item_index": index, "item_index": index,
"query_text": item["query"], "query_text": item["query"],
"gold_chunk_ids": item.get("gold_chunk_ids") or [], "gold_chunk_ids": item.get("gold_chunk_ids") or [],
@ -152,7 +152,7 @@ class EvaluationService:
return questions, has_gold_chunks, has_gold_answers return questions, has_gold_chunks, has_gold_answers
async def upload_dataset( async def upload_dataset(
self, db_id: str, file_content: bytes, filename: str, name: str, description: str, created_by: str self, kb_id: str, file_content: bytes, filename: str, name: str, description: str, created_by: str
) -> dict[str, Any]: ) -> dict[str, Any]:
try: try:
questions, has_gold_chunks, has_gold_answers = self._parse_jsonl_questions(file_content) questions, has_gold_chunks, has_gold_answers = self._parse_jsonl_questions(file_content)
@ -162,7 +162,7 @@ class EvaluationService:
row = await self.eval_repo.create_dataset_with_items( row = await self.eval_repo.create_dataset_with_items(
{ {
"dataset_id": dataset_id, "dataset_id": dataset_id,
"db_id": db_id, "kb_id": kb_id,
"name": dataset_name, "name": dataset_name,
"description": description, "description": description,
"item_count": len(questions), "item_count": len(questions),
@ -176,16 +176,16 @@ class EvaluationService:
}, },
"created_by": created_by, "created_by": created_by,
}, },
self._build_dataset_items(dataset_id, db_id, questions), self._build_dataset_items(dataset_id, kb_id, questions),
) )
return self._dataset_to_dict(row) return self._dataset_to_dict(row)
except Exception as e: except Exception as e:
logger.error(f"上传评估数据集失败: {e}") logger.error(f"上传评估数据集失败: {e}")
raise raise
async def list_datasets(self, db_id: str) -> list[dict[str, Any]]: async def list_datasets(self, kb_id: str) -> list[dict[str, Any]]:
try: try:
rows = await self.eval_repo.list_datasets(db_id) rows = await self.eval_repo.list_datasets(kb_id)
for row in rows: for row in rows:
await self._sync_dataset_build_metadata(row) await self._sync_dataset_build_metadata(row)
return [self._dataset_to_dict(row) for row in rows] return [self._dataset_to_dict(row) for row in rows]
@ -194,11 +194,11 @@ class EvaluationService:
raise raise
async def get_dataset_detail( async def get_dataset_detail(
self, db_id: str, dataset_id: str, page: int = 1, page_size: int = 10 self, kb_id: str, dataset_id: str, page: int = 1, page_size: int = 10
) -> dict[str, Any]: ) -> dict[str, Any]:
try: try:
row = await self.eval_repo.get_dataset(dataset_id) row = await self.eval_repo.get_dataset(dataset_id)
if row is None or row.db_id != db_id: if row is None or row.kb_id != kb_id:
raise ValueError("Dataset not found") raise ValueError("Dataset not found")
if (row.build_metadata or {}).get("status", "completed") != "completed": if (row.build_metadata or {}).get("status", "completed") != "completed":
raise ValueError("Dataset is not ready") raise ValueError("Dataset is not ready")
@ -250,7 +250,7 @@ class EvaluationService:
async def generate_dataset( async def generate_dataset(
self, self,
db_id: str, kb_id: str,
name: str, name: str,
description: str, description: str,
count: int, count: int,
@ -269,7 +269,7 @@ class EvaluationService:
if generation_mode not in {"vector", "graph_enhanced"}: if generation_mode not in {"vector", "graph_enhanced"}:
raise ValueError("不支持的评估基准生成方式") raise ValueError("不支持的评估基准生成方式")
if generation_mode == "graph_enhanced": if generation_mode == "graph_enhanced":
indexed_count = await self.chunk_repo.count_graph_indexed_by_db_id(db_id) indexed_count = await self.chunk_repo.count_graph_indexed_by_kb_id(kb_id)
if indexed_count <= 0: if indexed_count <= 0:
raise ValueError("当前知识库尚未完成图索引，无法使用图增强构建") raise ValueError("当前知识库尚未完成图索引，无法使用图增强构建")
build_metadata = { build_metadata = {
@ -288,7 +288,7 @@ class EvaluationService:
await self.eval_repo.create_dataset( await self.eval_repo.create_dataset(
{ {
"dataset_id": dataset_id, "dataset_id": dataset_id,
"db_id": db_id, "kb_id": kb_id,
"name": name, "name": name,
"description": description, "description": description,
"item_count": 0, "item_count": 0,
@ -303,7 +303,7 @@ class EvaluationService:
task_type="dataset_generation", task_type="dataset_generation",
payload={ payload={
"dataset_id": dataset_id, "dataset_id": dataset_id,
"db_id": db_id, "kb_id": kb_id,
"created_by": created_by, "created_by": created_by,
"name": name, "name": name,
"description": description, "description": description,
@ -333,7 +333,7 @@ class EvaluationService:
payload = task.payload if task else {} payload = task.payload if task else {}
dataset_id = payload.get("dataset_id") dataset_id = payload.get("dataset_id")
db_id = payload.get("db_id") kb_id = payload.get("kb_id")
count = int(payload.get("count", 10)) count = int(payload.get("count", 10))
neighbors_count = int(payload.get("neighbors_count", 1)) neighbors_count = int(payload.get("neighbors_count", 1))
concurrency_count = normalize_generation_concurrency_count(payload.get("concurrency_count")) concurrency_count = normalize_generation_concurrency_count(payload.get("concurrency_count"))
@ -366,7 +366,7 @@ class EvaluationService:
) )
try: try:
kb_instance = await knowledge_base.aget_kb(db_id) kb_instance = await knowledge_base.aget_kb(kb_id)
if not kb_instance: if not kb_instance:
await report_progress(100, "知识库不存在") await report_progress(100, "知识库不存在")
raise ValueError("Knowledge Base not found") raise ValueError("Knowledge Base not found")
@ -378,7 +378,7 @@ class EvaluationService:
try: try:
async for item in iter_generated_benchmark_items( async for item in iter_generated_benchmark_items(
kb_instance=kb_instance, kb_instance=kb_instance,
db_id=db_id, kb_id=kb_id,
count=count, count=count,
neighbors_count=neighbors_count, neighbors_count=neighbors_count,
llm_model_spec=llm_model_spec, llm_model_spec=llm_model_spec,
@ -397,7 +397,7 @@ class EvaluationService:
if not questions: if not questions:
raise ValueError("未生成有效评估题目") raise ValueError("未生成有效评估题目")
await self.eval_repo.add_dataset_items(self._build_dataset_items(dataset_id, db_id, questions)) await self.eval_repo.add_dataset_items(self._build_dataset_items(dataset_id, kb_id, questions))
await self.eval_repo.update_dataset(dataset_id, {"item_count": len(questions)}) await self.eval_repo.update_dataset(dataset_id, {"item_count": len(questions)})
await self._update_dataset_build_metadata( await self._update_dataset_build_metadata(
dataset_id, dataset_id,
@ -419,26 +419,26 @@ class EvaluationService:
raise raise
async def run_evaluation( async def run_evaluation(
self, db_id: str, dataset_id: str, model_config: dict[str, Any] = None, created_by: str = "system" self, kb_id: str, dataset_id: str, model_config: dict[str, Any] = None, created_by: str = "system"
) -> str: ) -> str:
try: try:
run_id = f"run_{uuid.uuid4().hex[:8]}" run_id = f"run_{uuid.uuid4().hex[:8]}"
dataset_row = await self.eval_repo.get_dataset(dataset_id) dataset_row = await self.eval_repo.get_dataset(dataset_id)
if dataset_row is None or dataset_row.db_id != db_id: if dataset_row is None or dataset_row.kb_id != kb_id:
raise ValueError("Dataset not found") raise ValueError("Dataset not found")
if (dataset_row.build_metadata or {}).get("status", "completed") != "completed": if (dataset_row.build_metadata or {}).get("status", "completed") != "completed":
raise ValueError("Dataset is not ready") raise ValueError("Dataset is not ready")
retrieval_config = {} retrieval_config = {}
try: try:
kb_row = await self.kb_repo.get_by_id(db_id) kb_row = await self.kb_repo.get_by_kb_id(kb_id)
query_params = (kb_row.query_params if kb_row else None) or {} query_params = (kb_row.query_params if kb_row else None) or {}
retrieval_config = query_params.get("options", {}) if isinstance(query_params, dict) else {} retrieval_config = query_params.get("options", {}) if isinstance(query_params, dict) else {}
if not retrieval_config: if not retrieval_config:
kb_instance = await knowledge_base.aget_kb(db_id) kb_instance = await knowledge_base.aget_kb(kb_id)
if kb_instance: if kb_instance:
retrieval_config = kb_instance._get_default_query_params(db_id).get("options", {}) retrieval_config = kb_instance._get_default_query_params(kb_id).get("options", {})
logger.info(f"从知识库 {db_id} 加载检索配置: {list(retrieval_config.keys())}") logger.info(f"从知识库 {kb_id} 加载检索配置: {list(retrieval_config.keys())}")
except Exception as e: except Exception as e:
logger.error(f"获取知识库检索配置失败: {e}") logger.error(f"获取知识库检索配置失败: {e}")
@ -448,7 +448,7 @@ class EvaluationService:
await self.eval_repo.create_run( await self.eval_repo.create_run(
{ {
"run_id": run_id, "run_id": run_id,
"db_id": db_id, "kb_id": kb_id,
"dataset_id": dataset_id, "dataset_id": dataset_id,
"status": "running", "status": "running",
"retrieval_config": retrieval_config, "retrieval_config": retrieval_config,
@ -467,7 +467,7 @@ class EvaluationService:
task_type="rag_evaluation", task_type="rag_evaluation",
payload={ payload={
"run_id": run_id, "run_id": run_id,
"db_id": db_id, "kb_id": kb_id,
"dataset_id": dataset_id, "dataset_id": dataset_id,
"retrieval_config": retrieval_config, "retrieval_config": retrieval_config,
"created_by": created_by, "created_by": created_by,
@ -487,21 +487,21 @@ class EvaluationService:
payload = task.payload payload = task.payload
run_id = payload["run_id"] run_id = payload["run_id"]
db_id = payload["db_id"] kb_id = payload["kb_id"]
dataset_id = payload["dataset_id"] dataset_id = payload["dataset_id"]
retrieval_config = payload["retrieval_config"] retrieval_config = payload["retrieval_config"]
await context.set_progress(5, "加载评估数据集") await context.set_progress(5, "加载评估数据集")
dataset_row = await self.eval_repo.get_dataset(dataset_id) dataset_row = await self.eval_repo.get_dataset(dataset_id)
if dataset_row is None or dataset_row.db_id != db_id: if dataset_row is None or dataset_row.kb_id != kb_id:
raise ValueError("Dataset not found") raise ValueError("Dataset not found")
dataset_items = await self.eval_repo.list_all_dataset_items(dataset_id) dataset_items = await self.eval_repo.list_all_dataset_items(dataset_id)
if not dataset_items: if not dataset_items:
raise ValueError("Dataset has no items") raise ValueError("Dataset has no items")
kb_instance = await knowledge_base.aget_kb(db_id) kb_instance = await knowledge_base.aget_kb(kb_id)
if not kb_instance: if not kb_instance:
raise ValueError(f"Knowledge Base {db_id} not found") raise ValueError(f"Knowledge Base {kb_id} not found")
judge_llm = None judge_llm = None
if dataset_row.has_gold_answers: if dataset_row.has_gold_answers:
@ -544,7 +544,7 @@ class EvaluationService:
} }
question_result = await evaluate_question( question_result = await evaluate_question(
kb_instance=kb_instance, kb_instance=kb_instance,
db_id=db_id, kb_id=kb_id,
question_data=question_data, question_data=question_data,
retrieval_config=retrieval_config, retrieval_config=retrieval_config,
has_gold_chunks=dataset_row.has_gold_chunks, has_gold_chunks=dataset_row.has_gold_chunks,
@ -595,9 +595,9 @@ class EvaluationService:
await context.set_message(f"Error: {str(e)}") await context.set_message(f"Error: {str(e)}")
raise raise
async def list_runs(self, db_id: str) -> list[dict[str, Any]]: async def list_runs(self, kb_id: str) -> list[dict[str, Any]]:
try: try:
rows = await self.eval_repo.list_runs(db_id) rows = await self.eval_repo.list_runs(kb_id)
running_run_ids = {row.run_id for row in rows if row.status == "running"} running_run_ids = {row.run_id for row in rows if row.status == "running"}
task_by_run_id = {} task_by_run_id = {}
if running_run_ids: if running_run_ids:
@ -635,12 +635,12 @@ class EvaluationService:
raise raise
async def get_run_results( async def get_run_results(
self, db_id: str, run_id: str, page: int = 1, page_size: int = 20, error_only: bool = False self, kb_id: str, run_id: str, page: int = 1, page_size: int = 20, error_only: bool = False
) -> dict[str, Any]: ) -> dict[str, Any]:
if not re.match(r"^run_[a-f0-9]{8}$", run_id): if not re.match(r"^run_[a-f0-9]{8}$", run_id):
raise ValueError("Invalid run_id format") raise ValueError("Invalid run_id format")
row = await self.eval_repo.get_run(run_id) row = await self.eval_repo.get_run(run_id)
if row is None or row.db_id != db_id: if row is None or row.kb_id != kb_id:
task = await tasker.get_task(run_id) task = await tasker.get_task(run_id)
if task: if task:
return {"run_id": run_id, "status": task.status, "progress": task.progress, "message": task.message} return {"run_id": run_id, "status": task.status, "progress": task.progress, "message": task.message}
@ -686,11 +686,11 @@ class EvaluationService:
}, },
} }
async def delete_run(self, db_id: str, run_id: str) -> None: async def delete_run(self, kb_id: str, run_id: str) -> None:
if not re.match(r"^run_[a-f0-9]{8}$", run_id): if not re.match(r"^run_[a-f0-9]{8}$", run_id):
raise ValueError("Invalid run_id format") raise ValueError("Invalid run_id format")
row = await self.eval_repo.get_run(run_id) row = await self.eval_repo.get_run(run_id)
if row is None or row.db_id != db_id: if row is None or row.kb_id != kb_id:
raise ValueError("Run not found") raise ValueError("Run not found")
await self.eval_repo.delete_run(run_id) await self.eval_repo.delete_run(run_id)
logger.info(f"成功删除评估运行: {run_id}") logger.info(f"成功删除评估运行: {run_id}")

28
backend/package/yuxi/knowledge/graphs/graph_utils.py
View File

@ -16,12 +16,12 @@ def normalize_entity_name(text: str) -> str:
return " ".join(text.strip().lower().split()) return " ".join(text.strip().lower().split())
def compute_entity_id(db_id: str, normalized_name: str, label: str) -> str: def compute_entity_id(kb_id: str, normalized_name: str, label: str) -> str:
return hashstr(f"{db_id}:{normalized_name}:{label}", length=32) return hashstr(f"{kb_id}:{normalized_name}:{label}", length=32)
def compute_triple_id( def compute_triple_id(
db_id: str, kb_id: str,
source_normalized_name: str, source_normalized_name: str,
source_label: str, source_label: str,
relation_type: str, relation_type: str,
@ -29,17 +29,17 @@ def compute_triple_id(
target_label: str, target_label: str,
) -> str: ) -> str:
return hashstr( return hashstr(
f"{db_id}:{source_normalized_name}:{source_label}:{relation_type}:{target_normalized_name}:{target_label}", f"{kb_id}:{source_normalized_name}:{source_label}:{relation_type}:{target_normalized_name}:{target_label}",
length=32, length=32,
) )
def graph_entity_collection_name(db_id: str) -> str: def graph_entity_collection_name(kb_id: str) -> str:
return f"{db_id}_entity" return f"{kb_id}_entity"
def graph_triple_collection_name(db_id: str) -> str: def graph_triple_collection_name(kb_id: str) -> str:
return f"{db_id}_triple" return f"{kb_id}_triple"
def build_graph_payload(normalized_result: dict[str, Any]) -> dict[str, Any]: def build_graph_payload(normalized_result: dict[str, Any]) -> dict[str, Any]:
@ -99,7 +99,7 @@ def cypher_merge_chunk(db_label: str) -> str:
return f""" return f"""
MERGE (c:Chunk:MilvusKB:`{db_label}` {{chunk_id: $chunk_id}}) MERGE (c:Chunk:MilvusKB:`{db_label}` {{chunk_id: $chunk_id}})
SET c.file_id = $file_id, SET c.file_id = $file_id,
c.db_id = $db_id, c.kb_id = $kb_id,
c.chunk_index = $chunk_index, c.chunk_index = $chunk_index,
c.content_preview = $content_preview, c.content_preview = $content_preview,
c.start_char_pos = $start_char_pos, c.start_char_pos = $start_char_pos,
@ -112,14 +112,14 @@ def cypher_merge_entity_mention(db_label: str) -> str:
return f""" return f"""
MATCH (c:Chunk:MilvusKB:`{db_label}` {{chunk_id: $chunk_id}}) MATCH (c:Chunk:MilvusKB:`{db_label}` {{chunk_id: $chunk_id}})
MERGE (e:Entity:MilvusKB:`{db_label}` {{ MERGE (e:Entity:MilvusKB:`{db_label}` {{
db_id: $db_id, kb_id: $kb_id,
normalized_name: $normalized_name, normalized_name: $normalized_name,
label: $entity_label label: $entity_label
}}) }})
SET e.entity_id = $entity_id, SET e.entity_id = $entity_id,
e.name = $name, e.name = $name,
e.attributes = $attributes e.attributes = $attributes
MERGE (c)-[m:MENTIONS {{chunk_id: $chunk_id, file_id: $file_id, db_id: $db_id}}]->(e) MERGE (c)-[m:MENTIONS {{chunk_id: $chunk_id, file_id: $file_id, kb_id: $kb_id}}]->(e)
""" """
@ -127,17 +127,17 @@ def cypher_merge_relation(db_label: str) -> str:
"""MERGE 两个 Entity 之间的 RELATION 边。""" """MERGE 两个 Entity 之间的 RELATION 边。"""
return f""" return f"""
MATCH (source:Entity:MilvusKB:`{db_label}` {{ MATCH (source:Entity:MilvusKB:`{db_label}` {{
db_id: $db_id, kb_id: $kb_id,
normalized_name: $source_name, normalized_name: $source_name,
label: $source_label label: $source_label
}}) }})
MATCH (target:Entity:MilvusKB:`{db_label}` {{ MATCH (target:Entity:MilvusKB:`{db_label}` {{
db_id: $db_id, kb_id: $kb_id,
normalized_name: $target_name, normalized_name: $target_name,
label: $target_label label: $target_label
}}) }})
MERGE (source)-[r:RELATION {{ MERGE (source)-[r:RELATION {{
db_id: $db_id, kb_id: $kb_id,
chunk_id: $chunk_id, chunk_id: $chunk_id,
source_name: $source_name, source_name: $source_name,
target_name: $target_name, target_name: $target_name,

184
backend/package/yuxi/knowledge/graphs/milvus_graph_service.py
View File

@ -36,14 +36,14 @@ class MilvusGraphService:
def __init__( def __init__(
self, self,
*, *,
db_id: str | None = None, kb_id: str | None = None,
kb_repo: KnowledgeBaseRepository | None = None, kb_repo: KnowledgeBaseRepository | None = None,
chunk_repo: KnowledgeChunkRepository | None = None, chunk_repo: KnowledgeChunkRepository | None = None,
graph_repo: KnowledgeGraphRepository | None = None, graph_repo: KnowledgeGraphRepository | None = None,
graph_vector_store: MilvusGraphVectorStore | None = None, graph_vector_store: MilvusGraphVectorStore | None = None,
neo4j_connection: Neo4jConnectionManager | None = None, neo4j_connection: Neo4jConnectionManager | None = None,
): ):
self.db_id = db_id self.kb_id = kb_id
self.kb_repo = kb_repo or KnowledgeBaseRepository() self.kb_repo = kb_repo or KnowledgeBaseRepository()
self.chunk_repo = chunk_repo or KnowledgeChunkRepository() self.chunk_repo = chunk_repo or KnowledgeChunkRepository()
self.graph_repo = graph_repo or KnowledgeGraphRepository() self.graph_repo = graph_repo or KnowledgeGraphRepository()
@ -66,15 +66,15 @@ class MilvusGraphService:
def driver(self): def driver(self):
return self.connection.driver return self.connection.driver
async def get_status(self, db_id: str, *, tasker: Any = None) -> dict[str, Any]: async def get_status(self, kb_id: str, *, tasker: Any = None) -> dict[str, Any]:
kb = await self._get_milvus_kb(db_id) kb = await self._get_milvus_kb(kb_id)
params = dict(kb.additional_params or {}) params = dict(kb.additional_params or {})
config = params.get(GRAPH_CONFIG_KEY) or {} config = params.get(GRAPH_CONFIG_KEY) or {}
total_chunks, pending_chunks, indexed_chunks, graph_counts = await asyncio.gather( total_chunks, pending_chunks, indexed_chunks, graph_counts = await asyncio.gather(
self.chunk_repo.count_by_db_id(db_id), self.chunk_repo.count_by_kb_id(kb_id),
self.chunk_repo.count_graph_pending_by_db_id(db_id), self.chunk_repo.count_graph_pending_by_kb_id(kb_id),
self.chunk_repo.count_graph_indexed_by_db_id(db_id), self.chunk_repo.count_graph_indexed_by_kb_id(kb_id),
self.graph_repo.count_by_db_id(db_id), self.graph_repo.count_by_kb_id(kb_id),
) )
entity_count, relationship_count = graph_counts entity_count, relationship_count = graph_counts
@ -83,7 +83,7 @@ class MilvusGraphService:
if tasker is not None: if tasker is not None:
active_task = await tasker.find_task_by_payload( active_task = await tasker.find_task_by_payload(
task_type=GRAPH_TASK_TYPE, task_type=GRAPH_TASK_TYPE,
payload_match={"db_id": db_id}, payload_match={"kb_id": kb_id},
statuses={"pending", "running"}, statuses={"pending", "running"},
) )
if active_task: if active_task:
@ -92,7 +92,7 @@ class MilvusGraphService:
else: else:
failed_task = await tasker.find_task_by_payload( failed_task = await tasker.find_task_by_payload(
task_type=GRAPH_TASK_TYPE, task_type=GRAPH_TASK_TYPE,
payload_match={"db_id": db_id}, payload_match={"kb_id": kb_id},
statuses={"failed", "cancelled"}, statuses={"failed", "cancelled"},
) )
if failed_task: if failed_task:
@ -100,7 +100,7 @@ class MilvusGraphService:
build_task_progress = 0 build_task_progress = 0
return { return {
"db_id": db_id, "kb_id": kb_id,
"kb_type": kb.kb_type, "kb_type": kb.kb_type,
"configured": bool(config), "configured": bool(config),
"locked": bool(config.get("locked")), "locked": bool(config.get("locked")),
@ -116,12 +116,12 @@ class MilvusGraphService:
async def configure( async def configure(
self, self,
db_id: str, kb_id: str,
extractor_type: str, extractor_type: str,
extractor_options: dict[str, Any], extractor_options: dict[str, Any],
created_by: str, created_by: str,
) -> dict: ) -> dict:
kb = await self._get_milvus_kb(db_id) kb = await self._get_milvus_kb(kb_id)
additional_params = dict(kb.additional_params or {}) additional_params = dict(kb.additional_params or {})
existing_config = additional_params.get(GRAPH_CONFIG_KEY) or {} existing_config = additional_params.get(GRAPH_CONFIG_KEY) or {}
normalized_extractor_type = (extractor_type or "").lower() normalized_extractor_type = (extractor_type or "").lower()
@ -145,16 +145,16 @@ class MilvusGraphService:
config["updated_at"] = utc_isoformat() config["updated_at"] = utc_isoformat()
config["updated_by"] = created_by config["updated_by"] = created_by
additional_params[GRAPH_CONFIG_KEY] = config additional_params[GRAPH_CONFIG_KEY] = config
await self.kb_repo.update(db_id, {"additional_params": additional_params}) await self.kb_repo.update(kb_id, {"additional_params": additional_params})
return config return config
async def build_pending_chunks(self, db_id: str, *, batch_size: int, context=None) -> dict[str, Any]: async def build_pending_chunks(self, kb_id: str, *, batch_size: int, context=None) -> dict[str, Any]:
kb = await self._get_milvus_kb(db_id) kb = await self._get_milvus_kb(kb_id)
config = self._get_locked_config(kb.additional_params or {}) config = self._get_locked_config(kb.additional_params or {})
extractor_options = self._runtime_extractor_options(config) extractor_options = self._runtime_extractor_options(config)
extractor = GraphExtractorFactory.create(config["extractor_type"], extractor_options) extractor = GraphExtractorFactory.create(config["extractor_type"], extractor_options)
worker_count = self._get_worker_count(config) worker_count = self._get_worker_count(config)
total_pending = await self.chunk_repo.count_graph_pending_by_db_id(db_id) total_pending = await self.chunk_repo.count_graph_pending_by_kb_id(kb_id)
processed = 0 processed = 0
failed = 0 failed = 0
failed_chunk_ids: set[str] = set() failed_chunk_ids: set[str] = set()
@ -163,7 +163,7 @@ class MilvusGraphService:
while True: while True:
if context is not None: if context is not None:
await context.raise_if_cancelled() await context.raise_if_cancelled()
chunks = await self.chunk_repo.list_graph_pending_by_db_id(db_id, batch_size) chunks = await self.chunk_repo.list_graph_pending_by_kb_id(kb_id, batch_size)
unprocessed = [c for c in chunks if c.chunk_id not in failed_chunk_ids] unprocessed = [c for c in chunks if c.chunk_id not in failed_chunk_ids]
if not unprocessed: if not unprocessed:
break break
@ -182,23 +182,23 @@ class MilvusGraphService:
except asyncio.QueueEmpty: except asyncio.QueueEmpty:
return return
try: try:
extraction_result = await self._get_chunk_extraction_result(db_id, chunk, extractor) extraction_result = await self._get_chunk_extraction_result(kb_id, chunk, extractor)
async with write_lock: async with write_lock:
entities, triples = await asyncio.to_thread( entities, triples = await asyncio.to_thread(
self.write_chunk_graph, self.write_chunk_graph,
db_id, kb_id,
chunk, chunk,
extraction_result, extraction_result,
) )
await self.graph_repo.upsert_chunk_graph( await self.graph_repo.upsert_chunk_graph(
db_id=db_id, kb_id=kb_id,
file_id=chunk.file_id, file_id=chunk.file_id,
chunk_id=chunk.chunk_id, chunk_id=chunk.chunk_id,
entities=entities, entities=entities,
triples=triples, triples=triples,
) )
await self.graph_vector_store.insert_missing_graph_records( await self.graph_vector_store.insert_missing_graph_records(
db_id=db_id, kb_id=kb_id,
embedding_model_spec=kb.embedding_model_spec, embedding_model_spec=kb.embedding_model_spec,
entities=entities, entities=entities,
triples=triples, triples=triples,
@ -229,8 +229,8 @@ class MilvusGraphService:
await asyncio.gather(*workers, return_exceptions=True) await asyncio.gather(*workers, return_exceptions=True)
raise raise
remaining = await self.chunk_repo.count_graph_pending_by_db_id(db_id) remaining = await self.chunk_repo.count_graph_pending_by_kb_id(kb_id)
return {"db_id": db_id, "success": processed, "failed": failed, "remaining": remaining} return {"kb_id": kb_id, "success": processed, "failed": failed, "remaining": remaining}
@staticmethod @staticmethod
def _get_worker_count(config: dict[str, Any]) -> int: def _get_worker_count(config: dict[str, Any]) -> int:
@ -248,7 +248,7 @@ class MilvusGraphService:
options.pop("prompt", None) options.pop("prompt", None)
return options return options
async def _get_chunk_extraction_result(self, db_id: str, chunk, extractor: GraphExtractor) -> dict[str, Any]: async def _get_chunk_extraction_result(self, kb_id: str, chunk, extractor: GraphExtractor) -> dict[str, Any]:
extractor_type = extractor.extractor_type extractor_type = extractor.extractor_type
if chunk.extraction_result: if chunk.extraction_result:
return normalize_extraction_result(chunk.extraction_result, extractor_type) return normalize_extraction_result(chunk.extraction_result, extractor_type)
@ -256,7 +256,7 @@ class MilvusGraphService:
extraction_result = await extractor.extract( extraction_result = await extractor.extract(
chunk.content, chunk.content,
chunk_metadata={ chunk_metadata={
"db_id": db_id, "kb_id": kb_id,
"chunk_id": chunk.chunk_id, "chunk_id": chunk.chunk_id,
"file_id": chunk.file_id, "file_id": chunk.file_id,
"chunk_index": chunk.chunk_index, "chunk_index": chunk.chunk_index,
@ -268,22 +268,22 @@ class MilvusGraphService:
def write_chunk_graph( def write_chunk_graph(
self, self,
db_id: str, kb_id: str,
chunk, chunk,
normalized_result: dict[str, Any], normalized_result: dict[str, Any],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""将单个 chunk 的抽取结果写入 Neo4j。""" """将单个 chunk 的抽取结果写入 Neo4j。"""
label = safe_neo4j_label(db_id) label = safe_neo4j_label(kb_id)
graph_payload = build_graph_payload(normalized_result) graph_payload = build_graph_payload(normalized_result)
relation_extractor_type = graph_payload["metadata"].get("extractor_type", "unknown") relation_extractor_type = graph_payload["metadata"].get("extractor_type", "unknown")
entities = graph_payload["entities"] entities = graph_payload["entities"]
relations = graph_payload["relations"] relations = graph_payload["relations"]
entity_by_id = {entity["id"]: entity for entity in entities} entity_by_id = {entity["id"]: entity for entity in entities}
entity_records = self._build_entity_records(db_id, entities) entity_records = self._build_entity_records(kb_id, entities)
entity_record_by_local_id = { entity_record_by_local_id = {
entity["id"]: record for entity, record in zip(entities, entity_records, strict=True) entity["id"]: record for entity, record in zip(entities, entity_records, strict=True)
} }
triple_records = self._build_triple_records(db_id, relations, entity_record_by_local_id, graph_payload) triple_records = self._build_triple_records(kb_id, relations, entity_record_by_local_id, graph_payload)
content_preview = (chunk.content or "")[:300] content_preview = (chunk.content or "")[:300]
# 预构建 Cypher 模板（同一 chunk 内复用） # 预构建 Cypher 模板（同一 chunk 内复用）
@ -297,7 +297,7 @@ class MilvusGraphService:
merge_chunk_cypher, merge_chunk_cypher,
chunk_id=chunk.chunk_id, chunk_id=chunk.chunk_id,
file_id=chunk.file_id, file_id=chunk.file_id,
db_id=db_id, kb_id=kb_id,
chunk_index=chunk.chunk_index, chunk_index=chunk.chunk_index,
content_preview=content_preview, content_preview=content_preview,
start_char_pos=chunk.start_char_pos, start_char_pos=chunk.start_char_pos,
@ -311,7 +311,7 @@ class MilvusGraphService:
merge_entity_cypher, merge_entity_cypher,
chunk_id=chunk.chunk_id, chunk_id=chunk.chunk_id,
file_id=chunk.file_id, file_id=chunk.file_id,
db_id=db_id, kb_id=kb_id,
entity_id=entity_record["entity_id"], entity_id=entity_record["entity_id"],
normalized_name=normalize_entity_name(entity["text"]), normalized_name=normalize_entity_name(entity["text"]),
entity_label=entity.get("label") or "Entity", entity_label=entity.get("label") or "Entity",
@ -327,7 +327,7 @@ class MilvusGraphService:
target_record = entity_record_by_local_id[relation["target"]] target_record = entity_record_by_local_id[relation["target"]]
relation_type = relation.get("label") or "RELATED_TO" relation_type = relation.get("label") or "RELATED_TO"
triple_id = compute_triple_id( triple_id = compute_triple_id(
db_id, kb_id,
source_record["normalized_name"], source_record["normalized_name"],
source_record["label"], source_record["label"],
relation_type, relation_type,
@ -336,7 +336,7 @@ class MilvusGraphService:
) )
tx.run( tx.run(
merge_relation_cypher, merge_relation_cypher,
db_id=db_id, kb_id=kb_id,
chunk_id=chunk.chunk_id, chunk_id=chunk.chunk_id,
file_id=chunk.file_id, file_id=chunk.file_id,
source_name=normalize_entity_name(source["text"]), source_name=normalize_entity_name(source["text"]),
@ -352,16 +352,16 @@ class MilvusGraphService:
neo4j_write(self.driver, query) neo4j_write(self.driver, query)
return entity_records, triple_records return entity_records, triple_records
def _build_entity_records(self, db_id: str, entities: list[dict[str, Any]]) -> list[dict[str, Any]]: def _build_entity_records(self, kb_id: str, entities: list[dict[str, Any]]) -> list[dict[str, Any]]:
records = [] records = []
for entity in entities: for entity in entities:
label = entity.get("label") or "Entity" label = entity.get("label") or "Entity"
normalized_name = normalize_entity_name(entity["text"]) normalized_name = normalize_entity_name(entity["text"])
entity_id = compute_entity_id(db_id, normalized_name, label) entity_id = compute_entity_id(kb_id, normalized_name, label)
records.append( records.append(
{ {
"entity_id": entity_id, "entity_id": entity_id,
"db_id": db_id, "kb_id": kb_id,
"normalized_name": normalized_name, "normalized_name": normalized_name,
"label": label, "label": label,
"name": entity["text"], "name": entity["text"],
@ -373,7 +373,7 @@ class MilvusGraphService:
def _build_triple_records( def _build_triple_records(
self, self,
db_id: str, kb_id: str,
relations: list[dict[str, Any]], relations: list[dict[str, Any]],
entity_record_by_local_id: dict[str, dict[str, Any]], entity_record_by_local_id: dict[str, dict[str, Any]],
graph_payload: dict[str, Any], graph_payload: dict[str, Any],
@ -386,7 +386,7 @@ class MilvusGraphService:
target_record = entity_record_by_local_id[relation["target"]] target_record = entity_record_by_local_id[relation["target"]]
relation_type = relation.get("label") or "RELATED_TO" relation_type = relation.get("label") or "RELATED_TO"
triple_id = compute_triple_id( triple_id = compute_triple_id(
db_id, kb_id,
source_record["normalized_name"], source_record["normalized_name"],
source_record["label"], source_record["label"],
relation_type, relation_type,
@ -400,7 +400,7 @@ class MilvusGraphService:
records.append( records.append(
{ {
"triple_id": triple_id, "triple_id": triple_id,
"db_id": db_id, "kb_id": kb_id,
"source_entity_id": source_record["entity_id"], "source_entity_id": source_record["entity_id"],
"target_entity_id": target_record["entity_id"], "target_entity_id": target_record["entity_id"],
"relation_type": relation_type, "relation_type": relation_type,
@ -411,15 +411,15 @@ class MilvusGraphService:
) )
return records return records
async def reset(self, db_id: str, *, clear_extraction_result: bool, clear_config: bool) -> dict[str, Any]: async def reset(self, kb_id: str, *, clear_extraction_result: bool, clear_config: bool) -> dict[str, Any]:
kb = await self._get_milvus_kb(db_id) kb = await self._get_milvus_kb(kb_id)
await asyncio.to_thread(self.delete_graph, db_id) await asyncio.to_thread(self.delete_graph, kb_id)
await self.graph_repo.delete_by_db_id(db_id) await self.graph_repo.delete_by_kb_id(kb_id)
reset_chunks = await self.chunk_repo.reset_graph_state_by_db_id(db_id, clear_extraction_result) reset_chunks = await self.chunk_repo.reset_graph_state_by_kb_id(kb_id, clear_extraction_result)
if clear_config: if clear_config:
additional_params = dict(kb.additional_params or {}) additional_params = dict(kb.additional_params or {})
additional_params.pop(GRAPH_CONFIG_KEY, None) additional_params.pop(GRAPH_CONFIG_KEY, None)
await self.kb_repo.update(db_id, {"additional_params": additional_params}) await self.kb_repo.update(kb_id, {"additional_params": additional_params})
return { return {
"message": "图谱构建状态已重置", "message": "图谱构建状态已重置",
"status": "success", "status": "success",
@ -428,79 +428,79 @@ class MilvusGraphService:
"clear_config": clear_config, "clear_config": clear_config,
} }
def delete_graph(self, db_id: str) -> None: def delete_graph(self, kb_id: str) -> None:
label = safe_neo4j_label(db_id) label = safe_neo4j_label(kb_id)
def query(tx): def query(tx):
tx.run(f"MATCH (n:MilvusKB:`{label}`) DETACH DELETE n") tx.run(f"MATCH (n:MilvusKB:`{label}`) DETACH DELETE n")
neo4j_write(self.driver, query) neo4j_write(self.driver, query)
self.graph_vector_store.drop_graph_collections(db_id) self.graph_vector_store.drop_graph_collections(kb_id)
async def delete_file_graph(self, db_id: str, file_id: str) -> None: async def delete_file_graph(self, kb_id: str, file_id: str) -> None:
orphan_entity_ids, orphan_triple_ids = await self.graph_repo.delete_file_references(file_id) orphan_entity_ids, orphan_triple_ids = await self.graph_repo.delete_file_references(file_id)
await self.graph_vector_store.delete_graph_records( await self.graph_vector_store.delete_graph_records(
db_id, kb_id,
entity_ids=orphan_entity_ids, entity_ids=orphan_entity_ids,
triple_ids=orphan_triple_ids, triple_ids=orphan_triple_ids,
) )
await asyncio.to_thread(self._delete_file_graph_from_neo4j, db_id, file_id) await asyncio.to_thread(self._delete_file_graph_from_neo4j, kb_id, file_id)
def _delete_file_graph_from_neo4j(self, db_id: str, file_id: str) -> None: def _delete_file_graph_from_neo4j(self, kb_id: str, file_id: str) -> None:
label = safe_neo4j_label(db_id) label = safe_neo4j_label(kb_id)
def query(tx): def query(tx):
tx.run( tx.run(
f""" f"""
MATCH (:Chunk:MilvusKB:`{label}`)-[m:MENTIONS {{db_id: $db_id, file_id: $file_id}}]-> MATCH (:Chunk:MilvusKB:`{label}`)-[m:MENTIONS {{kb_id: $kb_id, file_id: $file_id}}]->
(:Entity:MilvusKB:`{label}`) (:Entity:MilvusKB:`{label}`)
DELETE m DELETE m
""", """,
db_id=db_id, kb_id=kb_id,
file_id=file_id, file_id=file_id,
) )
tx.run( tx.run(
f""" f"""
MATCH (:Entity:MilvusKB:`{label}`)-[r:RELATION {{db_id: $db_id, file_id: $file_id}}]-> MATCH (:Entity:MilvusKB:`{label}`)-[r:RELATION {{kb_id: $kb_id, file_id: $file_id}}]->
(:Entity:MilvusKB:`{label}`) (:Entity:MilvusKB:`{label}`)
DELETE r DELETE r
""", """,
db_id=db_id, kb_id=kb_id,
file_id=file_id, file_id=file_id,
) )
tx.run( tx.run(
f""" f"""
MATCH (c:Chunk:MilvusKB:`{label}` {{db_id: $db_id, file_id: $file_id}}) MATCH (c:Chunk:MilvusKB:`{label}` {{kb_id: $kb_id, file_id: $file_id}})
DETACH DELETE c DETACH DELETE c
""", """,
db_id=db_id, kb_id=kb_id,
file_id=file_id, file_id=file_id,
) )
tx.run( tx.run(
f""" f"""
MATCH (e:Entity:MilvusKB:`{label}` {{db_id: $db_id}}) MATCH (e:Entity:MilvusKB:`{label}` {{kb_id: $kb_id}})
WHERE NOT ()-[:MENTIONS]->(e) WHERE NOT ()-[:MENTIONS]->(e)
DETACH DELETE e DETACH DELETE e
""", """,
db_id=db_id, kb_id=kb_id,
) )
neo4j_write(self.driver, query) neo4j_write(self.driver, query)
async def query_nodes( async def query_nodes(
self, self,
db_id: str | None = None, kb_id: str | None = None,
*, *,
keyword: str = "", keyword: str = "",
max_depth: int = 1, max_depth: int = 1,
max_nodes: int = 50, max_nodes: int = 50,
exclude_chunk: bool = False, exclude_chunk: bool = False,
) -> dict[str, Any]: ) -> dict[str, Any]:
effective_db_id = db_id or self.db_id effective_kb_id = kb_id or self.kb_id
if not effective_db_id: if not effective_kb_id:
return {"nodes": [], "edges": []} return {"nodes": [], "edges": []}
label = safe_neo4j_label(effective_db_id) label = safe_neo4j_label(effective_kb_id)
limit = max_nodes limit = max_nodes
try: try:
with self.driver.session() as session: with self.driver.session() as session:
@ -509,21 +509,21 @@ class MilvusGraphService:
keyword=keyword, keyword=keyword,
limit=limit, limit=limit,
) )
return self._process_query_result(result, limit, effective_db_id, exclude_chunk) return self._process_query_result(result, limit, effective_kb_id, exclude_chunk)
except Exception as e: except Exception as e:
logger.error(f"Milvus graph query failed: {e}") logger.error(f"Milvus graph query failed: {e}")
return {"nodes": [], "edges": []} return {"nodes": [], "edges": []}
async def query_seed_subgraph( async def query_seed_subgraph(
self, self,
db_id: str, kb_id: str,
*, *,
entity_ids: list[str], entity_ids: list[str],
max_nodes: int, max_nodes: int,
) -> dict[str, Any]: ) -> dict[str, Any]:
if not entity_ids: if not entity_ids:
return {"nodes": [], "edges": []} return {"nodes": [], "edges": []}
label = safe_neo4j_label(db_id) label = safe_neo4j_label(kb_id)
cypher = f""" cypher = f"""
MATCH (seed:Entity:MilvusKB:`{label}`) MATCH (seed:Entity:MilvusKB:`{label}`)
WHERE seed.entity_id IN $entity_ids WHERE seed.entity_id IN $entity_ids
@ -546,14 +546,14 @@ class MilvusGraphService:
).single() ).single()
if not record: if not record:
return {"nodes": [], "edges": []} return {"nodes": [], "edges": []}
return self._process_subgraph_record(record, max_nodes, db_id) return self._process_subgraph_record(record, max_nodes, kb_id)
except Exception as e: except Exception as e:
logger.error(f"Milvus seed subgraph query failed: {e}") logger.error(f"Milvus seed subgraph query failed: {e}")
return {"nodes": [], "edges": []} return {"nodes": [], "edges": []}
async def query_and_rank_chunks_by_ppr( async def query_and_rank_chunks_by_ppr(
self, self,
db_id: str, kb_id: str,
seed_weights: dict[str, float], seed_weights: dict[str, float],
*, *,
max_nodes: int, max_nodes: int,
@ -563,7 +563,7 @@ class MilvusGraphService:
if not seed_weights: if not seed_weights:
return [] return []
subgraph = await self.query_seed_subgraph( subgraph = await self.query_seed_subgraph(
db_id, kb_id,
entity_ids=list(seed_weights.keys()), entity_ids=list(seed_weights.keys()),
max_nodes=max_nodes, max_nodes=max_nodes,
) )
@ -622,32 +622,32 @@ class MilvusGraphService:
) )
return ranked[:top_k] return ranked[:top_k]
async def get_labels(self, db_id: str | None = None) -> list[str]: async def get_labels(self, kb_id: str | None = None) -> list[str]:
effective_db_id = db_id or self.db_id effective_kb_id = kb_id or self.kb_id
if not effective_db_id: if not effective_kb_id:
return [] return []
label = safe_neo4j_label(effective_db_id) label = safe_neo4j_label(effective_kb_id)
cypher = f""" cypher = f"""
MATCH (n:MilvusKB:`{label}`) MATCH (n:MilvusKB:`{label}`)
UNWIND labels(n) AS node_label UNWIND labels(n) AS node_label
WITH DISTINCT node_label WITH DISTINCT node_label
WHERE node_label <> 'MilvusKB' AND node_label <> $db_id WHERE node_label <> 'MilvusKB' AND node_label <> $kb_id
RETURN node_label RETURN node_label
ORDER BY node_label ORDER BY node_label
""" """
try: try:
records = neo4j_read(self.driver, cypher, db_id=effective_db_id) records = neo4j_read(self.driver, cypher, kb_id=effective_kb_id)
return [record["node_label"] for record in records] return [record["node_label"] for record in records]
except Exception as e: except Exception as e:
logger.error(f"Failed to get Milvus graph labels: {e}") logger.error(f"Failed to get Milvus graph labels: {e}")
return [] return []
async def get_stats(self, db_id: str | None = None) -> dict[str, Any]: async def get_stats(self, kb_id: str | None = None) -> dict[str, Any]:
effective_db_id = db_id or self.db_id effective_kb_id = kb_id or self.kb_id
if not effective_db_id: if not effective_kb_id:
return {"total_nodes": 0, "total_edges": 0, "entity_types": []} return {"total_nodes": 0, "total_edges": 0, "entity_types": []}
label = safe_neo4j_label(effective_db_id) label = safe_neo4j_label(effective_kb_id)
stats_cypher = f""" stats_cypher = f"""
MATCH (n:MilvusKB:`{label}`) MATCH (n:MilvusKB:`{label}`)
@ -674,10 +674,10 @@ class MilvusGraphService:
logger.error(f"Failed to get Milvus graph stats: {e}") logger.error(f"Failed to get Milvus graph stats: {e}")
return {"total_nodes": 0, "total_edges": 0, "entity_types": []} return {"total_nodes": 0, "total_edges": 0, "entity_types": []}
async def _get_milvus_kb(self, db_id: str): async def _get_milvus_kb(self, kb_id: str):
kb = await self.kb_repo.get_by_id(db_id) kb = await self.kb_repo.get_by_kb_id(kb_id)
if kb is None: if kb is None:
raise ValueError(f"知识库 {db_id} 不存在") raise ValueError(f"知识库 {kb_id} 不存在")
if (kb.kb_type or "").lower() != "milvus": if (kb.kb_type or "").lower() != "milvus":
raise ValueError("仅 Milvus 知识库支持独立图谱构建") raise ValueError("仅 Milvus 知识库支持独立图谱构建")
return kb return kb
@ -737,7 +737,7 @@ class MilvusGraphService:
LIMIT {limit * 10} LIMIT {limit * 10}
""" """
def _process_query_result(self, result, limit: int, db_id: str, exclude_chunk: bool = False) -> dict[str, Any]: def _process_query_result(self, result, limit: int, kb_id: str, exclude_chunk: bool = False) -> dict[str, Any]:
nodes = [] nodes = []
edges = [] edges = []
node_ids = set() node_ids = set()
@ -748,7 +748,7 @@ class MilvusGraphService:
raw_node = record.get(key) raw_node = record.get(key)
if raw_node is None: if raw_node is None:
continue continue
node = self._normalize_node(raw_node, db_id) node = self._normalize_node(raw_node, kb_id)
if not node or node["id"] in node_ids: if not node or node["id"] in node_ids:
continue continue
if exclude_chunk and node.get("type") == "Chunk": if exclude_chunk and node.get("type") == "Chunk":
@ -766,14 +766,14 @@ class MilvusGraphService:
return {"nodes": nodes[:limit], "edges": edges[: limit * 2]} return {"nodes": nodes[:limit], "edges": edges[: limit * 2]}
def _process_subgraph_record(self, record: Any, limit: int, db_id: str) -> dict[str, Any]: def _process_subgraph_record(self, record: Any, limit: int, kb_id: str) -> dict[str, Any]:
nodes = [] nodes = []
edges = [] edges = []
node_ids = set() node_ids = set()
edge_ids = set() edge_ids = set()
for raw_node in record.get("nodes") or []: for raw_node in record.get("nodes") or []:
node = self._normalize_node(raw_node, db_id) node = self._normalize_node(raw_node, kb_id)
if not node or node["id"] in node_ids: if not node or node["id"] in node_ids:
continue continue
nodes.append(node) nodes.append(node)
@ -792,7 +792,7 @@ class MilvusGraphService:
return {"nodes": nodes, "edges": edges} return {"nodes": nodes, "edges": edges}
def _normalize_node(self, raw_node: Any, db_id: str | None = None) -> dict[str, Any]: def _normalize_node(self, raw_node: Any, kb_id: str | None = None) -> dict[str, Any]:
if hasattr(raw_node, "element_id"): if hasattr(raw_node, "element_id"):
node_id = raw_node.element_id node_id = raw_node.element_id
labels = list(raw_node.labels) labels = list(raw_node.labels)
@ -804,8 +804,8 @@ class MilvusGraphService:
else: else:
return {} return {}
effective_db_id = db_id or self.db_id effective_kb_id = kb_id or self.kb_id
db_label = properties.get("db_id") or effective_db_id db_label = properties.get("kb_id") or effective_kb_id
filtered_labels = [label for label in labels if label not in {"MilvusKB", db_label}] filtered_labels = [label for label in labels if label not in {"MilvusKB", db_label}]
entity_type = "Chunk" if "Chunk" in labels else properties.get("label", "Entity") entity_type = "Chunk" if "Chunk" in labels else properties.get("label", "Entity")
name = properties.get("name") or properties.get("content_preview") or properties.get("chunk_id") or "Unknown" name = properties.get("name") or properties.get("content_preview") or properties.get("chunk_id") or "Unknown"

32
backend/package/yuxi/knowledge/graphs/milvus_graph_vector_store.py
View File

@ -45,7 +45,7 @@ class MilvusGraphVectorStore:
async def insert_missing_graph_records( async def insert_missing_graph_records(
self, self,
*, *,
db_id: str, kb_id: str,
embedding_model_spec: str, embedding_model_spec: str,
entities: list[dict[str, Any]], entities: list[dict[str, Any]],
triples: list[dict[str, Any]], triples: list[dict[str, Any]],
@ -57,8 +57,8 @@ class MilvusGraphVectorStore:
if not embedding_info or embedding_info.model_type != "embedding": if not embedding_info or embedding_info.model_type != "embedding":
raise ValueError(f"Unsupported embedding model: {embedding_model_spec}") raise ValueError(f"Unsupported embedding model: {embedding_model_spec}")
entity_collection = self._get_or_create_entity_collection(db_id, embedding_info) entity_collection = self._get_or_create_entity_collection(kb_id, embedding_info)
triple_collection = self._get_or_create_triple_collection(db_id, embedding_info) triple_collection = self._get_or_create_triple_collection(kb_id, embedding_info)
entity_ids = [entity["entity_id"] for entity in entities] entity_ids = [entity["entity_id"] for entity in entities]
triple_ids = [triple["triple_id"] for triple in triples] triple_ids = [triple["triple_id"] for triple in triples]
@ -83,24 +83,24 @@ class MilvusGraphVectorStore:
if missing_triples: if missing_triples:
await asyncio.to_thread(self._insert_triples, triple_collection, missing_triples, triple_embeddings) await asyncio.to_thread(self._insert_triples, triple_collection, missing_triples, triple_embeddings)
async def delete_graph_records(self, db_id: str, *, entity_ids: list[str], triple_ids: list[str]) -> None: async def delete_graph_records(self, kb_id: str, *, entity_ids: list[str], triple_ids: list[str]) -> None:
tasks = [] tasks = []
if entity_ids: if entity_ids:
tasks.append(asyncio.to_thread(self._delete_ids, graph_entity_collection_name(db_id), entity_ids)) tasks.append(asyncio.to_thread(self._delete_ids, graph_entity_collection_name(kb_id), entity_ids))
if triple_ids: if triple_ids:
tasks.append(asyncio.to_thread(self._delete_ids, graph_triple_collection_name(db_id), triple_ids)) tasks.append(asyncio.to_thread(self._delete_ids, graph_triple_collection_name(kb_id), triple_ids))
if tasks: if tasks:
await asyncio.gather(*tasks) await asyncio.gather(*tasks)
async def search_entities( async def search_entities(
self, self,
*, *,
db_id: str, kb_id: str,
query_text: str, query_text: str,
embedding_model_spec: str, embedding_model_spec: str,
top_k: int, top_k: int,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
collection_name = graph_entity_collection_name(db_id) collection_name = graph_entity_collection_name(kb_id)
if not utility.has_collection(collection_name, using=self.connection_alias): if not utility.has_collection(collection_name, using=self.connection_alias):
return [] return []
return await self._search_graph_collection( return await self._search_graph_collection(
@ -114,12 +114,12 @@ class MilvusGraphVectorStore:
async def search_triples( async def search_triples(
self, self,
*, *,
db_id: str, kb_id: str,
query_text: str, query_text: str,
embedding_model_spec: str, embedding_model_spec: str,
top_k: int, top_k: int,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
collection_name = graph_triple_collection_name(db_id) collection_name = graph_triple_collection_name(kb_id)
if not utility.has_collection(collection_name, using=self.connection_alias): if not utility.has_collection(collection_name, using=self.connection_alias):
return [] return []
return await self._search_graph_collection( return await self._search_graph_collection(
@ -130,8 +130,8 @@ class MilvusGraphVectorStore:
output_fields=["id", "content", "source_id", "target_id"], output_fields=["id", "content", "source_id", "target_id"],
) )
def drop_graph_collections(self, db_id: str) -> None: def drop_graph_collections(self, kb_id: str) -> None:
for collection_name in [graph_entity_collection_name(db_id), graph_triple_collection_name(db_id)]: for collection_name in [graph_entity_collection_name(kb_id), graph_triple_collection_name(kb_id)]:
try: try:
if utility.has_collection(collection_name, using=self.connection_alias): if utility.has_collection(collection_name, using=self.connection_alias):
utility.drop_collection(collection_name, using=self.connection_alias) utility.drop_collection(collection_name, using=self.connection_alias)
@ -195,8 +195,8 @@ class MilvusGraphVectorStore:
records.append(record) records.append(record)
return records return records
def _get_or_create_entity_collection(self, db_id: str, embedding_info: Any) -> Collection: def _get_or_create_entity_collection(self, kb_id: str, embedding_info: Any) -> Collection:
collection_name = graph_entity_collection_name(db_id) collection_name = graph_entity_collection_name(kb_id)
fields = [ fields = [
FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True), FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True),
FieldSchema( FieldSchema(
@ -211,8 +211,8 @@ class MilvusGraphVectorStore:
] ]
return self._get_or_create_collection(collection_name, fields, embedding_info) return self._get_or_create_collection(collection_name, fields, embedding_info)
def _get_or_create_triple_collection(self, db_id: str, embedding_info: Any) -> Collection: def _get_or_create_triple_collection(self, kb_id: str, embedding_info: Any) -> Collection:
collection_name = graph_triple_collection_name(db_id) collection_name = graph_triple_collection_name(kb_id)
fields = [ fields = [
FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True), FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True),
FieldSchema( FieldSchema(

18
backend/package/yuxi/knowledge/implementations/dify.py
View File

@ -63,18 +63,18 @@ class DifyKB(ReadOnlyConnectors):
raise ValueError("Dify api_url 必须以 /v1 结尾") raise ValueError("Dify api_url 必须以 /v1 结尾")
return params return params
async def aquery(self, query_text: str, db_id: str, agent_call: bool = False, **kwargs) -> list[dict]: async def aquery(self, query_text: str, kb_id: str, agent_call: bool = False, **kwargs) -> list[dict]:
del agent_call del agent_call
metadata = self.databases_meta.get(db_id, {}).get("metadata", {}) or {} metadata = self.databases_meta.get(kb_id, {}).get("metadata", {}) or {}
api_url = str(metadata.get("dify_api_url") or "").strip() api_url = str(metadata.get("dify_api_url") or "").strip()
token = str(metadata.get("dify_token") or "").strip() token = str(metadata.get("dify_token") or "").strip()
dataset_id = str(metadata.get("dify_dataset_id") or "").strip() dataset_id = str(metadata.get("dify_dataset_id") or "").strip()
if not api_url or not token or not dataset_id: if not api_url or not token or not dataset_id:
logger.error(f"Dify config incomplete for db_id={db_id}") logger.error(f"Dify config incomplete for kb_id={kb_id}")
return [] return []
query_params = self._get_query_params(db_id) query_params = self._get_query_params(kb_id)
merged = {**query_params, **kwargs} merged = {**query_params, **kwargs}
search_mode = str(merged.get("search_mode", "vector")).lower() search_mode = str(merged.get("search_mode", "vector")).lower()
@ -109,7 +109,7 @@ class DifyKB(ReadOnlyConnectors):
try: try:
response_json = await self._request_dify(client_payload=payload, request_url=request_url, headers=headers) response_json = await self._request_dify(client_payload=payload, request_url=request_url, headers=headers)
except Exception as e: # noqa: BLE001 except Exception as e: # noqa: BLE001
logger.error(f"Dify query failed for db_id={db_id}: {e}, {traceback.format_exc()}") logger.error(f"Dify query failed for kb_id={kb_id}: {e}, {traceback.format_exc()}")
# 一些 Dify 部署版本对 retrieval_model 兼容性较差，失败时降级为仅 query 请求重试一次 # 一些 Dify 部署版本对 retrieval_model 兼容性较差，失败时降级为仅 query 请求重试一次
try: try:
response_json = await self._request_dify( response_json = await self._request_dify(
@ -117,10 +117,10 @@ class DifyKB(ReadOnlyConnectors):
request_url=request_url, request_url=request_url,
headers=headers, headers=headers,
) )
logger.warning(f"Dify query fallback to query-only succeeded for db_id={db_id}") logger.warning(f"Dify query fallback to query-only succeeded for kb_id={kb_id}")
except Exception as fallback_error: # noqa: BLE001 except Exception as fallback_error: # noqa: BLE001
logger.error( logger.error(
f"Dify query fallback failed for db_id={db_id}: {fallback_error}, {traceback.format_exc()}" f"Dify query fallback failed for kb_id={kb_id}: {fallback_error}, {traceback.format_exc()}"
) )
return [] return []
@ -172,8 +172,8 @@ class DifyKB(ReadOnlyConnectors):
raise e raise e
return response.json() return response.json()
def get_query_params_config(self, db_id: str, **kwargs) -> dict: def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
del db_id, kwargs del kb_id, kwargs
options = [ options = [
{ {
"key": "search_mode", "key": "search_mode",

162
backend/package/yuxi/knowledge/implementations/milvus.py
View File

@ -266,7 +266,7 @@ class MilvusKB(KnowledgeBase):
# 连接名称 # 连接名称
self.connection_alias = f"milvus_{hashstr(work_dir, 6)}" self.connection_alias = f"milvus_{hashstr(work_dir, 6)}"
# 存储集合映射 {db_id: Collection} # 存储集合映射 {kb_id: Collection}
self.collections: dict[str, Any] = {} self.collections: dict[str, Any] = {}
# 元数据锁 # 元数据锁
@ -297,22 +297,22 @@ class MilvusKB(KnowledgeBase):
logger.error(f"Failed to connect to Milvus: {e}") logger.error(f"Failed to connect to Milvus: {e}")
raise raise
async def _create_kb_instance(self, db_id: str, kb_config: dict) -> Any: async def _create_kb_instance(self, kb_id: str, kb_config: dict) -> Any:
"""创建 Milvus 集合""" """创建 Milvus 集合"""
logger.info(f"Creating Milvus collection for {db_id}") logger.info(f"Creating Milvus collection for {kb_id}")
if not (metadata := self.databases_meta.get(db_id)): if not (metadata := self.databases_meta.get(kb_id)):
raise ValueError(f"Database {db_id} not found") raise ValueError(f"Database {kb_id} not found")
embedding_model_spec = metadata.get("embedding_model_spec") embedding_model_spec = metadata.get("embedding_model_spec")
if not embedding_model_spec: if not embedding_model_spec:
raise ValueError(f"Embedding model spec not found for database {db_id}") raise ValueError(f"Embedding model spec not found for database {kb_id}")
embedding_info = model_cache.get_model_info(embedding_model_spec) embedding_info = model_cache.get_model_info(embedding_model_spec)
if not embedding_info or embedding_info.model_type != "embedding": if not embedding_info or embedding_info.model_type != "embedding":
raise ValueError(f"Unsupported embedding model: {embedding_model_spec}") raise ValueError(f"Unsupported embedding model: {embedding_model_spec}")
collection_name = db_id collection_name = kb_id
try: try:
# 检查集合是否存在 # 检查集合是否存在
@ -329,18 +329,18 @@ class MilvusKB(KnowledgeBase):
f"expected='{expected_model}', found_in_description='{description}'" f"expected='{expected_model}', found_in_description='{description}'"
) )
utility.drop_collection(collection_name, using=self.connection_alias) utility.drop_collection(collection_name, using=self.connection_alias)
return self._create_new_collection(collection_name, embedding_info, db_id) return self._create_new_collection(collection_name, embedding_info, kb_id)
if not self._collection_supports_bm25(collection): if not self._collection_supports_bm25(collection):
logger.warning(f"Collection {collection_name} schema does not support BM25, recreating") logger.warning(f"Collection {collection_name} schema does not support BM25, recreating")
utility.drop_collection(collection_name, using=self.connection_alias) utility.drop_collection(collection_name, using=self.connection_alias)
return self._create_new_collection(collection_name, embedding_info, db_id) return self._create_new_collection(collection_name, embedding_info, kb_id)
logger.info(f"Retrieved existing collection: {collection_name}") logger.info(f"Retrieved existing collection: {collection_name}")
return collection return collection
else: else:
logger.info(f"Collection {collection_name} not found, creating new one") logger.info(f"Collection {collection_name} not found, creating new one")
return self._create_new_collection(collection_name, embedding_info, db_id) return self._create_new_collection(collection_name, embedding_info, kb_id)
except (connections.MilvusException, RuntimeError) as e: except (connections.MilvusException, RuntimeError) as e:
logger.error(f"Error checking collection {collection_name}: {e}") logger.error(f"Error checking collection {collection_name}: {e}")
@ -350,7 +350,7 @@ class MilvusKB(KnowledgeBase):
logger.debug(f"Traceback: {traceback.format_exc()}") logger.debug(f"Traceback: {traceback.format_exc()}")
raise raise
def _create_new_collection(self, collection_name: str, embedding_info: Any, db_id: str) -> Collection: def _create_new_collection(self, collection_name: str, embedding_info: Any, kb_id: str) -> Collection:
"""创建新的 Milvus 集合""" """创建新的 Milvus 集合"""
embedding_dim = embedding_info.dimension or 1024 embedding_dim = embedding_info.dimension or 1024
model_name = embedding_info.model_id model_name = embedding_info.model_id
@ -380,7 +380,7 @@ class MilvusKB(KnowledgeBase):
schema = CollectionSchema( schema = CollectionSchema(
fields=fields, fields=fields,
description=f"Knowledge base collection for {db_id} using {model_name}", description=f"Knowledge base collection for {kb_id} using {model_name}",
functions=[bm25_function], functions=[bm25_function],
) )
@ -439,24 +439,24 @@ class MilvusKB(KnowledgeBase):
method = model.batch_encode if sync else model.abatch_encode method = model.batch_encode if sync else model.abatch_encode
return partial(method, batch_size=batch_size) return partial(method, batch_size=batch_size)
async def _get_milvus_collection(self, db_id: str): async def _get_milvus_collection(self, kb_id: str):
"""获取或创建 Milvus 集合""" """获取或创建 Milvus 集合"""
if db_id in self.collections: if kb_id in self.collections:
return self.collections[db_id] return self.collections[kb_id]
if db_id not in self.databases_meta: if kb_id not in self.databases_meta:
return None return None
try: try:
# 创建集合 # 创建集合
collection = await self._create_kb_instance(db_id, {}) collection = await self._create_kb_instance(kb_id, {})
await self._initialize_kb_instance(collection) await self._initialize_kb_instance(collection)
self.collections[db_id] = collection self.collections[kb_id] = collection
return collection return collection
except Exception as e: except Exception as e:
logger.error(f"Failed to create Milvus collection for {db_id}: {e}") logger.error(f"Failed to create Milvus collection for {kb_id}: {e}")
logger.error(f"Traceback: {traceback.format_exc()}") logger.error(f"Traceback: {traceback.format_exc()}")
return None return None
@ -464,12 +464,12 @@ class MilvusKB(KnowledgeBase):
"""将文本分割成块""" """将文本分割成块"""
return chunk_markdown(text, file_id, filename, params) return chunk_markdown(text, file_id, filename, params)
def _build_chunk_pg_records(self, db_id: str, chunks: list[dict]) -> list[dict[str, Any]]: def _build_chunk_pg_records(self, kb_id: str, chunks: list[dict]) -> list[dict[str, Any]]:
return [ return [
{ {
"chunk_id": chunk["chunk_id"], "chunk_id": chunk["chunk_id"],
"file_id": chunk["file_id"], "file_id": chunk["file_id"],
"db_id": db_id, "kb_id": kb_id,
"chunk_index": chunk["chunk_index"], "chunk_index": chunk["chunk_index"],
"content": chunk["content"], "content": chunk["content"],
"start_char_pos": chunk.get("start_char_pos"), "start_char_pos": chunk.get("start_char_pos"),
@ -486,7 +486,7 @@ class MilvusKB(KnowledgeBase):
async def _insert_chunks_to_stores( async def _insert_chunks_to_stores(
self, self,
db_id: str, kb_id: str,
file_id: str, file_id: str,
collection: Collection, collection: Collection,
chunks: list[dict], chunks: list[dict],
@ -508,7 +508,7 @@ class MilvusKB(KnowledgeBase):
def _insert_milvus_records(): def _insert_milvus_records():
collection.insert(entities) collection.insert(entities)
pg_task = chunk_repo.batch_upsert(self._build_chunk_pg_records(db_id, chunks)) pg_task = chunk_repo.batch_upsert(self._build_chunk_pg_records(kb_id, chunks))
milvus_task = asyncio.to_thread(_insert_milvus_records) milvus_task = asyncio.to_thread(_insert_milvus_records)
results = await asyncio.gather(pg_task, milvus_task, return_exceptions=True) results = await asyncio.gather(pg_task, milvus_task, return_exceptions=True)
errors = [result for result in results if isinstance(result, Exception)] errors = [result for result in results if isinstance(result, Exception)]
@ -545,7 +545,7 @@ class MilvusKB(KnowledgeBase):
return "未知来源" return "未知来源"
return self.files_meta.get(file_id, {}).get("filename") or "未知来源" return self.files_meta.get(file_id, {}).get("filename") or "未知来源"
def _build_file_name_expr(self, db_id: str, file_name: str | None) -> str | None: def _build_file_name_expr(self, kb_id: str, file_name: str | None) -> str | None:
if not file_name: if not file_name:
return None return None
@ -553,7 +553,7 @@ class MilvusKB(KnowledgeBase):
matched_file_ids = [ matched_file_ids = [
file_id file_id
for file_id, file_meta in self.files_meta.items() for file_id, file_meta in self.files_meta.items()
if file_meta.get("database_id") == db_id and file_name_pattern in (file_meta.get("filename") or "") if file_meta.get("kb_id") == kb_id and file_name_pattern in (file_meta.get("filename") or "")
] ]
if not matched_file_ids: if not matched_file_ids:
return 'file_id == "__no_matching_file__"' return 'file_id == "__no_matching_file__"'
@ -564,13 +564,13 @@ class MilvusKB(KnowledgeBase):
return f'file_id in ["{joined_ids}"]' return f'file_id in ["{joined_ids}"]'
async def index_file( async def index_file(
self, db_id: str, file_id: str, operator_id: str | None = None, params: dict | None = None self, kb_id: str, file_id: str, operator_id: str | None = None, params: dict | None = None
) -> dict: ) -> dict:
""" """
Index parsed file (Status: INDEXING -> INDEXED/ERROR_INDEXING) Index parsed file (Status: INDEXING -> INDEXED/ERROR_INDEXING)
Args: Args:
db_id: Database ID kb_id: Database ID
file_id: File ID file_id: File ID
operator_id: ID of the user performing the operation operator_id: ID of the user performing the operation
params: Override processing params to apply during indexing (merged on top of stored params) params: Override processing params to apply during indexing (merged on top of stored params)
@ -578,15 +578,15 @@ class MilvusKB(KnowledgeBase):
Returns: Returns:
Updated file metadata Updated file metadata
""" """
if db_id not in self.databases_meta: if kb_id not in self.databases_meta:
raise ValueError(f"Database {db_id} not found") raise ValueError(f"Database {kb_id} not found")
# Get/Create collection # Get/Create collection
collection = await self._get_milvus_collection(db_id) collection = await self._get_milvus_collection(kb_id)
if not collection: if not collection:
raise ValueError(f"Failed to get Milvus collection for {db_id}") raise ValueError(f"Failed to get Milvus collection for {kb_id}")
embedding_model_spec = self.databases_meta[db_id].get("embedding_model_spec") embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
embedding_function = self._get_embedding_function(embedding_model_spec) embedding_function = self._get_embedding_function(embedding_model_spec)
# Get file meta # Get file meta
@ -627,7 +627,7 @@ class MilvusKB(KnowledgeBase):
# 将传入的 params 作为 request_params，确保用户指定的参数始终覆盖存储的参数 # 将传入的 params 作为 request_params，确保用户指定的参数始终覆盖存储的参数
params = resolve_processing_params( params = resolve_processing_params(
kb_additional_params=self.databases_meta.get(db_id, {}).get("metadata"), kb_additional_params=self.databases_meta.get(kb_id, {}).get("metadata"),
file_processing_params=file_meta.get("processing_params"), file_processing_params=file_meta.get("processing_params"),
request_params=params, request_params=params,
) )
@ -656,8 +656,8 @@ class MilvusKB(KnowledgeBase):
embeddings = await embedding_function(texts) embeddings = await embedding_function(texts)
# Clean up existing chunks if any (for re-indexing) # Clean up existing chunks if any (for re-indexing)
await self.delete_file_chunks_only(db_id, file_id) await self.delete_file_chunks_only(kb_id, file_id)
await self._insert_chunks_to_stores(db_id, file_id, collection, chunks, embeddings) await self._insert_chunks_to_stores(kb_id, file_id, collection, chunks, embeddings)
logger.info(f"Indexed file {file_id} into Milvus") logger.info(f"Indexed file {file_id} into Milvus")
@ -685,16 +685,16 @@ class MilvusKB(KnowledgeBase):
# Remove from processing queue # Remove from processing queue
self._remove_from_processing_queue(file_id) self._remove_from_processing_queue(file_id)
async def update_content(self, db_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]: async def update_content(self, kb_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]:
"""更新内容 - 根据file_ids重新解析文件并更新向量库""" """更新内容 - 根据file_ids重新解析文件并更新向量库"""
if db_id not in self.databases_meta: if kb_id not in self.databases_meta:
raise ValueError(f"Database {db_id} not found") raise ValueError(f"Database {kb_id} not found")
collection = await self._get_milvus_collection(db_id) collection = await self._get_milvus_collection(kb_id)
if not collection: if not collection:
raise ValueError(f"Failed to get Milvus collection for {db_id}") raise ValueError(f"Failed to get Milvus collection for {kb_id}")
embedding_model_spec = self.databases_meta[db_id].get("embedding_model_spec") embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
embedding_function = self._get_embedding_function(embedding_model_spec) embedding_function = self._get_embedding_function(embedding_model_spec)
# 处理默认参数 # 处理默认参数
@ -724,7 +724,7 @@ class MilvusKB(KnowledgeBase):
# 更新状态为处理中 # 更新状态为处理中
async with self._metadata_lock: async with self._metadata_lock:
resolved_params = resolve_processing_params( resolved_params = resolve_processing_params(
kb_additional_params=self.databases_meta.get(db_id, {}).get("metadata"), kb_additional_params=self.databases_meta.get(kb_id, {}).get("metadata"),
file_processing_params=self.files_meta[file_id].get("processing_params"), file_processing_params=self.files_meta[file_id].get("processing_params"),
request_params=params, request_params=params,
) )
@ -733,7 +733,7 @@ class MilvusKB(KnowledgeBase):
await self._persist_file(file_id) await self._persist_file(file_id)
# 重新解析文件为 markdown # 重新解析文件为 markdown
parse_params = {**resolved_params, "image_bucket": "public", "image_prefix": f"{db_id}/kb-images"} 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) markdown_content = await Parser.aparse(source=file_path, params=parse_params)
# 重新生成 chunks # 重新生成 chunks
@ -741,12 +741,12 @@ class MilvusKB(KnowledgeBase):
logger.info(f"Split {filename} into {len(chunks)} chunks") logger.info(f"Split {filename} into {len(chunks)} chunks")
# 先删除现有 chunks，保留文件元数据 # 先删除现有 chunks，保留文件元数据
await self.delete_file_chunks_only(db_id, file_id) await self.delete_file_chunks_only(kb_id, file_id)
if chunks: if chunks:
texts = [chunk["content"] for chunk in chunks] texts = [chunk["content"] for chunk in chunks]
embeddings = await embedding_function(texts) embeddings = await embedding_function(texts)
await self._insert_chunks_to_stores(db_id, file_id, collection, chunks, embeddings) await self._insert_chunks_to_stores(kb_id, file_id, collection, chunks, embeddings)
logger.info(f"Updated file {file_path} in Milvus. Done.") logger.info(f"Updated file {file_path} in Milvus. Done.")
@ -804,13 +804,13 @@ class MilvusKB(KnowledgeBase):
chunk["distance"] = hit.distance chunk["distance"] = hit.distance
return chunk return chunk
async def aquery(self, query_text: str, db_id: str, agent_call: bool = False, **kwargs) -> list[dict]: async def aquery(self, query_text: str, kb_id: str, agent_call: bool = False, **kwargs) -> list[dict]:
"""异步查询知识库""" """异步查询知识库"""
collection = await self._get_milvus_collection(db_id) collection = await self._get_milvus_collection(kb_id)
if not collection: if not collection:
raise ValueError(f"Database {db_id} not found") raise ValueError(f"Database {kb_id} not found")
query_params = self._get_query_params(db_id) query_params = self._get_query_params(kb_id)
# 合并查询参数：kwargs（临时参数）优先级高于 query_params（持久化参数） # 合并查询参数：kwargs（临时参数）优先级高于 query_params（持久化参数）
# 这样允许用户在单次查询中临时覆盖持久化配置 # 这样允许用户在单次查询中临时覆盖持久化配置
merged_kwargs = {**query_params, **kwargs} merged_kwargs = {**query_params, **kwargs}
@ -835,14 +835,14 @@ class MilvusKB(KnowledgeBase):
else: else:
recall_top_k = final_top_k recall_top_k = final_top_k
file_expr = self._build_file_name_expr(db_id, merged_kwargs.get("file_name")) file_expr = self._build_file_name_expr(kb_id, merged_kwargs.get("file_name"))
if file_expr: if file_expr:
logger.debug(f"Using filter expression: {file_expr}") logger.debug(f"Using filter expression: {file_expr}")
output_fields = ["content", "chunk_id", "file_id", "chunk_index"] output_fields = ["content", "chunk_id", "file_id", "chunk_index"]
retrieved_chunks: list[dict] = [] retrieved_chunks: list[dict] = []
if search_mode == "vector": if search_mode == "vector":
embedding_model_spec = self.databases_meta[db_id].get("embedding_model_spec") embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
embedding_function = self._get_embedding_function(embedding_model_spec, sync=True) embedding_function = self._get_embedding_function(embedding_model_spec, sync=True)
query_embedding = embedding_function([query_text]) query_embedding = embedding_function([query_text])
@ -895,7 +895,7 @@ class MilvusKB(KnowledgeBase):
logger.debug(f"Milvus BM25 query response: {len(retrieved_chunks)} chunks found") logger.debug(f"Milvus BM25 query response: {len(retrieved_chunks)} chunks found")
else: else:
embedding_model_spec = self.databases_meta[db_id].get("embedding_model_spec") embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
embedding_function = self._get_embedding_function(embedding_model_spec, sync=True) embedding_function = self._get_embedding_function(embedding_model_spec, sync=True)
query_embedding = embedding_function([query_text]) query_embedding = embedding_function([query_text])
bm25_top_k = int(merged_kwargs.get("bm25_top_k", recall_top_k)) bm25_top_k = int(merged_kwargs.get("bm25_top_k", recall_top_k))
@ -939,7 +939,7 @@ class MilvusKB(KnowledgeBase):
logger.debug(f"Milvus hybrid query response: {len(retrieved_chunks)} chunks found") logger.debug(f"Milvus hybrid query response: {len(retrieved_chunks)} chunks found")
if use_graph_retrieval: if use_graph_retrieval:
graph_chunks = await self._retrieve_graph_chunks(query_text, db_id, retrieved_chunks, merged_kwargs) graph_chunks = await self._retrieve_graph_chunks(query_text, kb_id, retrieved_chunks, merged_kwargs)
if graph_chunks: if graph_chunks:
graph_weight = float(merged_kwargs.get("graph_weight", 1.0)) graph_weight = float(merged_kwargs.get("graph_weight", 1.0))
retrieved_chunks = self._fuse_chunk_rankings(retrieved_chunks, graph_chunks, graph_weight) retrieved_chunks = self._fuse_chunk_rankings(retrieved_chunks, graph_chunks, graph_weight)
@ -974,7 +974,7 @@ class MilvusKB(KnowledgeBase):
key=lambda item: item.get("rerank_score", item.get("score", 0.0)), reverse=True key=lambda item: item.get("rerank_score", item.get("score", 0.0)), reverse=True
) )
elapsed = time.time() - rerank_start elapsed = time.time() - rerank_start
logger.info(f"Reranking completed for {db_id} in {elapsed:.3f}s with model {reranker_model}") logger.info(f"Reranking completed for {kb_id} in {elapsed:.3f}s with model {reranker_model}")
finally: finally:
await reranker.aclose() await reranker.aclose()
@ -991,7 +991,7 @@ class MilvusKB(KnowledgeBase):
async def _retrieve_graph_chunks( async def _retrieve_graph_chunks(
self, self,
query_text: str, query_text: str,
db_id: str, kb_id: str,
base_chunks: list[dict], base_chunks: list[dict],
query_params: dict[str, Any], query_params: dict[str, Any],
) -> list[dict]: ) -> list[dict]:
@ -999,7 +999,7 @@ class MilvusKB(KnowledgeBase):
from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService
from yuxi.knowledge.graphs.milvus_graph_vector_store import MilvusGraphVectorStore from yuxi.knowledge.graphs.milvus_graph_vector_store import MilvusGraphVectorStore
embedding_model_spec = self.databases_meta[db_id].get("embedding_model_spec") embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
if not embedding_model_spec: if not embedding_model_spec:
return [] return []
@ -1011,25 +1011,25 @@ class MilvusKB(KnowledgeBase):
vector_store = MilvusGraphVectorStore() vector_store = MilvusGraphVectorStore()
entity_hits, triple_hits = await asyncio.gather( entity_hits, triple_hits = await asyncio.gather(
vector_store.search_entities( vector_store.search_entities(
db_id=db_id, kb_id=kb_id,
query_text=query_text, query_text=query_text,
embedding_model_spec=embedding_model_spec, embedding_model_spec=embedding_model_spec,
top_k=entity_top_k, top_k=entity_top_k,
), ),
vector_store.search_triples( vector_store.search_triples(
db_id=db_id, kb_id=kb_id,
query_text=query_text, query_text=query_text,
embedding_model_spec=embedding_model_spec, embedding_model_spec=embedding_model_spec,
top_k=triple_top_k, top_k=triple_top_k,
), ),
) )
seed_weights = await self._build_graph_seed_weights(db_id, base_chunks, entity_hits, triple_hits) seed_weights = await self._build_graph_seed_weights(kb_id, base_chunks, entity_hits, triple_hits)
if not seed_weights: if not seed_weights:
return [] return []
graph_service = MilvusGraphService() graph_service = MilvusGraphService()
graph_scores = await graph_service.query_and_rank_chunks_by_ppr( graph_scores = await graph_service.query_and_rank_chunks_by_ppr(
db_id, kb_id,
seed_weights, seed_weights,
max_nodes=graph_max_nodes, max_nodes=graph_max_nodes,
top_k=graph_top_k, top_k=graph_top_k,
@ -1045,12 +1045,12 @@ class MilvusKB(KnowledgeBase):
for chunk in chunks for chunk in chunks
] ]
except Exception as exc: # noqa: BLE001 except Exception as exc: # noqa: BLE001
logger.error(f"Graph retrieval failed for {db_id}: {exc}") logger.error(f"Graph retrieval failed for {kb_id}: {exc}")
return [] return []
async def _build_graph_seed_weights( async def _build_graph_seed_weights(
self, self,
db_id: str, kb_id: str,
base_chunks: list[dict], base_chunks: list[dict],
entity_hits: list[dict[str, Any]], entity_hits: list[dict[str, Any]],
triple_hits: list[dict[str, Any]], triple_hits: list[dict[str, Any]],
@ -1129,18 +1129,18 @@ class MilvusKB(KnowledgeBase):
return sorted(fused.values(), key=lambda item: item.get("fusion_score", 0.0), reverse=True) return sorted(fused.values(), key=lambda item: item.get("fusion_score", 0.0), reverse=True)
async def delete_file_chunks_only(self, db_id: str, file_id: str) -> None: async def delete_file_chunks_only(self, kb_id: str, file_id: str) -> None:
"""仅删除文件的chunks数据，保留元数据（用于更新操作）""" """仅删除文件的chunks数据，保留元数据（用于更新操作）"""
chunk_repo = KnowledgeChunkRepository() chunk_repo = KnowledgeChunkRepository()
if await chunk_repo.count_graph_indexed_by_file_id(file_id): if await chunk_repo.count_graph_indexed_by_file_id(file_id):
from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService
try: try:
await MilvusGraphService().delete_file_graph(db_id, file_id) await MilvusGraphService().delete_file_graph(kb_id, file_id)
except Exception as e: except Exception as e:
logger.error(f"Failed to delete graph data for file {file_id}: {e}") logger.error(f"Failed to delete graph data for file {file_id}: {e}")
await chunk_repo.delete_by_file_id(file_id) await chunk_repo.delete_by_file_id(file_id)
collection = await self._get_milvus_collection(db_id) collection = await self._get_milvus_collection(kb_id)
if collection: if collection:
# 先查询文件是否存在，避免不必要的删除操作 # 先查询文件是否存在，避免不必要的删除操作
@ -1150,10 +1150,10 @@ class MilvusKB(KnowledgeBase):
logger.error(f"Error checking file existence in Milvus: {e}") logger.error(f"Error checking file existence in Milvus: {e}")
# 注意：这里不删除 files_meta[file_id]，保留元数据用于后续操作 # 注意：这里不删除 files_meta[file_id]，保留元数据用于后续操作
async def delete_file(self, db_id: str, file_id: str) -> None: async def delete_file(self, kb_id: str, file_id: str) -> None:
"""删除文件（包括元数据）""" """删除文件（包括元数据）"""
# 先删除 Milvus 中的 chunks 数据 # 先删除 Milvus 中的 chunks 数据
await self.delete_file_chunks_only(db_id, file_id) await self.delete_file_chunks_only(kb_id, file_id)
# 使用锁确保元数据操作的原子性 # 使用锁确保元数据操作的原子性
async with self._metadata_lock: async with self._metadata_lock:
@ -1163,14 +1163,14 @@ class MilvusKB(KnowledgeBase):
await KnowledgeFileRepository().delete(file_id) await KnowledgeFileRepository().delete(file_id)
async def get_file_basic_info(self, db_id: str, file_id: str) -> dict: async def get_file_basic_info(self, kb_id: str, file_id: str) -> dict:
"""获取文件基本信息（仅元数据）""" """获取文件基本信息（仅元数据）"""
if file_id not in self.files_meta: if file_id not in self.files_meta:
raise Exception(f"File not found: {file_id}") raise Exception(f"File not found: {file_id}")
return {"meta": self.files_meta[file_id]} return {"meta": self.files_meta[file_id]}
async def get_file_content(self, db_id: str, file_id: str) -> dict: async def get_file_content(self, kb_id: str, file_id: str) -> dict:
"""获取文件内容信息（chunks和lines）""" """获取文件内容信息（chunks和lines）"""
if file_id not in self.files_meta: if file_id not in self.files_meta:
raise Exception(f"File not found: {file_id}") raise Exception(f"File not found: {file_id}")
@ -1211,37 +1211,37 @@ class MilvusKB(KnowledgeBase):
return content_info return content_info
async def get_file_info(self, db_id: str, file_id: str) -> dict: async def get_file_info(self, kb_id: str, file_id: str) -> dict:
"""获取文件完整信息（基本信息+内容信息）""" """获取文件完整信息（基本信息+内容信息）"""
if file_id not in self.files_meta: if file_id not in self.files_meta:
raise Exception(f"File not found: {file_id}") raise Exception(f"File not found: {file_id}")
# 合并基本信息和内容信息 # 合并基本信息和内容信息
basic_info = await self.get_file_basic_info(db_id, file_id) basic_info = await self.get_file_basic_info(kb_id, file_id)
content_info = await self.get_file_content(db_id, file_id) content_info = await self.get_file_content(kb_id, file_id)
return {**basic_info, **content_info} return {**basic_info, **content_info}
def delete_database(self, db_id: str) -> dict: def delete_database(self, kb_id: str) -> dict:
"""删除数据库，同时清除Milvus中的集合""" """删除数据库，同时清除Milvus中的集合"""
# Drop Milvus collection # Drop Milvus collection
try: try:
if utility.has_collection(db_id, using=self.connection_alias): if utility.has_collection(kb_id, using=self.connection_alias):
utility.drop_collection(db_id, using=self.connection_alias) utility.drop_collection(kb_id, using=self.connection_alias)
logger.info(f"Dropped Milvus collection for {db_id}") logger.info(f"Dropped Milvus collection for {kb_id}")
else: else:
logger.info(f"Milvus collection {db_id} does not exist, skipping") logger.info(f"Milvus collection {kb_id} does not exist, skipping")
except Exception as e: except Exception as e:
logger.error(f"Failed to drop Milvus collection {db_id}: {e}") logger.error(f"Failed to drop Milvus collection {kb_id}: {e}")
from yuxi.knowledge.graphs.milvus_graph_vector_store import MilvusGraphVectorStore from yuxi.knowledge.graphs.milvus_graph_vector_store import MilvusGraphVectorStore
MilvusGraphVectorStore().drop_graph_collections(db_id) MilvusGraphVectorStore().drop_graph_collections(kb_id)
# Call base method to delete local files and metadata # Call base method to delete local files and metadata
return super().delete_database(db_id) return super().delete_database(kb_id)
def get_query_params_config(self, db_id: str, **kwargs) -> dict: def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
"""获取 Milvus 知识库的查询参数配置""" """获取 Milvus 知识库的查询参数配置"""
return {"type": "milvus", "options": _retrieval_config_options()} return {"type": "milvus", "options": _retrieval_config_options()}

32
backend/package/yuxi/knowledge/implementations/notion.py
View File

@ -227,11 +227,11 @@ class NotionKB(ReadOnlyConnectors):
"notion_version": notion_version, "notion_version": notion_version,
} }
async def aquery(self, query_text: str, db_id: str, agent_call: bool = False, **kwargs) -> list[dict]: async def aquery(self, query_text: str, kb_id: str, agent_call: bool = False, **kwargs) -> list[dict]:
del agent_call del agent_call
try: try:
token, data_source_id, notion_version = self._get_connection_config(db_id) token, data_source_id, notion_version = self._get_connection_config(kb_id)
query_params = self._get_query_params(db_id) query_params = self._get_query_params(kb_id)
merged = {**query_params, **kwargs} merged = {**query_params, **kwargs}
final_top_k = min(max(int(merged.get("final_top_k", 10)), 1), 50) final_top_k = min(max(int(merged.get("final_top_k", 10)), 1), 50)
max_scan_pages = min(max(int(merged.get("max_scan_pages", 100)), 10), 1000) max_scan_pages = min(max(int(merged.get("max_scan_pages", 100)), 10), 1000)
@ -270,7 +270,7 @@ class NotionKB(ReadOnlyConnectors):
return sorted(results, key=lambda item: item.get("score", 0.0), reverse=True)[:final_top_k] return sorted(results, key=lambda item: item.get("score", 0.0), reverse=True)[:final_top_k]
except (NotionAPIError, httpx.HTTPError, ValueError) as exc: except (NotionAPIError, httpx.HTTPError, ValueError) as exc:
logger.error(f"Notion query failed for db_id={db_id}: {exc}, {traceback.format_exc()}") logger.error(f"Notion query failed for kb_id={kb_id}: {exc}, {traceback.format_exc()}")
return [] return []
async def _search_candidate_pages( async def _search_candidate_pages(
@ -372,13 +372,13 @@ class NotionKB(ReadOnlyConnectors):
}, },
} }
async def open_file_content(self, db_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict: async def open_file_content(self, kb_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict:
content = await self._read_page_markdown(db_id, file_id) content = await self._read_page_markdown(kb_id, file_id)
return self._build_open_file_window(content, offset=offset, limit=limit) return self._build_open_file_window(content, offset=offset, limit=limit)
async def find_file_content( async def find_file_content(
self, self,
db_id: str, kb_id: str,
file_id: str, file_id: str,
patterns: list[str], patterns: list[str],
*, *,
@ -387,7 +387,7 @@ class NotionKB(ReadOnlyConnectors):
max_windows: int = 5, max_windows: int = 5,
window_size: int = 80, window_size: int = 80,
) -> dict: ) -> dict:
content = await self._read_page_markdown(db_id, file_id) content = await self._read_page_markdown(kb_id, file_id)
return self._build_find_file_windows( return self._build_find_file_windows(
content, content,
patterns=patterns, patterns=patterns,
@ -397,9 +397,9 @@ class NotionKB(ReadOnlyConnectors):
window_size=window_size, window_size=window_size,
) )
async def _read_page_markdown(self, db_id: str, page_id: str) -> str: async def _read_page_markdown(self, kb_id: str, page_id: str) -> str:
token, data_source_id, notion_version = self._get_connection_config(db_id) token, data_source_id, notion_version = self._get_connection_config(kb_id)
cache_key = (db_id, page_id, data_source_id, notion_version) cache_key = (kb_id, page_id, data_source_id, notion_version)
cached = self._get_cached_page_markdown(cache_key) cached = self._get_cached_page_markdown(cache_key)
if cached is not None: if cached is not None:
return cached return cached
@ -430,15 +430,15 @@ class NotionKB(ReadOnlyConnectors):
self._page_markdown_cache.pop(oldest_key, None) self._page_markdown_cache.pop(oldest_key, None)
self._page_markdown_cache[cache_key] = (time.monotonic(), content) self._page_markdown_cache[cache_key] = (time.monotonic(), content)
def _get_connection_config(self, db_id: str) -> tuple[str, str, str]: def _get_connection_config(self, kb_id: str) -> tuple[str, str, str]:
metadata = self.databases_meta.get(db_id, {}).get("metadata", {}) or {} metadata = self.databases_meta.get(kb_id, {}).get("metadata", {}) or {}
token = str( token = str(
metadata.get("notion_token") or os.getenv("NOTION_TOKEN") or os.getenv("NOTION_API_KEY") or "" metadata.get("notion_token") or os.getenv("NOTION_TOKEN") or os.getenv("NOTION_API_KEY") or ""
).strip() ).strip()
data_source_id = str(metadata.get("notion_data_source_id") or "").strip() data_source_id = str(metadata.get("notion_data_source_id") or "").strip()
notion_version = str(metadata.get("notion_version") or NOTION_DEFAULT_VERSION).strip() or NOTION_DEFAULT_VERSION notion_version = str(metadata.get("notion_version") or NOTION_DEFAULT_VERSION).strip() or NOTION_DEFAULT_VERSION
if not token or not data_source_id: if not token or not data_source_id:
raise ValueError(f"Notion config incomplete for db_id={db_id}") raise ValueError(f"Notion config incomplete for kb_id={kb_id}")
return token, data_source_id, notion_version return token, data_source_id, notion_version
async def _page_to_markdown( async def _page_to_markdown(
@ -679,8 +679,8 @@ class NotionKB(ReadOnlyConnectors):
def _preview(content: str, line_count: int) -> str: def _preview(content: str, line_count: int) -> str:
return "\n".join(content.splitlines()[:line_count]) return "\n".join(content.splitlines()[:line_count])
def get_query_params_config(self, db_id: str, **kwargs) -> dict: def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
del db_id, kwargs del kb_id, kwargs
return { return {
"type": "notion", "type": "notion",
"options": [ "options": [

46
backend/package/yuxi/knowledge/implementations/read_only_connectors.py
View File

@ -17,8 +17,8 @@ class ReadOnlyConnectors(KnowledgeBase):
def _readonly_error() -> ValueError: def _readonly_error() -> ValueError:
return ValueError("只读检索连接器不支持该操作") return ValueError("只读检索连接器不支持该操作")
async def _create_kb_instance(self, db_id: str, config: dict) -> Any: async def _create_kb_instance(self, kb_id: str, config: dict) -> Any:
del db_id, config del kb_id, config
return None return None
async def _initialize_kb_instance(self, instance: Any) -> None: async def _initialize_kb_instance(self, instance: Any) -> None:
@ -26,53 +26,53 @@ class ReadOnlyConnectors(KnowledgeBase):
return None return None
async def add_file_record( async def add_file_record(
self, db_id: str, item: str, params: dict | None = None, operator_id: str | None = None self, kb_id: str, item: str, params: dict | None = None, operator_id: str | None = None
) -> dict: ) -> dict:
raise self._readonly_error() raise self._readonly_error()
async def parse_file(self, db_id: str, file_id: str, operator_id: str | None = None) -> dict: async def parse_file(self, kb_id: str, file_id: str, operator_id: str | None = None) -> dict:
raise self._readonly_error() raise self._readonly_error()
async def update_file_params(self, db_id: str, file_id: str, params: dict, operator_id: str | None = None) -> None: async def update_file_params(self, kb_id: str, file_id: str, params: dict, operator_id: str | None = None) -> None:
raise self._readonly_error() raise self._readonly_error()
async def create_folder(self, db_id: str, folder_name: str, parent_id: str | None = None) -> dict: async def create_folder(self, kb_id: str, folder_name: str, parent_id: str | None = None) -> dict:
raise self._readonly_error() raise self._readonly_error()
async def move_file(self, db_id: str, file_id: str, new_parent_id: str | None) -> dict: async def move_file(self, kb_id: str, file_id: str, new_parent_id: str | None) -> dict:
raise self._readonly_error() raise self._readonly_error()
async def delete_folder(self, db_id: str, folder_id: str) -> None: async def delete_folder(self, kb_id: str, folder_id: str) -> None:
raise self._readonly_error() raise self._readonly_error()
async def index_file( async def index_file(
self, self,
db_id: str, kb_id: str,
file_id: str, file_id: str,
operator_id: str | None = None, operator_id: str | None = None,
params: dict | None = None, params: dict | None = None,
) -> dict: ) -> dict:
raise self._readonly_error() raise self._readonly_error()
async def update_content(self, db_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]: async def update_content(self, kb_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]:
raise self._readonly_error() raise self._readonly_error()
async def delete_file(self, db_id: str, file_id: str) -> None: async def delete_file(self, kb_id: str, file_id: str) -> None:
raise self._readonly_error() raise self._readonly_error()
async def get_file_basic_info(self, db_id: str, file_id: str) -> dict: async def get_file_basic_info(self, kb_id: str, file_id: str) -> dict:
raise self._readonly_error() raise self._readonly_error()
async def get_file_content(self, db_id: str, file_id: str) -> dict: async def get_file_content(self, kb_id: str, file_id: str) -> dict:
raise self._readonly_error() raise self._readonly_error()
async def open_file_content(self, db_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict: async def open_file_content(self, kb_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict:
del offset, limit del offset, limit
raise self._readonly_error() raise self._readonly_error()
async def find_file_content( async def find_file_content(
self, self,
db_id: str, kb_id: str,
file_id: str, file_id: str,
patterns: list[str], patterns: list[str],
*, *,
@ -81,26 +81,26 @@ class ReadOnlyConnectors(KnowledgeBase):
max_windows: int = 5, max_windows: int = 5,
window_size: int = 80, window_size: int = 80,
) -> dict: ) -> dict:
del db_id, file_id, patterns, use_regex, case_sensitive, max_windows, window_size del kb_id, file_id, patterns, use_regex, case_sensitive, max_windows, window_size
raise self._readonly_error() raise self._readonly_error()
async def get_file_info(self, db_id: str, file_id: str) -> dict: async def get_file_info(self, kb_id: str, file_id: str) -> dict:
raise self._readonly_error() raise self._readonly_error()
async def list_file_tree( async def list_file_tree(
self, self,
db_id: str, kb_id: str,
parent_id: str | None = None, parent_id: str | None = None,
recursive: bool = False, recursive: bool = False,
files_only: bool = False, files_only: bool = False,
) -> dict: ) -> dict:
del db_id, parent_id, recursive, files_only del kb_id, parent_id, recursive, files_only
raise ValueError("只读检索连接器不支持文件树预览") raise ValueError("只读检索连接器不支持文件树预览")
async def read_file_preview(self, db_id: str, file_id: str, variant: str = "parsed") -> dict: async def read_file_preview(self, kb_id: str, file_id: str, variant: str = "parsed") -> dict:
del db_id, file_id, variant del kb_id, file_id, variant
raise ValueError("只读检索连接器不支持文件预览") raise ValueError("只读检索连接器不支持文件预览")
async def get_file_download(self, db_id: str, file_id: str, variant: str = "original") -> dict: async def get_file_download(self, kb_id: str, file_id: str, variant: str = "original") -> dict:
del db_id, file_id, variant del kb_id, file_id, variant
raise ValueError("只读检索连接器不支持文件下载") raise ValueError("只读检索连接器不支持文件下载")

234
backend/package/yuxi/knowledge/manager.py
View File

@ -104,19 +104,19 @@ class KnowledgeBaseManager:
logger.info(f"Created {kb_type} knowledge base instance") logger.info(f"Created {kb_type} knowledge base instance")
return kb_instance return kb_instance
async def move_file(self, db_id: str, file_id: str, new_parent_id: str | None) -> dict: async def move_file(self, kb_id: str, file_id: str, new_parent_id: str | None) -> dict:
""" """
移动文件/文件夹 移动文件/文件夹
""" """
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.move_file(db_id, file_id, new_parent_id) return await kb_instance.move_file(kb_id, file_id, new_parent_id)
async def _get_kb_for_database(self, db_id: str) -> KnowledgeBase: async def _get_kb_for_database(self, kb_id: str) -> KnowledgeBase:
""" """
根据数据库ID获取对应的知识库实例 根据数据库ID获取对应的知识库实例
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
Returns: Returns:
知识库实例 知识库实例
@ -127,10 +127,10 @@ class KnowledgeBaseManager:
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
kb_repo = KnowledgeBaseRepository() kb_repo = KnowledgeBaseRepository()
kb = await kb_repo.get_by_id(db_id) kb = await kb_repo.get_by_kb_id(kb_id)
if kb is None: if kb is None:
raise KBNotFoundError(f"Database {db_id} not found") raise KBNotFoundError(f"Database {kb_id} not found")
kb_type = kb.kb_type or "milvus" kb_type = kb.kb_type or "milvus"
@ -143,16 +143,16 @@ class KnowledgeBaseManager:
# 统一的外部接口 # 统一的外部接口
# ============================================================================= # =============================================================================
async def aget_kb(self, db_id: str) -> KnowledgeBase: async def aget_kb(self, kb_id: str) -> KnowledgeBase:
"""异步获取知识库实例 """异步获取知识库实例
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
Returns: Returns:
知识库实例 知识库实例
""" """
return await self._get_kb_for_database(db_id) return await self._get_kb_for_database(kb_id)
def _normalize_share_config( def _normalize_share_config(
self, self,
@ -206,20 +206,20 @@ class KnowledgeBaseManager:
for row in rows: for row in rows:
kb_type = row.kb_type or "milvus" kb_type = row.kb_type or "milvus"
if not KnowledgeBaseFactory.is_type_supported(kb_type): if not KnowledgeBaseFactory.is_type_supported(kb_type):
logger.warning(f"Skip unsupported database: db_id={row.db_id}, kb_type={kb_type}") logger.warning(f"Skip unsupported database: kb_id={row.kb_id}, kb_type={kb_type}")
continue continue
kb_instance = self._get_or_create_kb_instance(kb_type) kb_instance = self._get_or_create_kb_instance(kb_type)
db_info = kb_instance.get_database_info(row.db_id, include_files=False) db_info = kb_instance.get_database_info(row.kb_id, include_files=False)
if not db_info and kb_type not in metadata_reloaded_types: if not db_info and kb_type not in metadata_reloaded_types:
try: try:
await kb_instance._load_metadata() await kb_instance._load_metadata()
metadata_reloaded_types.add(kb_type) metadata_reloaded_types.add(kb_type)
except Exception as e: except Exception as e:
logger.warning(f"Failed to reload metadata for kb_type={kb_type}: {e}") logger.warning(f"Failed to reload metadata for kb_type={kb_type}: {e}")
db_info = kb_instance.get_database_info(row.db_id, include_files=False) db_info = kb_instance.get_database_info(row.kb_id, include_files=False)
if not db_info: if not db_info:
logger.warning(f"Skip database due to missing metadata: db_id={row.db_id}, kb_type={kb_type}") logger.warning(f"Skip database due to missing metadata: kb_id={row.kb_id}, kb_type={kb_type}")
continue continue
# 补充 share_config 和 additional_params # 补充 share_config 和 additional_params
@ -258,12 +258,12 @@ class KnowledgeBaseManager:
return False return False
async def check_accessible(self, user: dict, db_id: str) -> bool: async def check_accessible(self, user: dict, kb_id: str) -> bool:
"""检查用户是否有权限访问数据库 """检查用户是否有权限访问数据库
Args: Args:
user: 用户信息字典 user: 用户信息字典
db_id: 数据库ID kb_id: 数据库ID
Returns: Returns:
bool: 是否有权限 bool: 是否有权限
@ -275,7 +275,7 @@ class KnowledgeBaseManager:
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
kb_repo = KnowledgeBaseRepository() kb_repo = KnowledgeBaseRepository()
kb = await kb_repo.get_by_id(db_id) kb = await kb_repo.get_by_kb_id(kb_id)
if kb is None: if kb is None:
return False return False
@ -344,10 +344,10 @@ class KnowledgeBaseManager:
return True return True
return False return False
async def create_folder(self, db_id: str, folder_name: str, parent_id: str = None) -> dict: async def create_folder(self, kb_id: str, folder_name: str, parent_id: str = None) -> dict:
"""Create a folder in the database.""" """Create a folder in the database."""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.create_folder(db_id, folder_name, parent_id) return await kb_instance.create_folder(kb_id, folder_name, parent_id)
async def create_database( async def create_database(
self, self,
@ -401,16 +401,16 @@ class KnowledgeBaseManager:
llm_model_spec=llm_model_spec, llm_model_spec=llm_model_spec,
**kwargs, **kwargs,
) )
db_id = db_info["db_id"] kb_id = db_info["kb_id"]
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
kb_repo = KnowledgeBaseRepository() kb_repo = KnowledgeBaseRepository()
updated = await kb_repo.update(db_id, {"share_config": share_config, "created_by": created_by}) updated = await kb_repo.update(kb_id, {"share_config": share_config, "created_by": created_by})
if updated is None: if updated is None:
await kb_repo.create( await kb_repo.create(
{ {
"db_id": db_id, "kb_id": kb_id,
"name": database_name, "name": database_name,
"description": description, "description": description,
"kb_type": kb_type, "kb_type": kb_type,
@ -422,76 +422,76 @@ class KnowledgeBaseManager:
} }
) )
logger.info(f"Created {kb_type} database: {database_name} ({db_id}) with {kwargs}") logger.info(f"Created {kb_type} database: {database_name} ({kb_id}) with {kwargs}")
db_info["share_config"] = share_config db_info["share_config"] = share_config
return db_info return db_info
async def delete_database(self, db_id: str) -> dict: async def delete_database(self, kb_id: str) -> dict:
"""删除数据库""" """删除数据库"""
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
try: try:
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
result = await kb_instance.delete_database(db_id) result = await kb_instance.delete_database(kb_id)
# 删除数据库记录 # 删除数据库记录
kb_repo = KnowledgeBaseRepository() kb_repo = KnowledgeBaseRepository()
await kb_repo.delete(db_id) await kb_repo.delete(kb_id)
return result return result
except KBNotFoundError as e: except KBNotFoundError as e:
logger.warning(f"Database {db_id} not found during deletion: {e}") logger.warning(f"Database {kb_id} not found during deletion: {e}")
return {"message": "删除成功"} return {"message": "删除成功"}
async def add_file_record( async def add_file_record(
self, db_id: str, item: str, params: dict | None = None, operator_id: str | None = None self, kb_id: str, item: str, params: dict | None = None, operator_id: str | None = None
) -> dict: ) -> dict:
"""Add file record to metadata""" """Add file record to metadata"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.add_file_record(db_id, item, params, operator_id) return await kb_instance.add_file_record(kb_id, item, params, operator_id)
async def parse_file(self, db_id: str, file_id: str, operator_id: str | None = None) -> dict: async def parse_file(self, kb_id: str, file_id: str, operator_id: str | None = None) -> dict:
"""Parse file to Markdown""" """Parse file to Markdown"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.parse_file(db_id, file_id, operator_id) return await kb_instance.parse_file(kb_id, file_id, operator_id)
async def index_file( async def index_file(
self, db_id: str, file_id: str, operator_id: str | None = None, params: dict | None = None self, kb_id: str, file_id: str, operator_id: str | None = None, params: dict | None = None
) -> dict: ) -> dict:
"""Index parsed file""" """Index parsed file"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.index_file(db_id, file_id, operator_id, params=params) return await kb_instance.index_file(kb_id, file_id, operator_id, params=params)
async def update_file_params(self, db_id: str, file_id: str, params: dict, operator_id: str | None = None) -> None: async def update_file_params(self, kb_id: str, file_id: str, params: dict, operator_id: str | None = None) -> None:
"""Update file processing params""" """Update file processing params"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
await kb_instance.update_file_params(db_id, file_id, params, operator_id) await kb_instance.update_file_params(kb_id, file_id, params, operator_id)
async def aquery(self, query_text: str, db_id: str, **kwargs) -> str: async def aquery(self, query_text: str, kb_id: str, **kwargs) -> str:
"""异步查询知识库""" """异步查询知识库"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.aquery(query_text, db_id, **kwargs) return await kb_instance.aquery(query_text, kb_id, **kwargs)
async def export_data(self, db_id: str, format: str = "zip", **kwargs) -> str: async def export_data(self, kb_id: str, format: str = "zip", **kwargs) -> str:
"""导出知识库数据""" """导出知识库数据"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.export_data(db_id, format=format, **kwargs) return await kb_instance.export_data(kb_id, format=format, **kwargs)
async def get_database_info(self, db_id: str) -> dict | None: async def get_database_info(self, kb_id: str) -> dict | None:
"""获取数据库详细信息""" """获取数据库详细信息"""
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
kb_repo = KnowledgeBaseRepository() kb_repo = KnowledgeBaseRepository()
kb = await kb_repo.get_by_id(db_id) kb = await kb_repo.get_by_kb_id(kb_id)
if kb is None: if kb is None:
return None return None
try: try:
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
db_info = kb_instance.get_database_info(db_id) db_info = kb_instance.get_database_info(kb_id)
except KBNotFoundError: except KBNotFoundError:
db_info = { db_info = {
"db_id": db_id, "kb_id": kb_id,
"name": kb.name, "name": kb.name,
"description": kb.description, "description": kb.description,
"kb_type": kb.kb_type, "kb_type": kb.kb_type,
@ -509,39 +509,39 @@ class KnowledgeBaseManager:
return db_info return db_info
async def delete_folder(self, db_id: str, folder_id: str) -> None: async def delete_folder(self, kb_id: str, folder_id: str) -> None:
"""递归删除文件夹""" """递归删除文件夹"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
await kb_instance.delete_folder(db_id, folder_id) await kb_instance.delete_folder(kb_id, folder_id)
async def delete_file(self, db_id: str, file_id: str) -> None: async def delete_file(self, kb_id: str, file_id: str) -> None:
"""删除文件""" """删除文件"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
await kb_instance.delete_file(db_id, file_id) await kb_instance.delete_file(kb_id, file_id)
async def update_content(self, db_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]: async def update_content(self, kb_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]:
"""更新内容（重新分块）""" """更新内容（重新分块）"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.update_content(db_id, file_ids, params or {}) return await kb_instance.update_content(kb_id, file_ids, params or {})
async def get_file_basic_info(self, db_id: str, file_id: str) -> dict: async def get_file_basic_info(self, kb_id: str, file_id: str) -> dict:
"""获取文件基本信息（仅元数据）""" """获取文件基本信息（仅元数据）"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.get_file_basic_info(db_id, file_id) return await kb_instance.get_file_basic_info(kb_id, file_id)
async def get_file_content(self, db_id: str, file_id: str) -> dict: async def get_file_content(self, kb_id: str, file_id: str) -> dict:
"""获取文件内容信息（chunks和lines）""" """获取文件内容信息（chunks和lines）"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.get_file_content(db_id, file_id) return await kb_instance.get_file_content(kb_id, file_id)
async def open_file_content(self, db_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict: async def open_file_content(self, kb_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict:
"""按行窗口打开文件解析后的 Markdown 内容""" """按行窗口打开文件解析后的 Markdown 内容"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.open_file_content(db_id, file_id, offset, limit) return await kb_instance.open_file_content(kb_id, file_id, offset, limit)
async def find_file_content( async def find_file_content(
self, self,
db_id: str, kb_id: str,
file_id: str, file_id: str,
patterns: list[str], patterns: list[str],
*, *,
@ -550,9 +550,9 @@ class KnowledgeBaseManager:
max_windows: int = 5, max_windows: int = 5,
window_size: int = 80, window_size: int = 80,
) -> dict: ) -> dict:
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.find_file_content( return await kb_instance.find_file_content(
db_id, kb_id,
file_id, file_id,
patterns, patterns,
use_regex=use_regex, use_regex=use_regex,
@ -561,40 +561,40 @@ class KnowledgeBaseManager:
window_size=window_size, window_size=window_size,
) )
async def get_file_info(self, db_id: str, file_id: str) -> dict: async def get_file_info(self, kb_id: str, file_id: str) -> dict:
"""获取文件完整信息（基本信息+内容信息）""" """获取文件完整信息（基本信息+内容信息）"""
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.get_file_info(db_id, file_id) return await kb_instance.get_file_info(kb_id, file_id)
async def list_file_tree( async def list_file_tree(
self, self,
db_id: str, kb_id: str,
parent_id: str | None = None, parent_id: str | None = None,
recursive: bool = False, recursive: bool = False,
files_only: bool = False, files_only: bool = False,
) -> dict: ) -> dict:
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.list_file_tree(db_id, parent_id, recursive, files_only) return await kb_instance.list_file_tree(kb_id, parent_id, recursive, files_only)
async def read_file_preview(self, db_id: str, file_id: str, variant: str = "parsed") -> dict: async def read_file_preview(self, kb_id: str, file_id: str, variant: str = "parsed") -> dict:
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.read_file_preview(db_id, file_id, variant) return await kb_instance.read_file_preview(kb_id, file_id, variant)
async def get_file_download(self, db_id: str, file_id: str, variant: str = "original") -> dict: async def get_file_download(self, kb_id: str, file_id: str, variant: str = "original") -> dict:
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.get_file_download(db_id, file_id, variant) return await kb_instance.get_file_download(kb_id, file_id, variant)
async def file_name_existed_in_db(self, db_id: str | None, file_name: str | None) -> bool: async def file_name_existed_in_db(self, kb_id: str | None, file_name: str | None) -> bool:
"""检查指定数据库中是否存在同名的文件""" """检查指定数据库中是否存在同名的文件"""
if not db_id or not file_name: if not kb_id or not file_name:
return False return False
try: try:
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
except KBNotFoundError: except KBNotFoundError:
return False return False
for file_info in kb_instance.files_meta.values(): for file_info in kb_instance.files_meta.values():
if file_info.get("database_id") != db_id: if file_info.get("kb_id") != kb_id:
continue continue
if file_info.get("status") == "failed": if file_info.get("status") == "failed":
continue continue
@ -603,13 +603,13 @@ class KnowledgeBaseManager:
return False return False
async def get_same_name_files(self, db_id: str, filename: str) -> list[dict]: async def get_same_name_files(self, kb_id: str, filename: str) -> list[dict]:
"""获取同一知识库中同名文件列表 """获取同一知识库中同名文件列表
基于原始文件名直接比较 基于原始文件名直接比较
返回基础信息：文件名、大小、上传时间 返回基础信息：文件名、大小、上传时间
Args: Args:
db_id: 数据库ID kb_id: 数据库ID
filename: 要检测的文件名（原始文件名） filename: 要检测的文件名（原始文件名）
Returns: Returns:
@ -619,16 +619,16 @@ class KnowledgeBaseManager:
- created_at: 上传时间 - created_at: 上传时间
- file_id: 文件ID（用于下载） - file_id: 文件ID（用于下载）
""" """
if not db_id or not filename: if not kb_id or not filename:
return [] return []
try: try:
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
except KBNotFoundError: except KBNotFoundError:
return [] return []
same_name_files = [] same_name_files = []
for file_id, file_info in kb_instance.files_meta.items(): for file_id, file_info in kb_instance.files_meta.items():
if file_info.get("database_id") != db_id: if file_info.get("kb_id") != kb_id:
continue continue
if file_info.get("status") == "failed": if file_info.get("status") == "failed":
continue continue
@ -651,18 +651,18 @@ class KnowledgeBaseManager:
same_name_files.sort(key=lambda x: x.get("created_at", ""), reverse=True) same_name_files.sort(key=lambda x: x.get("created_at", ""), reverse=True)
return same_name_files return same_name_files
async def file_existed_in_db(self, db_id: str | None, content_hash: str | None) -> bool: async def file_existed_in_db(self, kb_id: str | None, content_hash: str | None) -> bool:
"""检查指定数据库中是否存在相同内容哈希的文件""" """检查指定数据库中是否存在相同内容哈希的文件"""
if not db_id or not content_hash: if not kb_id or not content_hash:
return False return False
try: try:
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
except KBNotFoundError: except KBNotFoundError:
return False return False
for file_info in kb_instance.files_meta.values(): for file_info in kb_instance.files_meta.values():
if file_info.get("database_id") != db_id: if file_info.get("kb_id") != kb_id:
continue continue
if file_info.get("status") == "failed": if file_info.get("status") == "failed":
continue continue
@ -673,7 +673,7 @@ class KnowledgeBaseManager:
async def update_database( async def update_database(
self, self,
db_id: str, kb_id: str,
name: str, name: str,
description: str, description: str,
llm_model_spec: str | None = None, llm_model_spec: str | None = None,
@ -687,12 +687,12 @@ class KnowledgeBaseManager:
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
kb_repo = KnowledgeBaseRepository() kb_repo = KnowledgeBaseRepository()
kb = await kb_repo.get_by_id(db_id) kb = await kb_repo.get_by_kb_id(kb_id)
if kb is None: if kb is None:
raise ValueError(f"数据库 {db_id} 不存在") raise ValueError(f"数据库 {kb_id} 不存在")
kb_instance = await self._get_kb_for_database(db_id) kb_instance = await self._get_kb_for_database(kb_id)
kb_instance.update_database(db_id, name, description, llm_model_spec, update_llm_model_spec) kb_instance.update_database(kb_id, name, description, llm_model_spec, update_llm_model_spec)
update_data: dict = { update_data: dict = {
"name": name, "name": name,
@ -711,8 +711,8 @@ class KnowledgeBaseManager:
deep_merge(current_additional_params, additional_params) deep_merge(current_additional_params, additional_params)
) )
update_data["additional_params"] = merged_additional_params update_data["additional_params"] = merged_additional_params
if db_id in kb_instance.databases_meta: if kb_id in kb_instance.databases_meta:
kb_instance.databases_meta[db_id]["metadata"] = merged_additional_params kb_instance.databases_meta[kb_id]["metadata"] = merged_additional_params
if share_config is not None: if share_config is not None:
update_data["share_config"] = self._normalize_share_config( update_data["share_config"] = self._normalize_share_config(
@ -722,9 +722,9 @@ class KnowledgeBaseManager:
) )
# 保存到数据库 # 保存到数据库
await kb_repo.update(db_id, update_data) await kb_repo.update(kb_id, update_data)
return await self.get_database_info(db_id) return await self.get_database_info(kb_id)
def get_retrievers(self) -> dict[str, dict]: def get_retrievers(self) -> dict[str, dict]:
"""获取所有检索器""" """获取所有检索器"""
@ -830,7 +830,7 @@ class KnowledgeBaseManager:
# 从数据库获取所有已知的数据库ID # 从数据库获取所有已知的数据库ID
kb_repo = KnowledgeBaseRepository() kb_repo = KnowledgeBaseRepository()
rows = await kb_repo.get_all() rows = await kb_repo.get_all()
all_known_db_ids = {row.db_id for row in rows} all_known_kb_ids = {row.kb_id for row in rows}
# 找出存在于 Milvus 但不在 metadata 中的集合 # 找出存在于 Milvus 但不在 metadata 中的集合
# missing_collections = actual_collection_names - metadata_collection_names # missing_collections = actual_collection_names - metadata_collection_names
@ -842,7 +842,7 @@ class KnowledgeBaseManager:
# 检查集合是否属于已知数据库 # 检查集合是否属于已知数据库
is_known = False is_known = False
if collection_name in all_known_db_ids: if collection_name in all_known_kb_ids:
is_known = True is_known = True
# 如果是已知集合，跳过 # 如果是已知集合，跳过
@ -873,36 +873,36 @@ class KnowledgeBaseManager:
metadata_collection_names = set(milvus_kb.databases_meta.keys()) metadata_collection_names = set(milvus_kb.databases_meta.keys())
# 检查文件级别的不一致（针对已知的数据库） # 检查文件级别的不一致（针对已知的数据库）
for db_id in metadata_collection_names: for kb_id in metadata_collection_names:
try: try:
if utility.has_collection(db_id, using=milvus_kb.connection_alias): if utility.has_collection(kb_id, using=milvus_kb.connection_alias):
from pymilvus import Collection from pymilvus import Collection
collection = Collection(name=db_id, using=milvus_kb.connection_alias) collection = Collection(name=kb_id, using=milvus_kb.connection_alias)
actual_count = collection.num_entities actual_count = collection.num_entities
# 获取 metadata 中记录的文件数量 # 获取 metadata 中记录的文件数量
metadata_files_count = sum( metadata_files_count = sum(
1 for file_info in milvus_kb.files_meta.values() if file_info.get("database_id") == db_id 1 for file_info in milvus_kb.files_meta.values() if file_info.get("kb_id") == kb_id
) )
# 如果向量数据库中有数据但 metadata 中没有文件记录，可能存在文件缺失 # 如果向量数据库中有数据但 metadata 中没有文件记录，可能存在文件缺失
if actual_count > 0 and metadata_files_count == 0: if actual_count > 0 and metadata_files_count == 0:
inconsistencies["missing_files"].append( inconsistencies["missing_files"].append(
{ {
"database_id": db_id, "kb_id": kb_id,
"vector_count": actual_count, "vector_count": actual_count,
"metadata_files_count": metadata_files_count, "metadata_files_count": metadata_files_count,
"detected_at": utc_isoformat(), "detected_at": utc_isoformat(),
} }
) )
logger.warning( logger.warning(
f"发现数据库 {db_id} 在 Milvus 中有 {actual_count} 条向量数据，" f"发现数据库 {kb_id} 在 Milvus 中有 {actual_count} 条向量数据，"
"但 metadata 中没有文件记录" "但 metadata 中没有文件记录"
) )
except Exception as e: except Exception as e:
logger.debug(f"检查数据库 {db_id} 的文件一致性时出错: {e}") logger.debug(f"检查数据库 {kb_id} 的文件一致性时出错: {e}")
except Exception as e: except Exception as e:
logger.error(f"检测 Milvus 数据不一致时出错: {e}") logger.error(f"检测 Milvus 数据不一致时出错: {e}")
@ -933,7 +933,7 @@ class KnowledgeBaseManager:
logger.warning(f" 缺失文件记录数量: {len(milvus_files_missing)}") logger.warning(f" 缺失文件记录数量: {len(milvus_files_missing)}")
for file_info in milvus_files_missing: for file_info in milvus_files_missing:
logger.warning( logger.warning(
f" - 数据库: {file_info['database_id']}, 向量数: {file_info['vector_count']}, " f" - 数据库: {file_info['kb_id']}, 向量数: {file_info['vector_count']}, "
f"元数据文件数: {file_info['metadata_files_count']}" f"元数据文件数: {file_info['metadata_files_count']}"
) )

8
backend/package/yuxi/knowledge/utils/kb_utils.py
View File

@ -49,14 +49,14 @@ async def calculate_content_hash(data: bytes | bytearray) -> str:
return sha256.hexdigest() return sha256.hexdigest()
async def prepare_item_metadata(item: str, content_type: str, db_id: str, params: dict | None = None) -> dict: async def prepare_item_metadata(item: str, content_type: str, kb_id: str, params: dict | None = None) -> dict:
""" """
准备文件或URL的元数据，文件来源必须是 MinIO URL。 准备文件或URL的元数据，文件来源必须是 MinIO URL。
Args: Args:
item: MinIO URL 或 URL item: MinIO URL 或 URL
content_type: 内容类型 ("file" 或 "url") content_type: 内容类型 ("file" 或 "url")
db_id: 数据库ID kb_id: 数据库ID
params: 处理参数，可选 params: 处理参数，可选
""" """
# 检查是否有预处理信息 (针对 URL 转 HTML 文件的情况) # 检查是否有预处理信息 (针对 URL 转 HTML 文件的情况)
@ -81,7 +81,7 @@ async def prepare_item_metadata(item: str, content_type: str, db_id: str, params
file_id = f"file_{hashstr(item + str(time.time()), 6)}" file_id = f"file_{hashstr(item + str(time.time()), 6)}"
metadata = { metadata = {
"database_id": db_id, "kb_id": kb_id,
"filename": filename_display, "filename": filename_display,
"path": item_path, "path": item_path,
"file_type": file_type, "file_type": file_type,
@ -147,7 +147,7 @@ async def prepare_item_metadata(item: str, content_type: str, db_id: str, params
raise ValueError(f"Unsupported content_type: {content_type}") raise ValueError(f"Unsupported content_type: {content_type}")
metadata = { metadata = {
"database_id": db_id, "kb_id": kb_id,
"filename": filename_display, # 使用显示用的文件名 "filename": filename_display, # 使用显示用的文件名
"path": item_path, "path": item_path,
"file_type": file_type, "file_type": file_type,