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 os
import re
import secrets
import string
from abc import ABC, abstractmethod
from typing import Any
@ -83,10 +85,10 @@ class KnowledgeBase(ABC):
"""由 KnowledgeBaseManager 调用，同步加载元数据"""
# 过滤出当前 kb_type 的知识库
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:
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"),
"description": meta.get("description"),
"kb_type": meta.get("kb_type"),
@ -100,9 +102,9 @@ class KnowledgeBase(ABC):
# 过滤文件
self.files_meta = {}
for file_id, meta in files_meta.items():
if meta.get("database_id") in self.databases_meta:
db_id = meta.get("database_id")
kb_additional_params = self.databases_meta.get(db_id, {}).get("metadata") or {}
if meta.get("kb_id") in self.databases_meta:
kb_id = meta.get("kb_id")
kb_additional_params = self.databases_meta.get(kb_id, {}).get("metadata") or {}
normalized_meta = dict(meta)
normalized_meta["processing_params"] = resolve_processing_params(
kb_additional_params=kb_additional_params,
@ -182,12 +184,12 @@ class KnowledgeBase(ABC):
return params
@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:
db_id: 数据库ID
kb_id: 数据库ID
config: 配置信息
Returns:
@ -206,13 +208,13 @@ class KnowledgeBase(ABC):
pass
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:
"""
Add a file record to metadata (Status: UPLOADED)
Args:
db_id: Database ID
kb_id: Database ID
item: File path or URL
params: Parameters
operator_id: Operator ID who created the file
@ -226,9 +228,9 @@ class KnowledgeBase(ABC):
content_type = params.get("content_type", "file")
# 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"]
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(
kb_additional_params=kb_additional_params,
file_processing_params=metadata.get("processing_params"),
@ -262,12 +264,12 @@ class KnowledgeBase(ABC):
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)
Args:
db_id: Database ID
kb_id: Database ID
file_id: File ID
operator_id: ID of the user performing the operation
@ -317,7 +319,7 @@ class KnowledgeBase(ABC):
# Prepare params
params = file_meta.get("processing_params", {}) or {}
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(
source=file_path,
@ -325,7 +327,7 @@ class KnowledgeBase(ABC):
)
# 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
self.files_meta[file_id]["status"] = FileStatus.PARSED
@ -354,7 +356,7 @@ class KnowledgeBase(ABC):
# Remove from processing queue
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"""
if file_id not in self.files_meta:
raise ValueError(f"File {file_id} not found")
@ -364,7 +366,7 @@ class KnowledgeBase(ABC):
return
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}")
@ -395,7 +397,7 @@ class KnowledgeBase(ABC):
self.files_meta[file_id]["updated_by"] = operator_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"""
from yuxi.storage.minio import get_minio_client
@ -403,7 +405,7 @@ class KnowledgeBase(ABC):
bucket_name = minio_client.KB_BUCKETS["parsed"]
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")
# Return standard HTTP URL from UploadResult
@ -431,9 +433,9 @@ class KnowledgeBase(ABC):
content_bytes = await self._read_minio_bytes(file_path)
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)
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")
return file_meta
@ -450,7 +452,7 @@ class KnowledgeBase(ABC):
variants.append({"key": "parsed", "label": "MD", "supported": True})
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"))
variants = [] if is_dir else self._knowledge_preview_variants(file_meta)
preview_modes = [item["key"] for item in variants]
@ -464,11 +466,11 @@ class KnowledgeBase(ABC):
path = f"{path}/"
return {
"source": "knowledge",
"db_id": db_id,
"kb_id": kb_id,
"file_id": file_id,
"parent_id": file_meta.get("parent_id"),
"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,
"is_dir": is_dir,
"size": 0 if is_dir else file_meta.get("size") or 0,
@ -487,7 +489,7 @@ class KnowledgeBase(ABC):
def _list_knowledge_children(
self,
db_id: str,
kb_id: str,
parent_id: str | None,
*,
recursive: bool,
@ -496,16 +498,16 @@ class KnowledgeBase(ABC):
children = [
(file_id, meta)
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 = []
for file_id, meta in children:
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"):
entries.extend(
self._list_knowledge_children(
db_id,
kb_id,
file_id,
recursive=True,
files_only=files_only,
@ -515,20 +517,20 @@ class KnowledgeBase(ABC):
async def list_file_tree(
self,
db_id: str,
kb_id: str,
parent_id: str | None = None,
recursive: bool = False,
files_only: bool = False,
) -> dict:
if db_id not in self.databases_meta:
raise ValueError(f"Database {db_id} not found")
if kb_id not in self.databases_meta:
raise ValueError(f"Database {kb_id} not found")
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"):
raise ValueError("Parent is not a folder")
return {
"entries": self._list_knowledge_children(
db_id,
kb_id,
parent_id,
recursive=recursive,
files_only=files_only,
@ -536,10 +538,10 @@ class KnowledgeBase(ABC):
"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
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"):
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
response = {
"source": "knowledge",
"db_id": db_id,
"kb_id": kb_id,
"file_id": file_id,
"filename": filename,
"variant": variant,
@ -626,8 +628,8 @@ class KnowledgeBase(ABC):
"message": message,
}
async def get_file_download(self, db_id: str, file_id: str, variant: str = "original") -> dict:
file_meta = self._get_file_meta(db_id, file_id)
async def get_file_download(self, kb_id: str, file_id: str, variant: str = "original") -> dict:
file_meta = self._get_file_meta(kb_id, file_id)
if file_meta.get("is_folder"):
raise ValueError("Cannot download a folder")
if variant not in {"original", "parsed"}:
@ -675,7 +677,7 @@ class KnowledgeBase(ABC):
}
@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):
return retrieval_results
@ -705,14 +707,14 @@ class KnowledgeBase(ABC):
results.append(
SearchResultSchema(
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 ""),
content=str(chunk.get("content") or ""),
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
def _build_find_file_windows(
@ -770,21 +772,21 @@ class KnowledgeBase(ABC):
break
return FindOutputSchema(
resource_id="",
kb_id="",
file_id="",
semantic=False,
match_mode="regex" if use_regex else "keyword",
total_matches=len(matched_indexes),
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 内容"""
file_meta = self.files_meta.get(file_id)
if file_meta is None:
raise Exception(f"文件不存在: {file_id}")
if file_meta.get("database_id") != db_id:
raise Exception(f"文件 {file_id} 不属于知识库 {db_id}")
if file_meta.get("kb_id") != kb_id:
raise Exception(f"文件 {file_id} 不属于知识库 {kb_id}")
if file_meta.get("is_folder"):
raise Exception(f"文件 {file_id} 是文件夹")
@ -797,7 +799,7 @@ class KnowledgeBase(ABC):
async def find_file_content(
self,
db_id: str,
kb_id: str,
file_id: str,
patterns: list[str],
*,
@ -809,8 +811,8 @@ class KnowledgeBase(ABC):
file_meta = self.files_meta.get(file_id)
if file_meta is None:
raise Exception(f"文件不存在: {file_id}")
if file_meta.get("database_id") != db_id:
raise Exception(f"文件 {file_id} 不属于知识库 {db_id}")
if file_meta.get("kb_id") != kb_id:
raise Exception(f"文件 {file_id} 不属于知识库 {kb_id}")
if file_meta.get("is_folder"):
raise Exception(f"文件 {file_id} 是文件夹")
@ -829,12 +831,12 @@ class KnowledgeBase(ABC):
)
@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)
Args:
db_id: Database ID
kb_id: Database ID
file_id: File ID
operator_id: ID of the user performing the operation
@ -864,13 +866,15 @@ class KnowledgeBase(ABC):
Returns:
数据库信息字典
"""
from yuxi.utils import hashstr
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,
"description": description,
"kb_type": self.kb_type,
@ -878,32 +882,32 @@ class KnowledgeBase(ABC):
"llm_model_spec": llm_model_spec,
"metadata": kwargs,
"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)
# 返回数据库信息
db_dict = self.databases_meta[db_id].copy()
db_dict["db_id"] = db_id
db_dict = self.databases_meta[kb_id].copy()
db_dict["kb_id"] = kb_id
db_dict["files"] = {}
return db_dict
async def delete_database(self, db_id: str) -> dict:
async def delete_database(self, kb_id: str) -> dict:
"""
删除数据库
Args:
db_id: 数据库ID
kb_id: 数据库ID
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.repositories.knowledge_base_repository import KnowledgeBaseRepository
from yuxi.storage.minio import get_minio_client
@ -911,7 +915,7 @@ class KnowledgeBase(ABC):
minio_client = get_minio_client()
# 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:
file_path = self.files_meta[file_id].get("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}")
# 删除解析后的 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)
del self.files_meta[file_id]
# 2. 并行删除所有知识库 bucket 中该 db_id 下的文件
prefix = f"{db_id}/"
# 2. 并行删除所有知识库 bucket 中该 kb_id 下的文件
prefix = f"{kb_id}/"
cleanup_buckets = {
minio_client.KB_BUCKETS["parsed"],
minio_client.KB_BUCKETS["documents"],
@ -940,13 +944,13 @@ class KnowledgeBase(ABC):
await asyncio.gather(*cleanup_tasks)
# 3. 删除数据库记录
del self.databases_meta[db_id]
del self.databases_meta[kb_id]
kb_repo = KnowledgeBaseRepository()
await kb_repo.delete(db_id)
await kb_repo.delete(kb_id)
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):
import shutil
@ -957,7 +961,7 @@ class KnowledgeBase(ABC):
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."""
import uuid
@ -968,7 +972,7 @@ class KnowledgeBase(ABC):
"filename": folder_name,
"is_folder": True,
"parent_id": parent_id,
"database_id": db_id,
"kb_id": kb_id,
"created_at": utc_isoformat(),
"status": "done",
"path": folder_name,
@ -978,12 +982,12 @@ class KnowledgeBase(ABC):
return self.files_meta[folder_id]
@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重新解析文件并更新向量库
Args:
db_id: 数据库ID
kb_id: 数据库ID
file_ids: 文件ID列表
params: 处理参数
@ -993,13 +997,13 @@ class KnowledgeBase(ABC):
pass
@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:
query_text: 查询文本
db_id: 数据库ID
kb_id: 数据库ID
**kwargs: 查询参数
Returns:
@ -1008,12 +1012,12 @@ class KnowledgeBase(ABC):
pass
@abstractmethod
def get_query_params_config(self, db_id: str, **kwargs) -> dict:
def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
"""
获取知识库类型的查询参数配置
Args:
db_id: 数据库ID
kb_id: 数据库ID
**kwargs: 额外参数
Returns:
@ -1037,54 +1041,54 @@ class KnowledgeBase(ABC):
"""
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
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:
query_params_meta = self.databases_meta[db_id].get("query_params") or {}
if kb_id in self.databases_meta:
query_params_meta = self.databases_meta[kb_id].get("query_params") or {}
return query_params_meta.get("options", {})
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": {...}}"""
config = self.get_query_params_config(db_id)
config = self.get_query_params_config(kb_id)
defaults = {}
for opt in config.get("options", []):
if "default" in opt:
defaults[opt["key"]] = opt["default"]
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:
db_id: 数据库ID
kb_id: 数据库ID
include_files: 是否包含文件信息，默认为True
Returns:
数据库信息或None
"""
if db_id not in self.databases_meta:
if kb_id not in self.databases_meta:
return None
meta = self.databases_meta[db_id].copy()
meta["db_id"] = db_id
meta = self.databases_meta[kb_id].copy()
meta["kb_id"] = kb_id
# 检查并修复异常的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
# 仅在需要时加载文件详情
if include_files:
db_files = {}
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"))
db_files[file_id] = {
"file_id": file_id,
@ -1126,22 +1130,22 @@ class KnowledgeBase(ABC):
self._ensure_metadata_loaded()
databases = []
for db_id, meta in self.databases_meta.items():
for kb_id, meta in self.databases_meta.items():
# 检查并修复异常的processing状态
self._check_and_fix_processing_status(db_id)
self._check_and_fix_processing_status(kb_id)
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
# 仅在需要时加载文件详情
if include_files:
db_files = {}
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"))
db_files[file_id] = {
"file_id": file_id,
@ -1209,13 +1213,13 @@ class KnowledgeBase(ABC):
with cls._processing_lock:
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:
db_id: 数据库ID
kb_id: 数据库ID
"""
try:
status_changed = False
@ -1228,7 +1232,7 @@ class KnowledgeBase(ABC):
# 检查该数据库下所有中间状态的文件
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")
if current_status in intermediate_states:
@ -1248,44 +1252,44 @@ class KnowledgeBase(ABC):
# 如果有状态变更，保存元数据
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:
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.
Args:
db_id: Database ID
kb_id: Database ID
folder_id: Folder ID to delete
"""
# Find all children
children = [
fid
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:
child_meta = self.files_meta.get(child_id)
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:
await self.delete_file(db_id, child_id)
await self.delete_file(kb_id, child_id)
# Delete the folder itself
# We call delete_file which should handle the actual removal.
# 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.
Args:
db_id: Database ID
kb_id: Database ID
file_id: File/Folder ID to move
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")
meta = self.files_meta[file_id]
if meta.get("database_id") != db_id:
raise ValueError(f"File {file_id} does not belong to database {db_id}")
if meta.get("kb_id") != kb_id:
raise ValueError(f"File {file_id} does not belong to database {kb_id}")
# Basic cycle detection for folders
if meta.get("is_folder") and new_parent_id:
@ -1320,23 +1324,23 @@ class KnowledgeBase(ABC):
return meta
@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:
db_id: 数据库ID
kb_id: 数据库ID
file_id: 文件ID
"""
pass
@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:
db_id: 数据库ID
kb_id: 数据库ID
file_id: 文件ID
Returns:
@ -1345,12 +1349,12 @@ class KnowledgeBase(ABC):
pass
@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）
Args:
db_id: 数据库ID
kb_id: 数据库ID
file_id: 文件ID
Returns:
@ -1359,12 +1363,12 @@ class KnowledgeBase(ABC):
pass
@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:
db_id: 数据库ID
kb_id: 数据库ID
file_id: 文件ID
Returns:
@ -1374,7 +1378,7 @@ class KnowledgeBase(ABC):
def update_database(
self,
db_id: str,
kb_id: str,
name: str,
description: str,
llm_model_spec: str | None = None,
@ -1384,7 +1388,7 @@ class KnowledgeBase(ABC):
更新数据库
Args:
db_id: 数据库ID
kb_id: 数据库ID
name: 新名称
description: 新描述
llm_model_spec: LLM 模型 spec（可选）
@ -1392,15 +1396,15 @@ class KnowledgeBase(ABC):
Returns:
更新后的数据库信息
"""
if db_id not in self.databases_meta:
raise ValueError(f"数据库 {db_id} 不存在")
if kb_id not in self.databases_meta:
raise ValueError(f"数据库 {kb_id} 不存在")
self.databases_meta[db_id]["name"] = name
self.databases_meta[db_id]["description"] = description
self.databases_meta[kb_id]["name"] = name
self.databases_meta[kb_id]["description"] = description
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]:
"""
@ -1410,19 +1414,19 @@ class KnowledgeBase(ABC):
检索器字典
"""
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):
results = await self.aquery(query_text, db_id, agent_call=True, **kwargs)
return self.build_search_output(db_id, results)
results = await self.aquery(query_text, kb_id, agent_call=True, **kwargs)
return self.build_search_output(kb_id, results)
return retriever
retrievers[db_id] = {
retrievers[kb_id] = {
"name": meta["name"],
"description": meta["description"],
"retriever": make_retriever(db_id),
"retriever": make_retriever(kb_id),
"metadata": meta,
}
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]
self.databases_meta = {
kb.db_id: {
kb.kb_id: {
"name": kb.name,
"description": kb.description,
"kb_type": kb.kb_type,
"embedding_model_spec": kb.embedding_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),
"created_at": utc_isoformat(kb.created_at) if kb.created_at else utc_isoformat(),
}
@ -1451,11 +1455,11 @@ class KnowledgeBase(ABC):
self.files_meta = {}
for kb in databases:
kb_additional_params = self.databases_meta.get(kb.db_id, {}).get("metadata") or {}
for record in await file_repo.list_by_db_id(kb.db_id):
kb_additional_params = self.databases_meta.get(kb.kb_id, {}).get("metadata") or {}
for record in await file_repo.list_by_kb_id(kb.kb_id):
self.files_meta[record.file_id] = {
"file_id": record.file_id,
"database_id": record.db_id,
"kb_id": record.kb_id,
"parent_id": record.parent_id,
"filename": record.filename,
"file_type": record.file_type,
@ -1542,11 +1546,11 @@ class KnowledgeBase(ABC):
self._normalize_metadata_state()
for db_id, meta in self.databases_meta.items():
existing = await kb_repo.get_by_id(db_id)
for kb_id, meta in self.databases_meta.items():
existing = await kb_repo.get_by_kb_id(kb_id)
payload = {
"db_id": db_id,
"name": meta.get("name") or db_id,
"kb_id": kb_id,
"name": meta.get("name") or kb_id,
"description": meta.get("description"),
"kb_type": meta.get("kb_type") or self.kb_type,
"embedding_model_spec": meta.get("embedding_model_spec"),
@ -1558,13 +1562,13 @@ class KnowledgeBase(ABC):
await kb_repo.create(payload)
for file_id, meta in self.files_meta.items():
db_id = meta.get("database_id")
if not db_id:
kb_id = meta.get("kb_id")
if not kb_id:
continue
await file_repo.upsert(
file_id=file_id,
data={
"db_id": db_id,
"kb_id": kb_id,
"parent_id": meta.get("parent_id"),
"filename": meta.get("filename") or "",
"original_filename": meta.get("original_filename"),
@ -1594,14 +1598,14 @@ class KnowledgeBase(ABC):
return
meta = self.files_meta[file_id]
db_id = meta.get("database_id")
if not db_id:
kb_id = meta.get("kb_id")
if not kb_id:
return
await file_repo.upsert(
file_id=file_id,
data={
"db_id": db_id,
"kb_id": kb_id,
"parent_id": meta.get("parent_id"),
"filename": meta.get("filename") or "",
"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
kb_repo = KnowledgeBaseRepository()
if db_id not in self.databases_meta:
if kb_id not in self.databases_meta:
return
meta = self.databases_meta[db_id]
existing = await kb_repo.get_by_id(db_id)
meta = self.databases_meta[kb_id]
existing = await kb_repo.get_by_kb_id(kb_id)
payload = {
"db_id": db_id,
"name": meta.get("name") or db_id,
"kb_id": kb_id,
"name": meta.get("name") or kb_id,
"description": meta.get("description"),
"kb_type": meta.get("kb_type") or self.kb_type,
"embedding_model_spec": meta.get("embedding_model_spec"),
@ -1647,7 +1651,7 @@ class KnowledgeBase(ABC):
await kb_repo.create(payload)
else:
await kb_repo.update(
db_id,
kb_id,
{
"name": payload["name"],
"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
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 = []
for fid, finfo in kb_instance.files_meta.items():
if finfo.get("database_id") != db_id:
if finfo.get("kb_id") != kb_id:
continue
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", []):
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(
*, 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]]:
if neighbors_count <= 0:
return []
@ -86,7 +86,7 @@ async def select_neighbor_chunks_by_kb_query(
candidates = await kb_instance.aquery(
anchor_content,
db_id,
kb_id,
search_mode="vector",
final_top_k=neighbors_count + 3,
use_reranker=False,
@ -120,7 +120,7 @@ async def select_neighbor_chunks_by_kb_query(
async def select_graph_enhanced_chunks(
*,
db_id: str,
kb_id: str,
anchor_chunk: dict[str, Any],
chunks_by_id: dict[str, dict[str, Any]],
context_count: int,
@ -146,7 +146,7 @@ async def select_graph_enhanced_chunks(
seed_weights[entity_id] = 1.0
ranked_chunks = await graph_service.query_and_rank_chunks_by_ppr(
db_id,
kb_id,
seed_weights,
max_nodes=GRAPH_PPR_MAX_NODES,
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(
*,
kb_instance: Any,
db_id: str,
kb_id: str,
all_chunks: list[dict[str, Any]],
llm: Any,
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")
anchor_chunk = graph_anchor_chunks[random.randrange(len(graph_anchor_chunks))]
ctx_chunks = await select_graph_enhanced_chunks(
db_id=db_id,
kb_id=kb_id,
anchor_chunk=anchor_chunk,
chunks_by_id=chunks_by_id,
context_count=context_count,
@ -222,7 +222,7 @@ async def _generate_benchmark_item_once(
anchor_chunk = all_chunks[random.randrange(len(all_chunks))]
neighbor_chunks = await select_neighbor_chunks_by_kb_query(
kb_instance=kb_instance,
db_id=db_id,
kb_id=kb_id,
anchor_chunk=anchor_chunk,
neighbors_count=context_count - 1,
)
@ -254,7 +254,7 @@ async def _generate_benchmark_item_once(
async def iter_generated_benchmark_items(
*,
kb_instance: Any,
db_id: str,
kb_id: str,
count: int,
neighbors_count: int,
llm_model_spec: str | None,
@ -267,7 +267,7 @@ async def iter_generated_benchmark_items(
if progress_cb:
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:
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}
@ -310,7 +310,7 @@ async def iter_generated_benchmark_items(
try:
item = await _generate_benchmark_item_once(
kb_instance=kb_instance,
db_id=db_id,
kb_id=kb_id,
all_chunks=all_chunks,
llm=llm,
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(
*,
kb_instance: Any,
db_id: str,
kb_id: str,
question_data: dict[str, Any],
retrieval_config: dict[str, Any],
has_gold_chunks: bool,
@ -67,7 +67,7 @@ async def evaluate_question(
select_model_fn: Callable[..., Any],
) -> dict[str, Any]:
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 = await generate_answer_if_needed(
query=query,

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

@ -35,7 +35,7 @@ class EvaluationService:
"dataset_id": row.dataset_id,
"name": row.name,
"description": row.description,
"db_id": row.db_id,
"kb_id": row.kb_id,
"item_count": row.item_count,
"has_gold_chunks": row.has_gold_chunks,
"has_gold_answers": row.has_gold_answers,
@ -93,13 +93,13 @@ class EvaluationService:
row.build_metadata = metadata
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]]:
return [
{
"item_id": f"dataset_item_{uuid.uuid4().hex[:12]}",
"dataset_id": dataset_id,
"db_id": db_id,
"kb_id": kb_id,
"item_index": index,
"query_text": item["query"],
"gold_chunk_ids": item.get("gold_chunk_ids") or [],
@ -152,7 +152,7 @@ class EvaluationService:
return questions, has_gold_chunks, has_gold_answers
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]:
try:
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(
{
"dataset_id": dataset_id,
"db_id": db_id,
"kb_id": kb_id,
"name": dataset_name,
"description": description,
"item_count": len(questions),
@ -176,16 +176,16 @@ class EvaluationService:
},
"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)
except Exception as e:
logger.error(f"上传评估数据集失败: {e}")
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:
rows = await self.eval_repo.list_datasets(db_id)
rows = await self.eval_repo.list_datasets(kb_id)
for row in rows:
await self._sync_dataset_build_metadata(row)
return [self._dataset_to_dict(row) for row in rows]
@ -194,11 +194,11 @@ class EvaluationService:
raise
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]:
try:
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")
if (row.build_metadata or {}).get("status", "completed") != "completed":
raise ValueError("Dataset is not ready")
@ -250,7 +250,7 @@ class EvaluationService:
async def generate_dataset(
self,
db_id: str,
kb_id: str,
name: str,
description: str,
count: int,
@ -269,7 +269,7 @@ class EvaluationService:
if generation_mode not in {"vector", "graph_enhanced"}:
raise ValueError("不支持的评估基准生成方式")
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:
raise ValueError("当前知识库尚未完成图索引，无法使用图增强构建")
build_metadata = {
@ -288,7 +288,7 @@ class EvaluationService:
await self.eval_repo.create_dataset(
{
"dataset_id": dataset_id,
"db_id": db_id,
"kb_id": kb_id,
"name": name,
"description": description,
"item_count": 0,
@ -303,7 +303,7 @@ class EvaluationService:
task_type="dataset_generation",
payload={
"dataset_id": dataset_id,
"db_id": db_id,
"kb_id": kb_id,
"created_by": created_by,
"name": name,
"description": description,
@ -333,7 +333,7 @@ class EvaluationService:
payload = task.payload if task else {}
dataset_id = payload.get("dataset_id")
db_id = payload.get("db_id")
kb_id = payload.get("kb_id")
count = int(payload.get("count", 10))
neighbors_count = int(payload.get("neighbors_count", 1))
concurrency_count = normalize_generation_concurrency_count(payload.get("concurrency_count"))
@ -366,7 +366,7 @@ class EvaluationService:
)
try:
kb_instance = await knowledge_base.aget_kb(db_id)
kb_instance = await knowledge_base.aget_kb(kb_id)
if not kb_instance:
await report_progress(100, "知识库不存在")
raise ValueError("Knowledge Base not found")
@ -378,7 +378,7 @@ class EvaluationService:
try:
async for item in iter_generated_benchmark_items(
kb_instance=kb_instance,
db_id=db_id,
kb_id=kb_id,
count=count,
neighbors_count=neighbors_count,
llm_model_spec=llm_model_spec,
@ -397,7 +397,7 @@ class EvaluationService:
if not questions:
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._update_dataset_build_metadata(
dataset_id,
@ -419,26 +419,26 @@ class EvaluationService:
raise
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:
try:
run_id = f"run_{uuid.uuid4().hex[:8]}"
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")
if (dataset_row.build_metadata or {}).get("status", "completed") != "completed":
raise ValueError("Dataset is not ready")
retrieval_config = {}
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 {}
retrieval_config = query_params.get("options", {}) if isinstance(query_params, dict) else {}
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:
retrieval_config = kb_instance._get_default_query_params(db_id).get("options", {})
logger.info(f"从知识库 {db_id} 加载检索配置: {list(retrieval_config.keys())}")
retrieval_config = kb_instance._get_default_query_params(kb_id).get("options", {})
logger.info(f"从知识库 {kb_id} 加载检索配置: {list(retrieval_config.keys())}")
except Exception as e:
logger.error(f"获取知识库检索配置失败: {e}")
@ -448,7 +448,7 @@ class EvaluationService:
await self.eval_repo.create_run(
{
"run_id": run_id,
"db_id": db_id,
"kb_id": kb_id,
"dataset_id": dataset_id,
"status": "running",
"retrieval_config": retrieval_config,
@ -467,7 +467,7 @@ class EvaluationService:
task_type="rag_evaluation",
payload={
"run_id": run_id,
"db_id": db_id,
"kb_id": kb_id,
"dataset_id": dataset_id,
"retrieval_config": retrieval_config,
"created_by": created_by,
@ -487,21 +487,21 @@ class EvaluationService:
payload = task.payload
run_id = payload["run_id"]
db_id = payload["db_id"]
kb_id = payload["kb_id"]
dataset_id = payload["dataset_id"]
retrieval_config = payload["retrieval_config"]
await context.set_progress(5, "加载评估数据集")
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")
dataset_items = await self.eval_repo.list_all_dataset_items(dataset_id)
if not dataset_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:
raise ValueError(f"Knowledge Base {db_id} not found")
raise ValueError(f"Knowledge Base {kb_id} not found")
judge_llm = None
if dataset_row.has_gold_answers:
@ -544,7 +544,7 @@ class EvaluationService:
}
question_result = await evaluate_question(
kb_instance=kb_instance,
db_id=db_id,
kb_id=kb_id,
question_data=question_data,
retrieval_config=retrieval_config,
has_gold_chunks=dataset_row.has_gold_chunks,
@ -595,9 +595,9 @@ class EvaluationService:
await context.set_message(f"Error: {str(e)}")
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:
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"}
task_by_run_id = {}
if running_run_ids:
@ -635,12 +635,12 @@ class EvaluationService:
raise
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]:
if not re.match(r"^run_[a-f0-9]{8}$", run_id):
raise ValueError("Invalid run_id format")
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)
if task:
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):
raise ValueError("Invalid run_id format")
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")
await self.eval_repo.delete_run(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())
def compute_entity_id(db_id: str, normalized_name: str, label: str) -> str:
return hashstr(f"{db_id}:{normalized_name}:{label}", length=32)
def compute_entity_id(kb_id: str, normalized_name: str, label: str) -> str:
return hashstr(f"{kb_id}:{normalized_name}:{label}", length=32)
def compute_triple_id(
db_id: str,
kb_id: str,
source_normalized_name: str,
source_label: str,
relation_type: str,
@ -29,17 +29,17 @@ def compute_triple_id(
target_label: str,
) -> str:
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,
)
def graph_entity_collection_name(db_id: str) -> str:
return f"{db_id}_entity"
def graph_entity_collection_name(kb_id: str) -> str:
return f"{kb_id}_entity"
def graph_triple_collection_name(db_id: str) -> str:
return f"{db_id}_triple"
def graph_triple_collection_name(kb_id: str) -> str:
return f"{kb_id}_triple"
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"""
MERGE (c:Chunk:MilvusKB:`{db_label}` {{chunk_id: $chunk_id}})
SET c.file_id = $file_id,
c.db_id = $db_id,
c.kb_id = $kb_id,
c.chunk_index = $chunk_index,
c.content_preview = $content_preview,
c.start_char_pos = $start_char_pos,
@ -112,14 +112,14 @@ def cypher_merge_entity_mention(db_label: str) -> str:
return f"""
MATCH (c:Chunk:MilvusKB:`{db_label}` {{chunk_id: $chunk_id}})
MERGE (e:Entity:MilvusKB:`{db_label}` {{
db_id: $db_id,
kb_id: $kb_id,
normalized_name: $normalized_name,
label: $entity_label
}})
SET e.entity_id = $entity_id,
e.name = $name,
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 边。"""
return f"""
MATCH (source:Entity:MilvusKB:`{db_label}` {{
db_id: $db_id,
kb_id: $kb_id,
normalized_name: $source_name,
label: $source_label
}})
MATCH (target:Entity:MilvusKB:`{db_label}` {{
db_id: $db_id,
kb_id: $kb_id,
normalized_name: $target_name,
label: $target_label
}})
MERGE (source)-[r:RELATION {{
db_id: $db_id,
kb_id: $kb_id,
chunk_id: $chunk_id,
source_name: $source_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__(
self,
*,
db_id: str | None = None,
kb_id: str | None = None,
kb_repo: KnowledgeBaseRepository | None = None,
chunk_repo: KnowledgeChunkRepository | None = None,
graph_repo: KnowledgeGraphRepository | None = None,
graph_vector_store: MilvusGraphVectorStore | 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.chunk_repo = chunk_repo or KnowledgeChunkRepository()
self.graph_repo = graph_repo or KnowledgeGraphRepository()
@ -66,15 +66,15 @@ class MilvusGraphService:
def driver(self):
return self.connection.driver
async def get_status(self, db_id: str, *, tasker: Any = None) -> dict[str, Any]:
kb = await self._get_milvus_kb(db_id)
async def get_status(self, kb_id: str, *, tasker: Any = None) -> dict[str, Any]:
kb = await self._get_milvus_kb(kb_id)
params = dict(kb.additional_params or {})
config = params.get(GRAPH_CONFIG_KEY) or {}
total_chunks, pending_chunks, indexed_chunks, graph_counts = await asyncio.gather(
self.chunk_repo.count_by_db_id(db_id),
self.chunk_repo.count_graph_pending_by_db_id(db_id),
self.chunk_repo.count_graph_indexed_by_db_id(db_id),
self.graph_repo.count_by_db_id(db_id),
self.chunk_repo.count_by_kb_id(kb_id),
self.chunk_repo.count_graph_pending_by_kb_id(kb_id),
self.chunk_repo.count_graph_indexed_by_kb_id(kb_id),
self.graph_repo.count_by_kb_id(kb_id),
)
entity_count, relationship_count = graph_counts
@ -83,7 +83,7 @@ class MilvusGraphService:
if tasker is not None:
active_task = await tasker.find_task_by_payload(
task_type=GRAPH_TASK_TYPE,
payload_match={"db_id": db_id},
payload_match={"kb_id": kb_id},
statuses={"pending", "running"},
)
if active_task:
@ -92,7 +92,7 @@ class MilvusGraphService:
else:
failed_task = await tasker.find_task_by_payload(
task_type=GRAPH_TASK_TYPE,
payload_match={"db_id": db_id},
payload_match={"kb_id": kb_id},
statuses={"failed", "cancelled"},
)
if failed_task:
@ -100,7 +100,7 @@ class MilvusGraphService:
build_task_progress = 0
return {
"db_id": db_id,
"kb_id": kb_id,
"kb_type": kb.kb_type,
"configured": bool(config),
"locked": bool(config.get("locked")),
@ -116,12 +116,12 @@ class MilvusGraphService:
async def configure(
self,
db_id: str,
kb_id: str,
extractor_type: str,
extractor_options: dict[str, Any],
created_by: str,
) -> dict:
kb = await self._get_milvus_kb(db_id)
kb = await self._get_milvus_kb(kb_id)
additional_params = dict(kb.additional_params or {})
existing_config = additional_params.get(GRAPH_CONFIG_KEY) or {}
normalized_extractor_type = (extractor_type or "").lower()
@ -145,16 +145,16 @@ class MilvusGraphService:
config["updated_at"] = utc_isoformat()
config["updated_by"] = created_by
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
async def build_pending_chunks(self, db_id: str, *, batch_size: int, context=None) -> dict[str, Any]:
kb = await self._get_milvus_kb(db_id)
async def build_pending_chunks(self, kb_id: str, *, batch_size: int, context=None) -> dict[str, Any]:
kb = await self._get_milvus_kb(kb_id)
config = self._get_locked_config(kb.additional_params or {})
extractor_options = self._runtime_extractor_options(config)
extractor = GraphExtractorFactory.create(config["extractor_type"], extractor_options)
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
failed = 0
failed_chunk_ids: set[str] = set()
@ -163,7 +163,7 @@ class MilvusGraphService:
while True:
if context is not None:
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]
if not unprocessed:
break
@ -182,23 +182,23 @@ class MilvusGraphService:
except asyncio.QueueEmpty:
return
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:
entities, triples = await asyncio.to_thread(
self.write_chunk_graph,
db_id,
kb_id,
chunk,
extraction_result,
)
await self.graph_repo.upsert_chunk_graph(
db_id=db_id,
kb_id=kb_id,
file_id=chunk.file_id,
chunk_id=chunk.chunk_id,
entities=entities,
triples=triples,
)
await self.graph_vector_store.insert_missing_graph_records(
db_id=db_id,
kb_id=kb_id,
embedding_model_spec=kb.embedding_model_spec,
entities=entities,
triples=triples,
@ -229,8 +229,8 @@ class MilvusGraphService:
await asyncio.gather(*workers, return_exceptions=True)
raise
remaining = await self.chunk_repo.count_graph_pending_by_db_id(db_id)
return {"db_id": db_id, "success": processed, "failed": failed, "remaining": remaining}
remaining = await self.chunk_repo.count_graph_pending_by_kb_id(kb_id)
return {"kb_id": kb_id, "success": processed, "failed": failed, "remaining": remaining}
@staticmethod
def _get_worker_count(config: dict[str, Any]) -> int:
@ -248,7 +248,7 @@ class MilvusGraphService:
options.pop("prompt", None)
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
if chunk.extraction_result:
return normalize_extraction_result(chunk.extraction_result, extractor_type)
@ -256,7 +256,7 @@ class MilvusGraphService:
extraction_result = await extractor.extract(
chunk.content,
chunk_metadata={
"db_id": db_id,
"kb_id": kb_id,
"chunk_id": chunk.chunk_id,
"file_id": chunk.file_id,
"chunk_index": chunk.chunk_index,
@ -268,22 +268,22 @@ class MilvusGraphService:
def write_chunk_graph(
self,
db_id: str,
kb_id: str,
chunk,
normalized_result: dict[str, Any],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""将单个 chunk 的抽取结果写入 Neo4j。"""
label = safe_neo4j_label(db_id)
label = safe_neo4j_label(kb_id)
graph_payload = build_graph_payload(normalized_result)
relation_extractor_type = graph_payload["metadata"].get("extractor_type", "unknown")
entities = graph_payload["entities"]
relations = graph_payload["relations"]
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["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]
# 预构建 Cypher 模板（同一 chunk 内复用）
@ -297,7 +297,7 @@ class MilvusGraphService:
merge_chunk_cypher,
chunk_id=chunk.chunk_id,
file_id=chunk.file_id,
db_id=db_id,
kb_id=kb_id,
chunk_index=chunk.chunk_index,
content_preview=content_preview,
start_char_pos=chunk.start_char_pos,
@ -311,7 +311,7 @@ class MilvusGraphService:
merge_entity_cypher,
chunk_id=chunk.chunk_id,
file_id=chunk.file_id,
db_id=db_id,
kb_id=kb_id,
entity_id=entity_record["entity_id"],
normalized_name=normalize_entity_name(entity["text"]),
entity_label=entity.get("label") or "Entity",
@ -327,7 +327,7 @@ class MilvusGraphService:
target_record = entity_record_by_local_id[relation["target"]]
relation_type = relation.get("label") or "RELATED_TO"
triple_id = compute_triple_id(
db_id,
kb_id,
source_record["normalized_name"],
source_record["label"],
relation_type,
@ -336,7 +336,7 @@ class MilvusGraphService:
)
tx.run(
merge_relation_cypher,
db_id=db_id,
kb_id=kb_id,
chunk_id=chunk.chunk_id,
file_id=chunk.file_id,
source_name=normalize_entity_name(source["text"]),
@ -352,16 +352,16 @@ class MilvusGraphService:
neo4j_write(self.driver, query)
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 = []
for entity in entities:
label = entity.get("label") or "Entity"
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(
{
"entity_id": entity_id,
"db_id": db_id,
"kb_id": kb_id,
"normalized_name": normalized_name,
"label": label,
"name": entity["text"],
@ -373,7 +373,7 @@ class MilvusGraphService:
def _build_triple_records(
self,
db_id: str,
kb_id: str,
relations: list[dict[str, Any]],
entity_record_by_local_id: dict[str, dict[str, Any]],
graph_payload: dict[str, Any],
@ -386,7 +386,7 @@ class MilvusGraphService:
target_record = entity_record_by_local_id[relation["target"]]
relation_type = relation.get("label") or "RELATED_TO"
triple_id = compute_triple_id(
db_id,
kb_id,
source_record["normalized_name"],
source_record["label"],
relation_type,
@ -400,7 +400,7 @@ class MilvusGraphService:
records.append(
{
"triple_id": triple_id,
"db_id": db_id,
"kb_id": kb_id,
"source_entity_id": source_record["entity_id"],
"target_entity_id": target_record["entity_id"],
"relation_type": relation_type,
@ -411,15 +411,15 @@ class MilvusGraphService:
)
return records
async def reset(self, db_id: str, *, clear_extraction_result: bool, clear_config: bool) -> dict[str, Any]:
kb = await self._get_milvus_kb(db_id)
await asyncio.to_thread(self.delete_graph, db_id)
await self.graph_repo.delete_by_db_id(db_id)
reset_chunks = await self.chunk_repo.reset_graph_state_by_db_id(db_id, clear_extraction_result)
async def reset(self, kb_id: str, *, clear_extraction_result: bool, clear_config: bool) -> dict[str, Any]:
kb = await self._get_milvus_kb(kb_id)
await asyncio.to_thread(self.delete_graph, kb_id)
await self.graph_repo.delete_by_kb_id(kb_id)
reset_chunks = await self.chunk_repo.reset_graph_state_by_kb_id(kb_id, clear_extraction_result)
if clear_config:
additional_params = dict(kb.additional_params or {})
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 {
"message": "图谱构建状态已重置",
"status": "success",
@ -428,79 +428,79 @@ class MilvusGraphService:
"clear_config": clear_config,
}
def delete_graph(self, db_id: str) -> None:
label = safe_neo4j_label(db_id)
def delete_graph(self, kb_id: str) -> None:
label = safe_neo4j_label(kb_id)
def query(tx):
tx.run(f"MATCH (n:MilvusKB:`{label}`) DETACH DELETE n")
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)
await self.graph_vector_store.delete_graph_records(
db_id,
kb_id,
entity_ids=orphan_entity_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:
label = safe_neo4j_label(db_id)
def _delete_file_graph_from_neo4j(self, kb_id: str, file_id: str) -> None:
label = safe_neo4j_label(kb_id)
def query(tx):
tx.run(
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}`)
DELETE m
""",
db_id=db_id,
kb_id=kb_id,
file_id=file_id,
)
tx.run(
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}`)
DELETE r
""",
db_id=db_id,
kb_id=kb_id,
file_id=file_id,
)
tx.run(
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
""",
db_id=db_id,
kb_id=kb_id,
file_id=file_id,
)
tx.run(
f"""
MATCH (e:Entity:MilvusKB:`{label}` {{db_id: $db_id}})
MATCH (e:Entity:MilvusKB:`{label}` {{kb_id: $kb_id}})
WHERE NOT ()-[:MENTIONS]->(e)
DETACH DELETE e
""",
db_id=db_id,
kb_id=kb_id,
)
neo4j_write(self.driver, query)
async def query_nodes(
self,
db_id: str | None = None,
kb_id: str | None = None,
*,
keyword: str = "",
max_depth: int = 1,
max_nodes: int = 50,
exclude_chunk: bool = False,
) -> dict[str, Any]:
effective_db_id = db_id or self.db_id
if not effective_db_id:
effective_kb_id = kb_id or self.kb_id
if not effective_kb_id:
return {"nodes": [], "edges": []}
label = safe_neo4j_label(effective_db_id)
label = safe_neo4j_label(effective_kb_id)
limit = max_nodes
try:
with self.driver.session() as session:
@ -509,21 +509,21 @@ class MilvusGraphService:
keyword=keyword,
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:
logger.error(f"Milvus graph query failed: {e}")
return {"nodes": [], "edges": []}
async def query_seed_subgraph(
self,
db_id: str,
kb_id: str,
*,
entity_ids: list[str],
max_nodes: int,
) -> dict[str, Any]:
if not entity_ids:
return {"nodes": [], "edges": []}
label = safe_neo4j_label(db_id)
label = safe_neo4j_label(kb_id)
cypher = f"""
MATCH (seed:Entity:MilvusKB:`{label}`)
WHERE seed.entity_id IN $entity_ids
@ -546,14 +546,14 @@ class MilvusGraphService:
).single()
if not record:
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:
logger.error(f"Milvus seed subgraph query failed: {e}")
return {"nodes": [], "edges": []}
async def query_and_rank_chunks_by_ppr(
self,
db_id: str,
kb_id: str,
seed_weights: dict[str, float],
*,
max_nodes: int,
@ -563,7 +563,7 @@ class MilvusGraphService:
if not seed_weights:
return []
subgraph = await self.query_seed_subgraph(
db_id,
kb_id,
entity_ids=list(seed_weights.keys()),
max_nodes=max_nodes,
)
@ -622,32 +622,32 @@ class MilvusGraphService:
)
return ranked[:top_k]
async def get_labels(self, db_id: str | None = None) -> list[str]:
effective_db_id = db_id or self.db_id
if not effective_db_id:
async def get_labels(self, kb_id: str | None = None) -> list[str]:
effective_kb_id = kb_id or self.kb_id
if not effective_kb_id:
return []
label = safe_neo4j_label(effective_db_id)
label = safe_neo4j_label(effective_kb_id)
cypher = f"""
MATCH (n:MilvusKB:`{label}`)
UNWIND labels(n) AS 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
ORDER BY node_label
"""
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]
except Exception as e:
logger.error(f"Failed to get Milvus graph labels: {e}")
return []
async def get_stats(self, db_id: str | None = None) -> dict[str, Any]:
effective_db_id = db_id or self.db_id
if not effective_db_id:
async def get_stats(self, kb_id: str | None = None) -> dict[str, Any]:
effective_kb_id = kb_id or self.kb_id
if not effective_kb_id:
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"""
MATCH (n:MilvusKB:`{label}`)
@ -674,10 +674,10 @@ class MilvusGraphService:
logger.error(f"Failed to get Milvus graph stats: {e}")
return {"total_nodes": 0, "total_edges": 0, "entity_types": []}
async def _get_milvus_kb(self, db_id: str):
kb = await self.kb_repo.get_by_id(db_id)
async def _get_milvus_kb(self, kb_id: str):
kb = await self.kb_repo.get_by_kb_id(kb_id)
if kb is None:
raise ValueError(f"知识库 {db_id} 不存在")
raise ValueError(f"知识库 {kb_id} 不存在")
if (kb.kb_type or "").lower() != "milvus":
raise ValueError("仅 Milvus 知识库支持独立图谱构建")
return kb
@ -737,7 +737,7 @@ class MilvusGraphService:
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 = []
edges = []
node_ids = set()
@ -748,7 +748,7 @@ class MilvusGraphService:
raw_node = record.get(key)
if raw_node is None:
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:
continue
if exclude_chunk and node.get("type") == "Chunk":
@ -766,14 +766,14 @@ class MilvusGraphService:
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 = []
edges = []
node_ids = set()
edge_ids = set()
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:
continue
nodes.append(node)
@ -792,7 +792,7 @@ class MilvusGraphService:
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"):
node_id = raw_node.element_id
labels = list(raw_node.labels)
@ -804,8 +804,8 @@ class MilvusGraphService:
else:
return {}
effective_db_id = db_id or self.db_id
db_label = properties.get("db_id") or effective_db_id
effective_kb_id = kb_id or self.kb_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}]
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"

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(
self,
*,
db_id: str,
kb_id: str,
embedding_model_spec: str,
entities: 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":
raise ValueError(f"Unsupported embedding model: {embedding_model_spec}")
entity_collection = self._get_or_create_entity_collection(db_id, embedding_info)
triple_collection = self._get_or_create_triple_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(kb_id, embedding_info)
entity_ids = [entity["entity_id"] for entity in entities]
triple_ids = [triple["triple_id"] for triple in triples]
@ -83,24 +83,24 @@ class MilvusGraphVectorStore:
if missing_triples:
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 = []
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:
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:
await asyncio.gather(*tasks)
async def search_entities(
self,
*,
db_id: str,
kb_id: str,
query_text: str,
embedding_model_spec: str,
top_k: int,
) -> 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):
return []
return await self._search_graph_collection(
@ -114,12 +114,12 @@ class MilvusGraphVectorStore:
async def search_triples(
self,
*,
db_id: str,
kb_id: str,
query_text: str,
embedding_model_spec: str,
top_k: int,
) -> 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):
return []
return await self._search_graph_collection(
@ -130,8 +130,8 @@ class MilvusGraphVectorStore:
output_fields=["id", "content", "source_id", "target_id"],
)
def drop_graph_collections(self, db_id: str) -> None:
for collection_name in [graph_entity_collection_name(db_id), graph_triple_collection_name(db_id)]:
def drop_graph_collections(self, kb_id: str) -> None:
for collection_name in [graph_entity_collection_name(kb_id), graph_triple_collection_name(kb_id)]:
try:
if utility.has_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)
return records
def _get_or_create_entity_collection(self, db_id: str, embedding_info: Any) -> Collection:
collection_name = graph_entity_collection_name(db_id)
def _get_or_create_entity_collection(self, kb_id: str, embedding_info: Any) -> Collection:
collection_name = graph_entity_collection_name(kb_id)
fields = [
FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True),
FieldSchema(
@ -211,8 +211,8 @@ class MilvusGraphVectorStore:
]
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:
collection_name = graph_triple_collection_name(db_id)
def _get_or_create_triple_collection(self, kb_id: str, embedding_info: Any) -> Collection:
collection_name = graph_triple_collection_name(kb_id)
fields = [
FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True),
FieldSchema(

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

@ -63,18 +63,18 @@ class DifyKB(ReadOnlyConnectors):
raise ValueError("Dify api_url 必须以 /v1 结尾")
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
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()
token = str(metadata.get("dify_token") or "").strip()
dataset_id = str(metadata.get("dify_dataset_id") or "").strip()
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 []
query_params = self._get_query_params(db_id)
query_params = self._get_query_params(kb_id)
merged = {**query_params, **kwargs}
search_mode = str(merged.get("search_mode", "vector")).lower()
@ -109,7 +109,7 @@ class DifyKB(ReadOnlyConnectors):
try:
response_json = await self._request_dify(client_payload=payload, request_url=request_url, headers=headers)
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 请求重试一次
try:
response_json = await self._request_dify(
@ -117,10 +117,10 @@ class DifyKB(ReadOnlyConnectors):
request_url=request_url,
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
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 []
@ -172,8 +172,8 @@ class DifyKB(ReadOnlyConnectors):
raise e
return response.json()
def get_query_params_config(self, db_id: str, **kwargs) -> dict:
del db_id, kwargs
def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
del kb_id, kwargs
options = [
{
"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)}"
# 存储集合映射 {db_id: Collection}
# 存储集合映射 {kb_id: Collection}
self.collections: dict[str, Any] = {}
# 元数据锁
@ -297,22 +297,22 @@ class MilvusKB(KnowledgeBase):
logger.error(f"Failed to connect to Milvus: {e}")
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 集合"""
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)):
raise ValueError(f"Database {db_id} not found")
if not (metadata := self.databases_meta.get(kb_id)):
raise ValueError(f"Database {kb_id} not found")
embedding_model_spec = metadata.get("embedding_model_spec")
if not embedding_model_spec:
raise ValueError(f"Embedding model spec not found for database {db_id}")
raise ValueError(f"Embedding model spec not found for database {kb_id}")
embedding_info = model_cache.get_model_info(embedding_model_spec)
if not embedding_info or embedding_info.model_type != "embedding":
raise ValueError(f"Unsupported embedding model: {embedding_model_spec}")
collection_name = db_id
collection_name = kb_id
try:
# 检查集合是否存在
@ -329,18 +329,18 @@ class MilvusKB(KnowledgeBase):
f"expected='{expected_model}', found_in_description='{description}'"
)
utility.drop_collection(collection_name, using=self.connection_alias)
return self._create_new_collection(collection_name, embedding_info, db_id)
return self._create_new_collection(collection_name, embedding_info, kb_id)
if not self._collection_supports_bm25(collection):
logger.warning(f"Collection {collection_name} schema does not support BM25, recreating")
utility.drop_collection(collection_name, using=self.connection_alias)
return self._create_new_collection(collection_name, embedding_info, db_id)
return self._create_new_collection(collection_name, embedding_info, kb_id)
logger.info(f"Retrieved existing collection: {collection_name}")
return collection
else:
logger.info(f"Collection {collection_name} not found, creating new one")
return self._create_new_collection(collection_name, embedding_info, db_id)
return self._create_new_collection(collection_name, embedding_info, kb_id)
except (connections.MilvusException, RuntimeError) as e:
logger.error(f"Error checking collection {collection_name}: {e}")
@ -350,7 +350,7 @@ class MilvusKB(KnowledgeBase):
logger.debug(f"Traceback: {traceback.format_exc()}")
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 集合"""
embedding_dim = embedding_info.dimension or 1024
model_name = embedding_info.model_id
@ -380,7 +380,7 @@ class MilvusKB(KnowledgeBase):
schema = CollectionSchema(
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],
)
@ -439,24 +439,24 @@ class MilvusKB(KnowledgeBase):
method = model.batch_encode if sync else model.abatch_encode
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 集合"""
if db_id in self.collections:
return self.collections[db_id]
if kb_id in self.collections:
return self.collections[kb_id]
if db_id not in self.databases_meta:
if kb_id not in self.databases_meta:
return None
try:
# 创建集合
collection = await self._create_kb_instance(db_id, {})
collection = await self._create_kb_instance(kb_id, {})
await self._initialize_kb_instance(collection)
self.collections[db_id] = collection
self.collections[kb_id] = collection
return collection
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()}")
return None
@ -464,12 +464,12 @@ class MilvusKB(KnowledgeBase):
"""将文本分割成块"""
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 [
{
"chunk_id": chunk["chunk_id"],
"file_id": chunk["file_id"],
"db_id": db_id,
"kb_id": kb_id,
"chunk_index": chunk["chunk_index"],
"content": chunk["content"],
"start_char_pos": chunk.get("start_char_pos"),
@ -486,7 +486,7 @@ class MilvusKB(KnowledgeBase):
async def _insert_chunks_to_stores(
self,
db_id: str,
kb_id: str,
file_id: str,
collection: Collection,
chunks: list[dict],
@ -508,7 +508,7 @@ class MilvusKB(KnowledgeBase):
def _insert_milvus_records():
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)
results = await asyncio.gather(pg_task, milvus_task, return_exceptions=True)
errors = [result for result in results if isinstance(result, Exception)]
@ -545,7 +545,7 @@ class MilvusKB(KnowledgeBase):
return "未知来源"
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:
return None
@ -553,7 +553,7 @@ class MilvusKB(KnowledgeBase):
matched_file_ids = [
file_id
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:
return 'file_id == "__no_matching_file__"'
@ -564,13 +564,13 @@ class MilvusKB(KnowledgeBase):
return f'file_id in ["{joined_ids}"]'
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:
"""
Index parsed file (Status: INDEXING -> INDEXED/ERROR_INDEXING)
Args:
db_id: Database ID
kb_id: Database ID
file_id: File ID
operator_id: ID of the user performing the operation
params: Override processing params to apply during indexing (merged on top of stored params)
@ -578,15 +578,15 @@ class MilvusKB(KnowledgeBase):
Returns:
Updated file metadata
"""
if db_id not in self.databases_meta:
raise ValueError(f"Database {db_id} not found")
if kb_id not in self.databases_meta:
raise ValueError(f"Database {kb_id} not found")
# Get/Create collection
collection = await self._get_milvus_collection(db_id)
collection = await self._get_milvus_collection(kb_id)
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)
# Get file meta
@ -627,7 +627,7 @@ class MilvusKB(KnowledgeBase):
# 将传入的 params 作为 request_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"),
request_params=params,
)
@ -656,8 +656,8 @@ class MilvusKB(KnowledgeBase):
embeddings = await embedding_function(texts)
# Clean up existing chunks if any (for re-indexing)
await self.delete_file_chunks_only(db_id, file_id)
await self._insert_chunks_to_stores(db_id, file_id, collection, chunks, embeddings)
await self.delete_file_chunks_only(kb_id, file_id)
await self._insert_chunks_to_stores(kb_id, file_id, collection, chunks, embeddings)
logger.info(f"Indexed file {file_id} into Milvus")
@ -685,16 +685,16 @@ class MilvusKB(KnowledgeBase):
# Remove from processing queue
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重新解析文件并更新向量库"""
if db_id not in self.databases_meta:
raise ValueError(f"Database {db_id} not found")
if kb_id not in self.databases_meta:
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:
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)
# 处理默认参数
@ -724,7 +724,7 @@ class MilvusKB(KnowledgeBase):
# 更新状态为处理中
async with self._metadata_lock:
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"),
request_params=params,
)
@ -733,7 +733,7 @@ class MilvusKB(KnowledgeBase):
await self._persist_file(file_id)
# 重新解析文件为 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)
# 重新生成 chunks
@ -741,12 +741,12 @@ class MilvusKB(KnowledgeBase):
logger.info(f"Split {filename} into {len(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:
texts = [chunk["content"] for chunk in chunks]
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.")
@ -804,13 +804,13 @@ class MilvusKB(KnowledgeBase):
chunk["distance"] = hit.distance
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:
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（持久化参数）
# 这样允许用户在单次查询中临时覆盖持久化配置
merged_kwargs = {**query_params, **kwargs}
@ -835,14 +835,14 @@ class MilvusKB(KnowledgeBase):
else:
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:
logger.debug(f"Using filter expression: {file_expr}")
output_fields = ["content", "chunk_id", "file_id", "chunk_index"]
retrieved_chunks: list[dict] = []
if search_mode == "vector":
embedding_model_spec = self.databases_meta[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)
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")
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)
query_embedding = embedding_function([query_text])
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")
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:
graph_weight = float(merged_kwargs.get("graph_weight", 1.0))
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
)
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:
await reranker.aclose()
@ -991,7 +991,7 @@ class MilvusKB(KnowledgeBase):
async def _retrieve_graph_chunks(
self,
query_text: str,
db_id: str,
kb_id: str,
base_chunks: list[dict],
query_params: dict[str, Any],
) -> list[dict]:
@ -999,7 +999,7 @@ class MilvusKB(KnowledgeBase):
from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService
from yuxi.knowledge.graphs.milvus_graph_vector_store import MilvusGraphVectorStore
embedding_model_spec = self.databases_meta[db_id].get("embedding_model_spec")
embedding_model_spec = self.databases_meta[kb_id].get("embedding_model_spec")
if not embedding_model_spec:
return []
@ -1011,25 +1011,25 @@ class MilvusKB(KnowledgeBase):
vector_store = MilvusGraphVectorStore()
entity_hits, triple_hits = await asyncio.gather(
vector_store.search_entities(
db_id=db_id,
kb_id=kb_id,
query_text=query_text,
embedding_model_spec=embedding_model_spec,
top_k=entity_top_k,
),
vector_store.search_triples(
db_id=db_id,
kb_id=kb_id,
query_text=query_text,
embedding_model_spec=embedding_model_spec,
top_k=triple_top_k,
),
)
seed_weights = await self._build_graph_seed_weights(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:
return []
graph_service = MilvusGraphService()
graph_scores = await graph_service.query_and_rank_chunks_by_ppr(
db_id,
kb_id,
seed_weights,
max_nodes=graph_max_nodes,
top_k=graph_top_k,
@ -1045,12 +1045,12 @@ class MilvusKB(KnowledgeBase):
for chunk in chunks
]
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 []
async def _build_graph_seed_weights(
self,
db_id: str,
kb_id: str,
base_chunks: list[dict],
entity_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)
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数据，保留元数据（用于更新操作）"""
chunk_repo = KnowledgeChunkRepository()
if await chunk_repo.count_graph_indexed_by_file_id(file_id):
from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService
try:
await MilvusGraphService().delete_file_graph(db_id, file_id)
await MilvusGraphService().delete_file_graph(kb_id, file_id)
except Exception as e:
logger.error(f"Failed to delete graph data for file {file_id}: {e}")
await chunk_repo.delete_by_file_id(file_id)
collection = await self._get_milvus_collection(db_id)
collection = await self._get_milvus_collection(kb_id)
if collection:
# 先查询文件是否存在，避免不必要的删除操作
@ -1150,10 +1150,10 @@ class MilvusKB(KnowledgeBase):
logger.error(f"Error checking file existence in Milvus: {e}")
# 注意：这里不删除 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 数据
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:
@ -1163,14 +1163,14 @@ class MilvusKB(KnowledgeBase):
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:
raise Exception(f"File not found: {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）"""
if file_id not in self.files_meta:
raise Exception(f"File not found: {file_id}")
@ -1211,37 +1211,37 @@ class MilvusKB(KnowledgeBase):
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:
raise Exception(f"File not found: {file_id}")
# 合并基本信息和内容信息
basic_info = await self.get_file_basic_info(db_id, file_id)
content_info = await self.get_file_content(db_id, file_id)
basic_info = await self.get_file_basic_info(kb_id, file_id)
content_info = await self.get_file_content(kb_id, file_id)
return {**basic_info, **content_info}
def delete_database(self, db_id: str) -> dict:
def delete_database(self, kb_id: str) -> dict:
"""删除数据库，同时清除Milvus中的集合"""
# Drop Milvus collection
try:
if utility.has_collection(db_id, using=self.connection_alias):
utility.drop_collection(db_id, using=self.connection_alias)
logger.info(f"Dropped Milvus collection for {db_id}")
if utility.has_collection(kb_id, using=self.connection_alias):
utility.drop_collection(kb_id, using=self.connection_alias)
logger.info(f"Dropped Milvus collection for {kb_id}")
else:
logger.info(f"Milvus collection {db_id} does not exist, skipping")
logger.info(f"Milvus collection {kb_id} does not exist, skipping")
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
MilvusGraphVectorStore().drop_graph_collections(db_id)
MilvusGraphVectorStore().drop_graph_collections(kb_id)
# 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 知识库的查询参数配置"""
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,
}
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
try:
token, data_source_id, notion_version = self._get_connection_config(db_id)
query_params = self._get_query_params(db_id)
token, data_source_id, notion_version = self._get_connection_config(kb_id)
query_params = self._get_query_params(kb_id)
merged = {**query_params, **kwargs}
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)
@ -270,7 +270,7 @@ class NotionKB(ReadOnlyConnectors):
return sorted(results, key=lambda item: item.get("score", 0.0), reverse=True)[:final_top_k]
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 []
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:
content = await self._read_page_markdown(db_id, file_id)
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(kb_id, file_id)
return self._build_open_file_window(content, offset=offset, limit=limit)
async def find_file_content(
self,
db_id: str,
kb_id: str,
file_id: str,
patterns: list[str],
*,
@ -387,7 +387,7 @@ class NotionKB(ReadOnlyConnectors):
max_windows: int = 5,
window_size: int = 80,
) -> 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(
content,
patterns=patterns,
@ -397,9 +397,9 @@ class NotionKB(ReadOnlyConnectors):
window_size=window_size,
)
async def _read_page_markdown(self, db_id: str, page_id: str) -> str:
token, data_source_id, notion_version = self._get_connection_config(db_id)
cache_key = (db_id, page_id, data_source_id, notion_version)
async def _read_page_markdown(self, kb_id: str, page_id: str) -> str:
token, data_source_id, notion_version = self._get_connection_config(kb_id)
cache_key = (kb_id, page_id, data_source_id, notion_version)
cached = self._get_cached_page_markdown(cache_key)
if cached is not None:
return cached
@ -430,15 +430,15 @@ class NotionKB(ReadOnlyConnectors):
self._page_markdown_cache.pop(oldest_key, None)
self._page_markdown_cache[cache_key] = (time.monotonic(), content)
def _get_connection_config(self, db_id: str) -> tuple[str, str, str]:
metadata = self.databases_meta.get(db_id, {}).get("metadata", {}) or {}
def _get_connection_config(self, kb_id: str) -> tuple[str, str, str]:
metadata = self.databases_meta.get(kb_id, {}).get("metadata", {}) or {}
token = str(
metadata.get("notion_token") or os.getenv("NOTION_TOKEN") or os.getenv("NOTION_API_KEY") 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
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
async def _page_to_markdown(
@ -679,8 +679,8 @@ class NotionKB(ReadOnlyConnectors):
def _preview(content: str, line_count: int) -> str:
return "\n".join(content.splitlines()[:line_count])
def get_query_params_config(self, db_id: str, **kwargs) -> dict:
del db_id, kwargs
def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
del kb_id, kwargs
return {
"type": "notion",
"options": [

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

@ -17,8 +17,8 @@ class ReadOnlyConnectors(KnowledgeBase):
def _readonly_error() -> ValueError:
return ValueError("只读检索连接器不支持该操作")
async def _create_kb_instance(self, db_id: str, config: dict) -> Any:
del db_id, config
async def _create_kb_instance(self, kb_id: str, config: dict) -> Any:
del kb_id, config
return None
async def _initialize_kb_instance(self, instance: Any) -> None:
@ -26,53 +26,53 @@ class ReadOnlyConnectors(KnowledgeBase):
return None
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:
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()
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()
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()
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()
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()
async def index_file(
self,
db_id: str,
kb_id: str,
file_id: str,
operator_id: str | None = None,
params: dict | None = None,
) -> dict:
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()
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()
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()
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()
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
raise self._readonly_error()
async def find_file_content(
self,
db_id: str,
kb_id: str,
file_id: str,
patterns: list[str],
*,
@ -81,26 +81,26 @@ class ReadOnlyConnectors(KnowledgeBase):
max_windows: int = 5,
window_size: int = 80,
) -> 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()
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()
async def list_file_tree(
self,
db_id: str,
kb_id: str,
parent_id: str | None = None,
recursive: bool = False,
files_only: bool = False,
) -> dict:
del db_id, parent_id, recursive, files_only
del kb_id, parent_id, recursive, files_only
raise ValueError("只读检索连接器不支持文件树预览")
async def read_file_preview(self, db_id: str, file_id: str, variant: str = "parsed") -> dict:
del db_id, file_id, variant
async def read_file_preview(self, kb_id: str, file_id: str, variant: str = "parsed") -> dict:
del kb_id, file_id, variant
raise ValueError("只读检索连接器不支持文件预览")
async def get_file_download(self, db_id: str, file_id: str, variant: str = "original") -> dict:
del db_id, file_id, variant
async def get_file_download(self, kb_id: str, file_id: str, variant: str = "original") -> dict:
del kb_id, file_id, variant
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")
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)
return await kb_instance.move_file(db_id, file_id, new_parent_id)
kb_instance = await self._get_kb_for_database(kb_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获取对应的知识库实例
Args:
db_id: 数据库ID
kb_id: 数据库ID
Returns:
知识库实例
@ -127,10 +127,10 @@ class KnowledgeBaseManager:
from yuxi.repositories.knowledge_base_repository import 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:
raise KBNotFoundError(f"Database {db_id} not found")
raise KBNotFoundError(f"Database {kb_id} not found")
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:
db_id: 数据库ID
kb_id: 数据库ID
Returns:
知识库实例
"""
return await self._get_kb_for_database(db_id)
return await self._get_kb_for_database(kb_id)
def _normalize_share_config(
self,
@ -206,20 +206,20 @@ class KnowledgeBaseManager:
for row in rows:
kb_type = row.kb_type or "milvus"
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
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:
try:
await kb_instance._load_metadata()
metadata_reloaded_types.add(kb_type)
except Exception as 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:
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
# 补充 share_config 和 additional_params
@ -258,12 +258,12 @@ class KnowledgeBaseManager:
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:
user: 用户信息字典
db_id: 数据库ID
kb_id: 数据库ID
Returns:
bool: 是否有权限
@ -275,7 +275,7 @@ class KnowledgeBaseManager:
from yuxi.repositories.knowledge_base_repository import 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:
return False
@ -344,10 +344,10 @@ class KnowledgeBaseManager:
return True
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."""
kb_instance = await self._get_kb_for_database(db_id)
return await kb_instance.create_folder(db_id, folder_name, parent_id)
kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.create_folder(kb_id, folder_name, parent_id)
async def create_database(
self,
@ -401,16 +401,16 @@ class KnowledgeBaseManager:
llm_model_spec=llm_model_spec,
**kwargs,
)
db_id = db_info["db_id"]
kb_id = db_info["kb_id"]
from yuxi.repositories.knowledge_base_repository import 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:
await kb_repo.create(
{
"db_id": db_id,
"kb_id": kb_id,
"name": database_name,
"description": description,
"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
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
try:
kb_instance = await self._get_kb_for_database(db_id)
result = await kb_instance.delete_database(db_id)
kb_instance = await self._get_kb_for_database(kb_id)
result = await kb_instance.delete_database(kb_id)
# 删除数据库记录
kb_repo = KnowledgeBaseRepository()
await kb_repo.delete(db_id)
await kb_repo.delete(kb_id)
return result
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": "删除成功"}
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:
"""Add file record to metadata"""
kb_instance = await self._get_kb_for_database(db_id)
return await kb_instance.add_file_record(db_id, item, params, operator_id)
kb_instance = await self._get_kb_for_database(kb_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"""
kb_instance = await self._get_kb_for_database(db_id)
return await kb_instance.parse_file(db_id, file_id, operator_id)
kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.parse_file(kb_id, file_id, operator_id)
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:
"""Index parsed file"""
kb_instance = await self._get_kb_for_database(db_id)
return await kb_instance.index_file(db_id, file_id, operator_id, params=params)
kb_instance = await self._get_kb_for_database(kb_id)
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"""
kb_instance = await self._get_kb_for_database(db_id)
await kb_instance.update_file_params(db_id, file_id, params, operator_id)
kb_instance = await self._get_kb_for_database(kb_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)
return await kb_instance.aquery(query_text, db_id, **kwargs)
kb_instance = await self._get_kb_for_database(kb_id)
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)
return await kb_instance.export_data(db_id, format=format, **kwargs)
kb_instance = await self._get_kb_for_database(kb_id)
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
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:
return None
try:
kb_instance = await self._get_kb_for_database(db_id)
db_info = kb_instance.get_database_info(db_id)
kb_instance = await self._get_kb_for_database(kb_id)
db_info = kb_instance.get_database_info(kb_id)
except KBNotFoundError:
db_info = {
"db_id": db_id,
"kb_id": kb_id,
"name": kb.name,
"description": kb.description,
"kb_type": kb.kb_type,
@ -509,39 +509,39 @@ class KnowledgeBaseManager:
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)
await kb_instance.delete_folder(db_id, folder_id)
kb_instance = await self._get_kb_for_database(kb_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)
await kb_instance.delete_file(db_id, file_id)
kb_instance = await self._get_kb_for_database(kb_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)
return await kb_instance.update_content(db_id, file_ids, params or {})
kb_instance = await self._get_kb_for_database(kb_id)
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)
return await kb_instance.get_file_basic_info(db_id, file_id)
kb_instance = await self._get_kb_for_database(kb_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）"""
kb_instance = await self._get_kb_for_database(db_id)
return await kb_instance.get_file_content(db_id, file_id)
kb_instance = await self._get_kb_for_database(kb_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 内容"""
kb_instance = await self._get_kb_for_database(db_id)
return await kb_instance.open_file_content(db_id, file_id, offset, limit)
kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.open_file_content(kb_id, file_id, offset, limit)
async def find_file_content(
self,
db_id: str,
kb_id: str,
file_id: str,
patterns: list[str],
*,
@ -550,9 +550,9 @@ class KnowledgeBaseManager:
max_windows: int = 5,
window_size: int = 80,
) -> 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(
db_id,
kb_id,
file_id,
patterns,
use_regex=use_regex,
@ -561,40 +561,40 @@ class KnowledgeBaseManager:
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)
return await kb_instance.get_file_info(db_id, file_id)
kb_instance = await self._get_kb_for_database(kb_id)
return await kb_instance.get_file_info(kb_id, file_id)
async def list_file_tree(
self,
db_id: str,
kb_id: str,
parent_id: str | None = None,
recursive: bool = False,
files_only: bool = False,
) -> dict:
kb_instance = await self._get_kb_for_database(db_id)
return await kb_instance.list_file_tree(db_id, parent_id, recursive, files_only)
kb_instance = await self._get_kb_for_database(kb_id)
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:
kb_instance = await self._get_kb_for_database(db_id)
return await kb_instance.read_file_preview(db_id, file_id, variant)
async def read_file_preview(self, kb_id: str, file_id: str, variant: str = "parsed") -> dict:
kb_instance = await self._get_kb_for_database(kb_id)
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:
kb_instance = await self._get_kb_for_database(db_id)
return await kb_instance.get_file_download(db_id, file_id, variant)
async def get_file_download(self, kb_id: str, file_id: str, variant: str = "original") -> dict:
kb_instance = await self._get_kb_for_database(kb_id)
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
try:
kb_instance = await self._get_kb_for_database(db_id)
kb_instance = await self._get_kb_for_database(kb_id)
except KBNotFoundError:
return False
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
if file_info.get("status") == "failed":
continue
@ -603,13 +603,13 @@ class KnowledgeBaseManager:
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:
db_id: 数据库ID
kb_id: 数据库ID
filename: 要检测的文件名（原始文件名）
Returns:
@ -619,16 +619,16 @@ class KnowledgeBaseManager:
- created_at: 上传时间
- file_id: 文件ID（用于下载）
"""
if not db_id or not filename:
if not kb_id or not filename:
return []
try:
kb_instance = await self._get_kb_for_database(db_id)
kb_instance = await self._get_kb_for_database(kb_id)
except KBNotFoundError:
return []
same_name_files = []
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
if file_info.get("status") == "failed":
continue
@ -651,18 +651,18 @@ class KnowledgeBaseManager:
same_name_files.sort(key=lambda x: x.get("created_at", ""), reverse=True)
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
try:
kb_instance = await self._get_kb_for_database(db_id)
kb_instance = await self._get_kb_for_database(kb_id)
except KBNotFoundError:
return False
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
if file_info.get("status") == "failed":
continue
@ -673,7 +673,7 @@ class KnowledgeBaseManager:
async def update_database(
self,
db_id: str,
kb_id: str,
name: str,
description: str,
llm_model_spec: str | None = None,
@ -687,12 +687,12 @@ class KnowledgeBaseManager:
from yuxi.repositories.knowledge_base_repository import 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:
raise ValueError(f"数据库 {db_id} 不存在")
raise ValueError(f"数据库 {kb_id} 不存在")
kb_instance = await self._get_kb_for_database(db_id)
kb_instance.update_database(db_id, name, description, llm_model_spec, update_llm_model_spec)
kb_instance = await self._get_kb_for_database(kb_id)
kb_instance.update_database(kb_id, name, description, llm_model_spec, update_llm_model_spec)
update_data: dict = {
"name": name,
@ -711,8 +711,8 @@ class KnowledgeBaseManager:
deep_merge(current_additional_params, additional_params)
)
update_data["additional_params"] = merged_additional_params
if db_id in kb_instance.databases_meta:
kb_instance.databases_meta[db_id]["metadata"] = merged_additional_params
if kb_id in kb_instance.databases_meta:
kb_instance.databases_meta[kb_id]["metadata"] = merged_additional_params
if share_config is not None:
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]:
"""获取所有检索器"""
@ -830,7 +830,7 @@ class KnowledgeBaseManager:
# 从数据库获取所有已知的数据库ID
kb_repo = KnowledgeBaseRepository()
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 中的集合
# missing_collections = actual_collection_names - metadata_collection_names
@ -842,7 +842,7 @@ class KnowledgeBaseManager:
# 检查集合是否属于已知数据库
is_known = False
if collection_name in all_known_db_ids:
if collection_name in all_known_kb_ids:
is_known = True
# 如果是已知集合，跳过
@ -873,36 +873,36 @@ class KnowledgeBaseManager:
metadata_collection_names = set(milvus_kb.databases_meta.keys())
# 检查文件级别的不一致（针对已知的数据库）
for db_id in metadata_collection_names:
for kb_id in metadata_collection_names:
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
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
# 获取 metadata 中记录的文件数量
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 中没有文件记录，可能存在文件缺失
if actual_count > 0 and metadata_files_count == 0:
inconsistencies["missing_files"].append(
{
"database_id": db_id,
"kb_id": kb_id,
"vector_count": actual_count,
"metadata_files_count": metadata_files_count,
"detected_at": utc_isoformat(),
}
)
logger.warning(
f"发现数据库 {db_id} 在 Milvus 中有 {actual_count} 条向量数据，"
f"发现数据库 {kb_id} 在 Milvus 中有 {actual_count} 条向量数据，"
"但 metadata 中没有文件记录"
)
except Exception as e:
logger.debug(f"检查数据库 {db_id} 的文件一致性时出错: {e}")
logger.debug(f"检查数据库 {kb_id} 的文件一致性时出错: {e}")
except Exception as e:
logger.error(f"检测 Milvus 数据不一致时出错: {e}")
@ -933,7 +933,7 @@ class KnowledgeBaseManager:
logger.warning(f" 缺失文件记录数量: {len(milvus_files_missing)}")
for file_info in milvus_files_missing:
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']}"
)

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()
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。
Args:
item: MinIO URL 或 URL
content_type: 内容类型 ("file" 或 "url")
db_id: 数据库ID
kb_id: 数据库ID
params: 处理参数，可选
"""
# 检查是否有预处理信息 (针对 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)}"
metadata = {
"database_id": db_id,
"kb_id": kb_id,
"filename": filename_display,
"path": item_path,
"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}")
metadata = {
"database_id": db_id,
"kb_id": kb_id,
"filename": filename_display, # 使用显示用的文件名
"path": item_path,
"file_type": file_type,