fix csv display issue (microsoft#92)

1. fixed artifact paths 2. fixed displaying csv
mdwoicke · Dec 26, 2023 · 7baeb77 · 7baeb77
1 parent 304e5da
commit 7baeb77
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 13 deletions.
diff --git a/playground/UI/app.py b/playground/UI/app.py
@@ -55,10 +55,14 @@ def send_message_sync(msg: str) -> Round:
             ]:
                 continue
             elif atta.type == AttachmentType.artifact_paths:
-                artifact_paths = [item.replace("file://", "") for item in atta.content]
+                artifact_paths = atta.content
             else:
                 elements.append(
-                    cl.Text(name=atta.type.value, content=atta.content.encode(), display="inline"),
+                    cl.Text(
+                        name=atta.type.value,
+                        content=atta.content.encode(),
+                        display="inline",
+                    ),
                 )
         elements.append(
             cl.Text(
@@ -79,6 +83,25 @@ def send_message_sync(msg: str) -> Round:
         if len(artifact_paths) > 0:
             elements = []
             for path in artifact_paths:
-                image = cl.Image(name=path, display="inline", path=path, size="large")
-                elements.append(image)
+                # if path is image, display it
+                if path.endswith((".png", ".jpg", ".jpeg", ".gif")):
+                    image = cl.Image(
+                        name=path,
+                        display="inline",
+                        path=path,
+                        size="large",
+                    )
+                    elements.append(image)
+                elif path.endswith(".csv"):
+                    import pandas as pd
+
+                    data = pd.read_csv(path)
+                    row_count = len(data)
+                    table = cl.Text(
+                        name=path,
+                        content=f"There are {row_count} in the data. The top {min(row_count, 5)} rows are:\n"
+                        + data.head(n=5).to_markdown(),
+                        display="inline",
+                    )
+                    elements.append(table)
         await cl.Message(content=f"{post.message}", elements=elements).send()
diff --git a/taskweaver/code_interpreter/code_interpreter.py b/taskweaver/code_interpreter/code_interpreter.py
@@ -3,7 +3,7 @@
 
 from injector import inject
 
-from taskweaver.code_interpreter.code_executor import CodeExecutor, get_artifact_uri
+from taskweaver.code_interpreter.code_executor import CodeExecutor
 from taskweaver.code_interpreter.code_generator import CodeGenerator, format_code_revision_message
 from taskweaver.code_interpreter.code_generator.code_generator import format_output_revision_message
 from taskweaver.code_interpreter.code_verification import code_snippet_verification, format_code_correction_message
@@ -188,14 +188,10 @@ def reply(
             Attachment.create(
                 AttachmentType.artifact_paths,
                 [
-                    get_artifact_uri(
-                        execution_id=exec_result.execution_id,
-                        file=(
-                            a.file_name
-                            if os.path.isabs(a.file_name) or not self.config.use_local_uri
-                            else os.path.join(self.executor.execution_cwd, a.file_name)
-                        ),
-                        use_local_uri=self.config.use_local_uri,
+                    (
+                        a.file_name
+                        if os.path.isabs(a.file_name) or not self.config.use_local_uri
+                        else os.path.join(self.executor.execution_cwd, a.file_name)
                     )
                     for a in exec_result.artifact
                 ],