From 09c9bb5f91057c0cea052a086bcd4bc393765e63 Mon Sep 17 00:00:00 2001
From: David Brochart <david.brochart@gmail.com>
Date: Mon, 9 Sep 2024 09:39:29 +0200
Subject: [PATCH] Change shared model stream output schema

---
 .../fps_kernels/kernel_driver/driver.py       | 52 +++++++++----------
 1 file changed, 24 insertions(+), 28 deletions(-)

diff --git a/plugins/kernels/fps_kernels/kernel_driver/driver.py b/plugins/kernels/fps_kernels/kernel_driver/driver.py
index d1761b77..f375af95 100644
--- a/plugins/kernels/fps_kernels/kernel_driver/driver.py
+++ b/plugins/kernels/fps_kernels/kernel_driver/driver.py
@@ -228,24 +228,16 @@ async def _handle_outputs(self, outputs: Array, msg: Dict[str, Any]):
                     text = text[:-1]
                 if (not outputs) or (outputs[-1]["name"] != content["name"]):  # type: ignore
                     outputs.append(
-                        #Map(
-                        #    {
-                        #        "name": content["name"],
-                        #        "output_type": msg_type,
-                        #        "text": Array([content["text"]]),
-                        #    }
-                        #)
-                        {
-                            "name": content["name"],
-                            "output_type": msg_type,
-                            "text": [text],
-                        }
+                        Map(
+                            {
+                                "name": content["name"],
+                                "output_type": msg_type,
+                                "text": Array([content["text"]]),
+                            }
+                        )
                     )
                 else:
-                    #outputs[-1]["text"].append(content["text"])  # type: ignore
-                    last_output = outputs[-1]
-                    last_output["text"].append(text)  # type: ignore
-                    outputs[-1] = last_output
+                    outputs[-1]["text"].append(content["text"])  # type: ignore
         elif msg_type in ("display_data", "execute_result"):
             if "application/vnd.jupyter.ywidget-view+json" in content["data"]:
                 # this is a collaborative widget
@@ -258,21 +250,25 @@ async def _handle_outputs(self, outputs: Array, msg: Dict[str, Any]):
                         outputs.append(doc)
             else:
                 outputs.append(
-                    {
-                        "data": {"text/plain": [content["data"].get("text/plain", "")]},
-                        "execution_count": content["execution_count"],
-                        "metadata": {},
-                        "output_type": msg_type,
-                    }
+                    Map(
+                        {
+                            "data": {"text/plain": [content["data"].get("text/plain", "")]},
+                            "execution_count": content["execution_count"],
+                            "metadata": {},
+                            "output_type": msg_type,
+                        }
+                    )
                 )
         elif msg_type == "error":
             outputs.append(
-                {
-                    "ename": content["ename"],
-                    "evalue": content["evalue"],
-                    "output_type": "error",
-                    "traceback": content["traceback"],
-                }
+                Map(
+                    {
+                        "ename": content["ename"],
+                        "evalue": content["evalue"],
+                        "output_type": "error",
+                        "traceback": content["traceback"],
+                    }
+                )
             )