AI Assistant: restore Save to Macros in reloaded conversations, show saved macros, propose macros without running them

datalab/aiassistant/controller.py

@@ -93,6 +93,12 @@
     "  link so the user can opt in to persisting it. Always pick a short,\n"
     "  descriptive 'title' (it is shown in that link and used as the saved\n"
     "  macro's name).\n"
+    "- When the user asks you to write, create or give a macro (or a script),\n"
+    "  call 'create_and_run_macro' with 'autorun' set to false: the code is\n"
+    "  validated and offered through the 'Save to Macros' link, without being\n"
+    "  executed. Set 'autorun' to true (or omit it) only if the user also asks\n"
+    "  to run it. NEVER paste the code of a macro in your reply instead of\n"
+    "  calling the tool.\n"
     "- 'create_and_run_macro' returns the macro console output (stdout + stderr,\n"
     "  including Python tracebacks) and exit code. If exit_code != 0, READ the\n"
     "  output, fix the macro, and call the tool again.\n"

datalab/aiassistant/tools/builtin.py

@@ -589,6 +589,10 @@ def _tool_create_and_run_macro(
     }
     if not autorun:
         result["autorun"] = False
+        result["note"] = (
+            "Macro not executed. The user can save it to the Macro panel "
+            "with the 'Save to Macros' link."
+        )
         return result
 
     console = macropanel.console
@@ -1062,9 +1066,12 @@ def build_default_registry(expose_macro_tool: bool = True) -> ToolRegistry:
                     "the Macro panel; the user can opt in via a "
                     "'Save to Macros' link surfaced on the result, so always "
                     "pick a short, descriptive title (e.g. 'Generate "
-                    "super-Gaussian'). Use this for complex multi-step "
-                    "scripts; prefer 'apply_operation' for atomic actions. "
-                    "If 'ok' is False, read 'console_output' or "
+                    "super-Gaussian'). Set 'autorun' to false when the user "
+                    "asks you to write (not run) a macro: the code is then "
+                    "only validated and offered through the 'Save to Macros' "
+                    "link, without being executed. Use this for complex "
+                    "multi-step scripts; prefer 'apply_operation' for atomic "
+                    "actions. If 'ok' is False, read 'console_output' or "
                     "'validation_errors' and try again."
                 ),
                 parameters={
@@ -1075,7 +1082,16 @@ def build_default_registry(expose_macro_tool: bool = True) -> ToolRegistry:
                             "type": "string",
                             "description": "Full Python source code of the macro.",
                         },
