1257 lines
46 KiB
Python
1257 lines
46 KiB
Python
import asyncio
|
||
import os
|
||
from abc import ABC, abstractmethod
|
||
from typing import Any
|
||
|
||
from yuxi.knowledge.chunking.ragflow_like.presets import (
|
||
ensure_chunk_defaults_in_additional_params,
|
||
resolve_chunk_processing_params,
|
||
)
|
||
from yuxi.knowledge.utils import sanitize_processing_params
|
||
from yuxi.utils import logger
|
||
from yuxi.utils.datetime_utils import coerce_any_to_utc_datetime, utc_isoformat
|
||
|
||
|
||
class FileStatus:
|
||
UPLOADED = "uploaded"
|
||
PARSING = "parsing"
|
||
PARSED = "parsed"
|
||
ERROR_PARSING = "error_parsing"
|
||
INDEXING = "indexing"
|
||
INDEXED = "indexed"
|
||
ERROR_INDEXING = "error_indexing"
|
||
# Legacy status mapping
|
||
DONE = "done" # Map to INDEXED
|
||
FAILED = "failed" # Generic failure
|
||
|
||
|
||
class KnowledgeBaseException(Exception):
|
||
"""知识库统一异常基类"""
|
||
|
||
pass
|
||
|
||
|
||
class KBNotFoundError(KnowledgeBaseException):
|
||
"""知识库不存在错误"""
|
||
|
||
pass
|
||
|
||
|
||
class KBOperationError(KnowledgeBaseException):
|
||
"""知识库操作错误"""
|
||
|
||
pass
|
||
|
||
|
||
class KnowledgeBase(ABC):
|
||
"""知识库抽象基类，定义统一接口"""
|
||
|
||
# 类级别的处理队列，跟踪所有正在处理的文件
|
||
_processing_files = set()
|
||
_processing_lock = None
|
||
|
||
def __init__(self, work_dir: str):
|
||
"""
|
||
初始化知识库
|
||
|
||
Args:
|
||
work_dir: 工作目录
|
||
"""
|
||
import threading
|
||
|
||
self.work_dir = work_dir
|
||
self.databases_meta: dict[str, dict] = {}
|
||
self.files_meta: dict[str, dict] = {}
|
||
self.benchmarks_meta: dict[str, dict] = {}
|
||
self._metadata_loaded = False # 标记元数据是否已加载
|
||
|
||
# 初始化类级别的锁
|
||
if KnowledgeBase._processing_lock is None:
|
||
KnowledgeBase._processing_lock = threading.Lock()
|
||
|
||
os.makedirs(work_dir, exist_ok=True)
|
||
|
||
# 注意：不在 __init__ 中加载元数据，由 KnowledgeBaseManager 统一管理加载
|
||
|
||
def load_metadata(
|
||
self, global_databases_meta: dict[str, dict], files_meta: dict[str, dict], benchmarks_meta: dict[str, dict]
|
||
):
|
||
"""由 KnowledgeBaseManager 调用，同步加载元数据"""
|
||
# 过滤出当前 kb_type 的知识库
|
||
self.databases_meta = {}
|
||
for db_id, meta in global_databases_meta.items():
|
||
if meta.get("kb_type") == self.kb_type:
|
||
normalized_additional_params = ensure_chunk_defaults_in_additional_params(meta.get("additional_params"))
|
||
self.databases_meta[db_id] = {
|
||
"name": meta.get("name"),
|
||
"description": meta.get("description"),
|
||
"kb_type": meta.get("kb_type"),
|
||
"embed_info": meta.get("embed_info"),
|
||
"llm_info": meta.get("llm_info"),
|
||
"query_params": meta.get("query_params"),
|
||
"metadata": normalized_additional_params,
|
||
"created_at": meta.get("created_at"),
|
||
}
|
||
|
||
# 过滤文件
|
||
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 {}
|
||
normalized_meta = dict(meta)
|
||
normalized_meta["processing_params"] = resolve_chunk_processing_params(
|
||
kb_additional_params=kb_additional_params,
|
||
file_processing_params=meta.get("processing_params"),
|
||
)
|
||
self.files_meta[file_id] = normalized_meta
|
||
|
||
# 过滤评估基准
|
||
self.benchmarks_meta = {}
|
||
for kb_id, benchmarks in benchmarks_meta.items():
|
||
if kb_id in self.databases_meta:
|
||
self.benchmarks_meta[kb_id] = benchmarks
|
||
|
||
self._normalize_metadata_state()
|
||
self._metadata_loaded = True
|
||
logger.info(f"{self.kb_type}: 加载了 {len(self.databases_meta)} 个数据库的元数据")
|
||
|
||
def _ensure_metadata_loaded(self):
|
||
"""确保元数据已加载（延迟加载）"""
|
||
if not self._metadata_loaded:
|
||
logger.warning(f"{self.kb_type}: 元数据尚未加载，请确保 KnowledgeBaseManager 已调用 load_metadata()")
|
||
|
||
@staticmethod
|
||
def _normalize_timestamp(value: Any) -> str | None:
|
||
"""Convert persisted timestamps to a normalized UTC ISO string."""
|
||
try:
|
||
dt_value = coerce_any_to_utc_datetime(value)
|
||
except (TypeError, ValueError) as exc: # noqa: BLE001
|
||
logger.warning(f"Invalid timestamp encountered: {value!r} ({exc})")
|
||
return None
|
||
|
||
if not dt_value:
|
||
return None
|
||
return utc_isoformat(dt_value)
|
||
|
||
def _normalize_metadata_state(self) -> None:
|
||
"""Ensure in-memory metadata uses normalized timestamp formats."""
|
||
for meta in self.databases_meta.values():
|
||
if "created_at" in meta:
|
||
normalized = self._normalize_timestamp(meta.get("created_at"))
|
||
if normalized:
|
||
meta["created_at"] = normalized
|
||
|
||
for file_info in self.files_meta.values():
|
||
if "created_at" in file_info:
|
||
normalized = self._normalize_timestamp(file_info.get("created_at"))
|
||
if normalized:
|
||
file_info["created_at"] = normalized
|
||
|
||
for db_benchmarks in self.benchmarks_meta.values():
|
||
for b in db_benchmarks.values():
|
||
if "created_at" in b:
|
||
normalized = self._normalize_timestamp(b.get("created_at"))
|
||
if normalized:
|
||
b["created_at"] = normalized
|
||
if "updated_at" in b:
|
||
normalized = self._normalize_timestamp(b.get("updated_at"))
|
||
if normalized:
|
||
b["updated_at"] = normalized
|
||
|
||
@property
|
||
@abstractmethod
|
||
def kb_type(self) -> str:
|
||
"""知识库类型标识"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def _create_kb_instance(self, db_id: str, config: dict) -> Any:
|
||
"""
|
||
创建底层知识库实例
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
config: 配置信息
|
||
|
||
Returns:
|
||
底层知识库实例
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def _initialize_kb_instance(self, instance: Any) -> None:
|
||
"""
|
||
初始化底层知识库实例
|
||
|
||
Args:
|
||
instance: 底层知识库实例
|
||
"""
|
||
pass
|
||
|
||
async def add_file_record(
|
||
self, db_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
|
||
item: File path or URL
|
||
params: Parameters
|
||
operator_id: Operator ID who created the file
|
||
|
||
Returns:
|
||
File metadata record
|
||
"""
|
||
from yuxi.knowledge.utils.kb_utils import prepare_item_metadata
|
||
|
||
params = params or {}
|
||
content_type = params.get("content_type", "file")
|
||
|
||
# Prepare metadata
|
||
metadata = await prepare_item_metadata(item, content_type, db_id, params=params)
|
||
file_id = metadata["file_id"]
|
||
kb_additional_params = self.databases_meta.get(db_id, {}).get("metadata") or {}
|
||
metadata["processing_params"] = resolve_chunk_processing_params(
|
||
kb_additional_params=kb_additional_params,
|
||
file_processing_params=metadata.get("processing_params"),
|
||
)
|
||
|
||
# Initial status
|
||
metadata["status"] = FileStatus.UPLOADED
|
||
metadata["created_at"] = utc_isoformat()
|
||
if operator_id:
|
||
metadata["created_by"] = operator_id
|
||
|
||
# Save to metadata
|
||
self.files_meta[file_id] = metadata
|
||
await self._persist_file(file_id)
|
||
|
||
return metadata
|
||
|
||
async def parse_file(self, db_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
|
||
file_id: File ID
|
||
operator_id: ID of the user performing the operation
|
||
|
||
Returns:
|
||
Updated file metadata
|
||
"""
|
||
if file_id not in self.files_meta:
|
||
raise ValueError(f"File {file_id} not found")
|
||
|
||
file_meta = self.files_meta[file_id]
|
||
current_status = file_meta.get("status")
|
||
|
||
# Validate current status - only allow parsing from these states
|
||
allowed_statuses = {
|
||
FileStatus.UPLOADED,
|
||
FileStatus.ERROR_PARSING,
|
||
"failed", # Legacy status
|
||
}
|
||
|
||
if current_status not in allowed_statuses:
|
||
raise ValueError(
|
||
f"Cannot parse file with status '{current_status}'. "
|
||
f"File must be in one of these states: {', '.join(allowed_statuses)}"
|
||
)
|
||
|
||
file_path = file_meta.get("path")
|
||
if not file_path:
|
||
raise ValueError(f"File {file_id} has no valid path in metadata")
|
||
|
||
# Clear previous error if any
|
||
if "error" in file_meta:
|
||
self.files_meta[file_id].pop("error", None)
|
||
|
||
# Update status to PARSING and add to processing queue
|
||
self.files_meta[file_id]["status"] = FileStatus.PARSING
|
||
self.files_meta[file_id]["updated_at"] = utc_isoformat()
|
||
if operator_id:
|
||
self.files_meta[file_id]["updated_by"] = operator_id
|
||
await self._persist_file(file_id)
|
||
|
||
# Add to processing queue
|
||
self._add_to_processing_queue(file_id)
|
||
|
||
try:
|
||
from yuxi.plugins.parser.unified import Parser
|
||
|
||
# Prepare params
|
||
params = file_meta.get("processing_params", {}) or {}
|
||
params["image_bucket"] = "public"
|
||
params["image_prefix"] = f"{db_id}/kb-images"
|
||
|
||
markdown_content = await Parser.aparse(
|
||
source=file_path,
|
||
params=params,
|
||
)
|
||
|
||
# Save Markdown to MinIO
|
||
markdown_file_path = await self._save_markdown_to_minio(db_id, file_id, markdown_content)
|
||
|
||
# Update metadata
|
||
self.files_meta[file_id]["status"] = FileStatus.PARSED
|
||
self.files_meta[file_id]["markdown_file"] = markdown_file_path
|
||
self.files_meta[file_id]["updated_at"] = utc_isoformat()
|
||
if operator_id:
|
||
self.files_meta[file_id]["updated_by"] = operator_id
|
||
await self._persist_file(file_id)
|
||
|
||
return self.files_meta[file_id]
|
||
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
logger.error(f"Failed to parse file {file_id}: {error_msg}")
|
||
|
||
self.files_meta[file_id]["status"] = FileStatus.ERROR_PARSING
|
||
self.files_meta[file_id]["error"] = error_msg
|
||
self.files_meta[file_id]["updated_at"] = utc_isoformat()
|
||
if operator_id:
|
||
self.files_meta[file_id]["updated_by"] = operator_id
|
||
await self._persist_file(file_id)
|
||
|
||
raise
|
||
|
||
finally:
|
||
# 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:
|
||
"""Update file processing params"""
|
||
if file_id not in self.files_meta:
|
||
raise ValueError(f"File {file_id} not found")
|
||
|
||
# Skip if no params to update
|
||
if not params:
|
||
return
|
||
|
||
current_params = self.files_meta[file_id].get("processing_params", {}) or {}
|
||
kb_additional_params = self.databases_meta.get(db_id, {}).get("metadata") or {}
|
||
|
||
logger.debug(f"[update_file_params] file_id={file_id}, current_params={current_params}, new_params={params}")
|
||
|
||
current_params = resolve_chunk_processing_params(
|
||
kb_additional_params=kb_additional_params,
|
||
file_processing_params=current_params,
|
||
request_params=params,
|
||
)
|
||
|
||
self.files_meta[file_id]["processing_params"] = current_params
|
||
self.files_meta[file_id]["updated_at"] = utc_isoformat()
|
||
if operator_id:
|
||
self.files_meta[file_id]["updated_by"] = operator_id
|
||
|
||
logger.debug(f"[update_file_params] file_id={file_id}, updated_params={current_params}")
|
||
|
||
await self._persist_file(file_id)
|
||
|
||
async def _save_markdown_to_minio(self, db_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
|
||
|
||
minio_client = get_minio_client()
|
||
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"
|
||
data = content.encode("utf-8")
|
||
|
||
# Return standard HTTP URL from UploadResult
|
||
upload_result = await minio_client.aupload_file(
|
||
bucket_name=bucket_name,
|
||
object_name=object_name,
|
||
data=data,
|
||
)
|
||
|
||
return upload_result.url
|
||
|
||
async def _read_markdown_from_minio(self, file_path: str) -> str:
|
||
"""Read markdown content from MinIO"""
|
||
from yuxi.knowledge.utils.kb_utils import parse_minio_url
|
||
from yuxi.storage.minio import get_minio_client
|
||
|
||
if not file_path.startswith(("http://", "https://")):
|
||
raise ValueError(f"Invalid MinIO path format: {file_path}")
|
||
|
||
bucket_name, object_name = parse_minio_url(file_path)
|
||
minio_client = get_minio_client()
|
||
|
||
content_bytes = await minio_client.adownload_file(bucket_name, object_name)
|
||
return content_bytes.decode("utf-8")
|
||
|
||
@abstractmethod
|
||
async def index_file(self, db_id: str, file_id: str, operator_id: str | None = None) -> dict:
|
||
"""
|
||
Index parsed file (Status: INDEXING -> INDEXED/ERROR_INDEXING)
|
||
|
||
Args:
|
||
db_id: Database ID
|
||
file_id: File ID
|
||
operator_id: ID of the user performing the operation
|
||
|
||
Returns:
|
||
Updated file metadata
|
||
"""
|
||
pass
|
||
|
||
async def create_database(
|
||
self,
|
||
database_name: str,
|
||
description: str,
|
||
embed_info: dict | None = None,
|
||
llm_info: dict | None = None,
|
||
**kwargs,
|
||
) -> dict:
|
||
"""
|
||
创建数据库
|
||
|
||
Args:
|
||
database_name: 数据库名称
|
||
description: 数据库描述
|
||
embed_info: 嵌入模型信息
|
||
llm_info: LLM配置信息
|
||
**kwargs: 其他配置参数
|
||
|
||
Returns:
|
||
数据库信息字典
|
||
"""
|
||
from yuxi.utils import hashstr
|
||
|
||
kwargs = ensure_chunk_defaults_in_additional_params(kwargs)
|
||
|
||
# 从 kwargs 中获取 is_private 配置
|
||
is_private = kwargs.get("is_private", False)
|
||
prefix = "kb_private_" if is_private else "kb_"
|
||
db_id = f"{prefix}{hashstr(database_name, with_salt=True, length=32)}"
|
||
|
||
# 创建数据库记录
|
||
# 确保 Pydantic 模型被转换为字典，以便 JSON 序列化
|
||
embed_info_dump = embed_info.model_dump() if hasattr(embed_info, "model_dump") else embed_info
|
||
self.databases_meta[db_id] = {
|
||
"name": database_name,
|
||
"description": description,
|
||
"kb_type": self.kb_type,
|
||
"embed_info": embed_info_dump,
|
||
"llm_info": llm_info.model_dump() if hasattr(llm_info, "model_dump") else llm_info,
|
||
"metadata": kwargs,
|
||
"created_at": utc_isoformat(),
|
||
}
|
||
await self._persist_kb(db_id)
|
||
|
||
# 创建工作目录
|
||
working_dir = os.path.join(self.work_dir, db_id)
|
||
os.makedirs(working_dir, exist_ok=True)
|
||
|
||
# 返回数据库信息
|
||
db_dict = self.databases_meta[db_id].copy()
|
||
db_dict["db_id"] = db_id
|
||
db_dict["files"] = {}
|
||
|
||
return db_dict
|
||
|
||
async def delete_database(self, db_id: str) -> dict:
|
||
"""
|
||
删除数据库
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
|
||
Returns:
|
||
操作结果
|
||
"""
|
||
if db_id in self.databases_meta:
|
||
from yuxi.knowledge.utils.kb_utils import parse_minio_url
|
||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||
from yuxi.storage.minio import get_minio_client
|
||
|
||
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]
|
||
for file_id in files_to_delete:
|
||
file_path = self.files_meta[file_id].get("path")
|
||
if file_path and file_path.startswith(("http://", "https://")):
|
||
try:
|
||
bucket_name, object_name = parse_minio_url(file_path)
|
||
await minio_client.adelete_file(bucket_name, object_name)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to delete MinIO file {file_path}: {e}")
|
||
|
||
# 删除解析后的 markdown 文件
|
||
parsed_object = f"{db_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}/"
|
||
cleanup_buckets = {
|
||
minio_client.KB_BUCKETS["parsed"],
|
||
minio_client.KB_BUCKETS["documents"],
|
||
minio_client.KB_BUCKETS["images"],
|
||
}
|
||
cleanup_tasks = [
|
||
minio_client.adelete_objects_by_prefix(bucket_name, prefix) for bucket_name in cleanup_buckets
|
||
]
|
||
await asyncio.gather(*cleanup_tasks)
|
||
|
||
# 3. 删除数据库记录
|
||
del self.databases_meta[db_id]
|
||
kb_repo = KnowledgeBaseRepository()
|
||
await kb_repo.delete(db_id)
|
||
await self._save_metadata()
|
||
|
||
# 删除工作目录
|
||
working_dir = os.path.join(self.work_dir, db_id)
|
||
if os.path.exists(working_dir):
|
||
import shutil
|
||
|
||
try:
|
||
shutil.rmtree(working_dir)
|
||
except Exception as e:
|
||
logger.error(f"Error deleting working directory {working_dir}: {e}")
|
||
|
||
return {"message": "删除成功"}
|
||
|
||
async def create_folder(self, db_id: str, folder_name: str, parent_id: str | None = None) -> dict:
|
||
"""Create a folder in the database."""
|
||
import uuid
|
||
|
||
folder_id = f"folder-{uuid.uuid4()}"
|
||
|
||
self.files_meta[folder_id] = {
|
||
"file_id": folder_id,
|
||
"filename": folder_name,
|
||
"is_folder": True,
|
||
"parent_id": parent_id,
|
||
"database_id": db_id,
|
||
"created_at": utc_isoformat(),
|
||
"status": "done",
|
||
"path": folder_name,
|
||
"file_type": "folder",
|
||
}
|
||
await self._persist_file(folder_id)
|
||
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]:
|
||
"""
|
||
更新内容 - 根据file_ids重新解析文件并更新向量库
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
file_ids: 文件ID列表
|
||
params: 处理参数
|
||
|
||
Returns:
|
||
更新结果列表
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def aquery(self, query_text: str, db_id: str, **kwargs) -> list[dict]:
|
||
"""
|
||
异步查询知识库
|
||
|
||
Args:
|
||
query_text: 查询文本
|
||
db_id: 数据库ID
|
||
**kwargs: 查询参数
|
||
|
||
Returns:
|
||
一个包含字典的列表，每个字典代表一个检索到的文档块。
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
def get_query_params_config(self, db_id: str, **kwargs) -> dict:
|
||
"""
|
||
获取知识库类型的查询参数配置
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
**kwargs: 额外参数(如 reranker_names 等)
|
||
|
||
Returns:
|
||
dict: {
|
||
"type": "kb_type",
|
||
"options": [
|
||
{
|
||
"key": "param_name",
|
||
"label": "参数名称",
|
||
"type": "select|number|boolean",
|
||
"default": default_value,
|
||
"options": [...], # 对于 select 类型
|
||
"description": "参数描述",
|
||
"min": 1, # 对于 number 类型
|
||
"max": 100,
|
||
"step": 0.1
|
||
},
|
||
...
|
||
]
|
||
}
|
||
"""
|
||
pass
|
||
|
||
async def export_data(self, db_id: str, format: str = "zip", **kwargs) -> str:
|
||
pass
|
||
|
||
def query(self, query_text: str, db_id: str, **kwargs) -> list[dict]:
|
||
"""
|
||
同步查询知识库（兼容性方法）
|
||
|
||
Args:
|
||
query_text: 查询文本
|
||
db_id: 数据库ID
|
||
**kwargs: 查询参数
|
||
|
||
Returns:
|
||
一个包含字典的列表，每个字典代表一个检索到的文档块。
|
||
"""
|
||
import asyncio
|
||
|
||
logger.warning("query is deprecated, use aquery instead")
|
||
try:
|
||
loop = asyncio.get_running_loop()
|
||
return loop.run_until_complete(self.aquery(query_text, db_id, **kwargs))
|
||
except RuntimeError:
|
||
return asyncio.run(self.aquery(query_text, db_id, **kwargs))
|
||
|
||
def _get_query_params(self, db_id: str) -> dict:
|
||
"""从实例元数据中加载查询参数"""
|
||
if db_id in self.databases_meta:
|
||
query_params_meta = self.databases_meta[db_id].get("query_params") or {}
|
||
return query_params_meta.get("options", {})
|
||
return {}
|
||
|
||
def get_database_info(self, db_id: str, include_files: bool = True) -> dict | None:
|
||
"""
|
||
获取数据库详细信息
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
include_files: 是否包含文件信息，默认为True（保持向后兼容）
|
||
|
||
Returns:
|
||
数据库信息或None
|
||
"""
|
||
if db_id not in self.databases_meta:
|
||
return None
|
||
|
||
meta = self.databases_meta[db_id].copy()
|
||
meta["db_id"] = db_id
|
||
|
||
# 检查并修复异常的processing状态
|
||
self._check_and_fix_processing_status(db_id)
|
||
|
||
# 统计文件数量（始终计算，即使不加载文件详情）
|
||
db_file_count = sum(1 for file_info in self.files_meta.values() if file_info.get("database_id") == db_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:
|
||
created_at = self._normalize_timestamp(file_info.get("created_at"))
|
||
db_files[file_id] = {
|
||
"file_id": file_id,
|
||
"filename": file_info.get("filename", ""),
|
||
"path": file_info.get("path", ""),
|
||
"markdown_file": file_info.get("markdown_file", ""),
|
||
"type": file_info.get("file_type", ""),
|
||
"status": file_info.get("status", "done"),
|
||
"created_at": created_at,
|
||
"is_folder": file_info.get("is_folder", False),
|
||
"parent_id": file_info.get("parent_id", None),
|
||
}
|
||
|
||
# 按创建时间倒序排序文件列表
|
||
sorted_files = dict(
|
||
sorted(
|
||
db_files.items(),
|
||
key=lambda item: item[1].get("created_at") or "",
|
||
reverse=True,
|
||
)
|
||
)
|
||
|
||
meta["files"] = sorted_files
|
||
|
||
meta["status"] = "已连接"
|
||
return meta
|
||
|
||
def get_databases(self, include_files: bool = False) -> dict:
|
||
"""
|
||
获取所有数据库信息
|
||
|
||
Args:
|
||
include_files: 是否包含文件信息，默认False以减少响应大小
|
||
|
||
Returns:
|
||
数据库列表
|
||
"""
|
||
# 确保元数据已加载（延迟加载机制）
|
||
self._ensure_metadata_loaded()
|
||
|
||
databases = []
|
||
for db_id, meta in self.databases_meta.items():
|
||
# 检查并修复异常的processing状态
|
||
self._check_and_fix_processing_status(db_id)
|
||
|
||
db_dict = meta.copy()
|
||
db_dict["db_id"] = db_id
|
||
|
||
# 统计文件数量（始终计算，即使不加载文件详情）
|
||
db_file_count = sum(1 for file_info in self.files_meta.values() if file_info.get("database_id") == db_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:
|
||
created_at = self._normalize_timestamp(file_info.get("created_at"))
|
||
db_files[file_id] = {
|
||
"file_id": file_id,
|
||
"filename": file_info.get("filename", ""),
|
||
"path": file_info.get("path", ""),
|
||
"markdown_file": file_info.get("markdown_file", ""),
|
||
"type": file_info.get("file_type", ""),
|
||
"status": file_info.get("status", "done"),
|
||
"created_at": created_at,
|
||
"is_folder": file_info.get("is_folder", False),
|
||
"parent_id": file_info.get("parent_id", None),
|
||
}
|
||
|
||
# 按创建时间倒序排序文件列表
|
||
sorted_files = dict(
|
||
sorted(
|
||
db_files.items(),
|
||
key=lambda item: item[1].get("created_at") or "",
|
||
reverse=True,
|
||
)
|
||
)
|
||
|
||
db_dict["files"] = sorted_files
|
||
|
||
db_dict["status"] = "已连接"
|
||
databases.append(db_dict)
|
||
|
||
return {"databases": databases}
|
||
|
||
@classmethod
|
||
def _add_to_processing_queue(cls, file_id: str) -> None:
|
||
"""
|
||
将文件添加到处理队列
|
||
|
||
Args:
|
||
file_id: 文件ID
|
||
"""
|
||
with cls._processing_lock:
|
||
cls._processing_files.add(file_id)
|
||
logger.debug(f"Added file {file_id} to processing queue")
|
||
|
||
@classmethod
|
||
def _remove_from_processing_queue(cls, file_id: str) -> None:
|
||
"""
|
||
从处理队列中移除文件
|
||
|
||
Args:
|
||
file_id: 文件ID
|
||
"""
|
||
with cls._processing_lock:
|
||
cls._processing_files.discard(file_id)
|
||
logger.debug(f"Removed file {file_id} from processing queue")
|
||
|
||
@classmethod
|
||
def _is_file_in_processing_queue(cls, file_id: str) -> bool:
|
||
"""
|
||
检查文件是否在处理队列中
|
||
|
||
Args:
|
||
file_id: 文件ID
|
||
|
||
Returns:
|
||
bool: 文件是否在处理队列中
|
||
"""
|
||
with cls._processing_lock:
|
||
return file_id in cls._processing_files
|
||
|
||
def _check_and_fix_processing_status(self, db_id: str) -> None:
|
||
"""
|
||
检查并修复异常的处理中状态
|
||
如果文件状态为处理中但实际不在处理队列中，则修改为相应的错误状态
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
"""
|
||
try:
|
||
status_changed = False
|
||
|
||
# 定义需要检查的中间状态及其对应的错误状态
|
||
intermediate_states = {
|
||
FileStatus.PARSING: FileStatus.ERROR_PARSING,
|
||
FileStatus.INDEXING: FileStatus.ERROR_INDEXING,
|
||
"processing": "failed", # 兼容旧状态
|
||
}
|
||
|
||
# 检查该数据库下所有中间状态的文件
|
||
for file_id, file_info in self.files_meta.items():
|
||
if file_info.get("database_id") == db_id:
|
||
current_status = file_info.get("status")
|
||
|
||
if current_status in intermediate_states:
|
||
# 检查文件是否真的在处理队列中
|
||
if not self._is_file_in_processing_queue(file_id):
|
||
error_status = intermediate_states[current_status]
|
||
logger.warning(
|
||
f"File {file_id} has {current_status} status but is not in processing queue, "
|
||
f"marking as {error_status}"
|
||
)
|
||
self.files_meta[file_id]["status"] = error_status
|
||
self.files_meta[file_id]["error"] = (
|
||
f"{current_status.capitalize()} interrupted - process not found in queue"
|
||
)
|
||
self.files_meta[file_id]["updated_at"] = utc_isoformat()
|
||
status_changed = True
|
||
|
||
# 如果有状态变更，保存元数据
|
||
if status_changed:
|
||
logger.info(f"Fixed interrupted processing status for database {db_id}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error checking processing status for database {db_id}: {e}")
|
||
|
||
async def delete_folder(self, db_id: str, folder_id: str) -> None:
|
||
"""
|
||
Recursively delete a folder and its content.
|
||
|
||
Args:
|
||
db_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
|
||
]
|
||
|
||
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)
|
||
else:
|
||
await self.delete_file(db_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)
|
||
|
||
async def move_file(self, db_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
|
||
file_id: File/Folder ID to move
|
||
new_parent_id: New parent folder ID (None for root)
|
||
|
||
Returns:
|
||
dict: Updated metadata
|
||
"""
|
||
if file_id not in self.files_meta:
|
||
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}")
|
||
|
||
# Basic cycle detection for folders
|
||
if meta.get("is_folder") and new_parent_id:
|
||
# Check if new_parent_id is a child of file_id (or is file_id itself)
|
||
if new_parent_id == file_id:
|
||
raise ValueError("Cannot move a folder into itself")
|
||
|
||
# Walk up the tree from new_parent_id
|
||
current = new_parent_id
|
||
while current:
|
||
parent_meta = self.files_meta.get(current)
|
||
if not parent_meta:
|
||
break # Should not happen if integrity is maintained
|
||
if current == file_id:
|
||
raise ValueError("Cannot move a folder into its own subfolder")
|
||
current = parent_meta.get("parent_id")
|
||
|
||
meta["parent_id"] = new_parent_id
|
||
await self._persist_file(file_id)
|
||
return meta
|
||
|
||
@abstractmethod
|
||
async def delete_file(self, db_id: str, file_id: str) -> None:
|
||
"""
|
||
删除文件
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
file_id: 文件ID
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def get_file_basic_info(self, db_id: str, file_id: str) -> dict:
|
||
"""
|
||
获取文件基本信息（仅元数据）
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
file_id: 文件ID
|
||
|
||
Returns:
|
||
dict: 包含文件基本信息的字典
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def get_file_content(self, db_id: str, file_id: str) -> dict:
|
||
"""
|
||
获取文件内容信息（chunks和lines）
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
file_id: 文件ID
|
||
|
||
Returns:
|
||
dict: 包含文件内容信息的字典
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def get_file_info(self, db_id: str, file_id: str) -> dict:
|
||
"""
|
||
获取文件完整信息（基本信息+内容信息）- 保持向后兼容
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
file_id: 文件ID
|
||
|
||
Returns:
|
||
dict: 包含文件信息和chunks的字典
|
||
"""
|
||
pass
|
||
|
||
def get_db_upload_path(self, db_id: str | None = None) -> str:
|
||
"""
|
||
获取数据库上传路径
|
||
|
||
Args:
|
||
db_id: 数据库ID，可选
|
||
|
||
Returns:
|
||
上传路径
|
||
"""
|
||
if db_id:
|
||
uploads_folder = os.path.join(self.work_dir, db_id, "uploads")
|
||
os.makedirs(uploads_folder, exist_ok=True)
|
||
return uploads_folder
|
||
|
||
general_uploads = os.path.join(self.work_dir, "uploads")
|
||
os.makedirs(general_uploads, exist_ok=True)
|
||
return general_uploads
|
||
|
||
def update_database(self, db_id: str, name: str, description: str, llm_info: dict = None) -> dict:
|
||
"""
|
||
更新数据库
|
||
|
||
Args:
|
||
db_id: 数据库ID
|
||
name: 新名称
|
||
description: 新描述
|
||
llm_info: LLM配置信息（可选，仅用于 LightRAG 类型知识库）
|
||
|
||
Returns:
|
||
更新后的数据库信息
|
||
"""
|
||
if db_id not in self.databases_meta:
|
||
raise ValueError(f"数据库 {db_id} 不存在")
|
||
|
||
self.databases_meta[db_id]["name"] = name
|
||
self.databases_meta[db_id]["description"] = description
|
||
|
||
# 如果提供了 llm_info，则更新（仅针对 LightRAG 类型）
|
||
if llm_info is not None:
|
||
self.databases_meta[db_id]["llm_info"] = llm_info
|
||
|
||
asyncio.create_task(self._save_metadata())
|
||
|
||
return self.get_database_info(db_id)
|
||
|
||
def get_retrievers(self) -> dict[str, dict]:
|
||
"""
|
||
获取所有检索器
|
||
|
||
Returns:
|
||
检索器字典
|
||
"""
|
||
retrievers = {}
|
||
for db_id, meta in self.databases_meta.items():
|
||
|
||
def make_retriever(db_id):
|
||
async def retriever(query_text, **kwargs):
|
||
return await self.aquery(query_text, db_id, agent_call=True, **kwargs)
|
||
|
||
return retriever
|
||
|
||
retrievers[db_id] = {
|
||
"name": meta["name"],
|
||
"description": meta["description"],
|
||
"retriever": make_retriever(db_id),
|
||
"metadata": meta,
|
||
}
|
||
return retrievers
|
||
|
||
async def _load_metadata(self) -> None:
|
||
from yuxi.repositories.evaluation_repository import EvaluationRepository
|
||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
kb_repo = KnowledgeBaseRepository()
|
||
file_repo = KnowledgeFileRepository()
|
||
eval_repo = EvaluationRepository()
|
||
|
||
databases = [kb for kb in await kb_repo.get_all() if kb.kb_type == self.kb_type]
|
||
self.databases_meta = {
|
||
kb.db_id: {
|
||
"name": kb.name,
|
||
"description": kb.description,
|
||
"kb_type": kb.kb_type,
|
||
"embed_info": kb.embed_info,
|
||
"llm_info": kb.llm_info,
|
||
"query_params": kb.query_params,
|
||
"metadata": ensure_chunk_defaults_in_additional_params(kb.additional_params),
|
||
"created_at": utc_isoformat(kb.created_at) if kb.created_at else utc_isoformat(),
|
||
}
|
||
for kb in databases
|
||
}
|
||
|
||
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):
|
||
self.files_meta[record.file_id] = {
|
||
"file_id": record.file_id,
|
||
"database_id": record.db_id,
|
||
"parent_id": record.parent_id,
|
||
"filename": record.filename,
|
||
"file_type": record.file_type,
|
||
"path": record.path,
|
||
"markdown_file": record.markdown_file,
|
||
"status": record.status,
|
||
"content_hash": record.content_hash,
|
||
"size": record.file_size,
|
||
"content_type": record.content_type,
|
||
"processing_params": sanitize_processing_params(
|
||
resolve_chunk_processing_params(
|
||
kb_additional_params=kb_additional_params,
|
||
file_processing_params=record.processing_params,
|
||
)
|
||
),
|
||
"is_folder": record.is_folder,
|
||
"error": record.error_message,
|
||
"created_by": record.created_by,
|
||
"updated_by": record.updated_by,
|
||
"created_at": utc_isoformat(record.created_at) if record.created_at else None,
|
||
"updated_at": utc_isoformat(record.updated_at) if record.updated_at else None,
|
||
"original_filename": record.original_filename,
|
||
"minio_url": record.minio_url,
|
||
}
|
||
|
||
self.benchmarks_meta = {}
|
||
for kb in databases:
|
||
benchmarks = await eval_repo.list_benchmarks(kb.db_id)
|
||
if not benchmarks:
|
||
continue
|
||
self.benchmarks_meta[kb.db_id] = {}
|
||
for bench in benchmarks:
|
||
self.benchmarks_meta[kb.db_id][bench.benchmark_id] = {
|
||
"id": bench.benchmark_id,
|
||
"benchmark_id": bench.benchmark_id,
|
||
"name": bench.name,
|
||
"description": bench.description,
|
||
"db_id": bench.db_id,
|
||
"question_count": bench.question_count,
|
||
"has_gold_chunks": bench.has_gold_chunks,
|
||
"has_gold_answers": bench.has_gold_answers,
|
||
"benchmark_file": bench.data_file_path,
|
||
"created_by": bench.created_by,
|
||
"created_at": utc_isoformat(bench.created_at) if bench.created_at else None,
|
||
"updated_at": utc_isoformat(bench.updated_at) if bench.updated_at else None,
|
||
}
|
||
|
||
logger.info(f"Loaded {self.kb_type} metadata from database for {len(self.databases_meta)} databases")
|
||
|
||
async def _save_metadata(self) -> None:
|
||
from yuxi.repositories.evaluation_repository import EvaluationRepository
|
||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
kb_repo = KnowledgeBaseRepository()
|
||
file_repo = KnowledgeFileRepository()
|
||
eval_repo = EvaluationRepository()
|
||
|
||
self._normalize_metadata_state()
|
||
|
||
for db_id, meta in self.databases_meta.items():
|
||
existing = await kb_repo.get_by_id(db_id)
|
||
payload = {
|
||
"db_id": db_id,
|
||
"name": meta.get("name") or db_id,
|
||
"description": meta.get("description"),
|
||
"kb_type": meta.get("kb_type") or self.kb_type,
|
||
"embed_info": meta.get("embed_info"),
|
||
"llm_info": meta.get("llm_info"),
|
||
"query_params": meta.get("query_params"),
|
||
"additional_params": meta.get("metadata") or {},
|
||
}
|
||
if existing is None:
|
||
await kb_repo.create(payload)
|
||
else:
|
||
await kb_repo.update(
|
||
db_id,
|
||
{
|
||
"name": payload["name"],
|
||
"description": payload["description"],
|
||
"kb_type": payload["kb_type"],
|
||
"embed_info": payload["embed_info"],
|
||
"llm_info": payload["llm_info"],
|
||
"query_params": payload["query_params"],
|
||
"additional_params": payload["additional_params"],
|
||
},
|
||
)
|
||
|
||
for file_id, meta in self.files_meta.items():
|
||
db_id = meta.get("database_id")
|
||
if not db_id:
|
||
continue
|
||
await file_repo.upsert(
|
||
file_id=file_id,
|
||
data={
|
||
"db_id": db_id,
|
||
"parent_id": meta.get("parent_id"),
|
||
"filename": meta.get("filename") or "",
|
||
"original_filename": meta.get("original_filename"),
|
||
"file_type": meta.get("file_type"),
|
||
"path": meta.get("path"),
|
||
"minio_url": meta.get("minio_url"),
|
||
"markdown_file": meta.get("markdown_file"),
|
||
"status": meta.get("status"),
|
||
"content_hash": meta.get("content_hash"),
|
||
"file_size": meta.get("size"),
|
||
"content_type": meta.get("content_type"),
|
||
"processing_params": sanitize_processing_params(meta.get("processing_params")),
|
||
"is_folder": meta.get("is_folder", False),
|
||
"error_message": meta.get("error"),
|
||
"created_by": str(meta.get("created_by")) if meta.get("created_by") else None,
|
||
"updated_by": str(meta.get("updated_by")) if meta.get("updated_by") else None,
|
||
},
|
||
)
|
||
|
||
for db_id, benchmarks in self.benchmarks_meta.items():
|
||
for benchmark_id, meta in benchmarks.items():
|
||
existing = await eval_repo.get_benchmark(benchmark_id)
|
||
payload = {
|
||
"benchmark_id": benchmark_id,
|
||
"db_id": db_id,
|
||
"name": meta.get("name") or benchmark_id,
|
||
"description": meta.get("description"),
|
||
"question_count": int(meta.get("question_count") or 0),
|
||
"has_gold_chunks": bool(meta.get("has_gold_chunks")),
|
||
"has_gold_answers": bool(meta.get("has_gold_answers")),
|
||
"data_file_path": meta.get("benchmark_file"),
|
||
"created_by": str(meta.get("created_by")) if meta.get("created_by") else None,
|
||
}
|
||
if existing is None:
|
||
await eval_repo.create_benchmark(payload)
|
||
|
||
async def _persist_file(self, file_id: str) -> None:
|
||
"""只保存单个文件到数据库，避免全量遍历"""
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
file_repo = KnowledgeFileRepository()
|
||
|
||
if file_id not in self.files_meta:
|
||
return
|
||
|
||
meta = self.files_meta[file_id]
|
||
db_id = meta.get("database_id")
|
||
if not db_id:
|
||
return
|
||
|
||
await file_repo.upsert(
|
||
file_id=file_id,
|
||
data={
|
||
"db_id": db_id,
|
||
"parent_id": meta.get("parent_id"),
|
||
"filename": meta.get("filename") or "",
|
||
"original_filename": meta.get("original_filename"),
|
||
"file_type": meta.get("file_type"),
|
||
"path": meta.get("path"),
|
||
"minio_url": meta.get("minio_url"),
|
||
"markdown_file": meta.get("markdown_file"),
|
||
"status": meta.get("status"),
|
||
"content_hash": meta.get("content_hash"),
|
||
"file_size": meta.get("size"),
|
||
"content_type": meta.get("content_type"),
|
||
"processing_params": sanitize_processing_params(meta.get("processing_params")),
|
||
"is_folder": meta.get("is_folder", False),
|
||
"error_message": meta.get("error"),
|
||
"created_by": str(meta.get("created_by")) if meta.get("created_by") else None,
|
||
"updated_by": str(meta.get("updated_by")) if meta.get("updated_by") else None,
|
||
},
|
||
)
|
||
|
||
async def _persist_kb(self, db_id: str) -> None:
|
||
"""只保存单个知识库到数据库，避免全量遍历"""
|
||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||
|
||
kb_repo = KnowledgeBaseRepository()
|
||
|
||
if db_id not in self.databases_meta:
|
||
return
|
||
|
||
meta = self.databases_meta[db_id]
|
||
existing = await kb_repo.get_by_id(db_id)
|
||
payload = {
|
||
"db_id": db_id,
|
||
"name": meta.get("name") or db_id,
|
||
"description": meta.get("description"),
|
||
"kb_type": meta.get("kb_type") or self.kb_type,
|
||
"embed_info": meta.get("embed_info"),
|
||
"llm_info": meta.get("llm_info"),
|
||
"query_params": meta.get("query_params"),
|
||
"additional_params": meta.get("metadata") or {},
|
||
}
|
||
|
||
if existing is None:
|
||
await kb_repo.create(payload)
|
||
else:
|
||
await kb_repo.update(
|
||
db_id,
|
||
{
|
||
"name": payload["name"],
|
||
"description": payload["description"],
|
||
"kb_type": payload["kb_type"],
|
||
"embed_info": payload["embed_info"],
|
||
"llm_info": payload["llm_info"],
|
||
"query_params": payload["query_params"],
|
||
"additional_params": payload["additional_params"],
|
||
},
|
||
)
|