mirror of
https://github.com/jihe520/MathModelAgent.git
synced 2026-10-21 05:59:54 +08:00
233 lines
9.6 KiB
Python
233 lines
9.6 KiB
Python
from app.tools.base_interpreter import BaseCodeInterpreter
|
|
from app.tools.notebook_serializer import NotebookSerializer
|
|
import jupyter_client
|
|
from app.utils.log_util import logger
|
|
import os
|
|
from app.services.redis_manager import redis_manager
|
|
from app.schemas.response import (
|
|
OutputItem,
|
|
ResultModel,
|
|
StdErrModel,
|
|
SystemMessage,
|
|
)
|
|
|
|
|
|
class LocalCodeInterpreter(BaseCodeInterpreter):
|
|
def __init__(
|
|
self,
|
|
task_id: str,
|
|
work_dir: str,
|
|
notebook_serializer: NotebookSerializer,
|
|
):
|
|
super().__init__(task_id, work_dir, notebook_serializer)
|
|
self.km, self.kc = None, None
|
|
self.interrupt_signal = False
|
|
|
|
async def initialize(self):
|
|
# 本地内核一般不需异步上传文件，直接切换目录即可
|
|
# 初始化 Jupyter 内核管理器和客户端
|
|
logger.info("初始化本地内核")
|
|
self.km, self.kc = jupyter_client.manager.start_new_kernel(
|
|
kernel_name="python3"
|
|
)
|
|
self._pre_execute_code()
|
|
|
|
def _pre_execute_code(self):
|
|
init_code = (
|
|
f"import os\n"
|
|
f"work_dir = r'{self.work_dir}'\n"
|
|
f"os.makedirs(work_dir, exist_ok=True)\n"
|
|
f"os.chdir(work_dir)\n"
|
|
f"print('当前工作目录:', os.getcwd())\n"
|
|
# f"import matplotlib.pyplot as plt\n"
|
|
# f"import matplotlib as mpl\n"
|
|
# # 更完整的中文字体配置
|
|
# f"plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei', 'Microsoft YaHei', 'WenQuanYi Micro Hei', 'PingFang SC', 'Hiragino Sans GB', 'Heiti SC', 'DejaVu Sans', 'sans-serif']\n"
|
|
# f"plt.rcParams['axes.unicode_minus'] = False\n"
|
|
# f"plt.rcParams['font.family'] = 'sans-serif'\n"
|
|
# f"mpl.rcParams['font.size'] = 12\n"
|
|
# f"mpl.rcParams['axes.labelsize'] = 12\n"
|
|
# f"mpl.rcParams['xtick.labelsize'] = 10\n"
|
|
# f"mpl.rcParams['ytick.labelsize'] = 10\n"
|
|
# # 设置DPI以获得更清晰的显示
|
|
)
|
|
self.execute_code_(init_code)
|
|
|
|
async def execute_code(self, code: str) -> tuple[str, bool, str]:
|
|
logger.info(f"执行代码: {code}")
|
|
# 添加代码到notebook
|
|
self.notebook_serializer.add_code_cell_to_notebook(code)
|
|
|
|
text_to_gpt: list[str] = []
|
|
content_to_display: list[OutputItem] | None = []
|
|
error_occurred: bool = False
|
|
error_message: str = ""
|
|
|
|
await redis_manager.publish_message(
|
|
self.task_id,
|
|
SystemMessage(content="开始执行代码"),
|
|
)
|
|
# 执行 Python 代码
|
|
logger.info("开始在本地执行代码...")
|
|
execution = self.execute_code_(code)
|
|
logger.info("代码执行完成，开始处理结果...")
|
|
|
|
await redis_manager.publish_message(
|
|
self.task_id,
|
|
SystemMessage(content="代码执行完成"),
|
|
)
|
|
|
|
for mark, out_str in execution:
|
|
if mark in ("stdout", "execute_result_text", "display_text"):
|
|
text_to_gpt.append(self._truncate_text(f"[{mark}]\n{out_str}"))
|
|
# 添加text到notebook
|
|
content_to_display.append(
|
|
ResultModel(type="result", format="text", msg=out_str)
|
|
)
|
|
self.notebook_serializer.add_code_cell_output_to_notebook(out_str)
|
|
|
|
elif mark in (
|
|
"execute_result_png",
|
|
"execute_result_jpeg",
|
|
"display_png",
|
|
"display_jpeg",
|
|
):
|
|
# TODO: 视觉模型解释图像
|
|
text_to_gpt.append(f"[{mark} 图片已生成，内容为 base64，未展示]")
|
|
|
|
# 添加image到notebook
|
|
if "png" in mark:
|
|
self.notebook_serializer.add_image_to_notebook(out_str, "image/png")
|
|
content_to_display.append(
|
|
ResultModel(type="result", format="png", msg=out_str)
|
|
)
|
|
else:
|
|
self.notebook_serializer.add_image_to_notebook(
|
|
out_str, "image/jpeg"
|
|
)
|
|
content_to_display.append(
|
|
ResultModel(type="result", format="jpeg", msg=out_str)
|
|
)
|
|
|
|
elif mark == "error":
|
|
error_occurred = True
|
|
error_message = self.delete_color_control_char(out_str)
|
|
error_message = self._truncate_text(error_message)
|
|
logger.error(f"执行错误: {error_message}")
|
|
text_to_gpt.append(error_message)
|
|
# 添加error到notebook
|
|
self.notebook_serializer.add_code_cell_error_to_notebook(out_str)
|
|
content_to_display.append(StdErrModel(msg=out_str))
|
|
|
|
logger.info(f"text_to_gpt: {text_to_gpt}")
|
|
combined_text = "\n".join(text_to_gpt)
|
|
|
|
await self._push_to_websocket(content_to_display)
|
|
|
|
return (
|
|
combined_text,
|
|
error_occurred,
|
|
error_message,
|
|
)
|
|
|
|
def execute_code_(self, code) -> list[tuple[str, str]]:
|
|
msg_id = self.kc.execute(code)
|
|
logger.info(f"执行代码: {code}")
|
|
# Get the output of the code
|
|
msg_list = []
|
|
while True:
|
|
try:
|
|
iopub_msg = self.kc.get_iopub_msg(timeout=1)
|
|
msg_list.append(iopub_msg)
|
|
if (
|
|
iopub_msg["msg_type"] == "status"
|
|
and iopub_msg["content"].get("execution_state") == "idle"
|
|
):
|
|
break
|
|
except:
|
|
if self.interrupt_signal:
|
|
self.km.interrupt_kernel()
|
|
self.interrupt_signal = False
|
|
continue
|
|
|
|
all_output: list[tuple[str, str]] = []
|
|
for iopub_msg in msg_list:
|
|
if iopub_msg["msg_type"] == "stream":
|
|
if iopub_msg["content"].get("name") == "stdout":
|
|
output = iopub_msg["content"]["text"]
|
|
all_output.append(("stdout", output))
|
|
elif iopub_msg["msg_type"] == "execute_result":
|
|
if "data" in iopub_msg["content"]:
|
|
if "text/plain" in iopub_msg["content"]["data"]:
|
|
output = iopub_msg["content"]["data"]["text/plain"]
|
|
all_output.append(("execute_result_text", output))
|
|
if "text/html" in iopub_msg["content"]["data"]:
|
|
output = iopub_msg["content"]["data"]["text/html"]
|
|
all_output.append(("execute_result_html", output))
|
|
if "image/png" in iopub_msg["content"]["data"]:
|
|
output = iopub_msg["content"]["data"]["image/png"]
|
|
all_output.append(("execute_result_png", output))
|
|
if "image/jpeg" in iopub_msg["content"]["data"]:
|
|
output = iopub_msg["content"]["data"]["image/jpeg"]
|
|
all_output.append(("execute_result_jpeg", output))
|
|
elif iopub_msg["msg_type"] == "display_data":
|
|
if "data" in iopub_msg["content"]:
|
|
if "text/plain" in iopub_msg["content"]["data"]:
|
|
output = iopub_msg["content"]["data"]["text/plain"]
|
|
all_output.append(("display_text", output))
|
|
if "text/html" in iopub_msg["content"]["data"]:
|
|
output = iopub_msg["content"]["data"]["text/html"]
|
|
all_output.append(("display_html", output))
|
|
if "image/png" in iopub_msg["content"]["data"]:
|
|
output = iopub_msg["content"]["data"]["image/png"]
|
|
all_output.append(("display_png", output))
|
|
if "image/jpeg" in iopub_msg["content"]["data"]:
|
|
output = iopub_msg["content"]["data"]["image/jpeg"]
|
|
all_output.append(("display_jpeg", output))
|
|
elif iopub_msg["msg_type"] == "error":
|
|
# TODO: 正确返回格式
|
|
if "traceback" in iopub_msg["content"]:
|
|
output = "\n".join(iopub_msg["content"]["traceback"])
|
|
cleaned_output = self.delete_color_control_char(output)
|
|
all_output.append(("error", cleaned_output))
|
|
return all_output
|
|
|
|
async def get_created_images(self, section: str) -> list[str]:
|
|
"""获取新创建的图片列表"""
|
|
current_images = set()
|
|
files = os.listdir(self.work_dir)
|
|
for file in files:
|
|
if file.endswith((".png", ".jpg", ".jpeg")):
|
|
current_images.add(file)
|
|
|
|
# 计算新增的图片
|
|
new_images = current_images - self.last_created_images
|
|
|
|
# 更新last_created_images为当前的图片集合
|
|
self.last_created_images = current_images
|
|
|
|
logger.info(f"新创建的图片列表: {new_images}")
|
|
return list(new_images) # 最后转换为list返回
|
|
|
|
async def cleanup(self):
|
|
# 关闭内核
|
|
self.kc.shutdown()
|
|
logger.info("关闭内核")
|
|
self.km.shutdown_kernel()
|
|
|
|
def send_interrupt_signal(self):
|
|
self.interrupt_signal = True
|
|
|
|
def restart_jupyter_kernel(self):
|
|
"""Restart the Jupyter kernel and recreate the work directory."""
|
|
self.kc.shutdown()
|
|
self.km, self.kc = jupyter_client.manager.start_new_kernel(
|
|
kernel_name="python3"
|
|
)
|
|
self.interrupt_signal = False
|
|
self._create_work_dir()
|
|
|
|
def _create_work_dir(self):
|
|
"""Ensure the working directory exists after a restart."""
|
|
os.makedirs(self.work_dir, exist_ok=True)
|