Merge pull request #85 from cBioPortal/feat/clear-tools

feat(agent): clear_color_by and clear_render_controls tools

Author: hweej

packages/cell-explorer-agent/schema/app_config.schema.json

@@ -12,17 +12,38 @@
       "type": "string"
     },
     "colorBy": {
-      "type": "string",
-      "enum": [
-        "gene",
-        "category"
+      "anyOf": [
+        {
+          "type": "string",
+          "enum": [
+            "gene",
+            "category"
+          ]
+        },
+        {
+          "type": "null"
+        }
       ]
     },
     "gene": {
-      "type": "string"
+      "anyOf": [
+        {
+          "type": "string"
+        },
+        {
+          "type": "null"
+        }
+      ]
     },
     "category": {
-      "type": "string"
+      "anyOf": [
+        {
+          "type": "string"
+        },
+        {
+          "type": "null"
+        }
+      ]
     },
     "highlightedCategories": {
       "type": "array",
@@ -109,13 +130,27 @@
       "additionalProperties": false
     },
     "pointSize": {
-      "type": "number",
-      "exclusiveMinimum": 0
+      "anyOf": [
+        {
+          "type": "number",
+          "exclusiveMinimum": 0
+        },
+        {
+          "type": "null"
+        }
+      ]
     },
     "opacity": {
-      "type": "number",
-      "minimum": 0,
-      "maximum": 1
+      "anyOf": [
+        {
+          "type": "number",
+          "minimum": 0,
+          "maximum": 1
+        },
+        {
+          "type": "null"
+        }
+      ]
     },
     "summaryObsColumns": {
       "type": "array",

packages/cell-explorer-agent/src/cell_explorer_agent/prompt/system.py

@@ -65,6 +65,11 @@ def build_system_prompt(ctx: DatasetContext) -> str:
         "like 'in context', 'don't hide the rest', 'fade the others'), call "
         "set_selection_display_mode('dim') after the filter."
     )
+    lines.append(
+        "  - For 'stop coloring' / 'reset color' / 'undo my rendering changes' "
+        "use clear_color_by or clear_render_controls. These reset to defaults "
+        "without affecting filters."
+    )
     lines.append(
         "  - Do not pan/zoom (set_viewport) or change rendering (set_render_controls) "
         "unless the user explicitly asks. The user controls the camera and visual style; "

packages/cell-explorer-agent/src/cell_explorer_agent/tools/__init__.py

@@ -14,6 +14,7 @@
 from cell_explorer_agent.tools.data.schema import get_dataset_schema_tool
 from cell_explorer_agent.tools.registry import Tool, ToolCatalog, ToolKind
 from cell_explorer_agent.tools.ui_action.color import (
+    clear_color_by_tool,
     set_color_by_category_tool,
     set_color_by_gene_tool,
 )
@@ -26,7 +27,10 @@
     filter_by_gene_expression_tool,
 )
 from cell_explorer_agent.tools.ui_action.gene_label_column import set_gene_label_column_tool
-from cell_explorer_agent.tools.ui_action.render import set_render_controls_tool
+from cell_explorer_agent.tools.ui_action.render import (
+    clear_render_controls_tool,
+    set_render_controls_tool,
+)
 from cell_explorer_agent.tools.ui_action.selection_display_mode import (
     set_selection_display_mode_tool,
 )
@@ -56,6 +60,7 @@ def build_v1_catalog(z: ZarrAccess, *, config: AgentConfig) -> ToolCatalog:
     cat.register(set_embedding_tool(z))
     cat.register(set_color_by_gene_tool(z))
     cat.register(set_color_by_category_tool(z))
+    cat.register(clear_color_by_tool())
     cat.register(filter_by_ids_tool(z, filter_ids_max=config.filter_ids_max))
     cat.register(filter_by_gene_expression_tool(z))
     cat.register(set_selection_display_mode_tool())
@@ -70,6 +75,7 @@ def build_v1_catalog(z: ZarrAccess, *, config: AgentConfig) -> ToolCatalog:
         cat.register(set_summary_context_tool())
         cat.register(set_gene_label_column_tool(z))
         cat.register(set_render_controls_tool())
+        cat.register(clear_render_controls_tool())
 
     return cat
 

packages/cell-explorer-agent/src/cell_explorer_agent/tools/ui_action/color.py

@@ -30,6 +30,31 @@ async def run(gene: str) -> dict[str, Any]:
     )
 
 
