add

Author: jihe520

backend/app/core/agents.py

@@ -282,24 +282,17 @@ async def run(
         self,
         prompt: str,
         available_images: list[str] = None,
-        static_prefix: str = "/static/",
     ) -> str:
         """
         执行写作任务
         Args:
             prompt: 写作提示
             available_images: 可用的图片相对路径列表（如 20250420-173744-9f87792c/编号_分布.png）
-            static_prefix: 静态资源前缀
         """
         if available_images:
             self.available_images = available_images
             # 拼接成完整URL
-            image_list = "\n".join(
-                [
-                    f"- {static_prefix}{img if img.startswith('/') else '/' + img}"
-                    for img in available_images
-                ]
-            )
+            image_list = ",".join(available_images)
             image_prompt = f"\n可用的图片链接列表：\n{image_list}\n请在写作时适当引用这些图片链接。"
             prompt = prompt + image_prompt
 

backend/app/core/llm.py

@@ -5,6 +5,7 @@
 from app.schemas.response import AgentMessage, CoderMessage, WriterMessage
 from app.utils.enums import AgentType
 from app.utils.redis_manager import redis_manager
+import re
 
 
 class LLM:
@@ -79,6 +80,13 @@ async def send_message(self, agent_name, content, code=""):
         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)
         else:
             raise ValueError(f"无效的agent_name: {agent_name}")

backend/app/core/prompts.py

