jihe520__mathmodelagent/core/agents.py

import json
from core.LLM import BaseModel
from models.user_output import UserOutput
from tools.code_interpreter import CodeInterpreter
from utils.common_utils import get_current_files
from utils.enums import CompTemplate, FormatOutPut
from utils.logger import log
from utils.RichPrinter import RichPrinter
from utils.notebook_serializer import NotebookSerializer
from .functions import functions
from .prompt import *

class Agent:
    def __init__(
        self,
        model: BaseModel,
        max_chat_turns: int = 30,  # 单个agent最大对话轮次
        user_output: UserOutput = None,
    ) -> None:
        self.model = model
        self.chat_history: list[dict] = []  # 存储对话历史
        self.max_chat_turns = max_chat_turns  # 最大对话轮次
        self.current_chat_turns = 0  # 当前对话轮次计数器
        self.user_output = user_output

    def run(self, prompt: str, system_prompt: str) -> str:
        """
        执行agent的对话并返回结果和总结

        Args:
            prompt: 输入的提示

        Returns:
            str: 模型的响应
        """
        try:
            RichPrinter.agent_start(self.__class__.__name__)
            self.current_chat_turns = 0  # 重置对话轮次计数器

            # 更新对话历史
            self.append_chat_history({"role": "system", "content": system_prompt})
            self.append_chat_history({"role": "user", "content": prompt})

            # 获取历史消息用于本次对话
            response = self.model.chat(
                history=self.chat_history, agent_name=self.__class__.__name__
            )
            response_content = response.choices[0].message.content
            self.chat_history.append({"role": "assistant", "content": response_content})

            RichPrinter.agent_end(self.__class__.__name__)
            return response_content
        except Exception as e:
            error_msg = f"执行过程中遇到错误: {str(e)}"
            log.error(f"Agent执行失败: {str(e)}")
            return error_msg

    def append_chat_history(self, msg: dict) -> None:
        self.chat_history.append(msg)
        self.user_output.data_recorder.append_chat_history(
            msg, agent_name=self.__class__.__name__
        )


class ModelerAgent(Agent):  # 继承自Agent类而不是BaseModel
    def __init__(
        self,
        model: BaseModel,
        max_chat_turns: int = 30,  # 添加最大对话轮次限制
    ) -> None:
        super().__init__(model, max_chat_turns)
        self.system_prompt =MODELER_PROMPT

