mirror of
https://github.com/jihe520/MathModelAgent.git
synced 2026-10-02 10:55:49 +08:00
114 lines
4.0 KiB
Python
114 lines
4.0 KiB
Python
import nbformat
|
|
from nbformat import v4 as nbf
|
|
import ansi2html
|
|
import os
|
|
|
|
|
|
class NotebookSerializer:
|
|
def __init__(self, work_dir=None, notebook_name="notebook.ipynb"):
|
|
self.nb = nbf.new_notebook()
|
|
self.notebook_path = None
|
|
self.initialized = True
|
|
self.segmentation_output_content = {} # 保存coder_agent 在 jupyter 中执行的 output 结果内容
|
|
# {
|
|
# "eda": {
|
|
# }
|
|
# }
|
|
self.current_segmentation: str = ""
|
|
|
|
self.init_notebook(work_dir, notebook_name)
|
|
|
|
def init_notebook(self, work_dir=None, notebook_name="notebook.ipynb"):
|
|
"""初始化notebook路径
|
|
|
|
Args:
|
|
work_dir (str): jupyter工作目录路径
|
|
notebook_name (str): notebook文件名,默认为notebook.ipynb
|
|
"""
|
|
if work_dir:
|
|
# 确保使用jupyter工作目录
|
|
base, ext = os.path.splitext(notebook_name)
|
|
if ext.lower() != ".ipynb":
|
|
notebook_name += ".ipynb"
|
|
|
|
# 在jupyter工作目录下创建notebook文件
|
|
self.notebook_path = os.path.join(work_dir, notebook_name)
|
|
|
|
if os.path.exists(self.notebook_path):
|
|
raise FileExistsError(
|
|
f"文件 {self.notebook_path} 已存在。请选择其他文件名。"
|
|
)
|
|
|
|
def ansi_to_html(self, ansi_text):
|
|
converter = ansi2html.Ansi2HTMLConverter()
|
|
html_text = converter.convert(ansi_text)
|
|
return html_text
|
|
|
|
def write_to_notebook(self):
|
|
if self.notebook_path:
|
|
with open(self.notebook_path, "w", encoding="utf-8") as f:
|
|
nbformat.write(self.nb, f)
|
|
|
|
def add_code_cell_to_notebook(self, code):
|
|
code_cell = nbf.new_code_cell(source=code)
|
|
self.nb["cells"].append(code_cell)
|
|
self.write_to_notebook()
|
|
|
|
def add_code_cell_output_to_notebook(self, output):
|
|
"""添加代码单元格输出
|
|
|
|
Args:
|
|
output: 代码输出内容
|
|
"""
|
|
html_content = self.ansi_to_html(output)
|
|
if self.current_segmentation:
|
|
# 确保键存在
|
|
if self.current_segmentation not in self.segmentation_output_content:
|
|
self.segmentation_output_content[self.current_segmentation] = ""
|
|
self.segmentation_output_content[self.current_segmentation] += html_content
|
|
|
|
cell_output = nbf.new_output(
|
|
output_type="display_data", data={"text/html": html_content}
|
|
)
|
|
self.nb["cells"][-1]["outputs"].append(cell_output)
|
|
self.write_to_notebook()
|
|
|
|
def add_code_cell_error_to_notebook(self, error):
|
|
nbf_error_output = nbf.new_output(
|
|
output_type="error",
|
|
ename="Error",
|
|
evalue="Error message",
|
|
traceback=[error],
|
|
)
|
|
self.nb["cells"][-1]["outputs"].append(nbf_error_output)
|
|
self.write_to_notebook()
|
|
|
|
def add_image_to_notebook(self, image, mime_type):
|
|
image_output = nbf.new_output(
|
|
output_type="display_data", data={mime_type: image}
|
|
)
|
|
self.nb["cells"][-1]["outputs"].append(image_output)
|
|
self.write_to_notebook()
|
|
|
|
def add_markdown_to_notebook(self, content, title=None):
|
|
if title:
|
|
content = "##### " + title + ":\n" + content
|
|
markdown_cell = nbf.new_markdown_cell(content)
|
|
self.nb["cells"].append(markdown_cell)
|
|
self.write_to_notebook()
|
|
|
|
def add_markdown_segmentation_to_notebook(self, content, segmentation):
|
|
"""添加markdown分段并初始化对应的output内容存储
|
|
|
|
Args:
|
|
content: markdown内容
|
|
segmentation: 分段名称
|
|
"""
|
|
self.current_segmentation = segmentation
|
|
# 初始化该分段的output内容
|
|
self.segmentation_output_content[segmentation] = ""
|
|
self.add_markdown_to_notebook(content, segmentation)
|
|
|
|
def get_notebook_output_content(self, segmentation):
|
|
return self.segmentation_output_content[segmentation]
|