771 lines
28 KiB
Python
771 lines
28 KiB
Python
import json
|
||
import os
|
||
import tempfile
|
||
import shutil
|
||
import uuid
|
||
from typing import AsyncGenerator, Dict, List, Optional, Union, Any
|
||
from datetime import datetime
|
||
|
||
import uvicorn
|
||
from fastapi import FastAPI, HTTPException, Depends, Header, UploadFile, File
|
||
from fastapi.responses import StreamingResponse, HTMLResponse, FileResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from qwen_agent.llm.schema import ASSISTANT, FUNCTION
|
||
from pydantic import BaseModel, Field
|
||
|
||
# Import utility modules
|
||
from utils import (
|
||
# Models
|
||
Message, DatasetRequest, ChatRequest, ChatResponse, QueueTaskRequest, QueueTaskResponse,
|
||
QueueStatusResponse, TaskStatusResponse,
|
||
|
||
# File utilities
|
||
download_file, remove_file_or_directory, get_document_preview,
|
||
load_processed_files_log, save_processed_files_log, get_file_hash,
|
||
|
||
# Dataset management
|
||
download_dataset_files, generate_dataset_structure,
|
||
remove_dataset_directory, remove_dataset_directory_by_key,
|
||
|
||
# Project management
|
||
generate_project_readme, save_project_readme, get_project_status,
|
||
remove_project, list_projects, get_project_stats,
|
||
|
||
# Agent management
|
||
get_global_agent_manager, init_global_agent_manager
|
||
)
|
||
|
||
# Import modified_assistant
|
||
from modified_assistant import update_agent_llm
|
||
|
||
# Import queue manager
|
||
from task_queue.manager import queue_manager
|
||
from task_queue.integration_tasks import process_files_async, cleanup_project_async
|
||
from task_queue.task_status import task_status_store
|
||
import re
|
||
|
||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||
|
||
# Custom version for qwen-agent messages - keep this function as it's specific to this app
|
||
def get_content_from_messages(messages: List[dict], tool_response: bool = True) -> str:
|
||
"""Extract content from qwen-agent messages with special formatting"""
|
||
full_text = ''
|
||
content = []
|
||
TOOL_CALL_S = '[TOOL_CALL]'
|
||
TOOL_RESULT_S = '[TOOL_RESPONSE]'
|
||
THOUGHT_S = '[THINK]'
|
||
ANSWER_S = '[ANSWER]'
|
||
|
||
for msg in messages:
|
||
|
||
if msg['role'] == ASSISTANT:
|
||
if msg.get('reasoning_content'):
|
||
assert isinstance(msg['reasoning_content'], str), 'Now only supports text messages'
|
||
content.append(f'{THOUGHT_S}\n{msg["reasoning_content"]}')
|
||
if msg.get('content'):
|
||
assert isinstance(msg['content'], str), 'Now only supports text messages'
|
||
# 过滤掉流式输出中的不完整 tool_call 文本
|
||
content_text = msg["content"]
|
||
|
||
# 使用正则表达式替换不完整的 tool_call 模式为空字符串
|
||
|
||
# 匹配并替换不完整的 tool_call 模式
|
||
content_text = re.sub(r'<t?o?o?l?_?c?a?l?l?$', '', content_text)
|
||
# 只有在处理后内容不为空时才添加
|
||
if content_text.strip():
|
||
content.append(f'{ANSWER_S}\n{content_text}')
|
||
if msg.get('function_call'):
|
||
content_text = msg["function_call"]["arguments"]
|
||
content_text = re.sub(r'}\n<\/?t?o?o?l?_?c?a?l?l?$', '', content_text)
|
||
if content_text.strip():
|
||
content.append(f'{TOOL_CALL_S} {msg["function_call"]["name"]}\n{content_text}')
|
||
elif msg['role'] == FUNCTION:
|
||
if tool_response:
|
||
content.append(f'{TOOL_RESULT_S} {msg["name"]}\n{msg["content"]}')
|
||
else:
|
||
raise TypeError
|
||
|
||
if content:
|
||
full_text = '\n'.join(content)
|
||
|
||
return full_text
|
||
|
||
|
||
# Helper functions are now imported from utils module
|
||
|
||
|
||
|
||
|
||
|
||
# 全局助手管理器配置
|
||
max_cached_agents = int(os.getenv("MAX_CACHED_AGENTS", "20"))
|
||
|
||
# 初始化全局助手管理器
|
||
agent_manager = init_global_agent_manager(max_cached_agents=max_cached_agents)
|
||
|
||
app = FastAPI(title="Database Assistant API", version="1.0.0")
|
||
|
||
# 挂载public文件夹为静态文件服务
|
||
app.mount("/public", StaticFiles(directory="public"), name="static")
|
||
|
||
# 添加CORS中间件，支持前端页面
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"], # 在生产环境中应该设置为具体的前端域名
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
|
||
# Models are now imported from utils module
|
||
|
||
|
||
async def generate_stream_response(agent, messages, request) -> AsyncGenerator[str, None]:
|
||
"""生成流式响应"""
|
||
accumulated_content = ""
|
||
chunk_id = 0
|
||
try:
|
||
for response in agent.run(messages=messages):
|
||
previous_content = accumulated_content
|
||
accumulated_content = get_content_from_messages(response, tool_response=request.tool_response)
|
||
|
||
# 计算新增的内容
|
||
if accumulated_content.startswith(previous_content):
|
||
new_content = accumulated_content[len(previous_content):]
|
||
else:
|
||
new_content = accumulated_content
|
||
previous_content = ""
|
||
|
||
# 只有当有新内容时才发送chunk
|
||
if new_content:
|
||
chunk_id += 1
|
||
# 构造OpenAI格式的流式响应
|
||
chunk_data = {
|
||
"id": f"chatcmpl-{chunk_id}",
|
||
"object": "chat.completion.chunk",
|
||
"created": int(__import__('time').time()),
|
||
"model": request.model,
|
||
"choices": [{
|
||
"index": 0,
|
||
"delta": {
|
||
"content": new_content
|
||
},
|
||
"finish_reason": None
|
||
}]
|
||
}
|
||
|
||
yield f"data: {json.dumps(chunk_data, ensure_ascii=False)}\n\n"
|
||
|
||
# 发送最终完成标记
|
||
final_chunk = {
|
||
"id": f"chatcmpl-{chunk_id + 1}",
|
||
"object": "chat.completion.chunk",
|
||
"created": int(__import__('time').time()),
|
||
"model": request.model,
|
||
"choices": [{
|
||
"index": 0,
|
||
"delta": {},
|
||
"finish_reason": "stop"
|
||
}]
|
||
}
|
||
yield f"data: {json.dumps(final_chunk, ensure_ascii=False)}\n\n"
|
||
|
||
# 发送结束标记
|
||
yield "data: [DONE]\n\n"
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
error_details = traceback.format_exc()
|
||
print(f"Error in generate_stream_response: {str(e)}")
|
||
print(f"Full traceback: {error_details}")
|
||
|
||
error_data = {
|
||
"error": {
|
||
"message": f"Stream error: {str(e)}",
|
||
"type": "internal_error"
|
||
}
|
||
}
|
||
yield f"data: {json.dumps(error_data, ensure_ascii=False)}\n\n"
|
||
|
||
|
||
# Models are now imported from utils module
|
||
|
||
|
||
|
||
@app.post("/api/v1/files/process/async")
|
||
async def process_files_async_endpoint(request: QueueTaskRequest, authorization: Optional[str] = Header(None)):
|
||
"""
|
||
异步处理文件的队列版本API
|
||
与 /api/v1/files/process 功能相同，但使用队列异步处理
|
||
|
||
Args:
|
||
request: QueueTaskRequest containing dataset_id, files, system_prompt, mcp_settings, and queue options
|
||
authorization: Authorization header containing API key (Bearer <API_KEY>)
|
||
|
||
Returns:
|
||
QueueTaskResponse: Processing result with task ID for tracking
|
||
"""
|
||
try:
|
||
dataset_id = request.dataset_id
|
||
if not dataset_id:
|
||
raise HTTPException(status_code=400, detail="dataset_id is required")
|
||
|
||
# 估算处理时间（基于文件数量）
|
||
estimated_time = 0
|
||
if request.files:
|
||
total_files = sum(len(file_list) for file_list in request.files.values())
|
||
estimated_time = max(30, total_files * 10) # 每个文件预估10秒，最少30秒
|
||
|
||
# 提交异步任务
|
||
task_id = queue_manager.enqueue_multiple_files(
|
||
project_id=dataset_id,
|
||
file_paths=[],
|
||
original_filenames=[]
|
||
)
|
||
|
||
# 创建任务状态记录
|
||
import uuid
|
||
task_id = str(uuid.uuid4())
|
||
task_status_store.set_status(
|
||
task_id=task_id,
|
||
unique_id=dataset_id,
|
||
status="pending"
|
||
)
|
||
|
||
# 提交异步任务
|
||
task = process_files_async(
|
||
unique_id=dataset_id,
|
||
files=request.files,
|
||
system_prompt=request.system_prompt,
|
||
mcp_settings=request.mcp_settings,
|
||
task_id=task_id
|
||
)
|
||
|
||
return QueueTaskResponse(
|
||
success=True,
|
||
message=f"文件处理任务已提交到队列，项目ID: {dataset_id}",
|
||
unique_id=dataset_id,
|
||
task_id=task_id, # 使用我们自己的task_id
|
||
task_status="pending",
|
||
estimated_processing_time=estimated_time
|
||
)
|
||
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
print(f"Error submitting async file processing task: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
|
||
|
||
|
||
@app.get("/api/v1/task/{task_id}/status")
|
||
async def get_task_status(task_id: str):
|
||
"""获取任务状态 - 简单可靠"""
|
||
try:
|
||
status_data = task_status_store.get_status(task_id)
|
||
|
||
if not status_data:
|
||
return {
|
||
"success": False,
|
||
"message": "任务不存在或已过期",
|
||
"task_id": task_id,
|
||
"status": "not_found"
|
||
}
|
||
|
||
return {
|
||
"success": True,
|
||
"message": "任务状态获取成功",
|
||
"task_id": task_id,
|
||
"status": status_data["status"],
|
||
"unique_id": status_data["unique_id"],
|
||
"created_at": status_data["created_at"],
|
||
"updated_at": status_data["updated_at"],
|
||
"result": status_data.get("result"),
|
||
"error": status_data.get("error")
|
||
}
|
||
|
||
except Exception as e:
|
||
print(f"Error getting task status: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"获取任务状态失败: {str(e)}")
|
||
|
||
|
||
@app.delete("/api/v1/task/{task_id}")
|
||
async def delete_task(task_id: str):
|
||
"""删除任务记录"""
|
||
try:
|
||
success = task_status_store.delete_status(task_id)
|
||
if success:
|
||
return {
|
||
"success": True,
|
||
"message": f"任务记录已删除: {task_id}",
|
||
"task_id": task_id
|
||
}
|
||
else:
|
||
return {
|
||
"success": False,
|
||
"message": f"任务记录不存在: {task_id}",
|
||
"task_id": task_id
|
||
}
|
||
except Exception as e:
|
||
print(f"Error deleting task: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"删除任务记录失败: {str(e)}")
|
||
|
||
|
||
@app.get("/api/v1/tasks")
|
||
async def list_tasks(status: Optional[str] = None, dataset_id: Optional[str] = None, limit: int = 100):
|
||
"""列出任务，支持筛选"""
|
||
try:
|
||
if status or dataset_id:
|
||
# 使用搜索功能
|
||
tasks = task_status_store.search_tasks(status=status, unique_id=dataset_id, limit=limit)
|
||
else:
|
||
# 获取所有任务
|
||
all_tasks = task_status_store.list_all()
|
||
tasks = list(all_tasks.values())[:limit]
|
||
|
||
return {
|
||
"success": True,
|
||
"message": "任务列表获取成功",
|
||
"total_tasks": len(tasks),
|
||
"tasks": tasks,
|
||
"filters": {
|
||
"status": status,
|
||
"dataset_id": dataset_id,
|
||
"limit": limit
|
||
}
|
||
}
|
||
|
||
except Exception as e:
|
||
print(f"Error listing tasks: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"获取任务列表失败: {str(e)}")
|
||
|
||
|
||
@app.get("/api/v1/tasks/statistics")
|
||
async def get_task_statistics():
|
||
"""获取任务统计信息"""
|
||
try:
|
||
stats = task_status_store.get_statistics()
|
||
|
||
return {
|
||
"success": True,
|
||
"message": "统计信息获取成功",
|
||
"statistics": stats
|
||
}
|
||
|
||
except Exception as e:
|
||
print(f"Error getting statistics: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"获取统计信息失败: {str(e)}")
|
||
|
||
|
||
@app.post("/api/v1/tasks/cleanup")
|
||
async def cleanup_tasks(older_than_days: int = 7):
|
||
"""清理旧任务记录"""
|
||
try:
|
||
deleted_count = task_status_store.cleanup_old_tasks(older_than_days=older_than_days)
|
||
|
||
return {
|
||
"success": True,
|
||
"message": f"已清理 {deleted_count} 条旧任务记录",
|
||
"deleted_count": deleted_count,
|
||
"older_than_days": older_than_days
|
||
}
|
||
|
||
except Exception as e:
|
||
print(f"Error cleaning up tasks: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"清理任务记录失败: {str(e)}")
|
||
|
||
|
||
@app.get("/api/v1/projects/{dataset_id}/tasks")
|
||
async def get_project_tasks(dataset_id: str):
|
||
"""获取指定项目的所有任务"""
|
||
try:
|
||
tasks = task_status_store.get_by_unique_id(dataset_id)
|
||
|
||
return {
|
||
"success": True,
|
||
"message": "项目任务获取成功",
|
||
"dataset_id": dataset_id,
|
||
"total_tasks": len(tasks),
|
||
"tasks": tasks
|
||
}
|
||
|
||
except Exception as e:
|
||
print(f"Error getting project tasks: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"获取项目任务失败: {str(e)}")
|
||
|
||
|
||
@app.post("/api/v1/files/{dataset_id}/cleanup/async")
|
||
async def cleanup_project_async_endpoint(dataset_id: str, remove_all: bool = False):
|
||
"""异步清理项目文件"""
|
||
try:
|
||
task = cleanup_project_async(unique_id=dataset_id, remove_all=remove_all)
|
||
|
||
return {
|
||
"success": True,
|
||
"message": f"项目清理任务已提交到队列，项目ID: {dataset_id}",
|
||
"dataset_id": dataset_id,
|
||
"task_id": task.id,
|
||
"action": "remove_all" if remove_all else "cleanup_logs"
|
||
}
|
||
except Exception as e:
|
||
print(f"Error submitting cleanup task: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"提交清理任务失败: {str(e)}")
|
||
|
||
|
||
@app.post("/api/v1/chat/completions")
|
||
async def chat_completions(request: ChatRequest, authorization: Optional[str] = Header(None)):
|
||
"""
|
||
Chat completions API similar to OpenAI, supports both streaming and non-streaming
|
||
|
||
Args:
|
||
request: ChatRequest containing messages, model, dataset_ids (optional list), required bot_id, system_prompt, mcp_settings, and files
|
||
authorization: Authorization header containing API key (Bearer <API_KEY>)
|
||
|
||
Returns:
|
||
Union[ChatResponse, StreamingResponse]: Chat completion response or stream
|
||
|
||
Notes:
|
||
- dataset_ids: 可选参数，当提供时必须是项目ID列表（单个项目也使用数组格式）
|
||
- bot_id: 必需参数，机器人ID，用于创建项目目录
|
||
- 只有当提供 dataset_ids 时才会创建机器人项目目录：projects/robot/{bot_id}/
|
||
- 支持多知识库合并，自动处理文件夹重名冲突
|
||
|
||
Required Parameters:
|
||
- bot_id: str - 目标机器人项目ID
|
||
Optional Parameters:
|
||
- dataset_ids: List[str] - 源知识库项目ID列表（单个项目也使用数组格式）
|
||
|
||
Example:
|
||
{"bot_id": "my-bot-001"}
|
||
{"dataset_ids": ["project-123"], "bot_id": "my-bot-001"}
|
||
{"dataset_ids": ["project-123", "project-456"], "bot_id": "my-bot-002"}
|
||
"""
|
||
try:
|
||
# 从Authorization header中提取API key
|
||
api_key = None
|
||
if authorization:
|
||
# 移除 "Bearer " 前缀
|
||
if authorization.startswith("Bearer "):
|
||
api_key = authorization[7:]
|
||
else:
|
||
api_key = authorization
|
||
|
||
# 获取bot_id（必需参数）
|
||
bot_id = request.bot_id
|
||
if not bot_id:
|
||
raise HTTPException(status_code=400, detail="bot_id is required")
|
||
|
||
# 获取dataset_ids（可选参数，当提供时必须是数组）
|
||
dataset_ids_list = request.dataset_ids
|
||
project_dir = None
|
||
|
||
# 只有当提供了 dataset_ids 时才创建机器人目录并合并数据
|
||
if dataset_ids_list and len(dataset_ids_list) > 0:
|
||
from utils.multi_project_manager import create_robot_project
|
||
project_dir = create_robot_project(dataset_ids_list, bot_id)
|
||
|
||
# 收集额外参数作为 generate_cfg
|
||
exclude_fields = {'messages', 'model', 'model_server', 'dataset_ids', 'language', 'tool_response', 'system_prompt', 'mcp_settings' ,'stream', 'robot_type', 'bot_id'}
|
||
generate_cfg = {k: v for k, v in request.model_dump().items() if k not in exclude_fields}
|
||
|
||
# 从全局管理器获取或创建助手实例（配置读取逻辑已在agent_manager内部处理）
|
||
agent = await agent_manager.get_or_create_agent(
|
||
bot_id=bot_id,
|
||
project_dir=project_dir,
|
||
model_name=request.model,
|
||
api_key=api_key,
|
||
model_server=request.model_server,
|
||
generate_cfg=generate_cfg,
|
||
language=request.language,
|
||
system_prompt=request.system_prompt,
|
||
mcp_settings=request.mcp_settings,
|
||
robot_type=request.robot_type
|
||
)
|
||
# 构建包含项目信息的消息上下文
|
||
messages = []
|
||
for msg in request.messages:
|
||
if msg.role == "assistant":
|
||
# 对assistant消息进行[ANSWER]分割处理，只保留最后一段
|
||
content_parts = msg.content.split("[ANSWER]")
|
||
if content_parts:
|
||
# 取最后一段非空文本
|
||
last_part = content_parts[-1].strip()
|
||
messages.append({"role": msg.role, "content": last_part})
|
||
else:
|
||
messages.append({"role": msg.role, "content": msg.content})
|
||
else:
|
||
messages.append({"role": msg.role, "content": msg.content})
|
||
|
||
# 在最后一条消息的末尾追加回复语言
|
||
if messages and request.language:
|
||
language_map = {
|
||
'zh': '请用中文回复',
|
||
'en': 'Please reply in English',
|
||
'ja': '日本語で回答してください',
|
||
'jp': '日本語で回答してください'
|
||
}
|
||
language_instruction = language_map.get(request.language.lower(), '')
|
||
if language_instruction:
|
||
# 在最后一条消息末尾追加语言指令
|
||
messages[-1]['content'] = messages[-1]['content'] + f"\n\n{language_instruction}。"
|
||
|
||
# 根据stream参数决定返回流式还是非流式响应
|
||
if request.stream:
|
||
return StreamingResponse(
|
||
generate_stream_response(agent, messages, request),
|
||
media_type="text/event-stream",
|
||
headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}
|
||
)
|
||
else:
|
||
# 非流式响应
|
||
final_responses = agent.run_nonstream(messages)
|
||
|
||
if final_responses and len(final_responses) > 0:
|
||
# 使用 get_content_from_messages 处理响应，支持 tool_response 参数
|
||
content = get_content_from_messages(final_responses, tool_response=request.tool_response)
|
||
|
||
# 构造OpenAI格式的响应
|
||
return ChatResponse(
|
||
choices=[{
|
||
"index": 0,
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": content
|
||
},
|
||
"finish_reason": "stop"
|
||
}],
|
||
usage={
|
||
"prompt_tokens": sum(len(msg.content) for msg in request.messages),
|
||
"completion_tokens": len(content),
|
||
"total_tokens": sum(len(msg.content) for msg in request.messages) + len(content)
|
||
}
|
||
)
|
||
else:
|
||
raise HTTPException(status_code=500, detail="No response from agent")
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
error_details = traceback.format_exc()
|
||
print(f"Error in chat_completions: {str(e)}")
|
||
print(f"Full traceback: {error_details}")
|
||
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
|
||
|
||
|
||
@app.post("/api/v1/upload")
|
||
async def upload_file(file: UploadFile = File(...)):
|
||
"""
|
||
文件上传API接口，上传文件到 ./projects/uploads 目录
|
||
|
||
Args:
|
||
file: 上传的文件
|
||
|
||
Returns:
|
||
dict: 包含文件路径和文件名的响应
|
||
"""
|
||
try:
|
||
# 确保上传目录存在
|
||
upload_dir = os.path.join("projects", "uploads")
|
||
os.makedirs(upload_dir, exist_ok=True)
|
||
|
||
# 生成唯一文件名
|
||
file_extension = os.path.splitext(file.filename)[1] if file.filename else ""
|
||
unique_filename = f"{uuid.uuid4()}{file_extension}"
|
||
file_path = os.path.join(upload_dir, unique_filename)
|
||
|
||
# 保存文件
|
||
with open(file_path, "wb") as buffer:
|
||
shutil.copyfileobj(file.file, buffer)
|
||
|
||
return {
|
||
"success": True,
|
||
"message": "文件上传成功",
|
||
"filename": unique_filename,
|
||
"original_filename": file.filename,
|
||
"file_path": file_path
|
||
}
|
||
|
||
except Exception as e:
|
||
print(f"Error uploading file: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"文件上传失败: {str(e)}")
|
||
|
||
|
||
@app.get("/api/health")
|
||
async def health_check():
|
||
"""Health check endpoint"""
|
||
return {"message": "Database Assistant API is running"}
|
||
|
||
|
||
@app.get("/system/status")
|
||
async def system_status():
|
||
"""获取系统状态信息"""
|
||
# 获取助手缓存统计
|
||
cache_stats = agent_manager.get_cache_stats()
|
||
|
||
return {
|
||
"status": "running",
|
||
"storage_type": "File-Loaded Agent Manager",
|
||
"max_cached_agents": max_cached_agents,
|
||
"agent_cache": {
|
||
"total_cached_agents": cache_stats["total_cached_agents"],
|
||
"max_cached_agents": cache_stats["max_cached_agents"],
|
||
"cached_agents": cache_stats["agents"]
|
||
}
|
||
}
|
||
|
||
|
||
@app.post("/system/cleanup-cache")
|
||
async def cleanup_cache():
|
||
"""清理助手缓存"""
|
||
try:
|
||
# 清理助手实例缓存
|
||
cleared_count = agent_manager.clear_cache()
|
||
|
||
return {
|
||
"message": "缓存清理成功",
|
||
"cleared_agent_instances": cleared_count
|
||
}
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"缓存清理失败: {str(e)}")
|
||
|
||
|
||
@app.post("/system/cleanup-agent-cache")
|
||
async def cleanup_agent_cache():
|
||
"""仅清理助手实例缓存"""
|
||
try:
|
||
cleared_count = agent_manager.clear_cache()
|
||
return {
|
||
"message": "助手实例缓存清理成功",
|
||
"cleared_agent_instances": cleared_count
|
||
}
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"助手实例缓存清理失败: {str(e)}")
|
||
|
||
|
||
@app.get("/system/cached-projects")
|
||
async def get_cached_projects():
|
||
"""获取所有缓存的项目信息"""
|
||
try:
|
||
cache_stats = agent_manager.get_cache_stats()
|
||
|
||
return {
|
||
"cache_stats": cache_stats
|
||
}
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"获取缓存项目信息失败: {str(e)}")
|
||
|
||
|
||
@app.post("/system/remove-project-cache")
|
||
async def remove_project_cache(dataset_id: str):
|
||
"""移除特定项目的缓存"""
|
||
try:
|
||
removed_count = agent_manager.remove_cache_by_unique_id(dataset_id)
|
||
if removed_count > 0:
|
||
return {"message": f"项目缓存移除成功: {dataset_id}", "removed_count": removed_count}
|
||
else:
|
||
return {"message": f"未找到项目缓存: {dataset_id}", "removed_count": 0}
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"移除项目缓存失败: {str(e)}")
|
||
|
||
|
||
@app.get("/api/v1/files/{dataset_id}/status")
|
||
async def get_files_processing_status(dataset_id: str):
|
||
"""获取项目的文件处理状态"""
|
||
try:
|
||
# Load processed files log
|
||
processed_log = load_processed_files_log(dataset_id)
|
||
|
||
# Get project directory info
|
||
project_dir = os.path.join("projects", "data", dataset_id)
|
||
project_exists = os.path.exists(project_dir)
|
||
|
||
# Collect document.txt files
|
||
document_files = []
|
||
if project_exists:
|
||
for root, dirs, files in os.walk(project_dir):
|
||
for file in files:
|
||
if file == "document.txt":
|
||
document_files.append(os.path.join(root, file))
|
||
|
||
return {
|
||
"dataset_id": dataset_id,
|
||
"project_exists": project_exists,
|
||
"processed_files_count": len(processed_log),
|
||
"processed_files": processed_log,
|
||
"document_files_count": len(document_files),
|
||
"document_files": document_files,
|
||
"log_file_exists": os.path.exists(os.path.join("projects", "data", dataset_id, "processed_files.json"))
|
||
}
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"获取文件处理状态失败: {str(e)}")
|
||
|
||
|
||
@app.post("/api/v1/files/{dataset_id}/reset")
|
||
async def reset_files_processing(dataset_id: str):
|
||
"""重置项目的文件处理状态，删除处理日志和所有文件"""
|
||
try:
|
||
project_dir = os.path.join("projects", "data", dataset_id)
|
||
log_file = os.path.join("projects", "data", dataset_id, "processed_files.json")
|
||
|
||
# Load processed log to know what files to remove
|
||
processed_log = load_processed_files_log(dataset_id)
|
||
|
||
removed_files = []
|
||
# Remove all processed files and their dataset directories
|
||
for file_hash, file_info in processed_log.items():
|
||
# Remove local file in files directory
|
||
if 'local_path' in file_info:
|
||
if remove_file_or_directory(file_info['local_path']):
|
||
removed_files.append(file_info['local_path'])
|
||
|
||
# Handle new key-based structure first
|
||
if 'key' in file_info:
|
||
# Remove dataset directory by key
|
||
key = file_info['key']
|
||
if remove_dataset_directory_by_key(dataset_id, key):
|
||
removed_files.append(f"dataset/{key}")
|
||
elif 'filename' in file_info:
|
||
# Fallback to old filename-based structure
|
||
filename_without_ext = os.path.splitext(file_info['filename'])[0]
|
||
dataset_dir = os.path.join("projects", "data", dataset_id, "dataset", filename_without_ext)
|
||
if remove_file_or_directory(dataset_dir):
|
||
removed_files.append(dataset_dir)
|
||
|
||
# Also remove any specific dataset path if exists (fallback)
|
||
if 'dataset_path' in file_info:
|
||
if remove_file_or_directory(file_info['dataset_path']):
|
||
removed_files.append(file_info['dataset_path'])
|
||
|
||
# Remove the log file
|
||
if remove_file_or_directory(log_file):
|
||
removed_files.append(log_file)
|
||
|
||
# Remove the entire files directory
|
||
files_dir = os.path.join(project_dir, "files")
|
||
if remove_file_or_directory(files_dir):
|
||
removed_files.append(files_dir)
|
||
|
||
# Also remove the entire dataset directory (clean up any remaining files)
|
||
dataset_dir = os.path.join(project_dir, "dataset")
|
||
if remove_file_or_directory(dataset_dir):
|
||
removed_files.append(dataset_dir)
|
||
|
||
# Remove README.md if exists
|
||
readme_file = os.path.join(project_dir, "README.md")
|
||
if remove_file_or_directory(readme_file):
|
||
removed_files.append(readme_file)
|
||
|
||
return {
|
||
"message": f"文件处理状态重置成功: {dataset_id}",
|
||
"removed_files_count": len(removed_files),
|
||
"removed_files": removed_files
|
||
}
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"重置文件处理状态失败: {str(e)}")
|
||
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
uvicorn.run(app, host="0.0.0.0", port=8001)
|