feat(cli): introduce visualize tool for rich data display

- Implement `visualize` tool in core to support tables, charts, and diffs. - Add `RichDataDisplay` UI component using Ink for rendering visualizations. - Integrate visualization support into `ToolResultDisplay`. - Update system prompts to encourage the use of the `visualize` tool for structured data. - Add `info` semantic color to themes. - Fix shell parser initialization in tests.
2026-05-13 21:32:56 -07:00 · 2026-02-10 19:53:38 -05:00
parent 4a48d7cf93
commit 545a23f120
18 changed files with 1133 additions and 20 deletions
@@ -0,0 +1,204 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
+import type { ToolInvocation, ToolResult } from './tools.js';
+import type { MessageBus } from '../confirmation-bus/message-bus.js';
+import * as fs from 'node:fs/promises';
+import * as path from 'node:path';
+import { ToolErrorType } from './tool-error.js';
+
+interface VisualizeParams {
+  data: unknown;
+  type?: 'table' | 'bar_chart' | 'pie_chart' | 'line_chart' | 'diff';
+  title?: string;
+  save_as?: string;
+  open?: boolean;
+}
+
+class VisualizeInvocation extends BaseToolInvocation<
+  VisualizeParams,
+  ToolResult
+> {
+  getDescription(): string {
+    const type = this.params.type ?? 'table';
+    const action = this.params.save_as
+      ? `and saving to ${this.params.save_as}`
+      : '';
+    return `Visualizing data as ${type} ${action}`;
+  }
+
+  async execute(
+    _signal: AbortSignal,
+    _updateOutput?: (output: string) => void,
+  ): Promise<ToolResult> {
+    const { data, type = 'table', title, save_as } = this.params;
+
+    if (
+      type === 'table' ||
+      type === 'bar_chart' ||
+      type === 'pie_chart' ||
+      type === 'line_chart'
+    ) {
+      if (!Array.isArray(data)) {
+        return {
+          llmContent:
+            'Error: data must be an array of objects for this visualization type.',
+          returnDisplay:
+            'Error: data must be an array of objects for this visualization type.',
+          error: {
+            message: 'Data must be an array',
+            type: ToolErrorType.INVALID_TOOL_PARAMS,
+          },
+        };
+      }
+    } else if (type === 'diff') {
+      if (typeof data !== 'object' || !data) {
+        return {
+          llmContent: 'Error: data must be an object for diff visualization.',
+          returnDisplay:
+            'Error: data must be an object for diff visualization.',
+          error: {
+            message: 'Data must be an object',
+            type: ToolErrorType.INVALID_TOOL_PARAMS,
+          },
+        };
+      }
+    }
+
+    let savedFilePath: string | undefined;
+
+    if (save_as) {
+      try {
+        const absolutePath = path.resolve(save_as);
+        let content = '';
+        if (save_as.endsWith('.json')) {
+          content = JSON.stringify(data, null, 2);
+        } else if (save_as.endsWith('.csv') && Array.isArray(data)) {
+          // Basic CSV conversion
+          if (data.length > 0 && typeof data[0] === 'object') {
+            const headers = Object.keys(data[0] as object).join(',');
+            const rows = data
+              .map((row) =>
+                Object.values(row as object)
+                  .map((v) => {
+                    const s = String(v);
+                    // Quote if contains comma or newline
+                    if (
+                      s.includes(',') ||
+                      s.includes('\n') ||
+                      s.includes('"')
+                    ) {
+                      return `"${s.replace(/"/g, '""')}"`;
+                    }
+                    return s;
+                  })
+                  .join(','),
+              )
+              .join('\n');
+            content = `${headers}\n${rows}`;
+          } else {
+            content = '';
+          }
+        } else {
+          content = JSON.stringify(data, null, 2);
+        }
+
+        await fs.writeFile(absolutePath, content);
+        savedFilePath = absolutePath;
+      } catch (e) {
+        return {
+          llmContent: `Error saving file: ${e}`,
+          returnDisplay: `Error saving file: ${e}`,
+          error: {
+            message: `Error saving file: ${e}`,
+            type: ToolErrorType.EXECUTION_FAILED,
+          },
+        };
+      }
+    }
+
+    // Infer columns for table
+    let columns;
+    if (
+      type === 'table' &&
+      Array.isArray(data) &&
+      data.length > 0 &&
+      typeof data[0] === 'object'
+    ) {
+      columns = Object.keys(data[0] as object).map((key) => ({
+        key,
+        label: key,
+      }));
+    }
+
+    return {
+      llmContent:
+        'Visualization rendered in CLI.' +
+        (savedFilePath ? ` Saved to ${savedFilePath}` : ''),
+      returnDisplay: {
+        type,
+        title,
+        data,
+        columns,
+        savedFilePath,
+      },
+    };
+  }
+}
+
+export class VisualizeTool extends BaseDeclarativeTool<
+  VisualizeParams,
+  ToolResult
+> {
+  constructor(messageBus: MessageBus) {
+    super(
+      'visualize',
+      'Visualize Data',
+      'Renders structured data as tables, charts, or diffs, and optionally saves it to a file.',
+      Kind.Other,
+      {
+        type: 'object',
+        properties: {
+          data: {
+            description:
+              'The structured data to visualize. Array of objects for tables/charts. For diffs, provide an object with {fileDiff: string}, {old: string, new: string}, or {oldContent: string, newContent: string}. Can also be a string containing a unified diff.',
+          },
+          type: {
+            type: 'string',
+            enum: ['table', 'bar_chart', 'pie_chart', 'line_chart', 'diff'],
+            description: 'The visualization type. Default: table.',
+          },
+          title: {
+            type: 'string',
+            description: 'A title for the visualization.',
+          },
+          save_as: {
+            type: 'string',
+            description:
+              'File path to save the data (e.g. data.csv, data.json).',
+          },
+        },
+        required: ['data'],
+      },
+      messageBus,
+    );
+  }
+
+  protected createInvocation(
+    params: VisualizeParams,
+    messageBus: MessageBus,
+    toolName?: string,
+    toolDisplayName?: string,
+  ): ToolInvocation<VisualizeParams, ToolResult> {
+    return new VisualizeInvocation(
+      params,
+      messageBus,
+      toolName ?? this.name,
+      toolDisplayName ?? this.displayName,
+    );
+  }
+}