@@ -75,8 +75,7 @@ def get_writer_prompt(
         skill：熟练掌握{format_output}排版,
         output：你需要按照要求的格式排版,只输出{format_output}排版的内容
         
-        1. 当你输入图像引用时候，你需要将用户输入的文件名称路径切换为相对路径
-        如用户输入文件路径image_name.png,你转化为../jupyter/image_name.png,就可正确引用显示
+        1. 当你输入图像引用时候，image_name.png,就可正确引用显示
         2. 你不需要输出markdown的这个```markdown格式，只需要输出markdown的内容
         3. 严格按照参考用户输入的格式模板以及**正确的编号顺序**
         4. 不需要询问用户

backend/app/schemas/response.py

@@ -72,6 +72,7 @@ class CoderMessage(AgentMessage):
     agent_type: AgentType = AgentType.CODER
     code: str | None = None
     code_results: list[OutputItem] | None = None
+    files: list[str] | None = None
 
 
 class WriterMessage(AgentMessage):

backend/app/tools/code_interpreter.py

@@ -17,6 +17,9 @@
 import asyncio
 from app.config.setting import settings
 import json
+import base64
+
+from app.utils.common_utils import get_current_files
 
 
 def delete_color_control_char(string):
@@ -284,12 +287,19 @@ async def execute_code(self, code: str) -> tuple[str, bool, str]:
                         f"[{item.format} 图片已生成，内容为 base64，未展示]"
                     )
 
+        combined_text = "\n".join(text_to_gpt)
+
+        # 在代码执行完成后，立即同步文件
+        try:
+            await self.download_all_files_from_sandbox()
+            logger.info("文件同步完成")
+        except Exception as e:
+            logger.error(f"文件同步失败: {str(e)}")
+
         # 保存到分段内容
         ## TODO: Base64 等图像需要优化
         await self._push_to_websocket(content_to_display)
 
-        combined_text = "\n".join(text_to_gpt)
-
         return (
             combined_text,
             error_occurred,
@@ -302,6 +312,7 @@ async def _push_to_websocket(self, content_to_display: list[OutputItem] | None):
         agent_msg = CoderMessage(
             agent_type=AgentType.CODER,
             code_results=content_to_display,
+            files=get_current_files(self.work_dir, "all"),
         )
         logger.debug(f"发送消息: {agent_msg.model_dump_json()}")
         await redis_manager.publish_message(
@@ -346,36 +357,71 @@ def get_code_output(self, section: str) -> str:
 
     async def cleanup(self):
         """清理资源并关闭沙箱"""
-
-        if not await self.sbx.is_running():
-            logger.warning("沙箱已经关闭了")
-
         try:
             if self.sbx:
-                ## TODO: 生成一个文件，下载一份文件
-                await self.download_all_files_from_sandbox()
-                await self.sbx.kill()
-                logger.info("成功关闭沙箱环境")
+                if await self.sbx.is_running():
+                    try:
+                        await self.download_all_files_from_sandbox()
+                    except Exception as e:
+                        logger.error(f"下载文件失败: {str(e)}")
+                    finally:
+                        await self.sbx.kill()
+                        logger.info("成功关闭沙箱环境")
+                else:
+                    logger.warning("沙箱已经关闭，跳过清理步骤")
         except Exception as e:
             logger.error(f"清理沙箱环境失败: {str(e)}")
-            raise
+            # 这里可以选择不抛出异常，因为这是清理步骤
 
     async def download_all_files_from_sandbox(self) -> None:
-        """从沙箱中下载所有文件"""
+        """从沙箱中下载所有文件并与本地同步"""
         try:
-            files = await self.sbx.files.list("/home/user")
-            for file in files:
-                content = await self.sbx.files.read(file.path)
-                output_path = os.path.join(self.work_dir, file.name)
-                # 修复：确保写入文件时 content 一定是 bytes 类型
-                if isinstance(content, str):
-                    content = content.encode("utf-8")
-                with open(output_path, "wb") as f:
-                    f.write(content)
-                logger.info(f"成功下载文件: {file.name}")
+            # 获取沙箱中的文件列表
+            sandbox_files = await self.sbx.files.list("/home/user")
+            sandbox_files_dict = {f.name: f for f in sandbox_files}
+
+            # 获取本地文件列表
+            local_files = set()
+            if os.path.exists(self.work_dir):
+                local_files = set(os.listdir(self.work_dir))
+
+            # 下载新文件或更新已修改的文件
+            for file in sandbox_files:
+                try:
+                    local_path = os.path.join(self.work_dir, file.name)
+                    should_download = True
+
+                    # 检查文件是否需要更新
+                    if file.name in local_files:
+                        # 这里可以添加文件修改时间或内容哈希的比较
+                        # 暂时简单处理，有同名文件就更新
+                        pass
+
+                    if should_download:
+                        content = await self.sbx.files.read(file.path)
+
+                        # 确保目标目录存在
+                        os.makedirs(self.work_dir, exist_ok=True)
+
+                        # 写入文件
+                        with open(local_path, "wb") as f:
+                            if isinstance(content, str):
+                                if "," in content:  # 处理data URL格式
+                                    content = content.split(",")[1]
+                                    content = base64.b64decode(content)
+                                else:
+                                    content = content.encode("utf-8")
+                            f.write(content)
+                        logger.info(f"同步文件: {file.name}")
+
+                except Exception as e:
+                    logger.error(f"同步文件 {file.name} 失败: {str(e)}")
+                    continue
+
+            logger.info("文件同步完成")
+
         except Exception as e:
-            logger.error(f"下载文件失败: {str(e)}")
-            # 这里不再 raise，防止 cleanup 阶段因沙箱已关闭而报错
+            logger.error(f"文件同步失败: {str(e)}")
 
     async def shutdown_sandbox(self):
         """关闭沙箱环境"""

frontend/src/components/Files.vue

@@ -1,4 +1,5 @@
 <script setup lang="ts">
+import { computed } from 'vue'
 import Tree from '@/components/Tree.vue'
 import {
   SidebarContent,
@@ -8,49 +9,26 @@ import {
   SidebarMenu,
 } from '@/components/ui/sidebar'
 import { File } from 'lucide-vue-next'
+import { useTaskStore } from '@/stores/task'
 
-// This is sample data.
-const data = {
-  changes: [
-    {
-      file: 'README.md',
-      state: 'M',
-    },
-    {
-      file: 'api/hello/route.ts',
-      state: 'U',
-    },
-    {
-      file: 'app/layout.tsx',
-      state: 'M',
-    },
-  ],
-  tree: [
-    [
-      'app',
-      [
-        'api',
-        ['hello', ['route.ts']],
-        'page.tsx',
-        'layout.tsx',
-        ['blog', ['page.tsx']],
-      ],
-    ],
-    [
-      'components',
-      ['ui', 'button.tsx', 'card.tsx'],
-      'header.tsx',
-      'footer.tsx',
-    ],
-    ['lib', ['util.ts']],
-    ['public', 'favicon.ico', 'vercel.svg'],
-    '.eslintrc.json',
-    '.gitignore',
-    'next.config.js',
-    'tailwind.config.js',
-    'package.json',
-    'README.md',
-  ],
+const taskStore = useTaskStore()
+
+// 从消息中提取最新的文件列表
+const files = taskStore.files
+
+// 将文件列表转换为树形结构
+const fileTree = computed(() => {
+  return files.map(file => file)
+})
+
+const handleFileClick = (file: string) => {
+  // 处理文件点击
+  console.log('File clicked:', file)
+}
+
+const handleFileDownload = (file: string) => {
+  // 处理文件下载
+  console.log('Download file:', file)
 }
 </script>
 
@@ -61,7 +39,11 @@ const data = {
         <SidebarGroupLabel>Files</SidebarGroupLabel>
         <SidebarGroupContent>
           <SidebarMenu>
-            <Tree v-for="(item, index) in data.tree" :key="index" :item="item" />
+            <div v-if="fileTree.length === 0" class="px-4 py-2 text-sm text-gray-500">
+              No files
+            </div>
+            <Tree v-else v-for="(item, index) in fileTree" :key="index" :item="item" @click="handleFileClick(item)"
+              @download="handleFileDownload(item)" />
           </SidebarMenu>
         </SidebarGroupContent>
       </SidebarGroup>