feat: create ToolGuidelines for better Agent orientation of when and how to use marimo tools, add guidelines to GetMarimoRules tool

marimo/_ai/_tools/base.py

@@ -17,6 +17,7 @@
 )
 
 from marimo import _loggers
+from marimo._ai._tools.types import ToolGuidelines
 from marimo._ai._tools.utils.exceptions import ToolExecutionError
 from marimo._config.config import CopilotMode
 from marimo._server.ai.tools.types import (
@@ -95,6 +96,7 @@ class ToolBase(Generic[ArgsT, OutT], ABC):
     # Override in subclass, or rely on fallbacks below
     name: str = ""
     description: str = ""
+    guidelines: Optional[ToolGuidelines] = None
     Args: type[ArgsT]
     Output: type[OutT]
     context: ToolContext
@@ -127,7 +129,15 @@ def __init__(self, context: ToolContext) -> None:
 
         # get description from class docstring
         if self.description == "":
-            self.description = (self.__class__.__doc__ or "").strip()
+            base_description = (self.__class__.__doc__ or "").strip()
+
+            # If guidelines exist, append them
+            if self.guidelines is not None:
+                self.description = self._format_with_guidelines(
+                    base_description, self.guidelines
+                )
+            else:
+                self.description = base_description
 
     async def __call__(self, args: ArgsT) -> OutT:
         """
@@ -241,6 +251,34 @@ def _coerce_args(self, args: Any) -> ArgsT:  # type: ignore[override]
             return args  # type: ignore[return-value]
         return parse_raw(args, self.Args)
 
+    def _format_with_guidelines(
+        self, description: str, guidelines: ToolGuidelines
+    ) -> str:
+        """Combine description with structured guidelines."""
+        parts = [description] if description else []
+
+        if guidelines.when_to_use:
+            parts.append("\n## When to use:")
+            parts.extend(f"- {item}" for item in guidelines.when_to_use)
+
+        if guidelines.avoid_if:
+            parts.append("\n## Avoid if:")
+            parts.extend(f"- {item}" for item in guidelines.avoid_if)
+
+        if guidelines.prerequisites:
+            parts.append("\n## Prerequisites:")
+            parts.extend(f"- {item}" for item in guidelines.prerequisites)
+
+        if guidelines.side_effects:
+            parts.append("\n## Side effects:")
+            parts.extend(f"- {item}" for item in guidelines.side_effects)
+
+        if guidelines.additional_info:
+            parts.append("\n## Additional info:")
+            parts.append(guidelines.additional_info)
+
+        return "\n".join(parts)
+
     # error defaults/hooks
     def _default_error_code(self) -> str:
         return "UNEXPECTED_ERROR"

marimo/_ai/_tools/tools/rules.py

@@ -7,7 +7,7 @@
 import marimo._utils.requests as requests
 from marimo import _loggers
 from marimo._ai._tools.base import ToolBase
-from marimo._ai._tools.types import EmptyArgs, SuccessResult
+from marimo._ai._tools.types import EmptyArgs, SuccessResult, ToolGuidelines
 
 LOGGER = _loggers.marimo_logger()
 
@@ -29,6 +29,15 @@ class GetMarimoRules(ToolBase[EmptyArgs, GetMarimoRulesOutput]):
         The content of the rules file.
     """
 
+    guidelines = ToolGuidelines(
+        when_to_use=[
+            "Before using other marimo mcp tools, reading a marimo notebook, or writing to a notebook ALWAYS use this first to understand how marimo works",
+        ],
+        avoid_if=[
+            "The rules have already been retrieved recently, as they rarely change",
+        ],
+    )
+
     def handle(self, args: EmptyArgs) -> GetMarimoRulesOutput:
         del args
 

marimo/_ai/_tools/types.py

@@ -21,3 +21,14 @@ class SuccessResult:
 @dataclass
 class EmptyArgs:
     pass
+
+
+@dataclass
+class ToolGuidelines:
+    """Structured guidance for AI assistants on when and how to use a tool."""
+
+    when_to_use: Optional[list[str]] = None
+    avoid_if: Optional[list[str]] = None
+    prerequisites: Optional[list[str]] = None
+    side_effects: Optional[list[str]] = None
+    additional_info: Optional[str] = None