2025-09-01 22:37:03 +08:00
|
|
|
|
import asyncio
|
|
|
|
|
|
import os
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.knowledge.base import KBNotFoundError, KnowledgeBase
|
2026-05-18 09:40:34 +08:00
|
|
|
|
from yuxi.knowledge.chunking.ragflow_like.presets import deep_merge
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.knowledge.factory import KnowledgeBaseFactory
|
|
|
|
|
|
from yuxi.storage.postgres.models_business import User
|
|
|
|
|
|
from yuxi.utils import logger
|
|
|
|
|
|
from yuxi.utils.datetime_utils import utc_isoformat
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-05-18 21:32:10 +08:00
|
|
|
|
DEFAULT_SHARE_CONFIG = {"access_level": "global", "department_ids": [], "user_uids": []}
|
|
|
|
|
|
ACCESS_LEVELS = {"global", "department", "user"}
|
|
|
|
|
|
|
|
|
|
|
|
|
2025-07-21 18:18:47 +08:00
|
|
|
|
class KnowledgeBaseManager:
|
|
|
|
|
|
"""
|
|
|
|
|
|
知识库管理器
|
|
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
统一管理多种类型的知识库实例,直接通过 Repository 访问数据库,不维护冗余缓存。
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
|
def __init__(self, work_dir: str):
|
|
|
|
|
|
"""
|
|
|
|
|
|
初始化知识库管理器
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
work_dir: 工作目录
|
|
|
|
|
|
"""
|
|
|
|
|
|
self.work_dir = work_dir
|
|
|
|
|
|
os.makedirs(work_dir, exist_ok=True)
|
|
|
|
|
|
|
|
|
|
|
|
# 知识库实例缓存 {kb_type: kb_instance}
|
2025-07-26 03:36:54 +08:00
|
|
|
|
self.kb_instances: dict[str, KnowledgeBase] = {}
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2025-07-27 01:02:14 +08:00
|
|
|
|
# 元数据锁
|
|
|
|
|
|
self._metadata_lock = asyncio.Lock()
|
|
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
async def initialize(self):
|
|
|
|
|
|
"""异步初始化"""
|
2025-07-21 18:18:47 +08:00
|
|
|
|
# 初始化已存在的知识库实例
|
|
|
|
|
|
self._initialize_existing_kbs()
|
2026-01-21 15:15:52 +08:00
|
|
|
|
logger.info("KnowledgeBaseManager initialized")
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
def _initialize_existing_kbs(self):
|
|
|
|
|
|
"""初始化已存在的知识库实例"""
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
2026-01-20 02:14:51 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
async def _async_init():
|
|
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
|
|
|
|
|
rows = await kb_repo.get_all()
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
kb_types_in_use = set()
|
|
|
|
|
|
for row in rows:
|
2026-05-17 13:06:25 +08:00
|
|
|
|
kb_type = row.kb_type or "milvus"
|
|
|
|
|
|
if KnowledgeBaseFactory.is_type_supported(kb_type):
|
|
|
|
|
|
kb_types_in_use.add(kb_type)
|
|
|
|
|
|
else:
|
|
|
|
|
|
logger.warning(f"Skip unsupported knowledge base type during initialization: {kb_type}")
|
2025-10-23 22:57:01 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
logger.info(f"[InitializeKB] 发现 {len(kb_types_in_use)} 种知识库类型: {kb_types_in_use}")
|
|
|
|
|
|
|
|
|
|
|
|
# 为每种使用中的知识库类型创建实例并加载元数据
|
|
|
|
|
|
for kb_type in kb_types_in_use:
|
2026-05-24 00:13:05 +08:00
|
|
|
|
if not KnowledgeBaseFactory.is_type_supported(kb_type):
|
|
|
|
|
|
logger.warning(f"[InitializeKB] Skip initialization for unsupported knowledge base type: {kb_type}")
|
|
|
|
|
|
continue
|
2025-10-13 15:08:54 +08:00
|
|
|
|
try:
|
2026-01-21 15:15:52 +08:00
|
|
|
|
kb_instance = self._get_or_create_kb_instance(kb_type)
|
|
|
|
|
|
# 让 KB 实例自行加载元数据
|
|
|
|
|
|
await kb_instance._load_metadata()
|
|
|
|
|
|
logger.info(f"[InitializeKB] {kb_type} 实例已初始化")
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
|
logger.error(f"Failed to initialize {kb_type} knowledge base: {e}")
|
|
|
|
|
|
import traceback
|
2025-10-13 15:08:54 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
logger.error(traceback.format_exc())
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
# 在事件循环中运行异步初始化
|
|
|
|
|
|
try:
|
|
|
|
|
|
loop = asyncio.get_running_loop()
|
|
|
|
|
|
loop.create_task(_async_init())
|
|
|
|
|
|
except RuntimeError:
|
|
|
|
|
|
asyncio.run(_async_init())
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
|
|
|
|
|
def _get_or_create_kb_instance(self, kb_type: str) -> KnowledgeBase:
|
|
|
|
|
|
"""
|
|
|
|
|
|
获取或创建知识库实例
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
kb_type: 知识库类型
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
知识库实例
|
|
|
|
|
|
"""
|
|
|
|
|
|
if kb_type in self.kb_instances:
|
|
|
|
|
|
return self.kb_instances[kb_type]
|
|
|
|
|
|
|
|
|
|
|
|
# 创建新的知识库实例
|
|
|
|
|
|
kb_work_dir = os.path.join(self.work_dir, f"{kb_type}_data")
|
|
|
|
|
|
kb_instance = KnowledgeBaseFactory.create(kb_type, kb_work_dir)
|
|
|
|
|
|
|
|
|
|
|
|
self.kb_instances[kb_type] = kb_instance
|
|
|
|
|
|
logger.info(f"Created {kb_type} knowledge base instance")
|
|
|
|
|
|
return kb_instance
|
|
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def move_file(self, kb_id: str, file_id: str, new_parent_id: str | None) -> dict:
|
2025-12-30 14:28:48 +08:00
|
|
|
|
"""
|
|
|
|
|
|
移动文件/文件夹
|
|
|
|
|
|
"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.move_file(kb_id, file_id, new_parent_id)
|
2025-12-30 14:28:48 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def _get_kb_for_database(self, kb_id: str) -> KnowledgeBase:
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""
|
|
|
|
|
|
根据数据库ID获取对应的知识库实例
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_id: 数据库ID
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
知识库实例
|
|
|
|
|
|
|
|
|
|
|
|
Raises:
|
|
|
|
|
|
KBNotFoundError: 数据库不存在或知识库类型不支持
|
|
|
|
|
|
"""
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
2026-01-21 15:15:52 +08:00
|
|
|
|
|
|
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb = await kb_repo.get_by_kb_id(kb_id)
|
2026-01-21 15:15:52 +08:00
|
|
|
|
|
|
|
|
|
|
if kb is None:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
raise KBNotFoundError(f"Database {kb_id} not found")
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-05-17 13:06:25 +08:00
|
|
|
|
kb_type = kb.kb_type or "milvus"
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
|
|
|
|
|
if not KnowledgeBaseFactory.is_type_supported(kb_type):
|
|
|
|
|
|
raise KBNotFoundError(f"Unsupported knowledge base type: {kb_type}")
|
|
|
|
|
|
|
|
|
|
|
|
return self._get_or_create_kb_instance(kb_type)
|
|
|
|
|
|
|
|
|
|
|
|
# =============================================================================
|
2026-05-17 13:06:25 +08:00
|
|
|
|
# 统一的外部接口
|
2025-07-21 18:18:47 +08:00
|
|
|
|
# =============================================================================
|
|
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def aget_kb(self, kb_id: str) -> KnowledgeBase:
|
2026-01-21 15:15:52 +08:00
|
|
|
|
"""异步获取知识库实例
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_id: 数据库ID
|
2026-01-21 15:15:52 +08:00
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
知识库实例
|
|
|
|
|
|
"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
return await self._get_kb_for_database(kb_id)
|
2026-01-21 15:15:52 +08:00
|
|
|
|
|
2026-05-18 21:32:10 +08:00
|
|
|
|
def _normalize_share_config(
|
|
|
|
|
|
self,
|
|
|
|
|
|
share_config: dict | None,
|
|
|
|
|
|
*,
|
|
|
|
|
|
user_uid: str | None = None,
|
|
|
|
|
|
department_id: int | str | None = None,
|
|
|
|
|
|
) -> dict:
|
|
|
|
|
|
config = share_config or {}
|
|
|
|
|
|
access_level = config.get("access_level") or "global"
|
|
|
|
|
|
if access_level not in ACCESS_LEVELS:
|
|
|
|
|
|
raise ValueError("无效的知识库权限等级")
|
|
|
|
|
|
|
|
|
|
|
|
if access_level == "global":
|
|
|
|
|
|
return DEFAULT_SHARE_CONFIG.copy()
|
|
|
|
|
|
|
|
|
|
|
|
if access_level == "department":
|
|
|
|
|
|
department_ids = self._normalize_department_ids(config.get("department_ids"))
|
|
|
|
|
|
if department_id is not None:
|
|
|
|
|
|
department_ids.append(int(department_id))
|
|
|
|
|
|
department_ids = sorted(set(department_ids))
|
|
|
|
|
|
if not department_ids:
|
|
|
|
|
|
raise ValueError("部门共享至少需要选择一个部门")
|
|
|
|
|
|
return {"access_level": "department", "department_ids": department_ids, "user_uids": []}
|
|
|
|
|
|
|
|
|
|
|
|
user_uids = self._normalize_user_uids(config.get("user_uids"))
|
|
|
|
|
|
if user_uid:
|
|
|
|
|
|
user_uids.append(str(user_uid))
|
|
|
|
|
|
user_uids = sorted(set(user_uids))
|
|
|
|
|
|
if not user_uids:
|
|
|
|
|
|
raise ValueError("指定人可访问至少需要选择一个用户")
|
|
|
|
|
|
return {"access_level": "user", "department_ids": [], "user_uids": user_uids}
|
|
|
|
|
|
|
|
|
|
|
|
def _normalize_department_ids(self, department_ids: list | None) -> list[int]:
|
|
|
|
|
|
normalized = []
|
|
|
|
|
|
for department_id in department_ids or []:
|
|
|
|
|
|
normalized.append(int(department_id))
|
|
|
|
|
|
return normalized
|
|
|
|
|
|
|
|
|
|
|
|
def _normalize_user_uids(self, user_uids: list | None) -> list[str]:
|
|
|
|
|
|
return [uid for uid in (str(uid).strip() for uid in user_uids or []) if uid]
|
|
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
async def get_databases(self) -> dict:
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""获取所有数据库信息"""
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
|
|
|
|
|
rows = await kb_repo.get_all()
|
|
|
|
|
|
all_databases = []
|
2026-02-25 19:09:53 +08:00
|
|
|
|
metadata_reloaded_types: set[str] = set()
|
2026-01-21 15:15:52 +08:00
|
|
|
|
for row in rows:
|
2026-05-17 13:06:25 +08:00
|
|
|
|
kb_type = row.kb_type or "milvus"
|
2026-05-24 00:13:05 +08:00
|
|
|
|
if not KnowledgeBaseFactory.is_type_supported(kb_type):
|
2026-05-21 19:32:25 +08:00
|
|
|
|
logger.warning(f"Skip unsupported database: kb_id={row.kb_id}, kb_type={kb_type}")
|
2026-05-24 00:13:05 +08:00
|
|
|
|
continue
|
2026-02-25 19:09:53 +08:00
|
|
|
|
kb_instance = self._get_or_create_kb_instance(kb_type)
|
2026-05-21 19:32:25 +08:00
|
|
|
|
db_info = kb_instance.get_database_info(row.kb_id, include_files=False)
|
2026-02-25 19:09:53 +08:00
|
|
|
|
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}")
|
2026-05-21 19:32:25 +08:00
|
|
|
|
db_info = kb_instance.get_database_info(row.kb_id, include_files=False)
|
2026-02-25 19:09:53 +08:00
|
|
|
|
|
|
|
|
|
|
if not db_info:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
logger.warning(f"Skip database due to missing metadata: kb_id={row.kb_id}, kb_type={kb_type}")
|
2026-02-25 19:09:53 +08:00
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
|
|
# 补充 share_config 和 additional_params
|
2026-05-18 21:32:10 +08:00
|
|
|
|
db_info["share_config"] = row.share_config or DEFAULT_SHARE_CONFIG.copy()
|
2026-05-18 09:40:34 +08:00
|
|
|
|
db_info["additional_params"] = kb_instance.normalize_additional_params(row.additional_params)
|
|
|
|
|
|
db_info["created_by"] = row.created_by
|
2026-02-25 19:09:53 +08:00
|
|
|
|
all_databases.append(db_info)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
return {"databases": all_databases}
|
|
|
|
|
|
|
2026-05-19 18:37:29 +08:00
|
|
|
|
@staticmethod
|
|
|
|
|
|
def _database_info_accessible(user: dict, db_info: dict) -> bool:
|
2026-01-20 02:14:51 +08:00
|
|
|
|
if user.get("role") == "superadmin":
|
|
|
|
|
|
return True
|
|
|
|
|
|
|
2026-05-18 21:32:10 +08:00
|
|
|
|
user_uid = str(user.get("uid") or "")
|
2026-05-19 18:37:29 +08:00
|
|
|
|
if user_uid and db_info.get("created_by") == user_uid:
|
2026-01-20 02:14:51 +08:00
|
|
|
|
return True
|
|
|
|
|
|
|
2026-05-19 18:37:29 +08:00
|
|
|
|
share_config = db_info.get("share_config") or DEFAULT_SHARE_CONFIG.copy()
|
2026-05-18 21:32:10 +08:00
|
|
|
|
access_level = share_config.get("access_level")
|
|
|
|
|
|
if access_level == "global":
|
|
|
|
|
|
return True
|
2026-01-20 02:14:51 +08:00
|
|
|
|
|
2026-05-18 21:32:10 +08:00
|
|
|
|
if access_level == "department":
|
|
|
|
|
|
user_department_id = user.get("department_id")
|
|
|
|
|
|
if user_department_id is None:
|
|
|
|
|
|
return False
|
|
|
|
|
|
try:
|
|
|
|
|
|
department_ids = [int(dept_id) for dept_id in share_config.get("department_ids") or []]
|
|
|
|
|
|
return int(user_department_id) in department_ids
|
|
|
|
|
|
except (ValueError, TypeError):
|
|
|
|
|
|
return False
|
2026-01-20 02:14:51 +08:00
|
|
|
|
|
2026-05-18 21:32:10 +08:00
|
|
|
|
if access_level == "user":
|
|
|
|
|
|
return bool(user_uid and user_uid in (share_config.get("user_uids") or []))
|
|
|
|
|
|
|
|
|
|
|
|
return False
|
2026-01-20 02:14:51 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def check_accessible(self, user: dict, kb_id: str) -> bool:
|
2026-05-19 18:37:29 +08:00
|
|
|
|
"""检查用户是否有权限访问数据库
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
user: 用户信息字典
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_id: 数据库ID
|
2026-05-19 18:37:29 +08:00
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
bool: 是否有权限
|
|
|
|
|
|
"""
|
|
|
|
|
|
# 超级管理员有权访问所有
|
|
|
|
|
|
if user.get("role") == "superadmin":
|
|
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
|
|
|
|
|
|
|
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb = await kb_repo.get_by_kb_id(kb_id)
|
2026-05-19 18:37:29 +08:00
|
|
|
|
if kb is None:
|
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
|
return self._database_info_accessible(
|
|
|
|
|
|
user,
|
|
|
|
|
|
{
|
|
|
|
|
|
"created_by": kb.created_by,
|
|
|
|
|
|
"share_config": kb.share_config,
|
|
|
|
|
|
},
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2026-05-17 22:44:05 +08:00
|
|
|
|
async def get_databases_by_uid(self, uid: str) -> dict:
|
|
|
|
|
|
"""根据 uid 获取知识库列表"""
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.user_repository import UserRepository
|
2026-03-04 04:06:48 +08:00
|
|
|
|
|
|
|
|
|
|
# 通过数据库获取用户信息
|
|
|
|
|
|
user_repo = UserRepository()
|
2026-05-17 22:44:05 +08:00
|
|
|
|
user: User | None = await user_repo.get_by_uid(uid)
|
2026-03-04 04:06:48 +08:00
|
|
|
|
if not user:
|
2026-05-17 22:44:05 +08:00
|
|
|
|
logger.warning(f"User not found: {uid}")
|
2026-03-04 04:06:48 +08:00
|
|
|
|
return {"databases": []}
|
|
|
|
|
|
return await self.get_databases_by_user(user)
|
|
|
|
|
|
|
2026-03-14 15:58:35 +08:00
|
|
|
|
async def get_databases_by_user(self, user: User | dict) -> dict:
|
2026-03-04 04:06:48 +08:00
|
|
|
|
"""根据用户权限获取知识库列表"""
|
|
|
|
|
|
|
2026-03-14 15:58:35 +08:00
|
|
|
|
# 构建用户信息字典(支持 User 对象或 dict)
|
|
|
|
|
|
if isinstance(user, dict):
|
|
|
|
|
|
user_info = user
|
|
|
|
|
|
else:
|
|
|
|
|
|
user_info = {
|
2026-05-18 21:32:10 +08:00
|
|
|
|
"uid": user.uid,
|
2026-03-14 15:58:35 +08:00
|
|
|
|
"role": user.role,
|
|
|
|
|
|
"department_id": user.department_id,
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
user_role = user_info.get("role")
|
|
|
|
|
|
user_dept = user_info.get("department_id")
|
|
|
|
|
|
logger.info(f"Getting databases for user with role {user_role} and department {user_dept}")
|
2026-01-20 02:14:51 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
all_databases = (await self.get_databases()).get("databases", [])
|
2026-01-20 02:14:51 +08:00
|
|
|
|
|
|
|
|
|
|
# 超级管理员可以看到所有知识库
|
2026-03-04 04:06:48 +08:00
|
|
|
|
if user_info.get("role") == "superadmin":
|
2026-01-20 02:14:51 +08:00
|
|
|
|
return {"databases": all_databases}
|
|
|
|
|
|
|
2026-05-19 18:37:29 +08:00
|
|
|
|
filtered_databases = [
|
|
|
|
|
|
database for database in all_databases if self._database_info_accessible(user_info, database)
|
|
|
|
|
|
]
|
2026-01-20 02:14:51 +08:00
|
|
|
|
|
|
|
|
|
|
return {"databases": filtered_databases}
|
|
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
async def database_name_exists(self, database_name: str) -> bool:
|
|
|
|
|
|
"""检查知识库名称是否已存在"""
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
|
|
|
|
|
from yuxi.storage.postgres.manager import pg_manager
|
2026-01-14 22:00:53 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
# 确保 pg_manager 已初始化
|
|
|
|
|
|
if not pg_manager._initialized:
|
|
|
|
|
|
pg_manager.initialize()
|
2026-01-14 22:00:53 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
|
|
|
|
|
rows = await kb_repo.get_all()
|
|
|
|
|
|
for row in rows:
|
|
|
|
|
|
if (row.name or "").lower() == database_name.lower():
|
2026-01-14 22:00:53 +08:00
|
|
|
|
return True
|
|
|
|
|
|
return False
|
|
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def create_folder(self, kb_id: str, folder_name: str, parent_id: str = None) -> dict:
|
2025-12-30 14:28:48 +08:00
|
|
|
|
"""Create a folder in the database."""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.create_folder(kb_id, folder_name, parent_id)
|
2025-12-30 14:28:48 +08:00
|
|
|
|
|
2025-09-01 22:37:03 +08:00
|
|
|
|
async def create_database(
|
2026-01-21 15:15:52 +08:00
|
|
|
|
self,
|
|
|
|
|
|
database_name: str,
|
|
|
|
|
|
description: str,
|
2026-05-17 13:06:25 +08:00
|
|
|
|
kb_type: str = "milvus",
|
|
|
|
|
|
embedding_model_spec: str | None = None,
|
|
|
|
|
|
llm_model_spec: str | None = None,
|
2026-01-21 15:15:52 +08:00
|
|
|
|
share_config: dict | None = None,
|
2026-05-18 09:40:34 +08:00
|
|
|
|
created_by: str | None = None,
|
2026-05-18 21:32:10 +08:00
|
|
|
|
created_by_department_id: int | str | None = None,
|
2026-01-21 15:15:52 +08:00
|
|
|
|
**kwargs,
|
2025-09-01 22:37:03 +08:00
|
|
|
|
) -> dict:
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""
|
|
|
|
|
|
创建数据库
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
database_name: 数据库名称
|
|
|
|
|
|
description: 数据库描述
|
2026-05-17 13:06:25 +08:00
|
|
|
|
kb_type: 知识库类型,默认为 milvus
|
|
|
|
|
|
embedding_model_spec: 嵌入模型 spec
|
|
|
|
|
|
llm_model_spec: LLM 模型 spec
|
2026-01-20 02:14:51 +08:00
|
|
|
|
share_config: 共享配置
|
2026-05-18 09:40:34 +08:00
|
|
|
|
created_by: 创建者 uid
|
2026-05-18 21:32:10 +08:00
|
|
|
|
created_by_department_id: 创建者部门 ID
|
2026-05-16 17:21:27 +08:00
|
|
|
|
**kwargs: 其他配置参数
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
数据库信息字典
|
|
|
|
|
|
"""
|
|
|
|
|
|
if not KnowledgeBaseFactory.is_type_supported(kb_type):
|
|
|
|
|
|
available_types = list(KnowledgeBaseFactory.get_available_types().keys())
|
2025-07-26 03:36:54 +08:00
|
|
|
|
raise ValueError(f"Unsupported knowledge base type: {kb_type}. Available types: {available_types}")
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-01-14 22:00:53 +08:00
|
|
|
|
# 检查名称是否已存在
|
2026-01-21 15:15:52 +08:00
|
|
|
|
if await self.database_name_exists(database_name):
|
2026-01-14 22:00:53 +08:00
|
|
|
|
raise ValueError(f"知识库名称 '{database_name}' 已存在,请使用其他名称")
|
2026-01-21 15:15:52 +08:00
|
|
|
|
|
2026-05-18 21:32:10 +08:00
|
|
|
|
share_config = self._normalize_share_config(
|
|
|
|
|
|
share_config,
|
|
|
|
|
|
user_uid=created_by,
|
|
|
|
|
|
department_id=created_by_department_id,
|
|
|
|
|
|
)
|
2026-01-14 22:00:53 +08:00
|
|
|
|
|
2025-07-21 18:18:47 +08:00
|
|
|
|
kb_instance = self._get_or_create_kb_instance(kb_type)
|
2026-05-18 09:40:34 +08:00
|
|
|
|
kwargs = kb_instance.normalize_additional_params(kwargs)
|
2026-05-17 13:06:25 +08:00
|
|
|
|
db_info = await kb_instance.create_database(
|
|
|
|
|
|
database_name,
|
|
|
|
|
|
description,
|
|
|
|
|
|
embedding_model_spec=embedding_model_spec,
|
|
|
|
|
|
llm_model_spec=llm_model_spec,
|
|
|
|
|
|
**kwargs,
|
|
|
|
|
|
)
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_id = db_info["kb_id"]
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
2026-01-21 15:15:52 +08:00
|
|
|
|
|
|
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
2026-05-21 19:32:25 +08:00
|
|
|
|
updated = await kb_repo.update(kb_id, {"share_config": share_config, "created_by": created_by})
|
2026-01-21 15:15:52 +08:00
|
|
|
|
if updated is None:
|
|
|
|
|
|
await kb_repo.create(
|
|
|
|
|
|
{
|
2026-05-21 19:32:25 +08:00
|
|
|
|
"kb_id": kb_id,
|
2026-01-21 15:15:52 +08:00
|
|
|
|
"name": database_name,
|
|
|
|
|
|
"description": description,
|
|
|
|
|
|
"kb_type": kb_type,
|
2026-05-17 13:06:25 +08:00
|
|
|
|
"embedding_model_spec": embedding_model_spec,
|
|
|
|
|
|
"llm_model_spec": db_info.get("llm_model_spec"),
|
2026-01-21 15:15:52 +08:00
|
|
|
|
"additional_params": kwargs.copy(),
|
|
|
|
|
|
"share_config": share_config,
|
2026-05-18 09:40:34 +08:00
|
|
|
|
"created_by": created_by,
|
2026-01-21 15:15:52 +08:00
|
|
|
|
}
|
|
|
|
|
|
)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
logger.info(f"Created {kb_type} database: {database_name} ({kb_id}) with {kwargs}")
|
2026-01-20 02:14:51 +08:00
|
|
|
|
db_info["share_config"] = share_config
|
2025-07-21 18:18:47 +08:00
|
|
|
|
return db_info
|
|
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def delete_database(self, kb_id: str) -> dict:
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""删除数据库"""
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
2026-01-21 15:15:52 +08:00
|
|
|
|
|
2025-07-21 18:18:47 +08:00
|
|
|
|
try:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
result = await kb_instance.delete_database(kb_id)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
# 删除数据库记录
|
|
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
2026-05-21 19:32:25 +08:00
|
|
|
|
await kb_repo.delete(kb_id)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
|
|
|
|
|
return result
|
|
|
|
|
|
except KBNotFoundError as e:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
logger.warning(f"Database {kb_id} not found during deletion: {e}")
|
2025-07-27 01:02:14 +08:00
|
|
|
|
return {"message": "删除成功"}
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-01-04 21:35:25 +08:00
|
|
|
|
async def add_file_record(
|
2026-05-21 19:32:25 +08:00
|
|
|
|
self, kb_id: str, item: str, params: dict | None = None, operator_id: str | None = None
|
2026-01-04 21:35:25 +08:00
|
|
|
|
) -> dict:
|
|
|
|
|
|
"""Add file record to metadata"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.add_file_record(kb_id, item, params, operator_id)
|
2026-01-04 21:35:25 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def parse_file(self, kb_id: str, file_id: str, operator_id: str | None = None) -> dict:
|
2026-01-04 21:35:25 +08:00
|
|
|
|
"""Parse file to Markdown"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.parse_file(kb_id, file_id, operator_id)
|
2026-01-04 21:35:25 +08:00
|
|
|
|
|
2026-04-27 02:24:52 +08:00
|
|
|
|
async def index_file(
|
2026-05-21 19:32:25 +08:00
|
|
|
|
self, kb_id: str, file_id: str, operator_id: str | None = None, params: dict | None = None
|
2026-04-27 02:24:52 +08:00
|
|
|
|
) -> dict:
|
2026-01-04 21:35:25 +08:00
|
|
|
|
"""Index parsed file"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.index_file(kb_id, file_id, operator_id, params=params)
|
2026-01-04 21:35:25 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def update_file_params(self, kb_id: str, file_id: str, params: dict, operator_id: str | None = None) -> None:
|
2026-01-04 21:35:25 +08:00
|
|
|
|
"""Update file processing params"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
await kb_instance.update_file_params(kb_id, file_id, params, operator_id)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def aquery(self, query_text: str, kb_id: str, **kwargs) -> str:
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""异步查询知识库"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.aquery(query_text, kb_id, **kwargs)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def export_data(self, kb_id: str, format: str = "zip", **kwargs) -> str:
|
2025-08-10 21:58:41 +08:00
|
|
|
|
"""导出知识库数据"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.export_data(kb_id, format=format, **kwargs)
|
2025-08-10 21:58:41 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def get_database_info(self, kb_id: str) -> dict | None:
|
2026-01-21 15:15:52 +08:00
|
|
|
|
"""获取数据库详细信息"""
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
2026-01-09 01:21:41 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb = await kb_repo.get_by_kb_id(kb_id)
|
2026-01-21 15:15:52 +08:00
|
|
|
|
if kb is None:
|
|
|
|
|
|
return None
|
2026-01-09 01:21:41 +08:00
|
|
|
|
|
2025-07-21 18:18:47 +08:00
|
|
|
|
try:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
db_info = kb_instance.get_database_info(kb_id)
|
2026-01-21 15:15:52 +08:00
|
|
|
|
except KBNotFoundError:
|
|
|
|
|
|
db_info = {
|
2026-05-21 19:32:25 +08:00
|
|
|
|
"kb_id": kb_id,
|
2026-01-21 15:15:52 +08:00
|
|
|
|
"name": kb.name,
|
|
|
|
|
|
"description": kb.description,
|
|
|
|
|
|
"kb_type": kb.kb_type,
|
|
|
|
|
|
"files": {},
|
|
|
|
|
|
"row_count": 0,
|
|
|
|
|
|
"status": "已连接",
|
|
|
|
|
|
}
|
2025-07-26 03:36:54 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
# 添加数据库中的附加字段
|
2026-05-18 09:40:34 +08:00
|
|
|
|
db_info["additional_params"] = kb_instance.normalize_additional_params(kb.additional_params)
|
2026-05-18 21:32:10 +08:00
|
|
|
|
db_info["share_config"] = kb.share_config or DEFAULT_SHARE_CONFIG.copy()
|
2026-01-21 15:15:52 +08:00
|
|
|
|
db_info["mindmap"] = kb.mindmap
|
|
|
|
|
|
db_info["sample_questions"] = kb.sample_questions or []
|
|
|
|
|
|
db_info["query_params"] = kb.query_params
|
2025-07-26 03:36:54 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
return db_info
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def delete_folder(self, kb_id: str, folder_id: str) -> None:
|
2025-12-30 14:28:48 +08:00
|
|
|
|
"""递归删除文件夹"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
await kb_instance.delete_folder(kb_id, folder_id)
|
2025-12-30 14:28:48 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def delete_file(self, kb_id: str, file_id: str) -> None:
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""删除文件"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
await kb_instance.delete_file(kb_id, file_id)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def update_content(self, kb_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]:
|
2025-11-12 09:14:07 +08:00
|
|
|
|
"""更新内容(重新分块)"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.update_content(kb_id, file_ids, params or {})
|
2025-11-12 09:14:07 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def get_file_basic_info(self, kb_id: str, file_id: str) -> dict:
|
2025-09-21 23:48:56 +08:00
|
|
|
|
"""获取文件基本信息(仅元数据)"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.get_file_basic_info(kb_id, file_id)
|
2025-09-21 23:48:56 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def get_file_content(self, kb_id: str, file_id: str) -> dict:
|
2025-09-21 23:48:56 +08:00
|
|
|
|
"""获取文件内容信息(chunks和lines)"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.get_file_content(kb_id, file_id)
|
2025-09-21 23:48:56 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def open_file_content(self, kb_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict:
|
2026-05-16 11:39:15 +08:00
|
|
|
|
"""按行窗口打开文件解析后的 Markdown 内容"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.open_file_content(kb_id, file_id, offset, limit)
|
2026-05-16 11:39:15 +08:00
|
|
|
|
|
2026-05-18 14:09:42 +08:00
|
|
|
|
async def find_file_content(
|
|
|
|
|
|
self,
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_id: str,
|
2026-05-18 14:09:42 +08:00
|
|
|
|
file_id: str,
|
|
|
|
|
|
patterns: list[str],
|
|
|
|
|
|
*,
|
|
|
|
|
|
use_regex: bool = False,
|
|
|
|
|
|
case_sensitive: bool = False,
|
|
|
|
|
|
max_windows: int = 5,
|
|
|
|
|
|
window_size: int = 80,
|
|
|
|
|
|
) -> dict:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
2026-05-18 14:09:42 +08:00
|
|
|
|
return await kb_instance.find_file_content(
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_id,
|
2026-05-18 14:09:42 +08:00
|
|
|
|
file_id,
|
|
|
|
|
|
patterns,
|
|
|
|
|
|
use_regex=use_regex,
|
|
|
|
|
|
case_sensitive=case_sensitive,
|
|
|
|
|
|
max_windows=max_windows,
|
|
|
|
|
|
window_size=window_size,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def get_file_info(self, kb_id: str, file_id: str) -> dict:
|
2026-05-17 13:06:25 +08:00
|
|
|
|
"""获取文件完整信息(基本信息+内容信息)"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.get_file_info(kb_id, file_id)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-05-17 14:28:50 +08:00
|
|
|
|
async def list_file_tree(
|
|
|
|
|
|
self,
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_id: str,
|
2026-05-17 14:28:50 +08:00
|
|
|
|
parent_id: str | None = None,
|
|
|
|
|
|
recursive: bool = False,
|
|
|
|
|
|
files_only: bool = False,
|
|
|
|
|
|
) -> dict:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
|
|
|
|
|
return await kb_instance.list_file_tree(kb_id, parent_id, recursive, files_only)
|
2026-05-17 14:28:50 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
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)
|
2026-05-17 14:28:50 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
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)
|
2026-05-17 14:28:50 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def file_name_existed_in_db(self, kb_id: str | None, file_name: str | None) -> bool:
|
2025-11-29 17:49:42 +08:00
|
|
|
|
"""检查指定数据库中是否存在同名的文件"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
if not kb_id or not file_name:
|
2025-11-29 17:49:42 +08:00
|
|
|
|
return False
|
|
|
|
|
|
try:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
2025-11-29 17:49:42 +08:00
|
|
|
|
except KBNotFoundError:
|
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
|
for file_info in kb_instance.files_meta.values():
|
2026-05-21 19:32:25 +08:00
|
|
|
|
if file_info.get("kb_id") != kb_id:
|
2025-11-29 17:49:42 +08:00
|
|
|
|
continue
|
|
|
|
|
|
if file_info.get("status") == "failed":
|
|
|
|
|
|
continue
|
|
|
|
|
|
if file_info.get("file_name") == file_name:
|
|
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
|
|
return False
|
|
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def get_same_name_files(self, kb_id: str, filename: str) -> list[dict]:
|
2025-12-07 23:38:20 +08:00
|
|
|
|
"""获取同一知识库中同名文件列表
|
|
|
|
|
|
基于原始文件名直接比较
|
|
|
|
|
|
返回基础信息:文件名、大小、上传时间
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_id: 数据库ID
|
2025-12-07 23:38:20 +08:00
|
|
|
|
filename: 要检测的文件名(原始文件名)
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
同名文件列表,每项包含:
|
|
|
|
|
|
- filename: 文件名
|
|
|
|
|
|
- size: 文件大小
|
|
|
|
|
|
- created_at: 上传时间
|
|
|
|
|
|
- file_id: 文件ID(用于下载)
|
|
|
|
|
|
"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
if not kb_id or not filename:
|
2025-12-07 23:38:20 +08:00
|
|
|
|
return []
|
|
|
|
|
|
try:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
2025-12-07 23:38:20 +08:00
|
|
|
|
except KBNotFoundError:
|
|
|
|
|
|
return []
|
|
|
|
|
|
|
|
|
|
|
|
same_name_files = []
|
|
|
|
|
|
for file_id, file_info in kb_instance.files_meta.items():
|
2026-05-21 19:32:25 +08:00
|
|
|
|
if file_info.get("kb_id") != kb_id:
|
2025-12-07 23:38:20 +08:00
|
|
|
|
continue
|
|
|
|
|
|
if file_info.get("status") == "failed":
|
|
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
|
|
# 直接比较文件名(现在就是原始文件名)
|
|
|
|
|
|
current_filename = file_info.get("filename", "")
|
|
|
|
|
|
|
|
|
|
|
|
if current_filename.lower() == filename.lower():
|
2025-12-09 16:21:18 +08:00
|
|
|
|
same_name_files.append(
|
|
|
|
|
|
{
|
|
|
|
|
|
"file_id": file_id,
|
|
|
|
|
|
"filename": current_filename,
|
|
|
|
|
|
"size": file_info.get("size", 0),
|
|
|
|
|
|
"created_at": file_info.get("created_at", ""),
|
|
|
|
|
|
"content_hash": file_info.get("content_hash", ""),
|
|
|
|
|
|
}
|
|
|
|
|
|
)
|
2025-12-07 23:38:20 +08:00
|
|
|
|
|
|
|
|
|
|
# 按上传时间降序排序
|
|
|
|
|
|
same_name_files.sort(key=lambda x: x.get("created_at", ""), reverse=True)
|
|
|
|
|
|
return same_name_files
|
|
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
async def file_existed_in_db(self, kb_id: str | None, content_hash: str | None) -> bool:
|
2025-10-11 21:52:48 +08:00
|
|
|
|
"""检查指定数据库中是否存在相同内容哈希的文件"""
|
2026-05-21 19:32:25 +08:00
|
|
|
|
if not kb_id or not content_hash:
|
2025-10-11 21:52:48 +08:00
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
|
try:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_instance = await self._get_kb_for_database(kb_id)
|
2025-10-11 21:52:48 +08:00
|
|
|
|
except KBNotFoundError:
|
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
|
for file_info in kb_instance.files_meta.values():
|
2026-05-21 19:32:25 +08:00
|
|
|
|
if file_info.get("kb_id") != kb_id:
|
2025-10-11 21:52:48 +08:00
|
|
|
|
continue
|
2025-10-25 16:01:50 +08:00
|
|
|
|
if file_info.get("status") == "failed":
|
|
|
|
|
|
continue
|
2025-10-11 21:52:48 +08:00
|
|
|
|
if file_info.get("content_hash") == content_hash:
|
|
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
|
|
return False
|
|
|
|
|
|
|
2025-11-30 18:48:33 +08:00
|
|
|
|
async def update_database(
|
2026-01-21 15:15:52 +08:00
|
|
|
|
self,
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb_id: str,
|
2026-01-21 15:15:52 +08:00
|
|
|
|
name: str,
|
|
|
|
|
|
description: str,
|
2026-05-17 13:06:25 +08:00
|
|
|
|
llm_model_spec: str | None = None,
|
|
|
|
|
|
update_llm_model_spec: bool = False,
|
2026-01-21 15:15:52 +08:00
|
|
|
|
additional_params: dict | None = None,
|
|
|
|
|
|
share_config: dict | None = None,
|
2026-05-18 21:32:10 +08:00
|
|
|
|
operator_uid: str | None = None,
|
|
|
|
|
|
operator_department_id: int | str | None = None,
|
2025-11-30 18:48:33 +08:00
|
|
|
|
) -> dict:
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""更新数据库"""
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-02-20 19:35:43 +08:00
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
2026-05-21 19:32:25 +08:00
|
|
|
|
kb = await kb_repo.get_by_kb_id(kb_id)
|
2026-02-20 19:35:43 +08:00
|
|
|
|
if kb is None:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
raise ValueError(f"数据库 {kb_id} 不存在")
|
2026-02-20 19:35:43 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
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)
|
2025-11-30 18:48:33 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
update_data: dict = {
|
|
|
|
|
|
"name": name,
|
|
|
|
|
|
"description": description,
|
|
|
|
|
|
}
|
2026-05-17 13:06:25 +08:00
|
|
|
|
if update_llm_model_spec:
|
|
|
|
|
|
update_data["llm_model_spec"] = llm_model_spec
|
2026-02-20 19:35:43 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
if additional_params is not None:
|
2026-05-17 13:06:25 +08:00
|
|
|
|
current_additional_params = kb.additional_params or {}
|
|
|
|
|
|
current_graph_config = current_additional_params.get("graph_build_config") or {}
|
|
|
|
|
|
if current_graph_config.get("locked") and "graph_build_config" in additional_params:
|
|
|
|
|
|
raise ValueError("图谱抽取配置已锁定,请使用图谱重置接口重新配置")
|
|
|
|
|
|
|
2026-05-18 09:40:34 +08:00
|
|
|
|
merged_additional_params = kb_instance.normalize_additional_params(
|
2026-05-17 13:06:25 +08:00
|
|
|
|
deep_merge(current_additional_params, additional_params)
|
2026-02-20 19:35:43 +08:00
|
|
|
|
)
|
|
|
|
|
|
update_data["additional_params"] = merged_additional_params
|
2026-05-21 19:32:25 +08:00
|
|
|
|
if kb_id in kb_instance.databases_meta:
|
|
|
|
|
|
kb_instance.databases_meta[kb_id]["metadata"] = merged_additional_params
|
2026-02-20 19:35:43 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
if share_config is not None:
|
2026-05-18 21:32:10 +08:00
|
|
|
|
update_data["share_config"] = self._normalize_share_config(
|
|
|
|
|
|
share_config,
|
|
|
|
|
|
user_uid=operator_uid,
|
|
|
|
|
|
department_id=operator_department_id,
|
|
|
|
|
|
)
|
2026-01-20 02:14:51 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
# 保存到数据库
|
2026-05-21 19:32:25 +08:00
|
|
|
|
await kb_repo.update(kb_id, update_data)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
return await self.get_database_info(kb_id)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2025-07-26 03:36:54 +08:00
|
|
|
|
def get_retrievers(self) -> dict[str, dict]:
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""获取所有检索器"""
|
|
|
|
|
|
all_retrievers = {}
|
|
|
|
|
|
|
|
|
|
|
|
# 收集所有知识库的检索器
|
|
|
|
|
|
for kb_instance in self.kb_instances.values():
|
|
|
|
|
|
retrievers = kb_instance.get_retrievers()
|
|
|
|
|
|
all_retrievers.update(retrievers)
|
|
|
|
|
|
|
|
|
|
|
|
return all_retrievers
|
|
|
|
|
|
|
|
|
|
|
|
# =============================================================================
|
|
|
|
|
|
# 管理器特有的方法
|
|
|
|
|
|
# =============================================================================
|
|
|
|
|
|
|
2025-07-26 03:36:54 +08:00
|
|
|
|
def get_supported_kb_types(self) -> dict[str, dict]:
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""获取支持的知识库类型"""
|
|
|
|
|
|
return KnowledgeBaseFactory.get_available_types()
|
|
|
|
|
|
|
2025-07-26 03:36:54 +08:00
|
|
|
|
def get_kb_instance_info(self) -> dict[str, dict]:
|
2025-07-21 18:18:47 +08:00
|
|
|
|
"""获取知识库实例信息"""
|
|
|
|
|
|
info = {}
|
|
|
|
|
|
for kb_type, kb_instance in self.kb_instances.items():
|
|
|
|
|
|
info[kb_type] = {
|
|
|
|
|
|
"work_dir": kb_instance.work_dir,
|
|
|
|
|
|
"database_count": len(kb_instance.databases_meta),
|
2025-09-01 22:37:03 +08:00
|
|
|
|
"file_count": len(kb_instance.files_meta),
|
2025-07-21 18:18:47 +08:00
|
|
|
|
}
|
|
|
|
|
|
return info
|
|
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
async def get_statistics(self) -> dict:
|
2025-07-22 17:29:38 +08:00
|
|
|
|
"""获取统计信息"""
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
|
|
|
|
|
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
2026-01-21 15:15:52 +08:00
|
|
|
|
|
|
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
|
|
|
|
|
rows = await kb_repo.get_all()
|
|
|
|
|
|
|
|
|
|
|
|
stats = {"total_databases": len(rows), "kb_types": {}, "total_files": 0}
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2025-07-22 17:29:38 +08:00
|
|
|
|
# 按知识库类型统计
|
2026-01-21 15:15:52 +08:00
|
|
|
|
for row in rows:
|
2026-05-17 13:06:25 +08:00
|
|
|
|
kb_type = row.kb_type or "milvus"
|
2025-07-22 17:29:38 +08:00
|
|
|
|
if kb_type not in stats["kb_types"]:
|
|
|
|
|
|
stats["kb_types"][kb_type] = 0
|
|
|
|
|
|
stats["kb_types"][kb_type] += 1
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
file_repo = KnowledgeFileRepository()
|
|
|
|
|
|
files = await file_repo.get_all()
|
|
|
|
|
|
stats["total_files"] = len(files)
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2025-07-22 17:29:38 +08:00
|
|
|
|
return stats
|
2025-07-21 18:18:47 +08:00
|
|
|
|
|
2025-10-23 22:57:01 +08:00
|
|
|
|
# =============================================================================
|
|
|
|
|
|
# 数据一致性检测方法
|
|
|
|
|
|
# =============================================================================
|
|
|
|
|
|
|
|
|
|
|
|
async def detect_data_inconsistencies(self) -> dict:
|
|
|
|
|
|
"""
|
|
|
|
|
|
检测向量数据库中存在但在 metadata 中缺失的数据
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
包含不一致信息的字典,按知识库类型分组
|
|
|
|
|
|
"""
|
|
|
|
|
|
inconsistencies = {
|
|
|
|
|
|
"milvus": {"missing_collections": [], "missing_files": []},
|
|
|
|
|
|
"total_missing_collections": 0,
|
2025-10-24 00:11:52 +08:00
|
|
|
|
"total_missing_files": 0,
|
2025-10-23 22:57:01 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
logger.info("开始检测向量数据库与元数据的一致性...")
|
|
|
|
|
|
|
|
|
|
|
|
# 检测 Milvus 数据不一致
|
|
|
|
|
|
if "milvus" in self.kb_instances:
|
|
|
|
|
|
try:
|
|
|
|
|
|
milvus_inconsistencies = await self._detect_milvus_inconsistencies()
|
|
|
|
|
|
inconsistencies["milvus"] = milvus_inconsistencies
|
|
|
|
|
|
inconsistencies["total_missing_collections"] += len(milvus_inconsistencies["missing_collections"])
|
|
|
|
|
|
inconsistencies["total_missing_files"] += len(milvus_inconsistencies["missing_files"])
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
|
logger.error(f"检测 Milvus 数据不一致时出错: {e}")
|
|
|
|
|
|
|
|
|
|
|
|
# 输出检测结果到日志
|
|
|
|
|
|
self._log_inconsistencies(inconsistencies)
|
|
|
|
|
|
|
|
|
|
|
|
return inconsistencies
|
|
|
|
|
|
|
|
|
|
|
|
async def _detect_milvus_inconsistencies(self) -> dict:
|
|
|
|
|
|
"""检测 Milvus 中的数据不一致"""
|
2026-03-17 10:16:44 +08:00
|
|
|
|
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
2026-01-21 15:15:52 +08:00
|
|
|
|
|
2025-10-23 22:57:01 +08:00
|
|
|
|
inconsistencies = {"missing_collections": [], "missing_files": []}
|
|
|
|
|
|
|
|
|
|
|
|
milvus_kb = self.kb_instances["milvus"]
|
|
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
|
from pymilvus import utility
|
|
|
|
|
|
|
|
|
|
|
|
# 获取 Milvus 中所有实际的集合
|
|
|
|
|
|
actual_collection_names = set(utility.list_collections(using=milvus_kb.connection_alias))
|
|
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
# 从数据库获取所有已知的数据库ID
|
|
|
|
|
|
kb_repo = KnowledgeBaseRepository()
|
|
|
|
|
|
rows = await kb_repo.get_all()
|
2026-05-21 19:32:25 +08:00
|
|
|
|
all_known_kb_ids = {row.kb_id for row in rows}
|
2026-01-21 15:15:52 +08:00
|
|
|
|
|
2025-10-23 22:57:01 +08:00
|
|
|
|
# 找出存在于 Milvus 但不在 metadata 中的集合
|
2026-01-21 15:15:52 +08:00
|
|
|
|
# missing_collections = actual_collection_names - metadata_collection_names
|
|
|
|
|
|
for collection_name in actual_collection_names:
|
2025-10-23 22:57:01 +08:00
|
|
|
|
# 跳过一些系统集合
|
|
|
|
|
|
if not collection_name.startswith("kb_"):
|
|
|
|
|
|
continue
|
|
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
# 检查集合是否属于已知数据库
|
|
|
|
|
|
is_known = False
|
|
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
if collection_name in all_known_kb_ids:
|
2026-01-21 15:15:52 +08:00
|
|
|
|
is_known = True
|
|
|
|
|
|
|
|
|
|
|
|
# 如果是已知集合,跳过
|
|
|
|
|
|
if is_known:
|
|
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
|
|
# 如果是未知集合,记录下来
|
2025-10-24 00:11:52 +08:00
|
|
|
|
collection_info = {"collection_name": collection_name, "detected_at": utc_isoformat()}
|
2025-10-23 22:57:01 +08:00
|
|
|
|
|
|
|
|
|
|
# 尝试获取集合的基本信息
|
|
|
|
|
|
try:
|
|
|
|
|
|
from pymilvus import Collection
|
2025-10-24 00:11:52 +08:00
|
|
|
|
|
2025-10-23 22:57:01 +08:00
|
|
|
|
collection = Collection(name=collection_name, using=milvus_kb.connection_alias)
|
|
|
|
|
|
collection_info["count"] = collection.num_entities
|
|
|
|
|
|
collection_info["description"] = collection.description
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
|
logger.warning(f"无法获取集合 {collection_name} 的详细信息: {e}")
|
|
|
|
|
|
collection_info["count"] = "unknown"
|
|
|
|
|
|
|
|
|
|
|
|
inconsistencies["missing_collections"].append(collection_info)
|
2025-10-24 00:11:52 +08:00
|
|
|
|
logger.warning(
|
|
|
|
|
|
f"发现 Milvus 中存在但 metadata 中缺失的集合: {collection_name} "
|
|
|
|
|
|
f"(实体数: {collection_info['count']})"
|
|
|
|
|
|
)
|
2025-10-23 22:57:01 +08:00
|
|
|
|
|
2026-01-21 15:15:52 +08:00
|
|
|
|
# 获取 metadata 中记录的数据库ID(仅 Milvus 类型,用于检查文件一致性)
|
|
|
|
|
|
metadata_collection_names = set(milvus_kb.databases_meta.keys())
|
|
|
|
|
|
|
2025-10-23 22:57:01 +08:00
|
|
|
|
# 检查文件级别的不一致(针对已知的数据库)
|
2026-05-21 19:32:25 +08:00
|
|
|
|
for kb_id in metadata_collection_names:
|
2025-10-23 22:57:01 +08:00
|
|
|
|
try:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
if utility.has_collection(kb_id, using=milvus_kb.connection_alias):
|
2025-10-23 22:57:01 +08:00
|
|
|
|
from pymilvus import Collection
|
2025-10-24 00:11:52 +08:00
|
|
|
|
|
2026-05-21 19:32:25 +08:00
|
|
|
|
collection = Collection(name=kb_id, using=milvus_kb.connection_alias)
|
2025-10-23 22:57:01 +08:00
|
|
|
|
actual_count = collection.num_entities
|
|
|
|
|
|
|
|
|
|
|
|
# 获取 metadata 中记录的文件数量
|
2025-10-24 00:11:52 +08:00
|
|
|
|
metadata_files_count = sum(
|
2026-05-21 19:32:25 +08:00
|
|
|
|
1 for file_info in milvus_kb.files_meta.values() if file_info.get("kb_id") == kb_id
|
2025-10-24 00:11:52 +08:00
|
|
|
|
)
|
2025-10-23 22:57:01 +08:00
|
|
|
|
|
|
|
|
|
|
# 如果向量数据库中有数据但 metadata 中没有文件记录,可能存在文件缺失
|
|
|
|
|
|
if actual_count > 0 and metadata_files_count == 0:
|
2025-10-24 00:11:52 +08:00
|
|
|
|
inconsistencies["missing_files"].append(
|
|
|
|
|
|
{
|
2026-05-21 19:32:25 +08:00
|
|
|
|
"kb_id": kb_id,
|
2025-10-24 00:11:52 +08:00
|
|
|
|
"vector_count": actual_count,
|
|
|
|
|
|
"metadata_files_count": metadata_files_count,
|
|
|
|
|
|
"detected_at": utc_isoformat(),
|
|
|
|
|
|
}
|
|
|
|
|
|
)
|
|
|
|
|
|
logger.warning(
|
2026-05-21 19:32:25 +08:00
|
|
|
|
f"发现数据库 {kb_id} 在 Milvus 中有 {actual_count} 条向量数据,"
|
2025-10-24 00:11:52 +08:00
|
|
|
|
"但 metadata 中没有文件记录"
|
|
|
|
|
|
)
|
2025-10-23 22:57:01 +08:00
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
2026-05-21 19:32:25 +08:00
|
|
|
|
logger.debug(f"检查数据库 {kb_id} 的文件一致性时出错: {e}")
|
2025-10-23 22:57:01 +08:00
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
|
logger.error(f"检测 Milvus 数据不一致时出错: {e}")
|
|
|
|
|
|
|
|
|
|
|
|
return inconsistencies
|
|
|
|
|
|
|
|
|
|
|
|
def _log_inconsistencies(self, inconsistencies: dict) -> None:
|
|
|
|
|
|
"""将不一致检测结果输出到日志"""
|
|
|
|
|
|
total_missing_collections = inconsistencies["total_missing_collections"]
|
|
|
|
|
|
total_missing_files = inconsistencies["total_missing_files"]
|
|
|
|
|
|
|
|
|
|
|
|
if total_missing_collections == 0 and total_missing_files == 0:
|
|
|
|
|
|
logger.info("数据一致性检测完成,未发现不一致情况")
|
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
|
|
logger.warning("=" * 80)
|
|
|
|
|
|
logger.warning("数据一致性检测完成,发现以下不一致情况:")
|
|
|
|
|
|
logger.warning("=" * 80)
|
|
|
|
|
|
|
|
|
|
|
|
# Milvus 不一致情况
|
|
|
|
|
|
milvus_missing = inconsistencies["milvus"]["missing_collections"]
|
|
|
|
|
|
milvus_files_missing = inconsistencies["milvus"]["missing_files"]
|
|
|
|
|
|
if milvus_missing or milvus_files_missing:
|
2025-10-24 00:11:52 +08:00
|
|
|
|
logger.warning("Milvus 不一致情况:")
|
2025-10-23 22:57:01 +08:00
|
|
|
|
logger.warning(f" 缺失集合数量: {len(milvus_missing)}")
|
|
|
|
|
|
for collection_info in milvus_missing:
|
|
|
|
|
|
logger.warning(f" - 集合: {collection_info['collection_name']}, 实体数: {collection_info['count']}")
|
|
|
|
|
|
logger.warning(f" 缺失文件记录数量: {len(milvus_files_missing)}")
|
|
|
|
|
|
for file_info in milvus_files_missing:
|
2025-10-24 00:11:52 +08:00
|
|
|
|
logger.warning(
|
2026-05-21 19:32:25 +08:00
|
|
|
|
f" - 数据库: {file_info['kb_id']}, 向量数: {file_info['vector_count']}, "
|
2025-10-24 00:11:52 +08:00
|
|
|
|
f"元数据文件数: {file_info['metadata_files_count']}"
|
|
|
|
|
|
)
|
2025-10-23 22:57:01 +08:00
|
|
|
|
|
|
|
|
|
|
logger.warning("=" * 80)
|
|
|
|
|
|
logger.warning(f"总计:缺失集合 {total_missing_collections} 个,缺失文件记录 {total_missing_files} 个")
|
|
|
|
|
|
logger.warning("建议:检查这些不一致的数据,必要时进行数据清理或元数据修复")
|
|
|
|
|
|
logger.warning("=" * 80)
|
|
|
|
|
|
|
|
|
|
|
|
async def manual_consistency_check(self) -> dict:
|
|
|
|
|
|
"""
|
|
|
|
|
|
手动触发数据一致性检测
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
检测结果字典
|
|
|
|
|
|
"""
|
|
|
|
|
|
logger.info("手动触发数据一致性检测...")
|
|
|
|
|
|
return await self.detect_data_inconsistencies()
|