499 lines
19 KiB
Python
499 lines
19 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Queue tasks for file processing integration.
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import time
|
||
import hashlib
|
||
import shutil
|
||
from typing import Dict, List, Optional, Any
|
||
|
||
from task_queue.config import huey
|
||
from task_queue.manager import queue_manager
|
||
from task_queue.task_status import task_status_store
|
||
from utils import download_dataset_files, save_processed_files_log, load_processed_files_log
|
||
from utils.dataset_manager import remove_dataset_directory_by_key
|
||
|
||
|
||
def scan_upload_folder(upload_dir: str) -> List[str]:
|
||
"""
|
||
扫描上传文件夹中的所有支持格式的文件
|
||
|
||
Args:
|
||
upload_dir: 上传文件夹路径
|
||
|
||
Returns:
|
||
List[str]: 支持的文件路径列表
|
||
"""
|
||
supported_extensions = {
|
||
# 文本文件
|
||
'.txt', '.md', '.rtf',
|
||
# 文档文件
|
||
'.doc', '.docx', '.pdf', '.odt',
|
||
# 表格文件
|
||
'.xls', '.xlsx', '.csv', '.ods',
|
||
# 演示文件
|
||
'.ppt', '.pptx', '.odp',
|
||
# 电子书
|
||
'.epub', '.mobi',
|
||
# 网页文件
|
||
'.html', '.htm',
|
||
# 配置文件
|
||
'.json', '.xml', '.yaml', '.yml',
|
||
# 代码文件
|
||
'.py', '.js', '.java', '.cpp', '.c', '.go', '.rs',
|
||
# 压缩文件
|
||
'.zip', '.rar', '.7z', '.tar', '.gz'
|
||
}
|
||
|
||
scanned_files = []
|
||
|
||
if not os.path.exists(upload_dir):
|
||
return scanned_files
|
||
|
||
for root, dirs, files in os.walk(upload_dir):
|
||
for file in files:
|
||
# 跳过隐藏文件和系统文件
|
||
if file.startswith('.') or file.startswith('~'):
|
||
continue
|
||
|
||
file_path = os.path.join(root, file)
|
||
file_extension = os.path.splitext(file)[1].lower()
|
||
|
||
# 检查文件扩展名是否支持
|
||
if file_extension in supported_extensions:
|
||
scanned_files.append(file_path)
|
||
else:
|
||
# 对于没有扩展名的文件，也尝试处理（可能是文本文件）
|
||
if not file_extension:
|
||
try:
|
||
# 尝试读取文件头部来判断是否为文本文件
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
f.read(1024) # 读取前1KB
|
||
scanned_files.append(file_path)
|
||
except (UnicodeDecodeError, PermissionError):
|
||
# 不是文本文件或无法读取，跳过
|
||
pass
|
||
|
||
return scanned_files
|
||
|
||
|
||
@huey.task()
|
||
def process_files_async(
|
||
dataset_id: str,
|
||
files: Optional[Dict[str, List[str]]] = None,
|
||
upload_folder: Optional[Dict[str, str]] = None,
|
||
task_id: Optional[str] = None
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
异步处理文件任务 - 与现有files/process API兼容
|
||
|
||
Args:
|
||
dataset_id: 项目唯一ID
|
||
files: 按key分组的文件路径字典
|
||
upload_folder: 上传文件夹字典，按组名组织文件夹，例如 {'group1': 'my_project1', 'group2': 'my_project2'}
|
||
task_id: 任务ID（用于状态跟踪）
|
||
|
||
Returns:
|
||
处理结果字典
|
||
"""
|
||
try:
|
||
print(f"开始异步处理文件任务，项目ID: {dataset_id}")
|
||
|
||
# 如果有task_id，设置初始状态
|
||
if task_id:
|
||
task_status_store.set_status(
|
||
task_id=task_id,
|
||
unique_id=dataset_id,
|
||
status="running"
|
||
)
|
||
|
||
# 确保项目目录存在
|
||
project_dir = os.path.join("projects", "data", dataset_id)
|
||
if not os.path.exists(project_dir):
|
||
os.makedirs(project_dir, exist_ok=True)
|
||
|
||
# 处理文件：使用按key分组格式
|
||
processed_files_by_key = {}
|
||
|
||
# 如果提供了upload_folder，扫描这些文件夹中的文件
|
||
if upload_folder and not files:
|
||
scanned_files_by_group = {}
|
||
total_scanned_files = 0
|
||
|
||
for group_name, folder_name in upload_folder.items():
|
||
# 安全性检查：防止路径遍历攻击
|
||
safe_folder_name = os.path.basename(folder_name)
|
||
upload_dir = os.path.join("projects", "uploads", safe_folder_name)
|
||
|
||
if os.path.exists(upload_dir):
|
||
scanned_files = scan_upload_folder(upload_dir)
|
||
if scanned_files:
|
||
scanned_files_by_group[group_name] = scanned_files
|
||
total_scanned_files += len(scanned_files)
|
||
print(f"从上传文件夹 '{safe_folder_name}' (组: {group_name}) 扫描到 {len(scanned_files)} 个文件")
|
||
else:
|
||
print(f"上传文件夹 '{safe_folder_name}' (组: {group_name}) 中没有找到支持的文件")
|
||
else:
|
||
print(f"上传文件夹不存在: {upload_dir} (组: {group_name})")
|
||
|
||
if scanned_files_by_group:
|
||
files = scanned_files_by_group
|
||
print(f"总共从 {len(scanned_files_by_group)} 个组扫描到 {total_scanned_files} 个文件")
|
||
else:
|
||
print(f"所有上传文件夹中都没有找到支持的文件")
|
||
|
||
if files:
|
||
# 使用请求中的文件（按key分组）
|
||
# 由于这是异步任务，需要同步调用
|
||
import asyncio
|
||
try:
|
||
loop = asyncio.get_event_loop()
|
||
except RuntimeError:
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
|
||
processed_files_by_key = loop.run_until_complete(download_dataset_files(dataset_id, files))
|
||
total_files = sum(len(files_list) for files_list in processed_files_by_key.values())
|
||
print(f"异步处理了 {total_files} 个数据集文件，涉及 {len(processed_files_by_key)} 个key，项目ID: {dataset_id}")
|
||
else:
|
||
print(f"请求中未提供文件，项目ID: {dataset_id}")
|
||
|
||
# 收集项目目录下所有的 document.txt 文件
|
||
document_files = []
|
||
for root, dirs, files_list in os.walk(project_dir):
|
||
for file in files_list:
|
||
if file == "document.txt":
|
||
document_files.append(os.path.join(root, file))
|
||
|
||
# 生成项目README.md文件
|
||
try:
|
||
from utils.project_manager import save_project_readme
|
||
save_project_readme(dataset_id)
|
||
print(f"已生成README.md文件，项目ID: {dataset_id}")
|
||
except Exception as e:
|
||
print(f"生成README.md失败，项目ID: {dataset_id}, 错误: {str(e)}")
|
||
# 不影响主要处理流程，继续执行
|
||
|
||
# 构建结果文件列表
|
||
result_files = []
|
||
for key in processed_files_by_key.keys():
|
||
# 添加对应的dataset document.txt路径
|
||
document_path = os.path.join("projects", "data", dataset_id, "dataset", key, "document.txt")
|
||
if os.path.exists(document_path):
|
||
result_files.append(document_path)
|
||
|
||
# 对于没有在processed_files_by_key中但存在的document.txt文件，也添加到结果中
|
||
existing_document_paths = set(result_files) # 避免重复
|
||
for doc_file in document_files:
|
||
if doc_file not in existing_document_paths:
|
||
result_files.append(doc_file)
|
||
|
||
result = {
|
||
"status": "success",
|
||
"message": f"成功异步处理了 {len(result_files)} 个文档文件，涉及 {len(processed_files_by_key)} 个key",
|
||
"dataset_id": dataset_id,
|
||
"processed_files": result_files,
|
||
"processed_files_by_key": processed_files_by_key,
|
||
"document_files": document_files,
|
||
"total_files_processed": sum(len(files_list) for files_list in processed_files_by_key.values()),
|
||
"processing_time": time.time()
|
||
}
|
||
|
||
# 更新任务状态为完成
|
||
if task_id:
|
||
task_status_store.update_status(
|
||
task_id=task_id,
|
||
status="completed",
|
||
result=result
|
||
)
|
||
|
||
print(f"异步文件处理任务完成: {dataset_id}")
|
||
return result
|
||
|
||
except Exception as e:
|
||
error_msg = f"异步处理文件时发生错误: {str(e)}"
|
||
print(error_msg)
|
||
|
||
# 更新任务状态为错误
|
||
if task_id:
|
||
task_status_store.update_status(
|
||
task_id=task_id,
|
||
status="failed",
|
||
error=error_msg
|
||
)
|
||
|
||
return {
|
||
"status": "error",
|
||
"message": error_msg,
|
||
"dataset_id": dataset_id,
|
||
"error": str(e)
|
||
}
|
||
|
||
|
||
@huey.task()
|
||
def process_files_incremental_async(
|
||
dataset_id: str,
|
||
files_to_add: Optional[Dict[str, List[str]]] = None,
|
||
files_to_remove: Optional[Dict[str, List[str]]] = None,
|
||
system_prompt: Optional[str] = None,
|
||
mcp_settings: Optional[List[Dict]] = None,
|
||
task_id: Optional[str] = None
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
增量处理文件任务 - 支持添加和删除文件
|
||
|
||
Args:
|
||
dataset_id: 项目唯一ID
|
||
files_to_add: 按key分组的要添加的文件路径字典
|
||
files_to_remove: 按key分组的要删除的文件路径字典
|
||
system_prompt: 系统提示词
|
||
mcp_settings: MCP设置
|
||
task_id: 任务ID（用于状态跟踪）
|
||
|
||
Returns:
|
||
处理结果字典
|
||
"""
|
||
try:
|
||
print(f"开始增量处理文件任务，项目ID: {dataset_id}")
|
||
|
||
# 如果有task_id，设置初始状态
|
||
if task_id:
|
||
task_status_store.set_status(
|
||
task_id=task_id,
|
||
unique_id=dataset_id,
|
||
status="running"
|
||
)
|
||
|
||
# 确保项目目录存在
|
||
project_dir = os.path.join("projects", "data", dataset_id)
|
||
if not os.path.exists(project_dir):
|
||
os.makedirs(project_dir, exist_ok=True)
|
||
|
||
# 加载现有的处理日志
|
||
processed_log = load_processed_files_log(dataset_id)
|
||
print(f"加载现有处理日志，包含 {len(processed_log)} 个文件记录")
|
||
|
||
removed_files = []
|
||
added_files = []
|
||
|
||
# 1. 处理删除操作
|
||
if files_to_remove:
|
||
print(f"开始处理删除操作，涉及 {len(files_to_remove)} 个key分组")
|
||
for key, file_list in files_to_remove.items():
|
||
if not file_list: # 如果文件列表为空，删除整个key分组
|
||
print(f"删除整个key分组: {key}")
|
||
if remove_dataset_directory_by_key(dataset_id, key):
|
||
removed_files.append(f"dataset/{key}")
|
||
|
||
# 从处理日志中移除该key的所有记录
|
||
keys_to_remove = [file_hash for file_hash, file_info in processed_log.items()
|
||
if file_info.get('key') == key]
|
||
for file_hash in keys_to_remove:
|
||
del processed_log[file_hash]
|
||
removed_files.append(f"log_entry:{file_hash}")
|
||
else:
|
||
# 删除特定文件
|
||
for file_path in file_list:
|
||
print(f"删除特定文件: {key}/{file_path}")
|
||
|
||
# 实际删除文件
|
||
filename = os.path.basename(file_path)
|
||
|
||
# 删除原始文件
|
||
source_file = os.path.join("projects", "data", dataset_id, "files", key, filename)
|
||
if os.path.exists(source_file):
|
||
os.remove(source_file)
|
||
removed_files.append(f"file:{key}/{filename}")
|
||
|
||
# 删除处理后的文件目录
|
||
processed_dir = os.path.join("projects", "data", dataset_id, "processed", key, filename)
|
||
if os.path.exists(processed_dir):
|
||
shutil.rmtree(processed_dir)
|
||
removed_files.append(f"processed:{key}/{filename}")
|
||
|
||
# 计算文件hash以在日志中查找
|
||
file_hash = hashlib.md5(file_path.encode('utf-8')).hexdigest()
|
||
|
||
# 从处理日志中移除
|
||
if file_hash in processed_log:
|
||
del processed_log[file_hash]
|
||
removed_files.append(f"log_entry:{file_hash}")
|
||
|
||
# 2. 处理添加操作
|
||
processed_files_by_key = {}
|
||
if files_to_add:
|
||
print(f"开始处理添加操作，涉及 {len(files_to_add)} 个key分组")
|
||
# 使用异步处理下载文件
|
||
import asyncio
|
||
try:
|
||
loop = asyncio.get_event_loop()
|
||
except RuntimeError:
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
|
||
processed_files_by_key = loop.run_until_complete(download_dataset_files(dataset_id, files_to_add, incremental_mode=True))
|
||
total_added_files = sum(len(files_list) for files_list in processed_files_by_key.values())
|
||
print(f"异步处理了 {total_added_files} 个数据集文件，涉及 {len(processed_files_by_key)} 个key，项目ID: {dataset_id}")
|
||
|
||
# 记录添加的文件
|
||
for key, files_list in processed_files_by_key.items():
|
||
for file_path in files_list:
|
||
added_files.append(f"{key}/{file_path}")
|
||
else:
|
||
print(f"请求中未提供要添加的文件，项目ID: {dataset_id}")
|
||
|
||
# 保存更新后的处理日志
|
||
save_processed_files_log(dataset_id, processed_log)
|
||
print(f"已更新处理日志，当前包含 {len(processed_log)} 个文件记录")
|
||
|
||
# 保存system_prompt和mcp_settings到项目目录（如果提供）
|
||
if system_prompt:
|
||
system_prompt_file = os.path.join(project_dir, "system_prompt.md")
|
||
with open(system_prompt_file, 'w', encoding='utf-8') as f:
|
||
f.write(system_prompt)
|
||
print(f"已保存system_prompt，项目ID: {dataset_id}")
|
||
|
||
if mcp_settings:
|
||
mcp_settings_file = os.path.join(project_dir, "mcp_settings.json")
|
||
with open(mcp_settings_file, 'w', encoding='utf-8') as f:
|
||
json.dump(mcp_settings, f, ensure_ascii=False, indent=2)
|
||
print(f"已保存mcp_settings，项目ID: {dataset_id}")
|
||
|
||
# 生成项目README.md文件
|
||
try:
|
||
from utils.project_manager import save_project_readme
|
||
save_project_readme(dataset_id)
|
||
print(f"已生成README.md文件，项目ID: {dataset_id}")
|
||
except Exception as e:
|
||
print(f"生成README.md失败，项目ID: {dataset_id}, 错误: {str(e)}")
|
||
# 不影响主要处理流程，继续执行
|
||
|
||
# 收集项目目录下所有的 document.txt 文件
|
||
document_files = []
|
||
for root, dirs, files_list in os.walk(project_dir):
|
||
for file in files_list:
|
||
if file == "document.txt":
|
||
document_files.append(os.path.join(root, file))
|
||
|
||
# 构建结果文件列表
|
||
result_files = []
|
||
for key in processed_files_by_key.keys():
|
||
# 添加对应的dataset document.txt路径
|
||
document_path = os.path.join("projects", "data", dataset_id, "dataset", key, "document.txt")
|
||
if os.path.exists(document_path):
|
||
result_files.append(document_path)
|
||
|
||
# 对于没有在processed_files_by_key中但存在的document.txt文件，也添加到结果中
|
||
existing_document_paths = set(result_files) # 避免重复
|
||
for doc_file in document_files:
|
||
if doc_file not in existing_document_paths:
|
||
result_files.append(doc_file)
|
||
|
||
result = {
|
||
"status": "success",
|
||
"message": f"增量处理完成 - 添加了 {len(added_files)} 个文件，删除了 {len(removed_files)} 个文件，最终保留 {len(result_files)} 个文档文件",
|
||
"dataset_id": dataset_id,
|
||
"removed_files": removed_files,
|
||
"added_files": added_files,
|
||
"processed_files": result_files,
|
||
"processed_files_by_key": processed_files_by_key,
|
||
"document_files": document_files,
|
||
"total_files_added": sum(len(files_list) for files_list in processed_files_by_key.values()),
|
||
"total_files_removed": len(removed_files),
|
||
"final_files_count": len(result_files),
|
||
"processing_time": time.time()
|
||
}
|
||
|
||
# 更新任务状态为完成
|
||
if task_id:
|
||
task_status_store.update_status(
|
||
task_id=task_id,
|
||
status="completed",
|
||
result=result
|
||
)
|
||
|
||
print(f"增量文件处理任务完成: {dataset_id}")
|
||
return result
|
||
|
||
except Exception as e:
|
||
error_msg = f"增量处理文件时发生错误: {str(e)}"
|
||
print(error_msg)
|
||
|
||
# 更新任务状态为错误
|
||
if task_id:
|
||
task_status_store.update_status(
|
||
task_id=task_id,
|
||
status="failed",
|
||
error=error_msg
|
||
)
|
||
|
||
return {
|
||
"status": "error",
|
||
"message": error_msg,
|
||
"dataset_id": dataset_id,
|
||
"error": str(e)
|
||
}
|
||
|
||
|
||
@huey.task()
|
||
def cleanup_project_async(
|
||
dataset_id: str,
|
||
remove_all: bool = False
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
异步清理项目文件
|
||
|
||
Args:
|
||
dataset_id: 项目唯一ID
|
||
remove_all: 是否删除整个项目目录
|
||
|
||
Returns:
|
||
清理结果字典
|
||
"""
|
||
try:
|
||
print(f"开始异步清理项目，项目ID: {dataset_id}")
|
||
|
||
project_dir = os.path.join("projects", "data", dataset_id)
|
||
removed_items = []
|
||
|
||
if remove_all and os.path.exists(project_dir):
|
||
import shutil
|
||
shutil.rmtree(project_dir)
|
||
removed_items.append(project_dir)
|
||
result = {
|
||
"status": "success",
|
||
"message": f"已删除整个项目目录: {project_dir}",
|
||
"dataset_id": dataset_id,
|
||
"removed_items": removed_items,
|
||
"action": "remove_all"
|
||
}
|
||
else:
|
||
# 只清理处理日志
|
||
log_file = os.path.join(project_dir, "processed_files.json")
|
||
if os.path.exists(log_file):
|
||
os.remove(log_file)
|
||
removed_items.append(log_file)
|
||
|
||
result = {
|
||
"status": "success",
|
||
"message": f"已清理项目处理日志，项目ID: {dataset_id}",
|
||
"dataset_id": dataset_id,
|
||
"removed_items": removed_items,
|
||
"action": "cleanup_logs"
|
||
}
|
||
|
||
print(f"异步清理任务完成: {dataset_id}")
|
||
return result
|
||
|
||
except Exception as e:
|
||
error_msg = f"异步清理项目时发生错误: {str(e)}"
|
||
print(error_msg)
|
||
return {
|
||
"status": "error",
|
||
"message": error_msg,
|
||
"dataset_id": dataset_id,
|
||
"error": str(e)
|
||
} |