MathModelAgent/models/user_input.py

from models.user_output import UserOutput
from utils.enums import CompTemplate, FormatOutPut
from core.LLM import BaseModel
from utils.logger import log
from utils.common_utils import simple_chat
import json
import os
from core.prompt import FORMAT_QUESTIONS_PROMPT

class UserInput:
    def __init__(
        self,
        comp_template: CompTemplate = CompTemplate.CHINA,
        format_output: FormatOutPut = FormatOutPut.Markdown,
        data_folder_path: str = "",
        bg_ques_all: str = "",
        model: BaseModel = None,  # TODO: 不同任务选择不同模型
        init_with_llm: bool = True,  # 新增参数控制是否使用LLM初始化问题
    ):
        self.comp_template: CompTemplate = comp_template  # 选择模板
        self.format_output: FormatOutPut = format_output  # 选择输出格式
        self.data_folder_path: str = (
            data_folder_path  # 数据集存放的文件夹 .\project\sample_data
        )
        self.bg_ques_all: str = bg_ques_all  # 用户输入的完整背景以及问题
        self.model: BaseModel = model  # 模型
        self.ques_count: int = 0  # 问题数量
        self.questions: dict[str, str | int] = {}  # 问题

        self.config_template = {}

        if init_with_llm:
            self.init_questions()

    def init_questions(self) -> None:
        """用户输入问题 使用LLM 格式化 questions"""
        # TODO:  "note": <补充说明,如果没有补充说明，请填 null>,

        history = [
            {
                "role": "system",
                "content": FORMAT_QUESTIONS_PROMPT,
            },
            {"role": "user", "content": self.bg_ques_all},
        ]
        json_str = simple_chat(self.model, history)
        json_str = json_str.replace("```json", "").replace("```", "").strip()
        log.info(f"json_str: {json_str}\n\n")
        # 检查返回的 JSON 字符串是否有效

        if not json_str:
            raise ValueError("返回的 JSON 字符串为空，请检查输入内容。")

        try:
            self.questions = json.loads(json_str)
            self.ques_count = self.questions["ques_count"]
        except json.JSONDecodeError as e:
            raise ValueError(f"JSON 解析错误: {e}")

    def set_questions_directly(self, questions: dict[str, str | int]) -> None:
        """直接设置questions,跳过LLM请求"""
        self.questions = questions
        self.ques_count = questions.get("ques_count", 0)

    def get_ques_count(self):
        return self.ques_count

    def get_comp_template(self):
        return self.comp_template

    def set_comp_template(self, comp_template):
        self.comp_template = comp_template

    def set_config_template(self, config_template):
        self.config_template = config_template

    def get_bg_ques_all(self):
        return self.bg_ques_all

    def get_format_output(self):
        return self.format_output

    def set_format_output(self, format_output):
        self.format_output = format_output

    def get_data_folder_path(self):
        return self.data_folder_path

    def set_data_folder_path(self, data_folder_path):
        self.data_folder_path = data_folder_path

    def get_questions(self):
        return self.questions

    def set_questions(self, questions):
        self.questions = questions

    def get_questions_quesx(self) -> dict[str, str]:
        """获取问题1,2,3...的键值对"""
        # 获取所有以 "ques" 开头的键值对
        questions_quesx = {
            key: value
            for key, value in self.questions.items()
            if key.startswith("ques") and key != "ques_count"
        }
        return questions_quesx

    def get_questions_quesx_keys(self) -> list[str]:
        """获取问题1,2...的键"""
        return list(self.get_questions_quesx().keys())

    def __str__(self):
        return f"choice: {self.comp_template}, format_out_put: {self.format_output}, data_folder_path: {self.data_folder_path}, bg_ques_all: {self.bg_ques_all}, model: {self.model}, ques_count: {self.ques_count}, questions: {self.questions}"

    def get_data_path(self) -> list[str]:
        """获取数据集完整路径"""
        data_folder_path = self.get_data_folder_path() # "" ./project/sample_data
        files = os.listdir(data_folder_path)
        full_paths = [os.path.abspath(os.path.join(data_folder_path, file)) for file in files]
        log.info(f"full_paths: {full_paths}")
        return full_paths

    def get_solution_flows(self):
        questions_quesx = self.get_questions_quesx()

        ques_flow = {
            key: {
                "coder_prompt": f"""
                        完成如下问题{value}
                    """,
            }
            for key, value in questions_quesx.items()
        }
        flows = {
            "eda": {
                # TODO ： 获取当前路径下的所有数据集
                "coder_prompt": f"""
                        对数据 {self.get_data_path()} 进行EDA分析(数据清洗,可视化),清洗后的数据保存当前目录下,**不需要复杂的模型**
                    """,
            },
            **ques_flow,
            "sensitivity_analysis": {
                "coder_prompt": """
                        根据上面建立的模型，选择一个模型，完成敏感性分析
                    """,
            },
        }

        return flows

    def get_writer_prompt(
        self, key: str, coder_response: str, notebook_serializer
    ) -> str:
        """根据不同的key生成对应的writer_prompt

        Args:
            key: 任务类型
            coder_response: 代码执行结果

        Returns:
            str: 生成的writer_prompt
        """

        notebook_output = notebook_serializer.get_notebook_output_content(key)
        # TODO: 结果{coder_response} 是否需要
        # TODO: 将当前产生的文件，路径发送给 writer_agent
        questions_quesx_keys = self.get_questions_quesx_keys()
        # TODO： 小标题编号
        # 题号最多6题
        quesx_writer_prompt = {
            key: f"""
                问题背景{self.get_questions()["background"]},不需要编写代码,代码手得到的结果{coder_response},{notebook_output},按照如下模板撰写：{self.config_template[key]}
            """
            for key in questions_quesx_keys
        }

        writer_prompt = {
            "eda": f"""
                问题背景{self.get_questions()["background"]},不需要编写代码,代码手得到的结果{coder_response},{notebook_output},按照如下模板撰写：{self.config_template["eda"]}
            """,
            **quesx_writer_prompt,
            "sensitivity_analysis": f"""
                问题背景{self.get_questions()["background"]},不需要编写代码,代码手得到的结果{coder_response},{notebook_output},按照如下模板撰写：{self.config_template["sensitivity_analysis"]}
            """,
        }

        if key in writer_prompt:
            return writer_prompt[key]
        else:
            raise ValueError(f"未知的任务类型: {key}")

    def get_write_flows(self, user_output: UserOutput):
        model_build_solve = user_output.get_model_build_solve()
        bg_ques_all = self.get_bg_ques_all()
        flows = {
            "firstPage": f"""问题背景{bg_ques_all},不需要编写代码,根据模型的求解的信息{model_build_solve}，按照如下模板撰写：{self.config_template["firstPage"]}，撰写标题，摘要，关键词""",
            "RepeatQues": f"""问题背景{bg_ques_all},不需要编写代码,根据模型的求解的信息{model_build_solve}，按照如下模板撰写：{self.config_template["RepeatQues"]}，撰写问题重述""",
            "analysisQues": f"""问题背景{bg_ques_all},不需要编写代码,根据模型的求解的信息{model_build_solve}，按照如下模板撰写：{self.config_template["analysisQues"]}，撰写问题分析""",
            "modelAssumption": f"""问题背景{bg_ques_all},不需要编写代码,根据模型的求解的信息{model_build_solve}，按照如下模板撰写：{self.config_template["modelAssumption"]}，撰写模型假设""",
            "symbol": f"""不需要编写代码,根据模型的求解的信息{model_build_solve}，按照如下模板撰写：{self.config_template["symbol"]}，撰写符号说明部分""",
            "judge": f"""不需要编写代码,根据模型的求解的信息{model_build_solve}，按照如下模板撰写：{self.config_template["judge"]}，撰写模型的评价部分""",
            # TODO: 修改参考文献插入方式
            "reference": f"""不需要编写代码,根据模型的求解的信息{model_build_solve}，可以生成参考文献,按照如下模板撰写：{self.config_template["reference"]}，撰写参考文献""",
        }
        return flows