mirror of
https://github.com/jihe520/MathModelAgent.git
synced 2026-10-02 02:44:56 +08:00
136 lines
4.0 KiB
Python
136 lines
4.0 KiB
Python
import gradio as gr
|
|
from config.config import Config
|
|
from utils.common_utils import create_work_directories, create_task_id, load_toml
|
|
from utils.logger import log
|
|
from utils.cli import get_ascii_banner
|
|
from utils.data_recorder import DataRecorder
|
|
from core.LLM import DeepSeekModel
|
|
from models.user_input import UserInput
|
|
from models.task import Task
|
|
from models.user_output import UserOutput
|
|
from gradio import ChatMessage
|
|
|
|
|
|
def bot(history: list, agent_name: str, message: str):
|
|
"""用于流式输出消息到chatbot"""
|
|
# 创建新消息
|
|
new_message = ChatMessage(
|
|
role="assistant",
|
|
content=message,
|
|
metadata={"title": agent_name, "status": "done"},
|
|
)
|
|
history.append(new_message)
|
|
return history
|
|
|
|
|
|
def task_start(template, output_format, dataset_path, question, history):
|
|
print(get_ascii_banner())
|
|
# === 初始化 ===
|
|
history = [] # 重置历史记录
|
|
|
|
# 创建消息回调函数
|
|
def message_callback(agent_name, message):
|
|
nonlocal history
|
|
history = bot(history, agent_name, message)
|
|
return history
|
|
|
|
try:
|
|
# 初始化日志和基本设置
|
|
log.set_console_level("WARNING")
|
|
task_id = create_task_id()
|
|
base_dir, dirs = create_work_directories(task_id)
|
|
try:
|
|
log.init(dirs["log"])
|
|
except ValueError:
|
|
pass
|
|
|
|
# 添加初始状态消息
|
|
history = bot(history, "系统", f"开始任务，ID: {task_id}")
|
|
|
|
config = Config(load_toml("config/config.toml"))
|
|
data_recorder = DataRecorder(dirs["log"])
|
|
|
|
# 初始化模型
|
|
deepseek_model = DeepSeekModel(
|
|
**config.get_model_config(),
|
|
data_recorder=data_recorder,
|
|
message_callback=message_callback,
|
|
)
|
|
|
|
# 创建用户输入
|
|
user_input = UserInput(
|
|
comp_template=template,
|
|
format_output=output_format,
|
|
data_folder_path=dataset_path,
|
|
bg_ques_all=question,
|
|
model=deepseek_model,
|
|
)
|
|
|
|
user_input.set_config_template(
|
|
config.get_config_template(user_input.comp_template)
|
|
)
|
|
|
|
# 创建任务
|
|
task = Task(
|
|
task_id=task_id,
|
|
base_dir=base_dir,
|
|
work_dirs=dirs,
|
|
llm=deepseek_model,
|
|
config=config,
|
|
)
|
|
|
|
# 运行任务
|
|
user_output: UserOutput = task.run(user_input, data_recorder)
|
|
user_output.save_result(ques_count=user_input.get_ques_count())
|
|
|
|
except Exception as e:
|
|
history = bot(history, "系统", f"错误: {str(e)}")
|
|
|
|
return history
|
|
|
|
|
|
# 创建 Blocks 界面
|
|
with gr.Blocks() as demo:
|
|
with gr.Row():
|
|
# 左侧输入部分
|
|
with gr.Column(scale=1):
|
|
template = gr.Dropdown(
|
|
["国赛", "美赛"],
|
|
label="模板",
|
|
info="选择你需要的模板(目前只支持中文)",
|
|
value="国赛",
|
|
)
|
|
output_format = gr.Dropdown(
|
|
["Markdown", "Latex"],
|
|
label="输出格式",
|
|
info="选择你需要的输出格式(目前只支持markdown)",
|
|
value="Markdown",
|
|
)
|
|
dataset_path = gr.Textbox(
|
|
label="输入数据集的相对文件夹路径：(推荐将数据集放在./project/sample_data 目录下)",
|
|
value="./project/sample_data",
|
|
)
|
|
question = gr.Textbox(
|
|
label="输入题目",
|
|
value="",
|
|
)
|
|
submit_btn = gr.Button("提交")
|
|
|
|
# 右侧聊天部分
|
|
with gr.Column(scale=1):
|
|
chatbot = gr.Chatbot(
|
|
value=[],
|
|
label="任务进度",
|
|
type="messages", # 指定使用 messages 类型
|
|
height=500,
|
|
)
|
|
|
|
submit_btn.click(
|
|
fn=task_start,
|
|
inputs=[template, output_format, dataset_path, question, chatbot],
|
|
outputs=chatbot,
|
|
queue=True,
|
|
)
|
|
|
|
demo.launch()
|