Wenjie Zhang
GitHub
commit
c2ae4c7d12
@ -8,7 +8,7 @@ from urllib.parse import quote, unquote
|
||||
|
||||
from fastapi import APIRouter, Body, Depends, File, HTTPException, Query, Request, UploadFile
|
||||
from fastapi.responses import FileResponse
|
||||
from starlette.responses import FileResponse as StarletteFileResponse
|
||||
from starlette.responses import StreamingResponse
|
||||
|
||||
from src.storage.db.models import User
|
||||
from server.utils.auth_middleware import get_admin_user
|
||||
@ -17,10 +17,47 @@ from src import config, knowledge_base
|
||||
from src.knowledge.indexing import SUPPORTED_FILE_EXTENSIONS, is_supported_file_extension, process_file_to_markdown
|
||||
from src.knowledge.utils import calculate_content_hash, merge_processing_params
|
||||
from src.models.embed import test_embedding_model_status, test_all_embedding_models_status
|
||||
from src.storage.minio.client import aupload_file_to_minio, get_minio_client
|
||||
from src.utils import hashstr, logger
|
||||
|
||||
knowledge = APIRouter(prefix="/knowledge", tags=["knowledge"])
|
||||
|
||||
media_types = {
|
||||
".pdf": "application/pdf",
|
||||
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
".doc": "application/msword",
|
||||
".txt": "text/plain",
|
||||
".md": "text/markdown",
|
||||
".json": "application/json",
|
||||
".csv": "text/csv",
|
||||
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
".xls": "application/vnd.ms-excel",
|
||||
".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
".ppt": "application/vnd.ms-powerpoint",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".bmp": "image/bmp",
|
||||
".svg": "image/svg+xml",
|
||||
".zip": "application/zip",
|
||||
".rar": "application/x-rar-compressed",
|
||||
".7z": "application/x-7z-compressed",
|
||||
".tar": "application/x-tar",
|
||||
".gz": "application/gzip",
|
||||
".html": "text/html",
|
||||
".htm": "text/html",
|
||||
".xml": "text/xml",
|
||||
".css": "text/css",
|
||||
".js": "application/javascript",
|
||||
".py": "text/x-python",
|
||||
".java": "text/x-java-source",
|
||||
".cpp": "text/x-c++src",
|
||||
".c": "text/x-csrc",
|
||||
".h": "text/x-chdr",
|
||||
".hpp": "text/x-c++hdr",
|
||||
}
|
||||
|
||||
# =============================================================================
|
||||
# === 数据库管理分组 ===
|
||||
# =============================================================================
|
||||
@ -179,12 +216,6 @@ async def export_database(
|
||||
if not os.path.exists(file_path):
|
||||
raise HTTPException(status_code=404, detail="Exported file not found.")
|
||||
|
||||
media_types = {
|
||||
"csv": "text/csv",
|
||||
"xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
"md": "text/markdown",
|
||||
"txt": "text/plain",
|
||||
}
|
||||
media_type = media_types.get(format, "application/octet-stream")
|
||||
|
||||
return FileResponse(path=file_path, filename=os.path.basename(file_path), media_type=media_type)
|
||||
@ -277,6 +308,7 @@ async def add_documents(
|
||||
|
||||
item_type = "URL" if content_type == "url" else "文件"
|
||||
failed_count = len([_p for _p in processed_items if _p.get("status") == "failed"])
|
||||
success_items = [_p for _p in processed_items if _p.get("status") == "done"]
|
||||
summary = {
|
||||
"db_id": db_id,
|
||||
"item_type": item_type,
|
||||
@ -284,6 +316,17 @@ async def add_documents(
|
||||
"failed": failed_count,
|
||||
}
|
||||
message = f"{item_type}处理完成，失败 {failed_count} 个" if failed_count else f"{item_type}处理完成"
|
||||
|
||||
for success_item in success_items:
|
||||
# 使用异步上传到minio的对应知识库,同名文件会被覆盖
|
||||
async with aiofiles.open(success_item["path"], "rb") as f:
|
||||
file_bytes = await f.read()
|
||||
# 上传的bucket名为ref-{refdb},refdb中的_替换为-
|
||||
refdb = db_id.replace("_", "-")
|
||||
url = await aupload_file_to_minio(
|
||||
f"ref-{refdb}", success_item["filename"], file_bytes, success_item["file_type"]
|
||||
)
|
||||
logger.info(f"上传文件成功: {url}")
|
||||
await context.set_result(summary | {"items": processed_items})
|
||||
await context.set_progress(100.0, message)
|
||||
return summary | {"items": processed_items}
|
||||
@ -354,6 +397,10 @@ async def delete_document(db_id: str, doc_id: str, current_user: User = Depends(
|
||||
"""删除文档"""
|
||||
logger.debug(f"DELETE document {doc_id} info in {db_id}")
|
||||
try:
|
||||
file_meta_info = await knowledge_base.get_file_basic_info(db_id, doc_id)
|
||||
file_name = file_meta_info.get("meta", {}).get("filename")
|
||||
minio_client = get_minio_client()
|
||||
await minio_client.adelete_file("ref-" + db_id.replace("_", "-"), file_name)
|
||||
await knowledge_base.delete_file(db_id, doc_id)
|
||||
return {"message": "删除成功"}
|
||||
except Exception as e:
|
||||
@ -473,24 +520,6 @@ async def download_document(db_id: str, doc_id: str, request: Request, current_u
|
||||
logger.debug(f"Download document {doc_id} from {db_id}")
|
||||
try:
|
||||
file_info = await knowledge_base.get_file_basic_info(db_id, doc_id)
|
||||
if not file_info:
|
||||
raise HTTPException(status_code=404, detail="File not found")
|
||||
|
||||
file_path = file_info.get("meta", {}).get("path")
|
||||
if not file_path:
|
||||
raise HTTPException(status_code=404, detail="File path not found in metadata")
|
||||
|
||||
# 安全检查：验证文件路径
|
||||
from src.knowledge.utils.kb_utils import validate_file_path
|
||||
|
||||
try:
|
||||
normalized_path = validate_file_path(file_path, db_id)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=403, detail=str(e))
|
||||
|
||||
if not os.path.exists(normalized_path):
|
||||
raise HTTPException(status_code=404, detail=f"File not found on disk: {file_info=}")
|
||||
|
||||
# 获取文件扩展名和MIME类型，解码URL编码的文件名
|
||||
filename = file_info.get("meta", {}).get("filename", "file")
|
||||
logger.debug(f"Original filename from database: {filename}")
|
||||
@ -505,46 +534,31 @@ async def download_document(db_id: str, doc_id: str, request: Request, current_u
|
||||
|
||||
_, ext = os.path.splitext(decoded_filename)
|
||||
|
||||
media_types = {
|
||||
".pdf": "application/pdf",
|
||||
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
".doc": "application/msword",
|
||||
".txt": "text/plain",
|
||||
".md": "text/markdown",
|
||||
".json": "application/json",
|
||||
".csv": "text/csv",
|
||||
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
".xls": "application/vnd.ms-excel",
|
||||
".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
".ppt": "application/vnd.ms-powerpoint",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".bmp": "image/bmp",
|
||||
".svg": "image/svg+xml",
|
||||
".zip": "application/zip",
|
||||
".rar": "application/x-rar-compressed",
|
||||
".7z": "application/x-7z-compressed",
|
||||
".tar": "application/x-tar",
|
||||
".gz": "application/gzip",
|
||||
".html": "text/html",
|
||||
".htm": "text/html",
|
||||
".xml": "text/xml",
|
||||
".css": "text/css",
|
||||
".js": "application/javascript",
|
||||
".py": "text/x-python",
|
||||
".java": "text/x-java-source",
|
||||
".cpp": "text/x-c++src",
|
||||
".c": "text/x-csrc",
|
||||
".h": "text/x-chdr",
|
||||
".hpp": "text/x-c++hdr",
|
||||
}
|
||||
media_type = media_types.get(ext.lower(), "application/octet-stream")
|
||||
|
||||
# 创建自定义FileResponse，避免文件名编码问题
|
||||
response = StarletteFileResponse(path=normalized_path, media_type=media_type)
|
||||
minio_client = get_minio_client()
|
||||
minio_response = await minio_client.adownload_response(
|
||||
bucket_name="ref-" + db_id.replace("_", "-"),
|
||||
object_name=filename,
|
||||
)
|
||||
|
||||
# 创建流式生成器
|
||||
async def minio_stream():
|
||||
try:
|
||||
while True:
|
||||
chunk = await asyncio.to_thread(minio_response.read, 8192)
|
||||
if not chunk:
|
||||
break
|
||||
yield chunk
|
||||
finally:
|
||||
minio_response.close()
|
||||
minio_response.release_conn()
|
||||
|
||||
# 创建StreamingResponse
|
||||
response = StreamingResponse(
|
||||
minio_stream(),
|
||||
media_type=media_type,
|
||||
)
|
||||
# 正确处理中文文件名的HTTP头部设置
|
||||
# HTTP头部只能包含ASCII字符，所以需要对中文文件名进行编码
|
||||
try:
|
||||
@ -1041,7 +1055,10 @@ async def upload_file(
|
||||
upload_dir = os.path.join(config.save_dir, "database", "uploads")
|
||||
|
||||
basename, ext = os.path.splitext(file.filename)
|
||||
filename = f"{basename}_{hashstr(basename, 4, with_salt=True)}{ext}".lower()
|
||||
# TODO:
|
||||
# 后续修改为遇到同名文件则在上传区域提示，是否删除旧文件，同时 filename name 也就不用添加 hash 了
|
||||
filename = f"{basename}_{hashstr(basename, 4, with_salt=True, salt='fixed_salt')}{ext}".lower()
|
||||
|
||||
file_path = os.path.join(upload_dir, filename)
|
||||
|
||||
# 在线程池中执行同步文件系统操作，避免阻塞事件循环
|
||||
@ -1049,15 +1066,13 @@ async def upload_file(
|
||||
|
||||
file_bytes = await file.read()
|
||||
|
||||
# 在线程池中执行计算密集型操作，避免阻塞事件循环
|
||||
content_hash = await asyncio.to_thread(calculate_content_hash, file_bytes)
|
||||
content_hash = await calculate_content_hash(file_bytes)
|
||||
|
||||
# 在线程池中执行同步数据库查询，避免阻塞事件循环
|
||||
file_exists = await asyncio.to_thread(knowledge_base.file_existed_in_db, db_id, content_hash)
|
||||
file_exists = await knowledge_base.file_existed_in_db(db_id, content_hash)
|
||||
if file_exists:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="数据库中已经存在了相同文件，File with the same content already exists in this database",
|
||||
detail="数据库中已经存在了相同内容文件，File with the same content already exists in this database",
|
||||
)
|
||||
|
||||
# 使用异步文件写入，避免阻塞事件循环
|
||||
|
||||
@ -173,7 +173,7 @@ class ChromaKB(KnowledgeBase):
|
||||
|
||||
return chunks
|
||||
|
||||
async def add_content(self, db_id: str, items: list[str], params: dict | None) -> list[dict]:
|
||||
async def add_content(self, db_id: str, items: list[str], params: dict | None = None) -> list[dict]:
|
||||
"""添加内容（文件/URL）"""
|
||||
if db_id not in self.databases_meta:
|
||||
raise ValueError(f"Database {db_id} not found")
|
||||
@ -187,7 +187,7 @@ class ChromaKB(KnowledgeBase):
|
||||
|
||||
for item in items:
|
||||
# 准备文件元数据
|
||||
metadata = prepare_item_metadata(item, content_type, db_id, params=params)
|
||||
metadata = await prepare_item_metadata(item, content_type, db_id, params=params)
|
||||
file_id = metadata["file_id"]
|
||||
filename = metadata["filename"]
|
||||
|
||||
|
||||
@ -226,7 +226,7 @@ class LightRagKB(KnowledgeBase):
|
||||
|
||||
for item in items:
|
||||
# 准备文件元数据
|
||||
metadata = prepare_item_metadata(item, content_type, db_id, params=params)
|
||||
metadata = await prepare_item_metadata(item, content_type, db_id, params=params)
|
||||
file_id = metadata["file_id"]
|
||||
item_path = metadata["path"]
|
||||
|
||||
|
||||
@ -226,7 +226,7 @@ class MilvusKB(KnowledgeBase):
|
||||
processed_items_info = []
|
||||
|
||||
for item in items:
|
||||
metadata = prepare_item_metadata(item, content_type, db_id, params=params)
|
||||
metadata = await prepare_item_metadata(item, content_type, db_id, params=params)
|
||||
file_id = metadata["file_id"]
|
||||
filename = metadata["filename"]
|
||||
|
||||
|
||||
@ -4,6 +4,7 @@ import re
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
import aiofiles
|
||||
from langchain_community.document_loaders import (
|
||||
CSVLoader,
|
||||
JSONLoader,
|
||||
@ -46,6 +47,47 @@ def is_supported_file_extension(file_name: str | os.PathLike[str]) -> bool:
|
||||
return Path(file_name).suffix.lower() in SUPPORTED_FILE_EXTENSIONS
|
||||
|
||||
|
||||
def _make_unique_columns(columns: list) -> list:
|
||||
"""
|
||||
处理重复的列名,给重复的列添加数字后缀
|
||||
|
||||
Args:
|
||||
columns: 原始列名列表
|
||||
|
||||
Returns:
|
||||
处理后的唯一列名列表
|
||||
|
||||
Example:
|
||||
['A', 'B', 'A', 'C', 'B'] -> ['A', 'B', 'A_2', 'C', 'B_2']
|
||||
"""
|
||||
if not columns:
|
||||
return columns
|
||||
|
||||
seen = {}
|
||||
unique_columns = []
|
||||
|
||||
for col in columns:
|
||||
# 处理 None 或空值
|
||||
if col is None or (isinstance(col, str) and not col.strip()):
|
||||
col = "Unnamed"
|
||||
|
||||
# 转换为字符串
|
||||
col_str = str(col)
|
||||
|
||||
# 如果这个列名已经出现过
|
||||
if col_str in seen:
|
||||
seen[col_str] += 1
|
||||
# 添加后缀
|
||||
unique_col = f"{col_str}_{seen[col_str]}"
|
||||
unique_columns.append(unique_col)
|
||||
else:
|
||||
# 第一次出现
|
||||
seen[col_str] = 1
|
||||
unique_columns.append(col_str)
|
||||
|
||||
return unique_columns
|
||||
|
||||
|
||||
def _extract_word_text(file_path: Path) -> str:
|
||||
"""
|
||||
Parse Word documents (.doc/.docx) into plain text.
|
||||
@ -332,20 +374,66 @@ async def process_file_to_markdown(file_path: str, params: dict | None = None) -
|
||||
elif file_ext in [".xls", ".xlsx"]:
|
||||
# 处理 Excel 文件
|
||||
import pandas as pd
|
||||
from openpyxl import load_workbook
|
||||
|
||||
# 读取所有工作表
|
||||
excel_file = pd.ExcelFile(file_path_obj)
|
||||
markdown_content = f"# {file_path_obj.name}\n\n"
|
||||
|
||||
for sheet_name in excel_file.sheet_names:
|
||||
df = pd.read_excel(file_path_obj, sheet_name=sheet_name)
|
||||
# 使用 openpyxl 加载工作簿以正确处理合并单元格
|
||||
wb = load_workbook(file_path_obj, data_only=True)
|
||||
|
||||
for sheet_name in wb.sheetnames:
|
||||
ws = wb[sheet_name]
|
||||
|
||||
# 先取消所有合并单元格，并填充值
|
||||
merged_ranges = list(ws.merged_cells.ranges)
|
||||
|
||||
for merged_range in merged_ranges:
|
||||
# 获取合并区域左上角单元格的值
|
||||
min_row, min_col, max_row, max_col = (
|
||||
merged_range.min_row,
|
||||
merged_range.min_col,
|
||||
merged_range.max_row,
|
||||
merged_range.max_col,
|
||||
)
|
||||
|
||||
# 获取左上角单元格的值
|
||||
top_left_value = ws.cell(row=min_row, column=min_col).value
|
||||
|
||||
# 取消合并
|
||||
ws.unmerge_cells(start_row=min_row, start_column=min_col, end_row=max_row, end_column=max_col)
|
||||
|
||||
# 在所有原合并单元格区域填充值
|
||||
for row in range(min_row, max_row + 1):
|
||||
for col in range(min_col, max_col + 1):
|
||||
ws.cell(row=row, column=col).value = top_left_value
|
||||
|
||||
# 转换为DataFrame
|
||||
data = []
|
||||
for row in ws.iter_rows(values_only=True):
|
||||
data.append(row)
|
||||
|
||||
# 第一行作为列名
|
||||
columns = data[0] if data else []
|
||||
df_data = data[1:] if len(data) > 1 else []
|
||||
|
||||
# 处理重复的列名,给重复的列添加后缀
|
||||
columns = _make_unique_columns(columns)
|
||||
|
||||
df = pd.DataFrame(df_data, columns=columns)
|
||||
|
||||
markdown_content += f"## {sheet_name}\n\n"
|
||||
|
||||
# 将每一行数据与表头组合成独立的表格
|
||||
for index, row in df.iterrows():
|
||||
# 创建包含表头和当前行的小表格
|
||||
row_df = pd.DataFrame([row], columns=df.columns)
|
||||
markdown_table = row_df.to_markdown(index=False)
|
||||
# 在最左列添加标题字段
|
||||
table_title = f"{file_path_obj.stem} - {sheet_name}" # 使用"文件名 - Sheet名"作为标题
|
||||
df.insert(0, "表格标题", table_title)
|
||||
|
||||
# 将每10行数据与表头组合成独立的表格
|
||||
chunk_size = 10
|
||||
for i in range(0, len(df), chunk_size):
|
||||
# 获取当前10行数据(或剩余不足10行的数据)
|
||||
chunk_df = df.iloc[i : i + chunk_size]
|
||||
markdown_content += f"### 数据行 {i + 1}-{min(i + chunk_size, len(df))}\n\n"
|
||||
markdown_table = chunk_df.to_markdown(index=False)
|
||||
markdown_content += f"{markdown_table}\n\n"
|
||||
|
||||
return markdown_content.strip()
|
||||
@ -354,8 +442,9 @@ async def process_file_to_markdown(file_path: str, params: dict | None = None) -
|
||||
# 处理 JSON 文件
|
||||
import json
|
||||
|
||||
with open(file_path_obj, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
async with aiofiles.open(file_path_obj, encoding="utf-8") as f:
|
||||
content = await f.read()
|
||||
data = json.loads(content)
|
||||
# 将 JSON 数据格式化为 markdown 代码块
|
||||
json_str = json.dumps(data, ensure_ascii=False, indent=2)
|
||||
return f"# {file_path_obj.name}\n\n```json\n{json_str}\n```"
|
||||
@ -430,7 +519,7 @@ def _process_zip_file(zip_path: str, db_id: str) -> dict:
|
||||
markdown_content = _replace_image_links(markdown_content, images_info)
|
||||
|
||||
# 4. 生成结果
|
||||
content_hash = calculate_content_hash(markdown_content.encode("utf-8"))
|
||||
content_hash = asyncio.run(calculate_content_hash(markdown_content.encode("utf-8")))
|
||||
|
||||
return {
|
||||
"markdown_content": markdown_content,
|
||||
@ -458,7 +547,7 @@ def _find_images_directory(zip_file: zipfile.ZipFile, md_file_path: str) -> str
|
||||
return None
|
||||
|
||||
|
||||
def _process_images(zip_file: zipfile.ZipFile, images_dir: str, db_id: str, md_file_path: str) -> list[dict]:
|
||||
async def _process_images(zip_file: zipfile.ZipFile, images_dir: str, db_id: str, md_file_path: str) -> list[dict]:
|
||||
"""处理图片：上传到MinIO并返回信息"""
|
||||
# 支持的图片格式
|
||||
SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}
|
||||
@ -475,7 +564,7 @@ def _process_images(zip_file: zipfile.ZipFile, images_dir: str, db_id: str, md_f
|
||||
image_names = [n for n in zip_file.namelist() if n.startswith(images_dir + "/")]
|
||||
|
||||
# 上传图片到MinIO
|
||||
minio_client = get_minio_client()
|
||||
minio_client = await get_minio_client()
|
||||
bucket_name = "kb-images"
|
||||
minio_client.ensure_bucket_exists(bucket_name)
|
||||
|
||||
|
||||
@ -351,7 +351,33 @@ class KnowledgeBaseManager:
|
||||
os.makedirs(general_uploads, exist_ok=True)
|
||||
return general_uploads
|
||||
|
||||
def file_existed_in_db(self, db_id: str | None, content_hash: str | None) -> bool:
|
||||
async def file_name_existed_in_db(self, db_id: str | None, file_name: str | None) -> bool:
|
||||
"""检查指定数据库中是否存在同名的文件"""
|
||||
if not db_id or not file_name:
|
||||
return False
|
||||
try:
|
||||
kb_instance = self._get_kb_for_database(db_id)
|
||||
except KBNotFoundError:
|
||||
return False
|
||||
|
||||
for file_info in kb_instance.files_meta.values():
|
||||
if file_info.get("database_id") != db_id:
|
||||
continue
|
||||
if file_info.get("status") == "failed":
|
||||
continue
|
||||
if file_info.get("file_name") == file_name:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
async def update_file(self, db_id: str, region_file_id: str, file_name: str, params: dict | None = None) -> dict:
|
||||
"""对单个文件执行更新"""
|
||||
kb_instance = self._get_kb_for_database(db_id)
|
||||
await kb_instance.delete_file(db_id, region_file_id)
|
||||
data_list = await kb_instance.add_content(db_id, [file_name], params or {})
|
||||
return data_list[0]
|
||||
|
||||
async def file_existed_in_db(self, db_id: str | None, content_hash: str | None) -> bool:
|
||||
"""检查指定数据库中是否存在相同内容哈希的文件"""
|
||||
if not db_id or not content_hash:
|
||||
return False
|
||||
|
||||
@ -3,6 +3,7 @@ import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import aiofiles
|
||||
from langchain_text_splitters import MarkdownTextSplitter
|
||||
|
||||
from src import config
|
||||
@ -89,7 +90,7 @@ def split_text_into_chunks(text: str, file_id: str, filename: str, params: dict
|
||||
chunks.append(
|
||||
{
|
||||
"id": f"{file_id}_chunk_{chunk_index}",
|
||||
"content": chunk_content.strip(),
|
||||
"content": chunk_content, # .strip(),
|
||||
"file_id": file_id,
|
||||
"filename": filename,
|
||||
"chunk_index": chunk_index,
|
||||
@ -102,7 +103,7 @@ def split_text_into_chunks(text: str, file_id: str, filename: str, params: dict
|
||||
return chunks
|
||||
|
||||
|
||||
def calculate_content_hash(data: bytes | bytearray | str | os.PathLike[str] | Path) -> str:
|
||||
async def calculate_content_hash(data: bytes | bytearray | str | os.PathLike[str] | Path) -> str:
|
||||
"""
|
||||
计算文件内容的 SHA-256 哈希值。
|
||||
|
||||
@ -120,15 +121,18 @@ def calculate_content_hash(data: bytes | bytearray | str | os.PathLike[str] | Pa
|
||||
|
||||
if isinstance(data, (str, os.PathLike, Path)):
|
||||
path = Path(data)
|
||||
with path.open("rb") as file_handle:
|
||||
for chunk in iter(lambda: file_handle.read(8192), b""):
|
||||
async with aiofiles.open(path, "rb") as file_handle:
|
||||
chunk = await file_handle.read(8192)
|
||||
while chunk:
|
||||
sha256.update(chunk)
|
||||
chunk = await file_handle.read(8192)
|
||||
|
||||
return sha256.hexdigest()
|
||||
|
||||
raise TypeError(f"Unsupported data type for hashing: {type(data)!r}")
|
||||
|
||||
|
||||
def prepare_item_metadata(item: str, content_type: str, db_id: str, params: dict | None = None) -> dict:
|
||||
async def prepare_item_metadata(item: str, content_type: str, db_id: str, params: dict | None = None) -> dict:
|
||||
"""
|
||||
准备文件或URL的元数据
|
||||
|
||||
@ -147,7 +151,7 @@ def prepare_item_metadata(item: str, content_type: str, db_id: str, params: dict
|
||||
content_hash = None
|
||||
try:
|
||||
if file_path.exists():
|
||||
content_hash = calculate_content_hash(file_path)
|
||||
content_hash = await calculate_content_hash(file_path)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning(f"Failed to calculate content hash for {file_path}: {exc}")
|
||||
else:
|
||||
|
||||
@ -3,11 +3,15 @@ MinIO 存储客户端
|
||||
简化的 MinIO 对象存储操作
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from datetime import timedelta
|
||||
from io import BytesIO
|
||||
|
||||
from urllib3 import BaseHTTPResponse
|
||||
|
||||
from minio import Minio
|
||||
from minio.error import S3Error
|
||||
from src.utils import logger
|
||||
@ -19,6 +23,10 @@ class StorageError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class StorageUploadError(StorageError):
|
||||
"""存储相关异常基类"""
|
||||
|
||||
|
||||
class UploadResult:
|
||||
"""简化的上传结果"""
|
||||
|
||||
@ -73,8 +81,8 @@ class MinIOClient:
|
||||
"""确保存储桶存在"""
|
||||
try:
|
||||
created = False
|
||||
if not self.client.bucket_exists(bucket_name):
|
||||
self.client.make_bucket(bucket_name)
|
||||
if not self.client.bucket_exists(bucket_name=bucket_name):
|
||||
self.client.make_bucket(bucket_name=bucket_name)
|
||||
created = True
|
||||
logger.info(f"存储桶 '{bucket_name}' 已创建")
|
||||
|
||||
@ -95,7 +103,7 @@ class MinIOClient:
|
||||
) -> UploadResult:
|
||||
"""上传文件到 MinIO"""
|
||||
try:
|
||||
self.ensure_bucket_exists(bucket_name)
|
||||
self.ensure_bucket_exists(bucket_name=bucket_name)
|
||||
|
||||
data_stream = BytesIO(data)
|
||||
result = self.client.put_object(
|
||||
@ -116,6 +124,14 @@ class MinIOClient:
|
||||
logger.error(error_msg)
|
||||
raise StorageError(error_msg)
|
||||
|
||||
async def aupload_file(
|
||||
self, bucket_name: str, object_name: str, data: bytes, content_type: str = "application/octet-stream"
|
||||
) -> UploadResult:
|
||||
result = await asyncio.to_thread(
|
||||
self.upload_file, bucket_name=bucket_name, object_name=object_name, data=data, content_type=content_type
|
||||
)
|
||||
return result
|
||||
|
||||
def upload_file_from_path(self, bucket_name: str, object_name: str, file_path: str) -> UploadResult:
|
||||
"""从文件路径上传文件"""
|
||||
try:
|
||||
@ -152,7 +168,7 @@ class MinIOClient:
|
||||
def download_file(self, bucket_name: str, object_name: str) -> bytes:
|
||||
"""下载文件"""
|
||||
try:
|
||||
response = self.client.get_object(bucket_name, object_name)
|
||||
response = self.client.get_object(bucket_name=bucket_name, object_name=object_name)
|
||||
data = response.read()
|
||||
response.close()
|
||||
logger.info(f"成功下载 '{object_name}' 从存储桶 '{bucket_name}'")
|
||||
@ -163,10 +179,46 @@ class MinIOClient:
|
||||
raise StorageError(f"对象 '{object_name}' 在存储桶 '{bucket_name}' 中不存在")
|
||||
raise StorageError(f"下载文件失败: {e}")
|
||||
|
||||
async def adownload_response(self, bucket_name: str, object_name: str) -> BaseHTTPResponse:
|
||||
"""异步下载文件"""
|
||||
try:
|
||||
response = await asyncio.to_thread(
|
||||
self.client.get_object,
|
||||
bucket_name=bucket_name,
|
||||
object_name=object_name,
|
||||
)
|
||||
return response
|
||||
|
||||
except S3Error as e:
|
||||
if "NoSuchKey" in str(e):
|
||||
raise StorageError(f"对象 '{object_name}' 在存储桶 '{bucket_name}' 中不存在")
|
||||
raise StorageError(f"下载文件失败: {e}")
|
||||
|
||||
async def adownload_file(self, bucket_name: str, object_name: str) -> bytes:
|
||||
"""异步下载文件"""
|
||||
try:
|
||||
response = await asyncio.to_thread(self.client.get_object, bucket_name=bucket_name, object_name=object_name)
|
||||
data = await asyncio.to_thread(response.read)
|
||||
response.close()
|
||||
logger.info(f"成功下载 '{object_name}' 从存储桶 '{bucket_name}'")
|
||||
return data
|
||||
|
||||
except S3Error as e:
|
||||
if "NoSuchKey" in str(e):
|
||||
raise StorageError(f"对象 '{object_name}' 在存储桶 '{bucket_name}' 中不存在")
|
||||
raise StorageError(f"下载文件失败: {e}")
|
||||
|
||||
def get_presigned_url(self, bucket_name: str, object_name: str, days=7) -> str:
|
||||
"""将minio放在内网访问，外部通过返回代理链接访问"""
|
||||
res_url = self.client.get_presigned_url(
|
||||
method="GET", bucket_name=bucket_name, object_name=object_name, expires=timedelta(days=days)
|
||||
)
|
||||
return res_url
|
||||
|
||||
def delete_file(self, bucket_name: str, object_name: str) -> bool:
|
||||
"""删除文件"""
|
||||
try:
|
||||
self.client.remove_object(bucket_name, object_name)
|
||||
self.client.remove_object(bucket_name=bucket_name, object_name=object_name)
|
||||
logger.info(f"成功删除 '{object_name}' 从存储桶 '{bucket_name}'")
|
||||
return True
|
||||
|
||||
@ -176,10 +228,19 @@ class MinIOClient:
|
||||
return False
|
||||
raise StorageError(f"删除文件失败: {e}")
|
||||
|
||||
async def adelete_file(self, bucket_name: str, object_name: str) -> bool:
|
||||
"""删除文件"""
|
||||
result = await asyncio.to_thread(
|
||||
self.delete_file,
|
||||
bucket_name=bucket_name,
|
||||
object_name=object_name,
|
||||
)
|
||||
return result
|
||||
|
||||
def file_exists(self, bucket_name: str, object_name: str) -> bool:
|
||||
"""检查文件是否存在"""
|
||||
try:
|
||||
self.client.stat_object(bucket_name, object_name)
|
||||
self.client.stat_object(bucket_name=bucket_name, object_name=object_name)
|
||||
return True
|
||||
except S3Error as e:
|
||||
if "NoSuchKey" in str(e):
|
||||
@ -210,7 +271,7 @@ class MinIOClient:
|
||||
}
|
||||
|
||||
try:
|
||||
self.client.set_bucket_policy(bucket_name, json.dumps(policy))
|
||||
self.client.set_bucket_policy(bucket_name=bucket_name, policy=json.dumps(policy))
|
||||
except S3Error as e:
|
||||
logger.warning(f"设置存储桶 '{bucket_name}' 公共读取策略失败: {e}")
|
||||
raise StorageError(f"无法设置存储桶公共访问策略: {e}")
|
||||
@ -241,7 +302,29 @@ def upload_image_to_minio(bucket_name: str, data: bytes, file_extension: str = "
|
||||
"""
|
||||
client = get_minio_client()
|
||||
file_name = f"{uuid.uuid4()}.{file_extension}"
|
||||
result = client.upload_file(
|
||||
client.upload_file(
|
||||
bucket_name=bucket_name, object_name=file_name, data=data, content_type=f"image/{file_extension}"
|
||||
)
|
||||
return result.url
|
||||
res_url = client.get_presigned_url(bucket_name=bucket_name, object_name=file_name, days=7)
|
||||
return res_url
|
||||
|
||||
|
||||
async def aupload_file_to_minio(bucket_name: str, file_name: str, data: bytes, file_extension: str) -> str:
|
||||
"""
|
||||
通过字节上传文件到 MinIO的异步接口，根据输入的file_extension确定文件格式，并返回资源url
|
||||
|
||||
Args:
|
||||
bucket_name: bucket_name
|
||||
file_name : filename
|
||||
data: 文件字节流
|
||||
file_extension: 输入的拓展名
|
||||
Returns:
|
||||
str: 文件访问 URL
|
||||
"""
|
||||
client = get_minio_client()
|
||||
# 根据扩展名猜测 content_type
|
||||
content_type = client._guess_content_type(file_extension)
|
||||
# 上传文件
|
||||
await client.aupload_file(bucket_name, file_name, data, content_type)
|
||||
res_url = client.get_presigned_url(bucket_name, file_name, days=7)
|
||||
return res_url
|
||||
|
||||
@ -26,7 +26,7 @@ def is_text_pdf(pdf_path):
|
||||
return text_ratio > 0.5
|
||||
|
||||
|
||||
def hashstr(input_string, length=None, with_salt=False):
|
||||
def hashstr(input_string, length=None, with_salt=False, salt=None):
|
||||
"""生成字符串的哈希值
|
||||
Args:
|
||||
input_string: 输入字符串
|
||||
@ -41,7 +41,8 @@ def hashstr(input_string, length=None, with_salt=False):
|
||||
encoded_string = str(input_string).encode("utf-8", errors="replace")
|
||||
|
||||
if with_salt:
|
||||
salt = str(time.time())
|
||||
if not salt:
|
||||
salt = str(time.time())
|
||||
encoded_string = (encoded_string.decode("utf-8") + salt).encode("utf-8")
|
||||
|
||||
hash = hashlib.md5(encoded_string).hexdigest()
|
||||
|
||||
@ -16,7 +16,7 @@ export function findOverlap(str1, str2) {
|
||||
let overlap = '';
|
||||
|
||||
// 从最长可能的重叠开始检查
|
||||
for (let i = maxOverlap; i > 0; i--) {
|
||||
for (let i = maxOverlap; i > 10; i--) {
|
||||
const endStr1 = str1.slice(-i);
|
||||
const startStr2 = str2.slice(0, i);
|
||||
|
||||
@ -63,7 +63,11 @@ export function mergeChunks(chunks) {
|
||||
|
||||
if (newContent.length > 0) {
|
||||
const startOffset = currentContent.length;
|
||||
currentContent += newContent;
|
||||
if (overlap.length > 0) {
|
||||
currentContent += newContent;
|
||||
} else {
|
||||
currentContent += `\n${newContent}`;
|
||||
}
|
||||
merged.push({
|
||||
...chunk,
|
||||
startOffset,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user