+def clear_color_by_tool() -> Tool:
+    async def run() -> dict[str, Any]:
+        # `colorBy: null` is the frontend's reset sentinel: applyConfig calls
+        # clearGene + clearObsColumn, and the canvas falls back to default gray.
+        payload: dict[str, Any] = {"colorBy": None}
+        validate_partial(payload)
+        return {"payload": payload}
+
+    return Tool(
+        name="clear_color_by",
+        kind="ui_action",
+        description=(
+            "Reset the embedding coloring to the default (gray points). Use "
+            "when the user says 'stop coloring' / 'clear the color' / 'reset "
+            "color'. Does not affect filters or rendering controls."
+        ),
+        args_schema={
+            "type": "object",
+            "properties": {},
+            "additionalProperties": False,
+        },
+        func=run,
+    )
+
+
 def set_color_by_category_tool(z: ZarrAccess) -> Tool:
     async def run(
         category: str,

packages/cell-explorer-agent/src/cell_explorer_agent/tools/ui_action/render.py

@@ -6,6 +6,31 @@
 from cell_explorer_agent.tools.registry import Tool
 
 
+def clear_render_controls_tool() -> Tool:
+    async def run() -> dict[str, Any]:
+        # `null` is the frontend's reset sentinel: applyConfig restores the
+        # store defaults (0.5 / 0.5).
+        payload: dict[str, Any] = {"pointSize": None, "opacity": None}
+        validate_partial(payload)
+        return {"payload": payload}
+
+    return Tool(
+        name="clear_render_controls",
+        kind="ui_action",
+        description=(
+            "Reset point_size and opacity to the frontend defaults (0.5 each). "
+            "Use when the user says 'reset rendering' / 'undo my point size and "
+            "opacity changes' / 'go back to default'."
+        ),
+        args_schema={
+            "type": "object",
+            "properties": {},
+            "additionalProperties": False,
+        },
+        func=run,
+    )
+
+
 def set_render_controls_tool() -> Tool:
     async def run(
         point_size: float | None = None,

packages/cell-explorer-agent/tests/test_tool_catalog_v1.py

@@ -18,6 +18,7 @@ async def test_v1_catalog_includes_all_tools(fake_zarr):
         "set_embedding",
         "set_color_by_gene",
         "set_color_by_category",
+        "clear_color_by",
         "filter_by_ids",
         "filter_by_gene_expression",
         "set_selection_display_mode",
@@ -29,6 +30,7 @@ async def test_v1_catalog_includes_all_tools(fake_zarr):
         "set_summary_context",
         "set_gene_label_column",
         "set_render_controls",
+        "clear_render_controls",
     }
 
 

packages/cell-explorer-agent/tests/test_tools_ui_action.py

@@ -1,5 +1,6 @@
 from cell_explorer_agent.tools.ui_action.embedding import set_embedding_tool
 from cell_explorer_agent.tools.ui_action.color import (
+    clear_color_by_tool,
     set_color_by_gene_tool,
     set_color_by_category_tool,
 )
@@ -17,7 +18,10 @@
 from cell_explorer_agent.tools.ui_action.viewport import set_viewport_tool
 from cell_explorer_agent.tools.ui_action.summary_context import set_summary_context_tool
 from cell_explorer_agent.tools.ui_action.gene_label_column import set_gene_label_column_tool
-from cell_explorer_agent.tools.ui_action.render import set_render_controls_tool
+from cell_explorer_agent.tools.ui_action.render import (
+    clear_render_controls_tool,
+    set_render_controls_tool,
+)
 from cell_explorer_agent.tools.ui_action.selection_display_mode import (
     set_selection_display_mode_tool,
 )
@@ -291,3 +295,17 @@ async def test_set_selection_display_mode_invalid_value():
     tool = set_selection_display_mode_tool()
     r = await tool.func(value="other")
     assert "error" in r
+
+
+async def test_clear_color_by_emits_null():
+    tool = clear_color_by_tool()
+    r = await tool.func()
+    assert r == {"payload": {"colorBy": None}}
+    assert tool.kind == "ui_action"
+
+
+async def test_clear_render_controls_emits_nulls():
+    tool = clear_render_controls_tool()
+    r = await tool.func()
+    assert r == {"payload": {"pointSize": None, "opacity": None}}
+    assert tool.kind == "ui_action"