Momento/memento-note/lib/ai/tools/chart.tool.ts

/**
 * Chart Tool for Notes
 * Allows AI to generate inline charts from note data
 */

import { tool } from 'ai'
import { z } from 'zod'
import { toolRegistry } from './registry'

// Chart generation tool - inserts chart as markdown code block
toolRegistry.register({
  name: 'generate_chart',
  description: 'Generate an inline chart from data and insert it into a note as markdown. The chart will be rendered directly in the note.',
  isInternal: true,
  buildTool: (ctx) =>
    tool({
      description: `Generate an inline chart from data and return it as a markdown code block that can be inserted into a note.

Available chart types:
- "bar": Vertical bar chart
- "horizontal-bar": Horizontal bar chart
- "line": Line chart for trends over time
- "area": Area chart (filled line)
- "pie": Pie/donut chart for proportions
- "radar": Radar chart for comparing multiple dimensions
- "funnel": Funnel chart for stages/conversions
- "gauge": Gauge/meter for a single value

Data format: Array of objects with "label" (string) and "value" (number).

IMPORTANT:
- Extract data from the note content when possible
- Keep labels short (max 20 characters)
- Round values to reasonable precision (max 2 decimal places)
- For time series, use short date formats (Jan, Feb, Mar or 2023, 2024)
- Maximum 12 data points for readability
- The output will be rendered as a visual chart in the note`,
      inputSchema: z.object({
        chartType: z.enum(['bar', 'horizontal-bar', 'line', 'area', 'pie', 'radar', 'funnel', 'gauge']).describe('Type of chart to generate'),
        title: z.string().optional().describe('Optional title for the chart'),
        data: z.array(z.object({
          label: z.string().describe('Label for the data point (e.g., month name, category)'),
          value: z.number().describe('Numeric value for the data point'),
        })).describe('Array of data points with label and value'),
        insertLocation: z.enum(['append', 'prepend', 'replace']).default('append').describe('Where to insert the chart in the note'),
        targetNoteId: z.string().optional().describe('Optional: specific note ID to update. If not provided, use the current note.'),
      }),
      execute: async ({ chartType, title, data, insertLocation, targetNoteId }) => {
        try {
          // Generate the markdown code block for the chart
          const chartData = data.map(d => `${d.label}: ${d.value}`).join('\n')
          const chartMarkdown = `\`\`\`chart
${chartType}${title ? `\n${title}` : ''}
${chartData}
\`\`\`\n`

          return {
            success: true,
            chartMarkdown,
            chartType,
            dataPointCount: data.length,
            message: `Chart generated with ${data.length} data points. Insert this markdown into the note.`,
          }
        } catch (e: any) {
          return { success: false, error: `Chart generation failed: ${e.message}` }
        }
      },
    }),
})

// Quick chart insert tool - directly updates a note with a chart
toolRegistry.register({
  name: 'insert_chart_in_note',
  description: 'Insert a chart directly into a note. Reads the note, extracts relevant data, generates a chart, and updates the note content.',
  isInternal: true,
  buildTool: (ctx) =>
    tool({
      description: `Insert a chart directly into a note. This tool will:
1. Read the current note content
2. Extract relevant data from the note (sales, metrics, comparisons, etc.)
3. Generate an appropriate chart type based on the data
4. Insert the chart markdown into the note

Use this when the user says "make a chart", "create a graph", "visualize this data", etc.

Choose the chart type based on the data:
- Use "bar" for comparing values across categories (default)
- Use "horizontal-bar" when labels are long
- Use "line" or "area" for time series or trends
- Use "pie" for showing proportions/percentages
- Use "radar" for comparing multiple attributes
- Use "funnel" for stages or conversion data
- Use "gauge" for progress/KPI (single value)`,
      inputSchema: z.object({
        noteId: z.string().describe('The note ID to read and update'),
        chartHint: z.string().optional().describe('Optional hint about what to chart (e.g., "sales by month", "comparison of products")'),
        insertLocation: z.enum(['append', 'prepend', 'before-section', 'after-section']).default('append').describe('Where to insert the chart'),
        sectionMarker: z.string().optional().describe('For before-section/after-section: the heading text to find (e.g., "## Sales Data")'),
      }),
      execute: async ({ noteId, chartHint, insertLocation, sectionMarker }) => {
        try {
          // We'll return the instructions for the AI to format the chart
          // The actual note update will be done by note_find_and_update or note_update
          return {
            success: true,
            instructions: 'Generate a chart markdown block using the generate_chart tool, then use note_update or note_find_and_update to insert it into the note.',
            noteId,
            chartHint,
            insertLocation,
            sectionMarker,
          }
        } catch (e: any) {
          return { success: false, error: `Failed to prepare chart insertion: ${e.message}` }
        }
      },
    }),
})