-                        "autorun": {"type": "boolean", "default": True},
+                        "autorun": {
+                            "type": "boolean",
+                            "default": True,
+                            "description": (
+                                "Run the macro (default true). Set to false "
+                                "to only validate the code and offer it to "
+                                "the user for saving to the Macro panel, "
+                                "without executing it."
+                            ),
+                        },
                         "timeout": {
                             "type": "number",
                             "default": 30.0,

datalab/aiassistant/widgets/chatpanel.py

@@ -7,9 +7,11 @@
 from __future__ import annotations
 
 import html
+import json
 import os.path as osp
+import textwrap
 import uuid
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from guidata.configtools import get_icon
 from guidata.widgets.dockable import DockableWidgetMixin
@@ -27,7 +29,10 @@
 from datalab.aiassistant.providers import get_provider
 from datalab.aiassistant.providers.base import TokenUsage
 from datalab.aiassistant.tools.builtin import build_default_registry
-from datalab.aiassistant.widgets.markdown import markdown_to_html
+from datalab.aiassistant.widgets.markdown import (
+    is_python_code_block,
+    markdown_to_html,
+)
 from datalab.aiassistant.widgets.toolconfirmdialog import ToolConfirmDialog
 from datalab.aiassistant.worker import AIWorker
 from datalab.config import Conf, _, get_config_path
@@ -45,6 +50,10 @@
 _INPUT_HISTORY_FILENAME = "input_history.txt"
 _INPUT_HISTORY_MAX = 500
 _CONVERSATIONS_MAX = 200
+_MACRO_TOOL_NAME = "create_and_run_macro"
+_DEFAULT_MACRO_TITLE = "AI Assistant macro"
+_SAVE_MACRO_LINK = "save-macro"
+_SHOW_MACRO_PANEL_LINK = "show-macro-panel"
 
 
 class _GuiBridge(QC.QObject):
@@ -118,11 +127,15 @@ def __init__(
             max_size=_INPUT_HISTORY_MAX,
         )
         self._current_conversation: Conversation | None = None
-        # Transient macros generated by the LLM via 'create_and_run_macro'.
-        # They are NOT added to the Macro panel by default; the user can
-        # save them on demand by clicking the 'Save to Macros' link
-        # surfaced after each macro tool call. Keyed by short uuid.
+        # Macros proposed by the LLM, either through 'create_and_run_macro'
+        # or as Python code blocks in its replies. They are NOT added to the
+        # Macro panel by default; the user can save them on demand by
+        # clicking the 'Save to Macros' link surfaced next to them. Keyed by
+        # short uuid.
         self._pending_macros: dict[str, tuple[str, str]] = {}
+        # Last user request, used to name the macros saved from the code
+        # blocks of the reply.
+        self._last_user_text: str | None = None
         self._setup_ui()
 
     # ------------------------------------------------------------------ UI
@@ -274,10 +287,30 @@ def _append_user(self, text: str) -> None:
     def _append_assistant(self, text: str) -> None:
         if not text:
             return
+        footer = (
+            self._code_block_save_link if Conf.ai_expose_macro_tool.get(True) else None
+        )
         self._append_html(
             f"<div style='margin-top:6px;color:#1a5fb4;'>"
             f"<b>{_('Assistant')}:</b>"
-            f"{markdown_to_html(text)}</div>"
+            f"{markdown_to_html(text, code_block_footer=footer)}</div>"
+        )
+
+    def _code_block_save_link(self, lang: str, code: str) -> str:
+        """Return a "Save to Macros" link for a Python code block of a reply.
+
+        The macro is named after the user request the reply answers.
+        """
+        if not is_python_code_block(lang, code):
+            return ""
+        title = (
+            derive_title(self._last_user_text)
+            if self._last_user_text and self._last_user_text.strip()
+            else _DEFAULT_MACRO_TITLE
+        )
+        return (
+            "<p style='margin:2px 0 4px 0;font-size:90%;'>"
+            f"{self._save_macro_link_html(title, textwrap.dedent(code))}</p>"
         )
 
     def _append_tool(self, name: str, ok: bool, summary: str) -> None:
@@ -288,50 +321,79 @@ def _append_tool(self, name: str, ok: bool, summary: str) -> None:
             f"{symbol} {html.escape(name)}: {html.escape(summary)}</div>"
         )
 
+    def _append_tool_result(
+        self, name: str, ok: bool, data: Any, error: str | None
+    ) -> None:
+        """Append the outcome of a tool call (live or reloaded).
+
+        For a macro proposed through 'create_and_run_macro', a "Save to
+        Macros" link is added so the user can opt in to persisting it in the
+        Macro panel (mirrors the DataLab-Web behaviour).
+        """
+        if ok:
+            summary = str(data)[:200] if data is not None else "ok"
+        else:
+            summary = error or "error"
+        self._append_tool(name, ok, summary)
+        if name == _MACRO_TOOL_NAME and isinstance(data, dict):
+            macro_payload = data.get("_macro")
+            if isinstance(macro_payload, dict):
+                title = str(macro_payload.get("title") or _DEFAULT_MACRO_TITLE)
+                code = str(macro_payload.get("code") or "")
+                if code:
+                    self._append_save_macro_link(title, code)
+
     def _append_system(self, text: str) -> None:
         self._append_html(
             f"<div style='margin-top:4px;color:#666;font-style:italic;'>"
             f"{html.escape(text)}</div>"
         )
 
