mirror of
https://github.com/jihe520/MathModelAgent.git
synced 2026-10-30 07:24:14 +08:00
274 lines
9.0 KiB
Python
274 lines
9.0 KiB
Python
from fastapi import APIRouter, BackgroundTasks, File, Form, UploadFile
|
|
from app.core.workflow import MathModelWorkFlow
|
|
from app.schemas.enums import CompTemplate, FormatOutPut
|
|
from app.utils.log_util import logger
|
|
from app.services.redis_manager import redis_manager
|
|
from app.schemas.request import Problem
|
|
from app.schemas.response import SystemMessage
|
|
from app.utils.common_utils import (
|
|
create_task_id,
|
|
create_work_dir,
|
|
get_current_files,
|
|
md_2_docx,
|
|
)
|
|
import os
|
|
import asyncio
|
|
from fastapi import HTTPException
|
|
from icecream import ic
|
|
from app.schemas.request import ExampleRequest
|
|
from pydantic import BaseModel
|
|
import litellm
|
|
from app.config.setting import settings
|
|
import requests
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
class ValidateApiKeyRequest(BaseModel):
|
|
api_key: str
|
|
base_url: str = "https://api.openai.com/v1"
|
|
model_id: str
|
|
|
|
|
|
class ValidateOpenalexEmailRequest(BaseModel):
|
|
email: str
|
|
|
|
|
|
class ValidateOpenalexEmailResponse(BaseModel):
|
|
valid: bool
|
|
message: str
|
|
|
|
|
|
class ValidateApiKeyResponse(BaseModel):
|
|
valid: bool
|
|
message: str
|
|
|
|
|
|
class SaveApiConfigRequest(BaseModel):
|
|
coordinator: dict
|
|
modeler: dict
|
|
coder: dict
|
|
writer: dict
|
|
openalex_email: str
|
|
|
|
|
|
@router.post("/save-api-config")
|
|
async def save_api_config(request: SaveApiConfigRequest):
|
|
"""
|
|
保存验证成功的 API 配置到 settings
|
|
"""
|
|
try:
|
|
# 更新各个模块的设置
|
|
if request.coordinator:
|
|
settings.COORDINATOR_API_KEY = request.coordinator.get("apiKey", "")
|
|
settings.COORDINATOR_MODEL = request.coordinator.get("modelId", "")
|
|
settings.COORDINATOR_BASE_URL = request.coordinator.get("baseUrl", "")
|
|
|
|
if request.modeler:
|
|
settings.MODELER_API_KEY = request.modeler.get("apiKey", "")
|
|
settings.MODELER_MODEL = request.modeler.get("modelId", "")
|
|
settings.MODELER_BASE_URL = request.modeler.get("baseUrl", "")
|
|
|
|
if request.coder:
|
|
settings.CODER_API_KEY = request.coder.get("apiKey", "")
|
|
settings.CODER_MODEL = request.coder.get("modelId", "")
|
|
settings.CODER_BASE_URL = request.coder.get("baseUrl", "")
|
|
|
|
if request.writer:
|
|
settings.WRITER_API_KEY = request.writer.get("apiKey", "")
|
|
settings.WRITER_MODEL = request.writer.get("modelId", "")
|
|
settings.WRITER_BASE_URL = request.writer.get("baseUrl", "")
|
|
|
|
if request.openalex_email:
|
|
settings.OPENALEX_EMAIL = request.openalex_email
|
|
|
|
return {"success": True, "message": "配置保存成功"}
|
|
except Exception as e:
|
|
logger.error(f"保存配置失败: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=f"保存配置失败: {str(e)}")
|
|
|
|
|
|
@router.post("/validate-api-key", response_model=ValidateApiKeyResponse)
|
|
async def validate_api_key(request: ValidateApiKeyRequest):
|
|
"""
|
|
验证 API Key 的有效性
|
|
"""
|
|
try:
|
|
# 使用 litellm 发送测试请求
|
|
await litellm.acompletion(
|
|
model=request.model_id,
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
max_tokens=1,
|
|
api_key=request.api_key,
|
|
base_url=request.base_url
|
|
if request.base_url != "https://api.openai.com/v1"
|
|
else None,
|
|
)
|
|
|
|
return ValidateApiKeyResponse(valid=True, message="✓ 模型 API 验证成功")
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
|
|
# 解析不同类型的错误
|
|
if "401" in error_msg or "Unauthorized" in error_msg:
|
|
return ValidateApiKeyResponse(valid=False, message="✗ API Key 无效或已过期")
|
|
elif "404" in error_msg or "Not Found" in error_msg:
|
|
return ValidateApiKeyResponse(
|
|
valid=False, message="✗ 模型 ID 不存在或 Base URL 错误"
|
|
)
|
|
elif "429" in error_msg or "rate limit" in error_msg.lower():
|
|
return ValidateApiKeyResponse(
|
|
valid=False, message="✗ 请求过于频繁，请稍后再试"
|
|
)
|
|
elif "403" in error_msg or "Forbidden" in error_msg:
|
|
return ValidateApiKeyResponse(
|
|
valid=False, message="✗ API 权限不足或账户余额不足"
|
|
)
|
|
else:
|
|
return ValidateApiKeyResponse(
|
|
valid=False, message=f"✗ 验证失败: {error_msg[:50]}..."
|
|
)
|
|
|
|
|
|
@router.post("/validate-openalex-email", response_model=ValidateOpenalexEmailResponse)
|
|
async def validate_openalex_email(request: ValidateOpenalexEmailRequest):
|
|
"""
|
|
验证 OpenAlex Email 的有效性
|
|
"""
|
|
try:
|
|
response = requests.get(
|
|
f"https://api.openalex.org/works?mailto={request.email}"
|
|
)
|
|
logger.debug(f"OpenAlex Email 验证响应: {response}")
|
|
response.raise_for_status()
|
|
return ValidateOpenalexEmailResponse(
|
|
valid=True, message="✓ OpenAlex Email 验证成功"
|
|
)
|
|
except Exception as e:
|
|
return ValidateOpenalexEmailResponse(
|
|
valid=False, message=f"✗ OpenAlex Email 验证失败: {str(e)}"
|
|
)
|
|
|
|
|
|
@router.post("/example")
|
|
async def exampleModeling(
|
|
example_request: ExampleRequest,
|
|
background_tasks: BackgroundTasks,
|
|
):
|
|
task_id = create_task_id()
|
|
work_dir = create_work_dir(task_id)
|
|
example_dir = os.path.join("app", "example", "example", example_request.source)
|
|
ic(example_dir)
|
|
with open(os.path.join(example_dir, "questions.txt"), "r", encoding="utf-8") as f:
|
|
ques_all = f.read()
|
|
|
|
current_files = get_current_files(example_dir, "data")
|
|
for file in current_files:
|
|
src_file = os.path.join(example_dir, file)
|
|
dst_file = os.path.join(work_dir, file)
|
|
with open(src_file, "rb") as src, open(dst_file, "wb") as dst:
|
|
dst.write(src.read())
|
|
# 存储任务ID
|
|
await redis_manager.set(f"task_id:{task_id}", task_id)
|
|
|
|
logger.info(f"Adding background task for task_id: {task_id}")
|
|
# 将任务添加到后台执行
|
|
background_tasks.add_task(
|
|
run_modeling_task_async,
|
|
task_id,
|
|
ques_all,
|
|
CompTemplate.CHINA,
|
|
FormatOutPut.Markdown,
|
|
)
|
|
return {"task_id": task_id, "status": "processing"}
|
|
|
|
|
|
@router.post("/modeling")
|
|
async def modeling(
|
|
background_tasks: BackgroundTasks,
|
|
ques_all: str = Form(...), # 从表单获取
|
|
comp_template: CompTemplate = Form(...), # 从表单获取
|
|
format_output: FormatOutPut = Form(...), # 从表单获取
|
|
files: list[UploadFile] = File(default=None),
|
|
):
|
|
task_id = create_task_id()
|
|
work_dir = create_work_dir(task_id)
|
|
|
|
# 如果有上传文件，保存文件
|
|
if files:
|
|
logger.info(f"开始处理上传的文件，工作目录: {work_dir}")
|
|
for file in files:
|
|
try:
|
|
data_file_path = os.path.join(work_dir, file.filename)
|
|
logger.info(f"保存文件: {file.filename} -> {data_file_path}")
|
|
|
|
# 确保文件名不为空
|
|
if not file.filename:
|
|
logger.warning("跳过空文件名")
|
|
continue
|
|
|
|
content = await file.read()
|
|
if not content:
|
|
logger.warning(f"文件 {file.filename} 内容为空")
|
|
continue
|
|
|
|
with open(data_file_path, "wb") as f:
|
|
f.write(content)
|
|
logger.info(f"成功保存文件: {data_file_path}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"保存文件 {file.filename} 失败: {str(e)}")
|
|
raise HTTPException(
|
|
status_code=500, detail=f"保存文件 {file.filename} 失败: {str(e)}"
|
|
)
|
|
else:
|
|
logger.warning("没有上传文件")
|
|
|
|
# 存储任务ID
|
|
await redis_manager.set(f"task_id:{task_id}", task_id)
|
|
|
|
logger.info(f"Adding background task for task_id: {task_id}")
|
|
# 将任务添加到后台执行
|
|
background_tasks.add_task(
|
|
run_modeling_task_async, task_id, ques_all, comp_template, format_output
|
|
)
|
|
return {"task_id": task_id, "status": "processing"}
|
|
|
|
|
|
async def run_modeling_task_async(
|
|
task_id: str,
|
|
ques_all: str,
|
|
comp_template: CompTemplate,
|
|
format_output: FormatOutPut,
|
|
):
|
|
logger.info(f"run modeling task for task_id: {task_id}")
|
|
|
|
problem = Problem(
|
|
task_id=task_id,
|
|
ques_all=ques_all,
|
|
comp_template=comp_template,
|
|
format_output=format_output,
|
|
)
|
|
|
|
# 发送任务开始状态
|
|
await redis_manager.publish_message(
|
|
task_id,
|
|
SystemMessage(content="任务开始处理"),
|
|
)
|
|
|
|
# 给一个短暂的延迟，确保 WebSocket 有机会连接
|
|
await asyncio.sleep(1)
|
|
|
|
# 创建任务并等待它完成
|
|
task = asyncio.create_task(MathModelWorkFlow().execute(problem))
|
|
# 设置超时时间（比如 300 分钟）
|
|
await asyncio.wait_for(task, timeout=3600 * 5)
|
|
|
|
# 发送任务完成状态
|
|
await redis_manager.publish_message(
|
|
task_id,
|
|
SystemMessage(content="任务处理完成", type="success"),
|
|
)
|
|
# 转换md为docx
|
|
md_2_docx(task_id)
|