From 0b31a79e359ed0cf8ca5f9f243902eeecf3b419f Mon Sep 17 00:00:00 2001
From: wuyiqunLu <wuyiqun.lu@landing.ai>
Date: Wed, 18 Sep 2024 12:01:47 +0800
Subject: [PATCH] feat: change the display on artifacts

---
 vision_agent/tools/meta_tools.py | 43 ++++++++++++++++++++++++++------
 vision_agent/utils/execute.py    |  8 +++---
 2 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/vision_agent/tools/meta_tools.py b/vision_agent/tools/meta_tools.py
index 477cbd30..bcdbefb6 100644
--- a/vision_agent/tools/meta_tools.py
+++ b/vision_agent/tools/meta_tools.py
@@ -1,4 +1,5 @@
 import difflib
+import json
 import os
 import pickle as pkl
 import re
@@ -70,8 +71,8 @@ def redisplay_results(execution: Execution) -> None:
             display({MimeType.TEXT_LATEX: result.latex}, raw=True)
         if result.json is not None:
             display({MimeType.APPLICATION_JSON: result.json}, raw=True)
-        if result.artifact_name is not None:
-            display({MimeType.TEXT_ARTIFACT_NAME: result.artifact_name}, raw=True)
+        if result.artifact is not None:
+            display({MimeType.APPLICATION_ARTIFACT: result.artifact}, raw=True)
         if result.extra is not None:
             display(result.extra, raw=True)
 
@@ -210,7 +211,14 @@ def create_code_artifact(artifacts: Artifacts, name: str) -> str:
         return_str = f"[Artifact {name} created]"
     print(return_str)
 
-    display({MimeType.TEXT_ARTIFACT_NAME: name}, raw=True)
+    display(
+        {
+            MimeType.APPLICATION_ARTIFACT: json.dumps(
+                {"name": name, "content": artifacts[name]}
+            )
+        },
+        raw=True,
+    )
     return return_str
 
 
@@ -294,7 +302,14 @@ def edit_code_artifact(
 
     artifacts[name] = "".join(edited_lines)
 
-    display({MimeType.TEXT_ARTIFACT_NAME: name}, raw=True)
+    display(
+        {
+            MimeType.APPLICATION_ARTIFACT: json.dumps(
+                {"name": name, "content": artifacts[name]}
+            )
+        },
+        raw=True,
+    )
     return open_code_artifact(artifacts, name, cur_line)
 
 
@@ -350,7 +365,10 @@ def detect_dogs(image_path: str):
     code_lines = code.splitlines(keepends=True)
     total_lines = len(code_lines)
 
-    display({MimeType.TEXT_ARTIFACT_NAME: name}, raw=True)
+    display(
+        {MimeType.APPLICATION_ARTIFACT: json.dumps({"name": name, "content": code})},
+        raw=True,
+    )
     return view_lines(code_lines, 0, total_lines, name, total_lines)
 
 
@@ -415,7 +433,10 @@ def detect_dogs(image_path: str):
     code_lines = code.splitlines(keepends=True)
     total_lines = len(code_lines)
 
-    display({MimeType.TEXT_ARTIFACT_NAME: name}, raw=True)
+    display(
+        {MimeType.APPLICATION_ARTIFACT: json.dumps({"name": name, "content": code})},
+        raw=True,
+    )
     return view_lines(code_lines, 0, total_lines, name, total_lines)
 
 
@@ -429,7 +450,6 @@ def write_media_artifact(artifacts: Artifacts, local_path: str) -> str:
     with open(local_path, "rb") as f:
         media = f.read()
     artifacts[Path(local_path).name] = media
-    display({MimeType.TEXT_ARTIFACT_NAME: Path(local_path).name}, raw=True)
     return f"[Media {Path(local_path).name} saved]"
 
 
@@ -596,7 +616,14 @@ def replacer(match: re.Match) -> str:
     diff = get_diff_with_prompts(name, code, new_code)
     print(diff)
 
-    display({MimeType.TEXT_ARTIFACT_NAME: name}, raw=True)
+    display(
+        {
+            MimeType.APPLICATION_ARTIFACT: json.dumps(
+                {"name": name, "content": new_code}
+            )
+        },
+        raw=True,
+    )
     return diff
 
 
diff --git a/vision_agent/utils/execute.py b/vision_agent/utils/execute.py
index cc01a89d..be29394a 100644
--- a/vision_agent/utils/execute.py
+++ b/vision_agent/utils/execute.py
@@ -56,7 +56,7 @@ class MimeType(str, Enum):
     TEXT_LATEX = "text/latex"
     APPLICATION_JSON = "application/json"
     APPLICATION_JAVASCRIPT = "application/javascript"
-    TEXT_ARTIFACT_NAME = "text/artifact/name"
+    APPLICATION_ARTIFACT = "application/artifact"
 
 
 class FileSerializer:
@@ -129,7 +129,7 @@ def __init__(self, is_main_result: bool, data: Dict[str, Any]):
         self.latex = data.pop(MimeType.TEXT_LATEX, None)
         self.json = data.pop(MimeType.APPLICATION_JSON, None)
         self.javascript = data.pop(MimeType.APPLICATION_JAVASCRIPT, None)
-        self.artifact_name = data.pop(MimeType.TEXT_ARTIFACT_NAME, None)
+        self.artifact = data.pop(MimeType.APPLICATION_ARTIFACT, None)
         self.extra = data
         # Only keeping the PNG representation if both PNG and JPEG are present
         if self.png and self.jpeg:
@@ -207,8 +207,8 @@ def formats(self) -> Iterable[str]:
             formats.append("javascript")
         if self.mp4:
             formats.append("mp4")
-        if self.artifact_name:
-            formats.append("artifact_name")
+        if self.artifact:
+            formats.append("artifact")
         if self.extra:
             formats.extend(iter(self.extra))
         return formats