-    def _append_save_macro_link(self, title: str, code: str) -> None:
-        """Append an inline "Save to Macros" link for a transient macro.
+    def _save_macro_link_html(self, title: str, code: str) -> str:
+        """Return an inline "Save to Macros" link for a proposed macro.
 
         The link uses the internal ``datalab:save-macro/<id>`` scheme,
         intercepted by :meth:`_on_anchor_clicked`. The macro source is
         stashed in ``self._pending_macros`` keyed by ``id``.
         """
         macro_id = uuid.uuid4().hex[:12]
         self._pending_macros[macro_id] = (title, code)
-        url = f"datalab:save-macro/{macro_id}"
+        url = f"datalab:{_SAVE_MACRO_LINK}/{macro_id}"
         save_label = _("Save to Macros")
-        self._append_html(
-            "<div style='margin-top:2px;margin-left:18px;font-size:90%;'>"
+        return (
             f"💾 <a href='{html.escape(url, quote=True)}'>"
             f"{html.escape(save_label)}</a> — "
             f"<i>{html.escape(title)}</i>"
-            "</div>"
+        )
+
+    def _append_save_macro_link(self, title: str, code: str) -> None:
+        """Append an inline "Save to Macros" link for a proposed macro."""
+        self._append_html(
+            "<div style='margin-top:2px;margin-left:18px;font-size:90%;'>"
+            f"{self._save_macro_link_html(title, code)}</div>"
         )
 
     def _on_anchor_clicked(self, url: QC.QUrl) -> None:
         """Dispatch internal ``datalab:`` links; forward others to the OS."""
         if url.scheme() == "datalab":
             path = url.path().lstrip("/")
-            if url.host() == "save-macro" or path.startswith("save-macro/"):
+            prefix = f"{_SAVE_MACRO_LINK}/"
+            if url.host() == _SAVE_MACRO_LINK or path.startswith(prefix):
                 # QUrl.path() for ``datalab:save-macro/<id>`` puts the
                 # ``save-macro`` segment in ``host`` on some platforms
                 # and in the leading path segment on others — handle both.
-                macro_id = (
-                    path.split("/", 1)[1] if path.startswith("save-macro/") else path
-                )
+                macro_id = path[len(prefix) :] if path.startswith(prefix) else path
                 self._save_macro(macro_id)
                 return
+            if _SHOW_MACRO_PANEL_LINK in (url.host(), path):
+                self.mainwindow.set_current_panel("macro")
+                return
         # External link — open in default browser (mirrors the previous
         # ``setOpenExternalLinks(True)`` behaviour).
         QG.QDesktopServices.openUrl(url)
 
     def _save_macro(self, macro_id: str) -> None:
-        """Persist a pending transient macro to the Macro panel."""
+        """Persist a pending macro to the Macro panel."""
         entry = self._pending_macros.pop(macro_id, None)
         if entry is None:
             self._append_system(
@@ -340,14 +402,21 @@ def _save_macro(self, macro_id: str) -> None:
             return
         title, code = entry
         try:
-            self.mainwindow.macropanel.add_macro_with_code(title, code)
+            self.mainwindow.macropanel.add_macro_with_code(title, code, source="ai")
         # pylint: disable-next=broad-exception-caught
         except BaseException as exc:  # noqa: BLE001
             self._append_system(_("Failed to save macro: %s") % exc)
             # Re-stash so the user can retry.
             self._pending_macros[macro_id] = entry
             return
-        self._append_system(_("Saved macro '%s' to the Macro panel.") % title)
+        # The Macro panel is tabbed with this one: offer to show it rather
+        # than switching automatically, which would hide the conversation.
+        saved = html.escape(_("Saved macro '%s' to the Macro panel.") % title)
+        show = html.escape(_("Show in Macro panel"))
+        self._append_html(
+            f"<div style='margin-top:4px;color:#666;font-style:italic;'>"
+            f"{saved} <a href='datalab:{_SHOW_MACRO_PANEL_LINK}'>{show}</a></div>"
+        )
 
     def _set_busy(self, busy: bool) -> None:
         self.send_button.setVisible(not busy)
@@ -527,6 +596,7 @@ def _on_new_conversation(self) -> None:
         if self._controller is not None:
             self._controller.reset()
         self._current_conversation = None
+        self._last_user_text = None
         self.history_view.clear()
         self._update_usage_badge(None, None)
         self._append_system(_("Conversation reset."))
@@ -564,6 +634,7 @@ def _on_send(self) -> None:
             self._current_conversation = Conversation.new()
             self._current_conversation.title = derive_title(text)
         self.input_edit.clear()
+        self._last_user_text = text
         self._append_user(text)
         self._set_busy(True)
         self._worker = AIWorker(self._controller, text, self)
@@ -573,25 +644,9 @@ def _on_send(self) -> None:
 
     def _on_turn_finished(self, result: TurnResult) -> None:
         for name, _args, tool_result in result.tool_executions:
-            if tool_result.ok:
-                summary = (
-                    str(tool_result.data)[:200]
-                    if tool_result.data is not None
-                    else "ok"
-                )
-            else:
-                summary = tool_result.error or "error"
-            self._append_tool(name, tool_result.ok, summary)
-            # When a transient macro was generated, surface a
-            # "Save to Macros" link so the user can opt in to persisting
-            # it in the Macro panel (mirrors the DataLab-Web behaviour).
-            if name == "create_and_run_macro" and isinstance(tool_result.data, dict):
-                macro_payload = tool_result.data.get("_macro")
-                if isinstance(macro_payload, dict):
-                    title = str(macro_payload.get("title") or "AI Assistant macro")
-                    code = str(macro_payload.get("code") or "")
-                    if code:
-                        self._append_save_macro_link(title, code)
+            self._append_tool_result(
+                name, tool_result.ok, tool_result.data, tool_result.error
+            )
         if result.cancelled:
             self._append_system(_("Tool execution cancelled by user."))
         if result.aborted:
@@ -647,9 +702,12 @@ def _load_conversation(self, conv_id: str) -> None:
 
     def _render_messages(self, messages: list[ChatMessage]) -> None:
         """Re-render persisted messages in the chat view."""
+        self._last_user_text = None
         for msg in messages:
             if msg.role == "user":
-                self._append_user(self._content_to_text(msg.content))
+                text = self._content_to_text(msg.content)
+                self._last_user_text = text
+                self._append_user(text)
             elif msg.role == "assistant":
                 if msg.content:
                     self._append_assistant(self._content_to_text(msg.content))
@@ -659,9 +717,27 @@ def _render_messages(self, messages: list[ChatMessage]) -> None:
                         % {"name": call.name, "args": call.arguments}
                     )
             elif msg.role == "tool":
-                content = self._content_to_text(msg.content)
-                summary = content[:200] if content else ""
-                self._append_tool(msg.name or "tool", True, summary)
+                ok, data, error = self._parse_tool_content(
+                    self._content_to_text(msg.content)
+                )
+                self._append_tool_result(msg.name or "tool", ok, data, error)
+
+    @staticmethod
+    def _parse_tool_content(content: str) -> tuple[bool, Any, str | None]:
+        """Decode a persisted tool message into ``(ok, data, error)``.
+
+        Tool messages hold the JSON produced by
+        :meth:`ToolResult.to_message_content`; anything else is shown as is.
+        """
+        try:
+            payload = json.loads(content)
+        except (TypeError, ValueError):
+            payload = None
+        if isinstance(payload, dict) and isinstance(payload.get("ok"), bool):
+            if payload["ok"]:
+                return True, payload.get("result"), None
+            return False, None, str(payload.get("error") or "error")
+        return True, content or None, None
 
     @staticmethod
     def _content_to_text(content) -> str: