fix bug

fix bug

Author: jihe520

backend/app/core/agents.py

@@ -34,7 +34,7 @@ def __init__(
         self.current_chat_turns = 0  # 当前对话轮次计数器
         self.user_output = user_output
 
-    async def run(self, prompt: str, system_prompt: str) -> str:
+    async def run(self, prompt: str, system_prompt: str, sub_title: str) -> str:
         """
         执行agent的对话并返回结果和总结
 
@@ -54,7 +54,9 @@ async def run(self, prompt: str, system_prompt: str) -> str:
 
             # 获取历史消息用于本次对话
             response = await self.model.chat(
-                history=self.chat_history, agent_name=self.__class__.__name__
+                history=self.chat_history,
+                agent_name=self.__class__.__name__,
+                sub_title=sub_title,
             )
             response_content = response.choices[0].message.content
             self.chat_history.append({"role": "assistant", "content": response_content})
@@ -140,7 +142,7 @@ async def run(self, prompt: str, subtask_title: str) -> CoderToWriter:
                 f"Failed to complete task after {self.max_retries} attempts. Last error: {last_error_message}"
             )
 
-            # try:
+        # try:
         while (
             not task_completed
             and retry_count < self.max_retries
@@ -189,6 +191,7 @@ async def run(self, prompt: str, subtask_title: str) -> CoderToWriter:
                     )
 
                     # 执行工具调用
+                    logger.warning("执行工具调用")
                     (
                         text_to_gpt,
                         error_occurred,
@@ -207,6 +210,8 @@ async def run(self, prompt: str, subtask_title: str) -> CoderToWriter:
                     )
 
                     if error_occurred:
+                        logger.warning("代码执行错误")
+
                         retry_count += 1
                         last_error_message = error_message
                         reflection_prompt = get_reflection_prompt(error_message, code)
@@ -224,6 +229,8 @@ async def run(self, prompt: str, subtask_title: str) -> CoderToWriter:
                     # 检查任务完成情况时也计入对话轮次
                     self.current_chat_turns += 1
                     # 使用所有执行结果生成检查提示
+                    logger.warning("判断是否完成")
+
                     completion_check_prompt = get_completion_check_prompt(
                         prompt, text_to_gpt
                     )
@@ -245,6 +252,7 @@ async def run(self, prompt: str, subtask_title: str) -> CoderToWriter:
                         hasattr(completion_response.choices[0].message, "tool_calls")
                         and completion_response.choices[0].message.tool_calls
                     ):
+                        logger.warning("没有调用工具，代表已经完成了")
                         task_completed = True
                         return completion_response.choices[0].message.content
 
@@ -282,21 +290,25 @@ async def run(
         self,
         prompt: str,
         available_images: list[str] = None,
+        sub_title: str = None,
     ) -> str:
         """
         执行写作任务
         Args:
             prompt: 写作提示
             available_images: 可用的图片相对路径列表（如 20250420-173744-9f87792c/编号_分布.png）
+            sub_title: 子任务标题
         """
+        logger.info(f"subtitle是:{sub_title}")
+
         if available_images:
             self.available_images = available_images
             # 拼接成完整URL
             image_list = ",".join(available_images)
             image_prompt = f"\n可用的图片链接列表：\n{image_list}\n请在写作时适当引用这些图片链接。"
             prompt = prompt + image_prompt
 
-        return await super().run(prompt, self.system_prompt)
+        return await super().run(prompt, self.system_prompt, sub_title)
 
     async def summarize(self) -> str:
         """

backend/app/core/llm.py

@@ -1,4 +1,5 @@
 import json
+from app.utils.common_utils import transform_link
 from openai import OpenAI
 from app.utils.log_util import logger
 import time
@@ -33,7 +34,10 @@ async def chat(
         retry_delay: float = 1.0,  # 添加重试延迟
         top_p: float | None = None,  # 添加top_p参数,
         agent_name: str = "NO_NAME",  # CoderAgent or WriterAgent
+        sub_title: str | None = None,
     ) -> str:
+        logger.info(f"subtitle是:{sub_title}")
+
         kwargs = {
             "model": self.model,
             "messages": history,
@@ -52,10 +56,11 @@ async def chat(
         for attempt in range(max_retries):
             try:
                 completion = self.client.chat.completions.create(**kwargs)
+                logger.info(f"API返回: {completion}")
                 if not completion or not hasattr(completion, "choices"):
                     raise ValueError("无效的API响应")
                 self.chat_count += 1
-                await self.analyse_completion(completion, agent_name)
+                await self.analyse_completion(completion, agent_name, sub_title)
                 return completion
             except json.JSONDecodeError:
                 logger.error(f"第{attempt + 1}次重试: API返回无效JSON")
@@ -65,7 +70,9 @@ async def chat(
                 logger.debug(f"请求参数: {kwargs}")
                 raise  # 如果所有重试都失败，则抛出异常
 
-    async def analyse_completion(self, completion, agent_name):
+    async def analyse_completion(self, completion, agent_name, sub_title):
+        logger.info(f"subtitle是:{sub_title}")
+
         code = ""
         if (
             hasattr(completion.choices[0].message, "tool_calls")
@@ -74,20 +81,22 @@ async def analyse_completion(self, completion, agent_name):
             tool_call = completion.choices[0].message.tool_calls[0]
             if tool_call.function.name == "execute_code":
                 code = json.loads(tool_call.function.arguments)["code"]
-        await self.send_message(agent_name, completion.choices[0].message.content, code)
+        (
+            await self.send_message(
+                agent_name, completion.choices[0].message.content, code, sub_title
+            )
+        )
 
-    async def send_message(self, agent_name, content, code=""):
+    async def send_message(self, agent_name, content, code="", sub_title=None):
+        logger.info(f"subtitle是:{sub_title}")
         if agent_name == "CoderAgent":
             agent_msg: CoderMessage = CoderMessage(content=content, code=code)
         elif agent_name == "WriterAgent":
-            # 判断content是否包含图片 xx.png,对其处理为    http://localhost:8000/static/20250428-200915-ebc154d4/512.jpg
-            if re.search(r"\.(png|jpg|jpeg|gif|bmp|webp)$", content):
-                content = re.sub(
-                    r"\.(png|jpg|jpeg|gif|bmp|webp)$",
-                    lambda match: f"http://localhost:8000/static/{self.task_id}/{match.group(0)}",
-                    content,
-                )
-            agent_msg: WriterMessage = WriterMessage(content=content)
+            # 处理 Markdown 格式的图片语法
+            content = transform_link(self.task_id, content)
+            agent_msg: WriterMessage = WriterMessage(
+                content=content, sub_title=sub_title
+            )
         else:
             raise ValueError(f"无效的agent_name: {agent_name}")
 

backend/app/core/workflow.py

@@ -126,6 +126,7 @@ async def execute(self, problem: Problem):
             writer_response = await writer_agent.run(
                 writer_prompt,
                 available_images=await e2b_code_interpreter.get_created_images(key),
+                sub_title=key,
             )
 
             await redis_manager.publish_message(
@@ -156,7 +157,7 @@ async def execute(self, problem: Problem):
                 comp_template=problem.comp_template,
                 format_output=problem.format_output,
             )
-            writer_response = await writer_agent.run(value)
+            writer_response = await writer_agent.run(prompt=value, sub_title=key)
             user_output.set_res(key, writer_response)
 
         logger.info(user_output.get_res())

backend/app/routers/modeling.py

@@ -6,7 +6,7 @@
 from app.utils.redis_manager import redis_manager
 from app.schemas.request import Problem
 from app.schemas.response import AgentMessage, AgentType, Message, SystemMessage
-from app.utils.common_utils import create_task_id, create_work_dir
+from app.utils.common_utils import create_task_id, create_work_dir, get_config_template
 import os
 import asyncio
 from fastapi import HTTPException
@@ -83,6 +83,13 @@ async def modeling(
     return {"task_id": task_id, "status": "processing"}
 
 
+@router.get("/writer_seque")
+async def get_writer_seque():
+    # 返回论文顺序
+    config_template: dict = get_config_template(CompTemplate.CHINA)
+    return list(config_template.keys())
+
+
 async def run_modeling_task_async(
     task_id: str,
     ques_all: str,

backend/app/schemas/response.py

@@ -77,6 +77,7 @@ class CoderMessage(AgentMessage):
 
 class WriterMessage(AgentMessage):
     agent_type: AgentType = AgentType.WRITER
+    sub_title: str | None = None
 
 
 # 所有可能的消息类型

backend/app/tools/code_interpreter.py

@@ -115,7 +115,7 @@ async def _pre_execute_code(self):
         )
         await self.execute_code(init_code)
 
-    def _truncate_text(self, text: str, max_length: int = 5000) -> str:
+    def _truncate_text(self, text: str, max_length: int = 2000) -> str:
         """截断文本，保留开头和结尾的重要信息"""
         if len(text) <= max_length:
             return text
@@ -264,6 +264,7 @@ async def execute_code(self, code: str) -> tuple[str, bool, str]:
                             msg=result._repr_javascript_(),
                         )
                     )
+
                     # 处理主要结果
                 # if result.is_main_result and result.text:
                 #     result_text = self._truncate_text(result.text)
@@ -278,15 +279,23 @@ async def execute_code(self, code: str) -> tuple[str, bool, str]:
         for item in content_to_display:
             if isinstance(item, dict):
                 if item.get("type") in ["stdout", "stderr", "error"]:
-                    text_to_gpt.append(item.get("content") or item.get("value") or "")
+                    text_to_gpt.append(
+                        self._truncate_text(
+                            item.get("content") or item.get("value") or ""
+                        )
+                    )
             elif isinstance(item, ResultModel):
                 if item.format in ["text", "html", "markdown", "json"]:
-                    text_to_gpt.append(f"[{item.format}]\n{item.msg}")
+                    text_to_gpt.append(
+                        self._truncate_text(f"[{item.format}]\n{item.msg}")
+                    )
                 elif item.format in ["png", "jpeg", "svg", "pdf"]:
                     text_to_gpt.append(
                         f"[{item.format} 图片已生成，内容为 base64，未展示]"
                     )
 
+        logger.info(f"text_to_gpt: {text_to_gpt}")
+
         combined_text = "\n".join(text_to_gpt)
 
         # 在代码执行完成后，立即同步文件

backend/app/utils/common_utils.py

@@ -2,9 +2,9 @@
 import datetime
 import hashlib
 import tomllib
-from app.core.llm import LLM
 from app.utils.enums import CompTemplate
 from app.utils.log_util import logger
+import re
 
 
 def create_task_id() -> str:
@@ -38,7 +38,8 @@ def get_work_dir(task_id: str) -> str:
         raise FileNotFoundError(f"工作目录不存在: {work_dir}")
 
 
-def get_config_template(comp_template: CompTemplate) -> dict:
+#  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"))
 
@@ -72,7 +73,7 @@ def get_current_files(folder_path: str, type: str = "all") -> list[str]:
         ]
 
 
-def simple_chat(model: LLM, history: list) -> str:
+def simple_chat(model, history: list) -> str:
     """
     Description of the function.
 
@@ -92,3 +93,13 @@ def simple_chat(model: LLM, history: list) -> str:
     completion = model.client.chat.completions.create(**kwargs)
 
     return completion.choices[0].message.content
+
+
+# 判断content是否包含图片 xx.png,对其处理为    ![filename](http://localhost:8000/static/20250428-200915-ebc154d4/filename.jpg)
+def transform_link(task_id: str, content: str):
+    content = re.sub(
+        r"!\[(.*?)\]\((.*?\.(?:png|jpg|jpeg|gif|bmp|webp))\)",
+        lambda match: f"![{match.group(1)}](http://localhost:8000/static/{task_id}/{match.group(2)})",
+        content,
+    )
+    return content

frontend/src/utils/response.ts

@@ -71,6 +71,7 @@ export interface CoderMessage extends AgentMessage {
 
 export interface WriterMessage extends AgentMessage {
   agent_type: 'WriterAgent';
+  sub_title?: string;
 }
 
 export type Message = SystemMessage | UserMessage | CoderMessage | WriterMessage;