mirror of
https://github.com/jihe520/MathModelAgent.git
synced 2026-09-19 08:09:48 +08:00
Sanjin
121 lines
3.7 KiB
Python
121 lines
3.7 KiB
Python
import os
|
|
import datetime
|
|
import hashlib
|
|
import tomllib
|
|
from app.schemas.enums import CompTemplate
|
|
from app.utils.log_util import logger
|
|
import re
|
|
import pypandoc
|
|
from app.config.setting import settings
|
|
from icecream import ic
|
|
|
|
|
|
def create_task_id() -> str:
|
|
"""生成任务ID"""
|
|
# 生成时间戳和随机hash
|
|
timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
|
|
random_hash = hashlib.md5(str(datetime.datetime.now()).encode()).hexdigest()[:8]
|
|
return f"{timestamp}-{random_hash}"
|
|
|
|
|
|
def create_work_dir(task_id: str) -> str:
|
|
# 设置主工作目录和子目录
|
|
work_dir = os.path.join("project", "work_dir", task_id)
|
|
|
|
try:
|
|
# 创建目录，如果目录已存在也不会报错
|
|
os.makedirs(work_dir, exist_ok=True)
|
|
return work_dir
|
|
except Exception as e:
|
|
# 捕获并记录创建目录时的异常
|
|
logger.error(f"创建工作目录失败: {str(e)}")
|
|
raise
|
|
|
|
|
|
def get_work_dir(task_id: str) -> str:
|
|
work_dir = os.path.join("project", "work_dir", task_id)
|
|
if os.path.exists(work_dir):
|
|
return work_dir
|
|
else:
|
|
logger.error(f"工作目录不存在: {work_dir}")
|
|
raise FileNotFoundError(f"工作目录不存在: {work_dir}")
|
|
|
|
|
|
# TODO: 是不是应该将 Prompt 写成一个 class
|
|
def get_config_template(comp_template: CompTemplate = CompTemplate.CHINA) -> dict:
|
|
if comp_template == CompTemplate.CHINA:
|
|
return load_toml(os.path.join("app", "config", "md_template.toml"))
|
|
|
|
|
|
def load_toml(path: str) -> dict:
|
|
with open(path, "rb") as f:
|
|
return tomllib.load(f)
|
|
|
|
|
|
def load_markdown(path: str) -> str:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return f.read()
|
|
|
|
|
|
def get_current_files(folder_path: str, type: str = "all") -> list[str]:
|
|
files = os.listdir(folder_path)
|
|
if type == "all":
|
|
return files
|
|
elif type == "md":
|
|
return [file for file in files if file.endswith(".md")]
|
|
elif type == "ipynb":
|
|
return [file for file in files if file.endswith(".ipynb")]
|
|
elif type == "data":
|
|
return [
|
|
file for file in files if file.endswith(".xlsx") or file.endswith(".csv")
|
|
]
|
|
elif type == "image":
|
|
return [
|
|
file for file in files if file.endswith(".png") or file.endswith(".jpg")
|
|
]
|
|
|
|
|
|
# 判断content是否包含图片 xx.png,对其处理为 
|
|
def transform_link(task_id: str, content: str):
|
|
content = re.sub(
|
|
r"!\[(.*?)\]\((.*?\.(?:png|jpg|jpeg|gif|bmp|webp))\)",
|
|
lambda match: f"})",
|
|
content,
|
|
)
|
|
return content
|
|
|
|
|
|
# TODO: fix 公式显示
|
|
def md_2_docx(task_id: str):
|
|
work_dir = get_work_dir(task_id)
|
|
md_path = os.path.join(work_dir, "res.md")
|
|
docx_path = os.path.join(work_dir, "res.docx")
|
|
|
|
extra_args = [
|
|
"--resource-path",
|
|
str(work_dir),
|
|
"--mathml", # MathML 格式公式
|
|
"--standalone",
|
|
]
|
|
|
|
pypandoc.convert_file(
|
|
source_file=md_path,
|
|
to="docx",
|
|
outputfile=docx_path,
|
|
format="markdown+tex_math_dollars",
|
|
extra_args=extra_args,
|
|
)
|
|
print(f"转换完成: {docx_path}")
|
|
logger.info(f"转换完成: {docx_path}")
|
|
|
|
|
|
def split_footnotes(text: str) -> tuple[str, list[tuple[str, str]]]:
|
|
main_text = re.sub(
|
|
r"\n\[\^\d+\]:.*?(?=\n\[\^|\n\n|\Z)", "", text, flags=re.DOTALL
|
|
).strip()
|
|
|
|
# 匹配脚注定义
|
|
footnotes = re.findall(r"\[\^(\d+)\]:\s*(.+?)(?=\n\[\^|\n\n|\Z)", text, re.DOTALL)
|
|
logger.info(f"main_text:{main_text} \n footnotes:{footnotes}")
|
|
return main_text, footnotes
|