Update Gradio_UI.py

Gradio_UI.py

@@ -123,6 +123,8 @@ def pull_messages_from_step(step_log: MemoryStep):
             content = str(args.get("answer", str(args))) if isinstance(args, dict) else str(args).strip()
 
             if used_code:
-                content = re.sub(r"```.*?\n", "", content)
-                content = re.sub(r"\s*<end_code>\s*", "", content).strip()
-                if not content.startswith("```python"):
+    content = re.sub(r"```.*?\n", "", content)
+    content = re.sub(r"\s*<end_code>\s*", "", content).strip()
+    if not content.startswith("```python"):
+        content = f"```python\n{content}\n```"
+