Backport PR jupyterlab#1155: Fix code output format in IPython (jupyt…

…erlab#1162)

Co-authored-by: Divyansh Choudhary <divyanshchoudhary99@gmail.com>

packages/jupyter-ai-magics/jupyter_ai_magics/magics.py

@@ -482,13 +482,7 @@ def display_output(self, output, display_format, md):
             lang_indicator = r"^```[a-zA-Z0-9]*\n"
             output = re.sub(lang_indicator, "", output)
             output = re.sub(r"\n```$", "", output)
-            new_cell_payload = dict(
-                source="set_next_input",
-                text=output,
-                replace=False,
-            )
-            ip = self.shell
-            ip.payload_manager.write_payload(new_cell_payload)
+            self.shell.set_next_input(output, replace=False)
             return HTML(
                 "AI generated code inserted below ⬇️", metadata=md
             )