# 代码强
class CoderAgent(Agent):  # 同样继承自Agent类
    def __init__(
        self,
        model: BaseModel,
        work_dir: str,  # 工作目录
        max_chat_turns: int = 60,  # TODO: 配置文件
        max_retries: int = 5,  # 重试次数
        user_output: UserOutput = None,
    ) -> None:
        super().__init__(model, max_chat_turns, user_output)
        self.work_dir = work_dir
        self.notebook_serializer = NotebookSerializer(work_dir)
        self.code_interpreter = CodeInterpreter(work_dir, self.notebook_serializer)

        self.max_retries = max_retries
        self.is_first_run = True
        self.created_images: list[str] = []  # 当前任务创建的图片列表
        self.system_prompt = CODER_PROMPT

        self.created_images: list[str] = []  # 当前求解创建的图片

    def run(self, prompt: str, subtask_title: str) -> str:
        RichPrinter.agent_start(self.__class__.__name__)
        self.created_images.clear()  # 清空上一次任务的图片列表
        # TODO: jupyter 的notebook 分节
        self.notebook_serializer.add_markdown_segmentation_to_notebook(
            content=prompt, segmentation=subtask_title
        )

        # 如果是第一次运行，则添加系统提示
        if self.is_first_run:
            self.is_first_run = False
            self.append_chat_history({"role": "system", "content": self.system_prompt})

        self.append_chat_history({"role": "user", "content": prompt})

        retry_count = 0
        last_error_message = ""
        task_completed = False

        if self.current_chat_turns >= self.max_chat_turns:
            raise Exception(
                f"Reached maximum number of chat turns ({self.max_chat_turns}). Task incomplete."
            )

        if retry_count >= self.max_retries:
            raise Exception(
                f"Failed to complete task after {self.max_retries} attempts. Last error: {last_error_message}"
            )

        try:
            while (
                not task_completed
                and retry_count < self.max_retries
                and self.current_chat_turns < self.max_chat_turns
            ):
                self.current_chat_turns += 1
                response = self.model.chat(
                    history=self.chat_history,
                    tools=functions,
                    tool_choice="auto",
                    agent_name=self.__class__.__name__,
                )

                if (
                    hasattr(response.choices[0].message, "tool_calls")
                    and response.choices[0].message.tool_calls
                ):
                    tool_call = response.choices[0].message.tool_calls[0]
                    tool_id = tool_call.id
                    # TODO: json JSON解析时遇到了无效的转义字符
                    if tool_call.function.name == "execute_code":
                        code = json.loads(tool_call.function.arguments)["code"]
                        full_content = response.choices[0].message.content

                        # 更新对话历史 - 添加助手的响应
                        self.append_chat_history(
                            {
                                "role": "assistant",
                                "content": full_content,
                                "tool_calls": [
                                    {
                                        "id": tool_id,
                                        "type": "function",
                                        "function": {
                                            "name": "execute_code",
                                            "arguments": json.dumps({"code": code}),
                                        },
                                    }
                                ],
                            }
                        )

                        # 执行代码并获取结果
                        (
                            text_to_gpt,
                            content_to_display,
                            error_occurred,
                            error_message,
                        ) = self.code_interpreter.jupyter_kernel.execute_code(code)

                        # 添加工具执行结果
                        self.append_chat_history(
                            {
                                "role": "tool",
                                "content": text_to_gpt,
                                "tool_call_id": tool_id,
                            }
                        )

                        if error_occurred:
                            retry_count += 1
                            last_error_message = error_message
                            reflection_prompt = get_reflection_prompt(error_message,code)

                            self.append_chat_history(
                                {"role": "user", "content": reflection_prompt}
                            )
                            continue

                        # 检查任务完成情况时也计入对话轮次
                        self.current_chat_turns += 1
                        completion_check_prompt = get_completion_check_prompt(prompt, text_to_gpt)
                        self.append_chat_history(
                            {"role": "user", "content": completion_check_prompt}
                        )

                        completion_response = self.model.chat(
                            history=self.chat_history,
                            tools=functions,
                            tool_choice="auto",
                            agent_name=self.__class__.__name__,
                        )

                        # # TODO: 压缩对话历史
                        # if completion_response.usage.prompt_tokens > self.max_tokens * 0.75:
                        #     # 保留系统消息
                        #     system_messages = [msg for msg in self.chat_history if msg["role"] == "system"]
                        #     recent_messages = self.chat_history[-5:]
                        #     self.chat_history = system_messages + recent_messages
                        # elif completion_response.usage.prompt_tokens > self.max_tokens:
                        #     raise Exception("token超出限制")

                        if not (
                            hasattr(
                                completion_response.choices[0].message, "tool_calls"
                            )
                            and completion_response.choices[0].message.tool_calls
                        ):
                            task_completed = True
                            RichPrinter.agent_end(self.__class__.__name__)
                            return completion_response.choices[0].message.content
                if retry_count >= self.max_retries:
                    return f"Failed to complete task after {self.max_retries} attempts. Last error: {last_error_message}"

                if self.current_chat_turns >= self.max_chat_turns:
                    return f"Reached maximum number of chat turns ({self.max_chat_turns}). Task incomplete."

            RichPrinter.agent_end(self.__class__.__name__)
            self._update_created_images()
            return response.choices[0].message.content
        except Exception as e:
            error_msg = f"执行过程中遇到错误: {str(e)}"
            log.error(f"Agent执行失败: {str(e)}")
            return error_msg

    def get_created_images(self) -> list[str]:
        """获取当前任务创建的图片列表"""
        return self.created_images

    def _update_created_images(self) -> None:
        """更新创建的图片列表"""
        current_images = set(get_current_files(self.work_dir, "image"))
        previous_images = set(self.created_images)
        self.created_images = list(current_images - previous_images)

    def get_notebook_serializer(self) -> NotebookSerializer:
        return self.notebook_serializer


# 长文本
# TODO: 并行 parallel
class WriterAgent(Agent):  # 同样继承自Agent类
    def __init__(
        self,
        model: BaseModel,
        max_chat_turns: int = 10,  # 添加最大对话轮次限制
        comp_template: CompTemplate = CompTemplate,
        format_output: FormatOutPut = FormatOutPut.Markdown,
        user_output: UserOutput = None,
    ) -> None:
        super().__init__(model, max_chat_turns, user_output)
        self.format_out_put = format_output
        self.comp_template = comp_template
        self.system_prompt = get_writer_prompt(format_output)
        self.available_images: list[str] = []

    def run(self, prompt: str, available_images: list[str] = None) -> str:
        """
        执行写作任务
        Args:
            prompt: 写作提示
            available_images: 可用的图片列表
        """
        if available_images:
            self.available_images = available_images
            image_list = "\n".join([f"- {img}" for img in available_images])
            image_prompt = (
                f"\n可用的图片列表：\n{image_list}\n请在写作时适当引用这些图片。"
            )
            prompt = prompt + image_prompt

        return super().run(prompt, self.system_prompt)

    def summarize(self) -> str:
        """
        总结对话内容
        """
        try:
            self.append_chat_history(
                {"role": "user", "content": "请简单总结以上完成什么任务取得什么结果:"}
            )
            # 获取历史消息用于本次对话
            response = self.model.chat(
                history=self.chat_history, agent_name=self.__class__.__name__
            )
            self.append_chat_history(
                {"role": "assistant", "content": response.choices[0].message.content}
            )
            return response.choices[0].message.content
        except Exception as e:
            log.error(f"总结生成失败: {str(e)}")
            # 返回一个基础总结，避免完全失败
            return "由于网络原因无法生成详细总结，但已完成主要任务处理。"