wehub-resource-sync
1654 lines
62 KiB
Python
1654 lines
62 KiB
Python
import asyncio
|
||
import mimetypes
|
||
import os
|
||
import re
|
||
import secrets
|
||
import string
|
||
from abc import ABC, abstractmethod
|
||
from typing import Any
|
||
|
||
from yuxi.knowledge.chunking.ragflow_like.presets import ensure_chunk_defaults_in_additional_params
|
||
from yuxi.knowledge.schemas import (
|
||
FindOutputSchema,
|
||
FindWindowSchema,
|
||
SearchOutputSchema,
|
||
SearchResultSchema,
|
||
)
|
||
from yuxi.knowledge.utils import resolve_processing_params, sanitize_processing_params
|
||
from yuxi.services.file_preview import (
|
||
MAX_BINARY_PREVIEW_SIZE_BYTES,
|
||
OfficePreviewConversionError,
|
||
convert_office_to_pdf,
|
||
detect_media_type,
|
||
is_binary_preview_type,
|
||
is_office_pdf_preview_file,
|
||
render_preview_payload,
|
||
render_preview_too_large_payload,
|
||
)
|
||
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"
|
||
|
||
|
||
INDEXED_STATS_STATUSES = {FileStatus.INDEXED, "done"}
|
||
|
||
|
||
def _should_repair_file_stats(file_meta: dict) -> bool:
|
||
status = file_meta.get("status")
|
||
return status is None or status in INDEXED_STATS_STATUSES
|
||
|
||
|
||
class KnowledgeBaseException(Exception):
|
||
"""知识库统一异常基类"""
|
||
|
||
pass
|
||
|
||
|
||
class KBNotFoundError(KnowledgeBaseException):
|
||
"""知识库不存在错误"""
|
||
|
||
pass
|
||
|
||
|
||
class KBOperationError(KnowledgeBaseException):
|
||
"""知识库操作错误"""
|
||
|
||
pass
|
||
|
||
|
||
class KnowledgeBase(ABC):
|
||
"""知识库抽象基类，定义统一接口"""
|
||
|
||
kb_type = ""
|
||
name = ""
|
||
description = ""
|
||
requires_embedding_model = True
|
||
supports_documents = True
|
||
apply_chunk_defaults = True
|
||
|
||
def __init__(self, work_dir: str):
|
||
"""
|
||
初始化知识库
|
||
|
||
Args:
|
||
work_dir: 工作目录
|
||
"""
|
||
self.work_dir = work_dir
|
||
self.databases_meta: dict[str, dict] = {}
|
||
self.benchmarks_meta: dict[str, dict] = {}
|
||
self._metadata_loaded = False # 标记元数据是否已加载
|
||
|
||
os.makedirs(work_dir, exist_ok=True)
|
||
|
||
# 注意：不在 __init__ 中加载元数据，由 KnowledgeBaseManager 统一管理加载
|
||
|
||
def load_metadata(
|
||
self,
|
||
global_databases_meta: dict[str, dict],
|
||
_unused_file_metadata: dict[str, dict],
|
||
benchmarks_meta: dict[str, dict],
|
||
):
|
||
"""由 KnowledgeBaseManager 调用，同步加载元数据"""
|
||
# 过滤出当前 kb_type 的知识库
|
||
self.databases_meta = {}
|
||
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[kb_id] = {
|
||
"name": meta.get("name"),
|
||
"description": meta.get("description"),
|
||
"kb_type": meta.get("kb_type"),
|
||
"embedding_model_spec": meta.get("embedding_model_spec"),
|
||
"llm_model_spec": meta.get("llm_model_spec"),
|
||
"query_params": meta.get("query_params"),
|
||
"metadata": normalized_additional_params,
|
||
"created_at": meta.get("created_at"),
|
||
}
|
||
|
||
del _unused_file_metadata
|
||
|
||
# 文件元数据以 PostgreSQL 为准，不在 KnowledgeBase 实例中缓存。
|
||
|
||
# 过滤评估基准
|
||
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 _file_record_to_meta(self, record: Any) -> dict:
|
||
kb_additional_params = self.databases_meta.get(record.kb_id, {}).get("metadata") or {}
|
||
return {
|
||
"file_id": record.file_id,
|
||
"kb_id": record.kb_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,
|
||
"chunk_count": int(getattr(record, "chunk_count", 0) or 0),
|
||
"token_count": int(getattr(record, "token_count", 0) or 0),
|
||
"content_type": record.content_type,
|
||
"processing_params": sanitize_processing_params(
|
||
resolve_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,
|
||
}
|
||
|
||
@staticmethod
|
||
def _file_meta_to_record_data(meta: dict) -> dict[str, Any]:
|
||
return {
|
||
"kb_id": meta.get("kb_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"),
|
||
"chunk_count": int(meta.get("chunk_count") or 0),
|
||
"token_count": int(meta.get("token_count") or 0),
|
||
"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 _load_file_meta(self, kb_id: str, file_id: str, *, refresh: bool = False) -> dict:
|
||
del refresh
|
||
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
record = await KnowledgeFileRepository().get_by_file_id(file_id)
|
||
if record is None or record.kb_id != kb_id:
|
||
raise ValueError(f"File {file_id} not found")
|
||
|
||
return self._file_record_to_meta(record)
|
||
|
||
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 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
|
||
|
||
@classmethod
|
||
def get_create_params_config(cls) -> dict[str, Any]:
|
||
"""获取创建知识库时的类型特定参数配置。"""
|
||
return {"options": []}
|
||
|
||
@classmethod
|
||
def validate_additional_params(cls, additional_params: dict | None) -> dict:
|
||
"""校验并规范化类型特定配置。"""
|
||
return dict(additional_params or {})
|
||
|
||
@classmethod
|
||
def normalize_additional_params(cls, additional_params: dict | None) -> dict:
|
||
"""规范化 additional_params，仅文档型知识库补充分块默认值。"""
|
||
params = cls.validate_additional_params(additional_params)
|
||
if cls.apply_chunk_defaults:
|
||
return ensure_chunk_defaults_in_additional_params(params)
|
||
return params
|
||
|
||
@abstractmethod
|
||
async def _create_kb_instance(self, kb_id: str, config: dict) -> Any:
|
||
"""
|
||
创建底层知识库实例
|
||
|
||
Args:
|
||
kb_id: 数据库ID
|
||
config: 配置信息
|
||
|
||
Returns:
|
||
底层知识库实例
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def _initialize_kb_instance(self, instance: Any) -> None:
|
||
"""
|
||
初始化底层知识库实例
|
||
|
||
Args:
|
||
instance: 底层知识库实例
|
||
"""
|
||
pass
|
||
|
||
async def add_file_record(
|
||
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:
|
||
kb_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, kb_id, params=params)
|
||
file_id = metadata["file_id"]
|
||
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"),
|
||
)
|
||
|
||
# Fallback: fetch file size from MinIO if not provided
|
||
if metadata.get("size") is None and content_type == "file":
|
||
try:
|
||
from yuxi.knowledge.utils.kb_utils import is_minio_url, parse_minio_url
|
||
from yuxi.storage.minio import get_minio_client
|
||
|
||
file_path = metadata.get("path") or item
|
||
if is_minio_url(file_path):
|
||
bucket_name, obj_name = parse_minio_url(file_path)
|
||
minio_client = get_minio_client()
|
||
file_size = await minio_client.astat_file(bucket_name, obj_name)
|
||
if file_size is not None:
|
||
metadata["size"] = file_size
|
||
except Exception as exc:
|
||
logger.warning(f"Failed to stat file size from MinIO for {item}: {exc}")
|
||
|
||
# Initial status
|
||
metadata["status"] = FileStatus.UPLOADED
|
||
metadata["created_at"] = utc_isoformat()
|
||
if operator_id:
|
||
metadata["created_by"] = operator_id
|
||
|
||
await self._persist_file_meta(file_id, metadata)
|
||
await self.refresh_database_stats(kb_id)
|
||
|
||
return metadata
|
||
|
||
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:
|
||
kb_id: Database ID
|
||
file_id: File ID
|
||
operator_id: ID of the user performing the operation
|
||
|
||
Returns:
|
||
Updated file metadata
|
||
"""
|
||
# Validate current status - only allow parsing from these states
|
||
allowed_statuses = {
|
||
FileStatus.UPLOADED,
|
||
FileStatus.ERROR_PARSING,
|
||
"failed", # Legacy status
|
||
}
|
||
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
file_repo = KnowledgeFileRepository()
|
||
claim_data = {"status": FileStatus.PARSING, "error_message": None}
|
||
if operator_id:
|
||
claim_data["updated_by"] = operator_id
|
||
claimed_record = await file_repo.update_fields_if_status(
|
||
kb_id=kb_id,
|
||
file_id=file_id,
|
||
allowed_statuses=allowed_statuses,
|
||
data=claim_data,
|
||
)
|
||
if claimed_record is None:
|
||
current_meta = await self._load_file_meta(kb_id, file_id)
|
||
current_status = current_meta.get("status")
|
||
raise ValueError(
|
||
f"Cannot parse file with status '{current_status}'. "
|
||
f"File must be in one of these states: {', '.join(allowed_statuses)}"
|
||
)
|
||
|
||
file_meta = self._file_record_to_meta(claimed_record)
|
||
file_path = file_meta.get("path")
|
||
if not file_path:
|
||
message = f"File {file_id} has no valid path in metadata"
|
||
update_data = {"status": FileStatus.ERROR_PARSING, "error_message": message}
|
||
if operator_id:
|
||
update_data["updated_by"] = operator_id
|
||
await file_repo.update_fields(file_id=file_id, kb_id=kb_id, data=update_data)
|
||
raise ValueError(message)
|
||
|
||
try:
|
||
from yuxi.knowledge.parser.unified import Parser
|
||
|
||
# Prepare params
|
||
params = file_meta.get("processing_params", {}) or {}
|
||
params["image_bucket"] = "public"
|
||
params["image_prefix"] = f"{kb_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(kb_id, file_id, markdown_content)
|
||
|
||
# Update metadata
|
||
file_meta["status"] = FileStatus.PARSED
|
||
file_meta["markdown_file"] = markdown_file_path
|
||
file_meta["error"] = None
|
||
file_meta["updated_at"] = utc_isoformat()
|
||
if operator_id:
|
||
file_meta["updated_by"] = operator_id
|
||
update_data = {
|
||
"status": FileStatus.PARSED,
|
||
"markdown_file": markdown_file_path,
|
||
"error_message": None,
|
||
}
|
||
if operator_id:
|
||
update_data["updated_by"] = operator_id
|
||
await file_repo.update_fields(file_id=file_id, kb_id=kb_id, data=update_data)
|
||
|
||
return file_meta
|
||
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
logger.error(f"Failed to parse file {file_id}: {error_msg}")
|
||
|
||
file_meta["status"] = FileStatus.ERROR_PARSING
|
||
file_meta["error"] = error_msg
|
||
file_meta["updated_at"] = utc_isoformat()
|
||
if operator_id:
|
||
file_meta["updated_by"] = operator_id
|
||
update_data = {"status": FileStatus.ERROR_PARSING, "error_message": error_msg}
|
||
if operator_id:
|
||
update_data["updated_by"] = operator_id
|
||
await file_repo.update_fields(file_id=file_id, kb_id=kb_id, data=update_data)
|
||
|
||
raise
|
||
|
||
async def update_file_params(self, kb_id: str, file_id: str, params: dict, operator_id: str | None = None) -> None:
|
||
"""Update file processing params"""
|
||
# Skip if no params to update
|
||
if not params:
|
||
return
|
||
|
||
file_meta = await self._load_file_meta(kb_id, file_id)
|
||
current_params = file_meta.get("processing_params", {}) 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}")
|
||
|
||
current_params = resolve_processing_params(
|
||
kb_additional_params=kb_additional_params,
|
||
file_processing_params=current_params,
|
||
request_params=params,
|
||
)
|
||
|
||
file_meta["processing_params"] = current_params
|
||
file_meta["updated_at"] = utc_isoformat()
|
||
if operator_id:
|
||
file_meta["updated_by"] = operator_id
|
||
|
||
logger.debug(f"[update_file_params] file_id={file_id}, updated_params={current_params}")
|
||
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
update_data = {"processing_params": sanitize_processing_params(current_params)}
|
||
if operator_id:
|
||
update_data["updated_by"] = operator_id
|
||
record = await KnowledgeFileRepository().update_fields(file_id=file_id, kb_id=kb_id, data=update_data)
|
||
if record is None:
|
||
raise ValueError(f"File {file_id} not found")
|
||
|
||
async def _mark_file_unparsed(self, kb_id: str, file_id: str, operator_id: str | None = None) -> None:
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
update_data = {"status": FileStatus.UPLOADED, "markdown_file": None, "error_message": None}
|
||
if operator_id:
|
||
update_data["updated_by"] = operator_id
|
||
record = await KnowledgeFileRepository().update_fields(file_id=file_id, kb_id=kb_id, data=update_data)
|
||
if record is None:
|
||
raise ValueError(f"File {file_id} not found")
|
||
|
||
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
|
||
|
||
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"{kb_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_minio_bytes(self, file_path: str) -> bytes:
|
||
from yuxi.knowledge.utils.kb_utils import is_minio_url, parse_minio_url
|
||
from yuxi.storage.minio import get_minio_client
|
||
|
||
if not file_path or not is_minio_url(file_path):
|
||
raise ValueError(f"Invalid MinIO path format: {file_path}")
|
||
|
||
bucket_name, object_name = parse_minio_url(file_path)
|
||
minio_client = get_minio_client()
|
||
return await minio_client.adownload_file(bucket_name, object_name)
|
||
|
||
async def _read_markdown_from_minio(self, file_path: str) -> str:
|
||
"""Read markdown content from MinIO"""
|
||
content_bytes = await self._read_minio_bytes(file_path)
|
||
return content_bytes.decode("utf-8")
|
||
|
||
async def _get_file_meta(self, kb_id: str, file_id: str) -> dict:
|
||
return await self._load_file_meta(kb_id, file_id)
|
||
|
||
@staticmethod
|
||
def _original_file_path(file_meta: dict) -> str | None:
|
||
return file_meta.get("minio_url") or file_meta.get("path")
|
||
|
||
def _knowledge_file_entry(self, kb_id: str, file_id: str, file_meta: dict) -> dict:
|
||
is_dir = bool(file_meta.get("is_folder"))
|
||
original_path = self._original_file_path(file_meta)
|
||
path = f"/{file_id}"
|
||
if is_dir:
|
||
path = f"{path}/"
|
||
return {
|
||
"source": "knowledge",
|
||
"kb_id": kb_id,
|
||
"file_id": file_id,
|
||
"parent_id": file_meta.get("parent_id"),
|
||
"path": path,
|
||
"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,
|
||
"modified_at": file_meta.get("updated_at") or file_meta.get("created_at") or "",
|
||
"readonly": True,
|
||
"status": file_meta.get("status", "done"),
|
||
"has_original_file": bool(original_path),
|
||
"has_parsed_markdown": bool(file_meta.get("markdown_file")),
|
||
}
|
||
|
||
def _sort_file_entries(self, entries: list[dict]) -> list[dict]:
|
||
return sorted(
|
||
entries,
|
||
key=lambda item: (not bool(item.get("is_dir")), str(item.get("name") or "").lower()),
|
||
)
|
||
|
||
async def _list_knowledge_children(
|
||
self,
|
||
kb_id: str,
|
||
parent_id: str | None,
|
||
*,
|
||
recursive: bool,
|
||
files_only: bool,
|
||
) -> list[dict]:
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
children = []
|
||
for record in await KnowledgeFileRepository().list_children(kb_id=kb_id, parent_id=parent_id):
|
||
meta = self._file_record_to_meta(record)
|
||
children.append((record.file_id, meta))
|
||
|
||
entries = []
|
||
for file_id, meta in children:
|
||
if not files_only or not meta.get("is_folder"):
|
||
entries.append(self._knowledge_file_entry(kb_id, file_id, meta))
|
||
if recursive and meta.get("is_folder"):
|
||
entries.extend(
|
||
await self._list_knowledge_children(
|
||
kb_id,
|
||
file_id,
|
||
recursive=True,
|
||
files_only=files_only,
|
||
)
|
||
)
|
||
return self._sort_file_entries(entries)
|
||
|
||
async def list_file_tree(
|
||
self,
|
||
kb_id: str,
|
||
parent_id: str | None = None,
|
||
recursive: bool = False,
|
||
files_only: bool = False,
|
||
) -> dict:
|
||
if kb_id not in self.databases_meta:
|
||
raise ValueError(f"Database {kb_id} not found")
|
||
if parent_id:
|
||
parent_meta = await self._get_file_meta(kb_id, parent_id)
|
||
if not parent_meta.get("is_folder"):
|
||
raise ValueError("Parent is not a folder")
|
||
return {
|
||
"entries": await self._list_knowledge_children(
|
||
kb_id,
|
||
parent_id,
|
||
recursive=recursive,
|
||
files_only=files_only,
|
||
),
|
||
"readonly": True,
|
||
}
|
||
|
||
@staticmethod
|
||
def _office_pdf_preview_path(kb_id: str, file_id: str) -> str:
|
||
return f"{kb_id}/preview/{file_id}.pdf"
|
||
|
||
async def _ensure_office_pdf_preview(self, kb_id: str, file_id: str, file_meta: dict) -> str:
|
||
from yuxi.storage.minio import get_minio_client
|
||
|
||
filename = file_meta.get("filename") or file_meta.get("original_filename") or file_id
|
||
if not is_office_pdf_preview_file(filename):
|
||
raise ValueError("当前文件类型不支持 PDF 预览")
|
||
|
||
minio_client = get_minio_client()
|
||
bucket_name = minio_client.KB_BUCKETS["parsed"]
|
||
object_name = self._office_pdf_preview_path(kb_id, file_id)
|
||
if await minio_client.astat_file(bucket_name, object_name) is not None:
|
||
return f"minio://{bucket_name}/{object_name}"
|
||
|
||
original_path = self._original_file_path(file_meta)
|
||
if not original_path:
|
||
raise ValueError("文件没有可转换的原始内容")
|
||
|
||
raw_content = await self._read_minio_bytes(original_path)
|
||
try:
|
||
pdf_content = await convert_office_to_pdf(filename, raw_content)
|
||
except OfficePreviewConversionError as exc:
|
||
raise ValueError(str(exc)) from exc
|
||
await minio_client.aupload_file(
|
||
bucket_name=bucket_name,
|
||
object_name=object_name,
|
||
data=pdf_content,
|
||
content_type="application/pdf",
|
||
)
|
||
return f"minio://{bucket_name}/{object_name}"
|
||
|
||
async def _get_minio_file_size(self, file_path: str) -> int | None:
|
||
from yuxi.knowledge.utils.kb_utils import is_minio_url, parse_minio_url
|
||
from yuxi.storage.minio import get_minio_client
|
||
|
||
if not file_path or not is_minio_url(file_path):
|
||
return None
|
||
bucket_name, object_name = parse_minio_url(file_path)
|
||
return await get_minio_client().astat_file(bucket_name, object_name)
|
||
|
||
async def read_file_preview(self, kb_id: str, file_id: str) -> dict:
|
||
file_meta = await self._get_file_meta(kb_id, file_id)
|
||
if file_meta.get("is_folder"):
|
||
raise ValueError("Cannot preview a folder")
|
||
|
||
filename = file_meta.get("filename") or file_meta.get("original_filename") or file_id
|
||
response = {
|
||
"source": "knowledge",
|
||
"kb_id": kb_id,
|
||
"file_id": file_id,
|
||
"filename": filename,
|
||
"readonly": True,
|
||
}
|
||
|
||
original_path = self._original_file_path(file_meta)
|
||
if not original_path:
|
||
return {
|
||
**response,
|
||
"content": None,
|
||
"preview_type": "unsupported",
|
||
"supported": False,
|
||
"message": "文件没有可预览的原始内容",
|
||
}
|
||
|
||
file_size = file_meta.get("size")
|
||
if file_size is None:
|
||
file_size = await self._get_minio_file_size(original_path)
|
||
if file_size is not None and int(file_size) > MAX_BINARY_PREVIEW_SIZE_BYTES:
|
||
return {**response, **render_preview_too_large_payload()}
|
||
|
||
if is_office_pdf_preview_file(filename):
|
||
preview_path = await self._ensure_office_pdf_preview(kb_id, file_id, file_meta)
|
||
stem = filename.rsplit(".", 1)[0] or file_id
|
||
return {
|
||
**response,
|
||
"content": await self._read_minio_bytes(preview_path),
|
||
"filename": f"{stem}.pdf",
|
||
"media_type": "application/pdf",
|
||
"preview_type": "pdf",
|
||
"supported": True,
|
||
"message": None,
|
||
"binary": True,
|
||
}
|
||
|
||
raw_content = await self._read_minio_bytes(original_path)
|
||
if len(raw_content) > MAX_BINARY_PREVIEW_SIZE_BYTES:
|
||
return {**response, **render_preview_too_large_payload()}
|
||
payload = render_preview_payload(filename, raw_content)
|
||
if is_binary_preview_type(payload["preview_type"]) and payload["supported"]:
|
||
return {
|
||
**response,
|
||
"content": raw_content,
|
||
"media_type": detect_media_type(filename, raw_content),
|
||
"preview_type": payload["preview_type"],
|
||
"supported": True,
|
||
"message": None,
|
||
"binary": True,
|
||
}
|
||
return {**response, **payload}
|
||
|
||
async def get_file_download(self, kb_id: str, file_id: str, variant: str = "original") -> dict:
|
||
file_meta = await 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"}:
|
||
raise ValueError("Unsupported download variant")
|
||
|
||
filename = file_meta.get("filename") or file_meta.get("original_filename") or file_id
|
||
if variant == "parsed":
|
||
markdown_file = file_meta.get("markdown_file")
|
||
if not markdown_file:
|
||
raise ValueError("文件尚未生成解析结果")
|
||
return {
|
||
"filename": f"{filename}.parsed.md",
|
||
"content": await self._read_minio_bytes(markdown_file),
|
||
"media_type": "text/markdown; charset=utf-8",
|
||
}
|
||
|
||
original_path = self._original_file_path(file_meta)
|
||
if not original_path:
|
||
raise ValueError("文件没有可下载的原始内容")
|
||
media_type = file_meta.get("content_type") or mimetypes.guess_type(filename)[0] or "application/octet-stream"
|
||
return {
|
||
"filename": filename,
|
||
"content": await self._read_minio_bytes(original_path),
|
||
"media_type": media_type,
|
||
}
|
||
|
||
def _build_open_file_window(self, content: str, *, offset: int = 0, limit: int = 800) -> dict[str, Any]:
|
||
lines = content.splitlines()
|
||
total_lines = len(lines)
|
||
start = min(max(int(offset), 0), total_lines)
|
||
window_size = min(max(int(limit), 1), 2000)
|
||
selected = lines[start : start + window_size]
|
||
end = start + len(selected)
|
||
|
||
return {
|
||
"start_line": start + 1 if selected else 0,
|
||
"end_line": end,
|
||
"total_lines": total_lines,
|
||
"offset": start,
|
||
"window_size": window_size,
|
||
"has_more_before": start > 0,
|
||
"has_more_after": end < total_lines,
|
||
"next_offset": end if end < total_lines else None,
|
||
"content": "\n".join(f"{start + idx + 1:6d}\t{line}" for idx, line in enumerate(selected)),
|
||
}
|
||
|
||
@staticmethod
|
||
def build_search_output(kb_id: str, retrieval_results: Any) -> dict[str, Any] | Any:
|
||
if not isinstance(retrieval_results, list):
|
||
return retrieval_results
|
||
|
||
results = []
|
||
for index, chunk in enumerate(retrieval_results):
|
||
if not isinstance(chunk, dict):
|
||
continue
|
||
|
||
metadata = chunk.get("metadata") if isinstance(chunk.get("metadata"), dict) else {}
|
||
metadata = {
|
||
key: value
|
||
for key, value in metadata.items()
|
||
if key not in {"filepath", "parsed_path", "path", "markdown_file"}
|
||
}
|
||
file_id = metadata.get("file_id") or chunk.get("file_id") or chunk.get("full_doc_id") or ""
|
||
chunk_id = metadata.get("chunk_id") or chunk.get("chunk_id") or chunk.get("id")
|
||
chunk_index = metadata.get("chunk_index")
|
||
if chunk_index is None:
|
||
chunk_index = chunk.get("chunk_index")
|
||
if chunk_index is not None:
|
||
metadata.setdefault("chunk_index", chunk_index)
|
||
if chunk.get("score") is not None:
|
||
metadata.setdefault("score", chunk.get("score"))
|
||
if chunk.get("distance") is not None:
|
||
metadata.setdefault("distance", chunk.get("distance"))
|
||
|
||
results.append(
|
||
SearchResultSchema(
|
||
id=str(chunk_id or f"{file_id}:{index + 1}"),
|
||
kb_id=str(kb_id),
|
||
file_id=str(file_id or ""),
|
||
content=str(chunk.get("content") or ""),
|
||
metadata=metadata,
|
||
)
|
||
)
|
||
|
||
return SearchOutputSchema(kb_id=str(kb_id), results=results).model_dump()
|
||
|
||
@staticmethod
|
||
def _build_find_file_windows(
|
||
content: str,
|
||
*,
|
||
patterns: list[str],
|
||
use_regex: bool = False,
|
||
case_sensitive: bool = False,
|
||
max_windows: int = 5,
|
||
window_size: int = 80,
|
||
) -> dict[str, Any]:
|
||
patterns = [pattern for pattern in patterns if pattern]
|
||
if not patterns:
|
||
raise ValueError("请提供至少一个 pattern")
|
||
|
||
lines = content.splitlines()
|
||
flags = 0 if case_sensitive else re.IGNORECASE
|
||
if use_regex:
|
||
matchers = [re.compile(pattern, flags) for pattern in patterns]
|
||
|
||
def line_matches(line: str) -> bool:
|
||
return any(matcher.search(line) for matcher in matchers)
|
||
|
||
else:
|
||
normalized_patterns = patterns if case_sensitive else [pattern.lower() for pattern in patterns]
|
||
|
||
def line_matches(line: str) -> bool:
|
||
haystack = line if case_sensitive else line.lower()
|
||
return any(pattern in haystack for pattern in normalized_patterns)
|
||
|
||
matched_indexes = [index for index, line in enumerate(lines) if line_matches(line)]
|
||
windows: list[FindWindowSchema] = []
|
||
covered_until = -1
|
||
normalized_window_size = min(max(int(window_size), 1), 200)
|
||
half_window = normalized_window_size // 2
|
||
|
||
for matched_index in matched_indexes:
|
||
if matched_index < covered_until:
|
||
continue
|
||
start = max(matched_index - half_window, 0)
|
||
end = min(start + normalized_window_size, len(lines))
|
||
start = max(end - normalized_window_size, 0)
|
||
matched_lines = [index + 1 for index in matched_indexes if start <= index < end]
|
||
selected = lines[start:end]
|
||
windows.append(
|
||
FindWindowSchema(
|
||
start_line=start + 1 if selected else 0,
|
||
end_line=end,
|
||
matched_lines=matched_lines,
|
||
content="\n".join(f"{start + idx + 1:6d}\t{line}" for idx, line in enumerate(selected)),
|
||
)
|
||
)
|
||
covered_until = end
|
||
if len(windows) >= max_windows:
|
||
break
|
||
|
||
return FindOutputSchema(
|
||
kb_id="",
|
||
file_id="",
|
||
semantic=False,
|
||
match_mode="regex" if use_regex else "keyword",
|
||
total_matches=len(matched_indexes),
|
||
windows=windows,
|
||
).model_dump(exclude={"kb_id", "file_id"})
|
||
|
||
async def open_file_content(self, kb_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict:
|
||
"""按行窗口打开文件解析后的 Markdown 内容"""
|
||
try:
|
||
file_meta = await self._load_file_meta(kb_id, file_id)
|
||
except ValueError as exc:
|
||
raise Exception(f"文件不存在: {file_id}") from exc
|
||
if file_meta.get("is_folder"):
|
||
raise Exception(f"文件 {file_id} 是文件夹")
|
||
|
||
markdown_file = file_meta.get("markdown_file")
|
||
if not markdown_file:
|
||
raise Exception(f"文件 {file_id} 没有解析后的 Markdown 内容")
|
||
|
||
content = await self._read_markdown_from_minio(markdown_file)
|
||
return self._build_open_file_window(content, offset=offset, limit=limit)
|
||
|
||
async def find_file_content(
|
||
self,
|
||
kb_id: str,
|
||
file_id: str,
|
||
patterns: list[str],
|
||
*,
|
||
use_regex: bool = False,
|
||
case_sensitive: bool = False,
|
||
max_windows: int = 5,
|
||
window_size: int = 80,
|
||
) -> dict:
|
||
try:
|
||
file_meta = await self._load_file_meta(kb_id, file_id)
|
||
except ValueError as exc:
|
||
raise Exception(f"文件不存在: {file_id}") from exc
|
||
if file_meta.get("is_folder"):
|
||
raise Exception(f"文件 {file_id} 是文件夹")
|
||
|
||
markdown_file = file_meta.get("markdown_file")
|
||
if not markdown_file:
|
||
raise Exception(f"文件 {file_id} 没有解析后的 Markdown 内容")
|
||
|
||
content = await self._read_markdown_from_minio(markdown_file)
|
||
return self._build_find_file_windows(
|
||
content,
|
||
patterns=patterns,
|
||
use_regex=use_regex,
|
||
case_sensitive=case_sensitive,
|
||
max_windows=max_windows,
|
||
window_size=window_size,
|
||
)
|
||
|
||
@abstractmethod
|
||
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:
|
||
kb_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,
|
||
embedding_model_spec: str | None = None,
|
||
llm_model_spec: str | None = None,
|
||
record_fields: dict[str, Any] | None = None,
|
||
**kwargs,
|
||
) -> dict:
|
||
"""
|
||
创建数据库
|
||
|
||
Args:
|
||
database_name: 数据库名称
|
||
description: 数据库描述
|
||
embedding_model_spec: 嵌入模型 spec
|
||
llm_model_spec: LLM 模型 spec
|
||
record_fields: 首次持久化知识库记录时由上层传入的受控业务字段
|
||
**kwargs: 其他配置参数
|
||
|
||
Returns:
|
||
数据库信息字典
|
||
"""
|
||
kwargs = self.normalize_additional_params(kwargs)
|
||
kwargs["stats"] = {"file_count": 0, "chunk_count": 0, "token_count": 0}
|
||
|
||
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[kb_id] = {
|
||
"name": database_name,
|
||
"description": description,
|
||
"kb_type": self.kb_type,
|
||
"embedding_model_spec": embedding_model_spec,
|
||
"llm_model_spec": llm_model_spec,
|
||
"metadata": kwargs,
|
||
"created_at": utc_isoformat(),
|
||
"query_params": self._get_default_query_params(kb_id),
|
||
}
|
||
await self._persist_kb(kb_id, record_fields=record_fields)
|
||
|
||
# 创建工作目录
|
||
working_dir = os.path.join(self.work_dir, kb_id)
|
||
os.makedirs(working_dir, exist_ok=True)
|
||
|
||
# 返回数据库信息
|
||
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, kb_id: str) -> dict:
|
||
"""
|
||
删除数据库
|
||
|
||
Args:
|
||
kb_id: 数据库ID
|
||
|
||
Returns:
|
||
操作结果
|
||
"""
|
||
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.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
from yuxi.storage.minio import get_minio_client
|
||
|
||
minio_client = get_minio_client()
|
||
file_repo = KnowledgeFileRepository()
|
||
|
||
# 1. 删除文件元数据中记录的 MinIO 文件
|
||
after_file_id = None
|
||
while True:
|
||
records = await file_repo.list_by_kb_id_after(kb_id, after_file_id=after_file_id, limit=500)
|
||
if not records:
|
||
break
|
||
after_file_id = records[-1].file_id
|
||
for record in records:
|
||
file_id = record.file_id
|
||
file_path = record.minio_url or record.path
|
||
if file_path and is_minio_url(file_path):
|
||
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"{kb_id}/parsed/{file_id}.md"
|
||
await minio_client.adelete_file(minio_client.KB_BUCKETS["parsed"], parsed_object)
|
||
|
||
# 2. 并行删除所有知识库 bucket 中该 kb_id 下的文件
|
||
prefix = f"{kb_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[kb_id]
|
||
await file_repo.delete_by_kb_id(kb_id)
|
||
kb_repo = KnowledgeBaseRepository()
|
||
await kb_repo.delete(kb_id)
|
||
|
||
# 删除工作目录
|
||
working_dir = os.path.join(self.work_dir, kb_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, kb_id: str, folder_name: str, parent_id: str | None = None) -> dict:
|
||
"""Create a folder in the database."""
|
||
import uuid
|
||
|
||
if parent_id:
|
||
parent_meta = await self._load_file_meta(kb_id, parent_id)
|
||
if not parent_meta.get("is_folder"):
|
||
raise ValueError("Parent is not a folder")
|
||
|
||
folder_id = f"folder-{uuid.uuid4()}"
|
||
|
||
folder_meta = {
|
||
"file_id": folder_id,
|
||
"filename": folder_name,
|
||
"is_folder": True,
|
||
"parent_id": parent_id,
|
||
"kb_id": kb_id,
|
||
"created_at": utc_isoformat(),
|
||
"status": "done",
|
||
"path": folder_name,
|
||
"file_type": "folder",
|
||
}
|
||
await self._persist_file_meta(folder_id, folder_meta)
|
||
return folder_meta
|
||
|
||
@abstractmethod
|
||
async def update_content(self, kb_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]:
|
||
"""
|
||
更新内容 - 根据file_ids重新解析文件并更新向量库
|
||
|
||
Args:
|
||
kb_id: 数据库ID
|
||
file_ids: 文件ID列表
|
||
params: 处理参数
|
||
|
||
Returns:
|
||
更新结果列表
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def aquery(self, query_text: str, kb_id: str, **kwargs) -> list[dict]:
|
||
"""
|
||
异步查询知识库
|
||
|
||
Args:
|
||
query_text: 查询文本
|
||
kb_id: 数据库ID
|
||
**kwargs: 查询参数
|
||
|
||
Returns:
|
||
一个包含字典的列表，每个字典代表一个检索到的文档块。
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
|
||
"""
|
||
获取知识库类型的查询参数配置
|
||
|
||
Args:
|
||
kb_id: 数据库ID
|
||
**kwargs: 额外参数
|
||
|
||
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, kb_id: str, format: str = "zip", **kwargs) -> str:
|
||
pass
|
||
|
||
def _get_query_params(self, kb_id: str) -> dict:
|
||
"""从实例元数据中加载查询参数"""
|
||
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, kb_id: str) -> dict[str, Any]:
|
||
"""从 get_query_params_config 中提取所有参数的默认值，返回 {"options": {...}}"""
|
||
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 _build_database_stats(self, kb_id: str) -> dict[str, int]:
|
||
del kb_id
|
||
return self._normalize_database_stats(None)
|
||
|
||
@staticmethod
|
||
def _normalize_database_stats(stats: dict | None) -> dict[str, int]:
|
||
normalized = {
|
||
"file_count": 0,
|
||
"folder_count": 0,
|
||
"row_count": 0,
|
||
"total_size": 0,
|
||
"chunk_count": 0,
|
||
"token_count": 0,
|
||
"pending_parse_count": 0,
|
||
"pending_index_count": 0,
|
||
"processing_count": 0,
|
||
}
|
||
if not isinstance(stats, dict):
|
||
return normalized
|
||
|
||
for key in normalized:
|
||
try:
|
||
normalized[key] = max(int(stats.get(key) or 0), 0)
|
||
except (TypeError, ValueError):
|
||
normalized[key] = 0
|
||
return normalized
|
||
|
||
def _get_database_stats(self, kb_id: str) -> dict[str, int]:
|
||
metadata = self.databases_meta.get(kb_id, {}).get("metadata") or {}
|
||
stats = metadata.get("stats") if isinstance(metadata, dict) else None
|
||
if isinstance(stats, dict):
|
||
return self._normalize_database_stats(stats)
|
||
return self._build_database_stats(kb_id)
|
||
|
||
def _set_database_stats(self, kb_id: str, stats: dict[str, int]) -> None:
|
||
if kb_id not in self.databases_meta:
|
||
raise ValueError(f"Database {kb_id} not found")
|
||
|
||
metadata = self.databases_meta[kb_id].setdefault("metadata", {})
|
||
metadata["stats"] = self._normalize_database_stats(stats)
|
||
|
||
async def refresh_database_stats(self, kb_id: str) -> dict[str, int]:
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
stats = await KnowledgeFileRepository().get_kb_file_stats(kb_id)
|
||
self._set_database_stats(kb_id, stats)
|
||
await self._persist_kb(kb_id)
|
||
return stats
|
||
|
||
async def repair_missing_file_stats(self, kb_id: str) -> dict:
|
||
if kb_id not in self.databases_meta:
|
||
raise ValueError(f"Database {kb_id} not found")
|
||
|
||
from yuxi.knowledge.chunking.ragflow_like.nlp import count_tokens
|
||
from yuxi.repositories.knowledge_chunk_repository import KnowledgeChunkRepository
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
chunk_repo = KnowledgeChunkRepository()
|
||
file_repo = KnowledgeFileRepository()
|
||
after_file_id = None
|
||
scanned_files = 0
|
||
scanned_indexed_files = 0
|
||
skipped_file_count = 0
|
||
scanned_token_files = 0
|
||
updated_files = 0
|
||
updated_chunk_files = 0
|
||
updated_token_files = 0
|
||
updated_size_files = 0
|
||
|
||
while True:
|
||
records = await file_repo.list_by_kb_id_after(
|
||
kb_id,
|
||
after_file_id=after_file_id,
|
||
limit=500,
|
||
files_only=True,
|
||
)
|
||
if not records:
|
||
break
|
||
after_file_id = records[-1].file_id
|
||
|
||
indexed_records = [record for record in records if _should_repair_file_stats({"status": record.status})]
|
||
indexed_file_ids = [record.file_id for record in indexed_records]
|
||
indexed_file_id_set = set(indexed_file_ids)
|
||
chunk_counts = await chunk_repo.count_by_file_ids(indexed_file_ids)
|
||
token_file_ids = [record.file_id for record in indexed_records if int(record.token_count or 0) <= 0]
|
||
token_counts = {file_id: 0 for file_id in token_file_ids}
|
||
for chunk in await chunk_repo.list_by_file_ids(token_file_ids):
|
||
token_counts[chunk.file_id] = token_counts.get(chunk.file_id, 0) + count_tokens(chunk.content or "")
|
||
size_updates = await self._fill_missing_file_sizes_for_records(records)
|
||
|
||
scanned_files += len(records)
|
||
scanned_indexed_files += len(indexed_file_ids)
|
||
skipped_file_count += len(records) - len(indexed_file_ids)
|
||
scanned_token_files += len(token_file_ids)
|
||
|
||
for record in records:
|
||
file_id = record.file_id
|
||
update_data: dict[str, Any] = {}
|
||
if file_id not in indexed_file_id_set:
|
||
if int(record.chunk_count or 0) != 0:
|
||
update_data["chunk_count"] = 0
|
||
updated_chunk_files += 1
|
||
if int(record.token_count or 0) != 0:
|
||
update_data["token_count"] = 0
|
||
updated_token_files += 1
|
||
else:
|
||
next_chunk_count = int(chunk_counts.get(file_id, 0))
|
||
if int(record.chunk_count or 0) != next_chunk_count:
|
||
update_data["chunk_count"] = next_chunk_count
|
||
updated_chunk_files += 1
|
||
|
||
if file_id in token_counts:
|
||
next_token_count = int(token_counts[file_id])
|
||
if record.token_count is None or int(record.token_count or 0) != next_token_count:
|
||
update_data["token_count"] = next_token_count
|
||
updated_token_files += 1
|
||
|
||
if file_id in size_updates:
|
||
update_data["file_size"] = size_updates[file_id]
|
||
updated_size_files += 1
|
||
|
||
if update_data:
|
||
updated_files += 1
|
||
await file_repo.update_fields(file_id=file_id, kb_id=kb_id, data=update_data)
|
||
|
||
stats = await self.refresh_database_stats(kb_id)
|
||
return {
|
||
"status": "success",
|
||
"stats": stats,
|
||
"scanned_files": scanned_files,
|
||
"scanned_indexed_files": scanned_indexed_files,
|
||
"skipped_unindexed_files": skipped_file_count,
|
||
"scanned_token_files": scanned_token_files,
|
||
"updated_files": updated_files,
|
||
"updated_chunk_files": updated_chunk_files,
|
||
"updated_token_files": updated_token_files,
|
||
"updated_size_files": updated_size_files,
|
||
}
|
||
|
||
def get_database_info(self, kb_id: str, include_files: bool = True) -> dict | None:
|
||
"""
|
||
获取数据库详细信息
|
||
|
||
Args:
|
||
kb_id: 数据库ID
|
||
include_files: 是否包含文件信息，默认为True
|
||
|
||
Returns:
|
||
数据库信息或None
|
||
"""
|
||
if kb_id not in self.databases_meta:
|
||
return None
|
||
|
||
meta = self.databases_meta[kb_id].copy()
|
||
meta["kb_id"] = kb_id
|
||
|
||
meta["stats"] = self._get_database_stats(kb_id)
|
||
meta["row_count"] = meta["stats"].get("row_count") or meta["stats"].get("file_count") or 0
|
||
|
||
if include_files:
|
||
meta["files"] = {}
|
||
meta["files_truncated"] = True
|
||
|
||
meta["status"] = "已连接"
|
||
return meta
|
||
|
||
def get_databases(self, include_files: bool = False) -> dict:
|
||
"""
|
||
获取所有数据库信息
|
||
|
||
Args:
|
||
include_files: 是否包含文件信息，默认False以减少响应大小
|
||
|
||
Returns:
|
||
数据库列表
|
||
"""
|
||
# 确保元数据已加载（延迟加载机制）
|
||
self._ensure_metadata_loaded()
|
||
|
||
databases = []
|
||
for kb_id, meta in self.databases_meta.items():
|
||
db_dict = meta.copy()
|
||
db_dict["kb_id"] = kb_id
|
||
db_dict["stats"] = self._get_database_stats(kb_id)
|
||
db_dict["row_count"] = db_dict["stats"].get("row_count") or db_dict["stats"].get("file_count") or 0
|
||
|
||
if include_files:
|
||
db_dict["files"] = {}
|
||
db_dict["files_truncated"] = True
|
||
|
||
db_dict["status"] = "已连接"
|
||
databases.append(db_dict)
|
||
|
||
return {"databases": databases}
|
||
|
||
async def delete_folder(self, kb_id: str, folder_id: str) -> None:
|
||
"""
|
||
Recursively delete a folder and its content.
|
||
|
||
Args:
|
||
kb_id: Database ID
|
||
folder_id: Folder ID to delete
|
||
"""
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
children = await KnowledgeFileRepository().list_children(kb_id=kb_id, parent_id=folder_id)
|
||
for child in children:
|
||
child_id = child.file_id
|
||
if child.is_folder:
|
||
await self.delete_folder(kb_id, child_id)
|
||
else:
|
||
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(kb_id, folder_id)
|
||
|
||
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:
|
||
kb_id: Database ID
|
||
file_id: File/Folder ID to move
|
||
new_parent_id: New parent folder ID (None for root)
|
||
|
||
Returns:
|
||
dict: Updated metadata
|
||
"""
|
||
meta = await self._load_file_meta(kb_id, file_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 = await self._load_file_meta(kb_id, current)
|
||
if current == new_parent_id and not parent_meta.get("is_folder"):
|
||
raise ValueError("Parent is not a folder")
|
||
if current == file_id:
|
||
raise ValueError("Cannot move a folder into its own subfolder")
|
||
current = parent_meta.get("parent_id")
|
||
elif new_parent_id:
|
||
parent_meta = await self._load_file_meta(kb_id, new_parent_id)
|
||
if not parent_meta.get("is_folder"):
|
||
raise ValueError("Parent is not a folder")
|
||
|
||
meta["parent_id"] = new_parent_id
|
||
await self._persist_file_meta(file_id, meta)
|
||
return meta
|
||
|
||
@abstractmethod
|
||
async def delete_file(self, kb_id: str, file_id: str) -> None:
|
||
"""
|
||
删除文件
|
||
|
||
Args:
|
||
kb_id: 数据库ID
|
||
file_id: 文件ID
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def get_file_basic_info(self, kb_id: str, file_id: str) -> dict:
|
||
"""
|
||
获取文件基本信息（仅元数据）
|
||
|
||
Args:
|
||
kb_id: 数据库ID
|
||
file_id: 文件ID
|
||
|
||
Returns:
|
||
dict: 包含文件基本信息的字典
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def get_file_content(self, kb_id: str, file_id: str) -> dict:
|
||
"""
|
||
获取文件内容信息（chunks和lines）
|
||
|
||
Args:
|
||
kb_id: 数据库ID
|
||
file_id: 文件ID
|
||
|
||
Returns:
|
||
dict: 包含文件内容信息的字典
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def get_file_info(self, kb_id: str, file_id: str) -> dict:
|
||
"""
|
||
获取文件完整信息（基本信息+内容信息）
|
||
|
||
Args:
|
||
kb_id: 数据库ID
|
||
file_id: 文件ID
|
||
|
||
Returns:
|
||
dict: 包含文件信息和chunks的字典
|
||
"""
|
||
pass
|
||
|
||
def update_database(
|
||
self,
|
||
kb_id: str,
|
||
name: str,
|
||
description: str,
|
||
llm_model_spec: str | None = None,
|
||
update_llm_model_spec: bool = False,
|
||
) -> dict:
|
||
"""
|
||
更新数据库
|
||
|
||
Args:
|
||
kb_id: 数据库ID
|
||
name: 新名称
|
||
description: 新描述
|
||
llm_model_spec: LLM 模型 spec（可选）
|
||
|
||
Returns:
|
||
更新后的数据库信息
|
||
"""
|
||
if kb_id not in self.databases_meta:
|
||
raise ValueError(f"数据库 {kb_id} 不存在")
|
||
|
||
self.databases_meta[kb_id]["name"] = name
|
||
self.databases_meta[kb_id]["description"] = description
|
||
if update_llm_model_spec:
|
||
self.databases_meta[kb_id]["llm_model_spec"] = llm_model_spec
|
||
|
||
return self.get_database_info(kb_id)
|
||
|
||
def get_retrievers(self) -> dict[str, dict]:
|
||
"""
|
||
获取所有检索器
|
||
|
||
Returns:
|
||
检索器字典
|
||
"""
|
||
retrievers = {}
|
||
for kb_id, meta in self.databases_meta.items():
|
||
|
||
def make_retriever(kb_id):
|
||
async def retriever(query_text, **kwargs):
|
||
results = await self.aquery(query_text, kb_id, agent_call=True, **kwargs)
|
||
return self.build_search_output(kb_id, results)
|
||
|
||
return retriever
|
||
|
||
retrievers[kb_id] = {
|
||
"name": meta["name"],
|
||
"description": meta["description"],
|
||
"retriever": make_retriever(kb_id),
|
||
"metadata": meta,
|
||
}
|
||
return retrievers
|
||
|
||
async def _load_metadata(self) -> None:
|
||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||
|
||
kb_repo = KnowledgeBaseRepository()
|
||
|
||
databases = [kb for kb in await kb_repo.get_all() if kb.kb_type == self.kb_type]
|
||
self.databases_meta = {
|
||
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.kb_id),
|
||
"metadata": self.normalize_additional_params(kb.additional_params),
|
||
"created_at": utc_isoformat(kb.created_at) if kb.created_at else utc_isoformat(),
|
||
}
|
||
for kb in databases
|
||
}
|
||
|
||
# 文件元数据量可能达到几十万级，启动阶段只加载 KB 级配置。
|
||
# 单文件操作按需查询 PostgreSQL，并在流程内通过局部变量传递。
|
||
|
||
self.benchmarks_meta = {}
|
||
self._normalize_metadata_state()
|
||
self._metadata_loaded = True
|
||
|
||
logger.info(f"Loaded {self.kb_type} metadata from database for {len(self.databases_meta)} databases")
|
||
|
||
async def _fill_missing_file_sizes_for_records(self, records: list[Any]) -> dict[str, int]:
|
||
"""为显式修复任务中的缺失 size 文件从 MinIO 补全大小信息。"""
|
||
from yuxi.knowledge.utils.kb_utils import is_minio_url, parse_minio_url
|
||
from yuxi.storage.minio import get_minio_client
|
||
|
||
candidates: list[tuple[str, str]] = []
|
||
for record in records:
|
||
if record.is_folder or record.file_size is not None:
|
||
continue
|
||
file_path = record.minio_url or record.path
|
||
if not file_path or not is_minio_url(file_path):
|
||
continue
|
||
candidates.append((record.file_id, file_path))
|
||
|
||
if not candidates:
|
||
return {}
|
||
|
||
minio_client = get_minio_client()
|
||
semaphore = asyncio.Semaphore(20)
|
||
|
||
async def _stat_file(file_id: str, file_path: str) -> tuple[str, int | None]:
|
||
bucket_name, obj_name = parse_minio_url(file_path)
|
||
try:
|
||
async with semaphore:
|
||
return file_id, await minio_client.astat_file(bucket_name, obj_name)
|
||
except Exception as exc:
|
||
logger.warning(f"Failed to fill size for {file_id}: {exc}")
|
||
return file_id, None
|
||
|
||
updates: dict[str, int] = {}
|
||
for offset in range(0, len(candidates), 100):
|
||
batch = candidates[offset : offset + 100]
|
||
for file_id, file_size in await asyncio.gather(
|
||
*(_stat_file(file_id, file_path) for file_id, file_path in batch)
|
||
):
|
||
if file_size is not None:
|
||
updates[file_id] = file_size
|
||
|
||
if updates:
|
||
logger.info(f"Filled {len(updates)}/{len(candidates)} missing file sizes from MinIO for {self.kb_type}")
|
||
return updates
|
||
|
||
async def _save_metadata(self) -> None:
|
||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||
|
||
kb_repo = KnowledgeBaseRepository()
|
||
|
||
self._normalize_metadata_state()
|
||
|
||
for kb_id, meta in self.databases_meta.items():
|
||
existing = await kb_repo.get_by_kb_id(kb_id)
|
||
payload = {
|
||
"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"),
|
||
"llm_model_spec": meta.get("llm_model_spec"),
|
||
"query_params": meta.get("query_params"),
|
||
"additional_params": meta.get("metadata") or {},
|
||
}
|
||
if existing is None:
|
||
await kb_repo.create(payload)
|
||
|
||
async def _persist_file_meta(self, file_id: str, meta: dict) -> None:
|
||
"""Persist one file metadata record without storing it on the KB instance."""
|
||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||
|
||
data = self._file_meta_to_record_data(meta)
|
||
if not data.get("kb_id"):
|
||
return
|
||
await KnowledgeFileRepository().upsert(file_id=file_id, data=data)
|
||
|
||
async def _persist_kb(self, kb_id: str, record_fields: dict[str, Any] | None = None) -> None:
|
||
"""只保存单个知识库到数据库，避免全量遍历"""
|
||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||
|
||
kb_repo = KnowledgeBaseRepository()
|
||
|
||
if kb_id not in self.databases_meta:
|
||
return
|
||
|
||
meta = self.databases_meta[kb_id]
|
||
existing = await kb_repo.get_by_kb_id(kb_id)
|
||
payload = {
|
||
"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"),
|
||
"llm_model_spec": meta.get("llm_model_spec"),
|
||
"query_params": meta.get("query_params"),
|
||
"additional_params": meta.get("metadata") or {},
|
||
}
|
||
if record_fields:
|
||
allowed_fields = {"share_config", "created_by"}
|
||
payload.update({key: value for key, value in record_fields.items() if key in allowed_fields})
|
||
|
||
if existing is None:
|
||
await kb_repo.create(payload)
|
||
else:
|
||
update_data = {
|
||
"name": payload["name"],
|
||
"description": payload["description"],
|
||
"kb_type": payload["kb_type"],
|
||
"embedding_model_spec": payload["embedding_model_spec"],
|
||
"llm_model_spec": payload["llm_model_spec"],
|
||
"query_params": payload["query_params"],
|
||
"additional_params": payload["additional_params"],
|
||
}
|
||
if record_fields:
|
||
update_data.update({key: payload[key] for key in allowed_fields if key in payload})
|
||
await kb_repo.update(kb_id, update_data)
|