@gram-ai/elements 1.18.4 → 1.18.6

package/src/lib/tools.ts

@@ -0,0 +1,210 @@
+import { JSONSchema7, ToolSet, type ToolCallOptions } from 'ai'
+import {
+  AssistantToolProps,
+  Tool,
+  makeAssistantTool,
+} from '@assistant-ui/react'
+import z from 'zod'
+import { FC } from 'react'
+
+/**
+ * Converts from assistant-ui tool format to the AI SDK tool shape
+ */
+export const toAISDKTools = (tools: Record<string, Tool>) => {
+  return Object.fromEntries(
+    Object.entries(tools).map(([name, tool]) => [
+      name,
+      {
+        ...(tool.description ? { description: tool.description } : undefined),
+        parameters: (tool.parameters instanceof z.ZodType
+          ? z.toJSONSchema(tool.parameters)
+          : tool.parameters) as JSONSchema7,
+      },
+    ])
+  )
+}
+
+/**
+ * Returns only frontend tools that are enabled
+ */
+export const getEnabledTools = (tools: Record<string, Tool>) => {
+  return Object.fromEntries(
+    Object.entries(tools).filter(
+      ([, tool]) => !tool.disabled && tool.type !== 'backend'
+    )
+  )
+}
+
+/**
+ * A frontend tool is a tool that is defined by the user and can be used in the chat.
+ */
+export type FrontendTool<TArgs extends Record<string, unknown>, TResult> = FC<
+  AssistantToolProps<TArgs, TResult>
+> & {
+  unstable_tool: AssistantToolProps<TArgs, TResult>
+}
+
+/**
+ * Module-level approval config that gets set by ElementsProvider at runtime.
+ * This allows defineFrontendTool to check approval status during execute.
+ */
+let approvalConfig: {
+  helpers: ApprovalHelpers
+  toolsRequiringApproval: Set<string>
+} | null = null
+
+/**
+ * Sets the approval configuration. Called by ElementsProvider.
+ */
+export function setFrontendToolApprovalConfig(
+  helpers: ApprovalHelpers,
+  toolsRequiringApproval: string[]
+): void {
+  approvalConfig = {
+    helpers,
+    toolsRequiringApproval: new Set(toolsRequiringApproval),
+  }
+}
+
+/**
+ * Clears the approval configuration. Called when ElementsProvider unmounts.
+ */
+export function clearFrontendToolApprovalConfig(): void {
+  approvalConfig = null
+}
+
+/**
+ * Make a frontend tool
+ */
+export const defineFrontendTool = <
+  TArgs extends Record<string, unknown>,
+  TResult,
+>(
+  tool: Tool,
+  name: string
+): FrontendTool<TArgs, TResult> => {
+  type ToolExecutionContext = Parameters<
+    NonNullable<Tool<Record<string, unknown>, void>['execute']>
+  >[1]
+  return makeAssistantTool({
+    ...tool,
+    execute: async (args: TArgs, context: ToolExecutionContext) => {
+      // Check if this tool requires approval at runtime
+      if (approvalConfig?.toolsRequiringApproval.has(name)) {
+        const { helpers } = approvalConfig
+        const toolCallId = context.toolCallId ?? ''
+
+        // Check if already approved (user chose "Approve always" previously)
+        if (!helpers.isToolApproved(name)) {
+          const approved = await helpers.requestApproval(name, toolCallId, args)
+
+          if (!approved) {
+            return {
+              content: [
+                {
+                  type: 'text',
+                  text: `Tool "${name}" execution was denied by the user. Please acknowledge this and continue without using this tool's result.`,
+                },
+              ],
+              isError: true,
+            } as TResult
+          }
+        }
+      }
+
+      return tool.execute?.(args, context)
+    },
+    toolName: name,
+  } as AssistantToolProps<TArgs, TResult>)
+}
+
+/**
+ * Helpers for requesting and tracking tool approval state.
+ */
+export interface ApprovalHelpers {
+  requestApproval: (
+    toolName: string,
+    toolCallId: string,
+    args: unknown
+  ) => Promise<boolean>
+  isToolApproved: (toolName: string) => boolean
+  whitelistTool: (toolName: string) => void
+}
+
+/**
+ * Wraps tools with approval logic based on the approval config.
+ */
+export function wrapToolsWithApproval(
+  tools: ToolSet,
+  toolsRequiringApproval: string[] | undefined,
+  approvalHelpers: ApprovalHelpers
+): ToolSet {
+  if (!toolsRequiringApproval || toolsRequiringApproval.length === 0) {
+    return tools
+  }
+
+  const approvalSet = new Set(toolsRequiringApproval)
+
+  return Object.fromEntries(
+    Object.entries(tools).map(([name, tool]) => {
+      if (!approvalSet.has(name)) {
+        return [name, tool]
+      }
+
+      const originalExecute = tool.execute
+      if (!originalExecute) {
+        return [name, tool]
+      }
+
+      return [
+        name,
+        {
+          ...tool,
+          execute: async (args: unknown, options?: ToolCallOptions) => {
+            const opts = (options ?? {}) as Parameters<
+              typeof originalExecute
+            >[1]
+            // Extract toolCallId from options
+            const toolCallId =
+              (opts as { toolCallId?: string }).toolCallId ?? ''
+
+            // Check if already approved (user chose "Approve always" previously)
+            if (approvalHelpers.isToolApproved(name)) {
+              return originalExecute(
+                args,
+                opts as Parameters<typeof originalExecute>[1]
+              )
+            }
+
+            // Request approval using the actual toolCallId from the stream
+            const approved = await approvalHelpers.requestApproval(
+              name,
+              toolCallId,
+              args
+            )
+
+            if (!approved) {
+              return {
+                content: [
+                  {
+                    type: 'text',
+                    text: `Tool "${name}" execution was denied by the user. Please acknowledge this and continue without using this tool's result.`,
+                  },
+                ],
+                isError: true,
+              }
+            }
+
+            // Note: Tool is marked as approved via the UI when user clicks "Approve always"
+            // (handled in tool-fallback.tsx via markToolApproved)
+
+            return originalExecute(
+              args,
+              opts as Parameters<typeof originalExecute>[1]
+            )
+          },
+        },
+      ]
+    })
+  ) as ToolSet
+}

package/src/lib/utils.ts

@@ -0,0 +1,16 @@
+import { clsx, type ClassValue } from 'clsx'
+import { twMerge } from 'tailwind-merge'
+
+export function cn(...inputs: ClassValue[]) {
+  return twMerge(clsx(inputs))
+}
+
+export function assertNever(value: unknown): never {
+  throw new Error(`Unexpected value: ${value}`)
+}
+
+export function assert(condition: unknown, message: string): asserts condition {
+  if (!condition) {
+    throw new Error(message)
+  }
+}

package/src/plugins/README.md

@@ -0,0 +1,49 @@
+# Plugin Development Guide
+
+This guide will help you create custom plugins for the Gram Elements library.
+
+## What are Plugins?
+
+Plugins enable you to add custom rendering capabilities to the Gram Elements library. They allow you to transform markdown code blocks with specific language identifiers into rich, interactive components.
+
+The typical plugin workflow is:
+
+1. **Extend System Prompt**: The plugin adds instructions to the system prompt, telling the LLM how to return data in a specific format
+2. **LLM Responds**: The LLM returns a code fence marked with your plugin's language identifier
+3. **Custom Rendering**: Your plugin's custom component renders the code block content
+
+## Plugin Interface
+
+A plugin is defined by the following TypeScript interface:
+
+```typescript
+interface Plugin {
+  // The language identifier for the code fence (e.g., "vega", "mermaid", "d3")
+  language: string
+
+  // Instructions for the LLM on how to use this plugin
+  prompt: string
+
+  // Your custom React component that renders the code block
+  SyntaxHighlighter: ComponentType<SyntaxHighlighterProps>
+
+  // Optional: Custom header component for the code block
+  CodeHeader?: ComponentType<CodeHeaderProps> | null
+
+  // Optional: Whether to override existing plugins with the same language
+  overrideExisting?: boolean
+}
+```
+
+## Support
+
+If you need help creating a plugin:
+
+- Check existing plugins for examples
+- Review the TypeScript types in `src/types/plugins.ts`
+- Open an issue on GitHub
+- Join our community Discord
+
+---
+
+Happy plugin building! 🚀

package/src/plugins/chart/component.tsx

@@ -0,0 +1,102 @@
+'use client'
+
+import { useDensity } from '@/hooks/useDensity'
+import { useRadius } from '@/hooks/useRadius'
+import { cn } from '@/lib/utils'
+import { useAssistantState } from '@assistant-ui/react'
+import { SyntaxHighlighterProps } from '@assistant-ui/react-markdown'
+import { AlertCircleIcon } from 'lucide-react'
+import { FC, useEffect, useMemo, useRef, useState } from 'react'
+import { parse, View, Warn } from 'vega'
+
+export const ChartRenderer: FC<SyntaxHighlighterProps> = ({ code }) => {
+  const message = useAssistantState(({ message }) => message)
+  const containerRef = useRef<HTMLDivElement>(null)
+  const viewRef = useRef<View | null>(null)
+  const [error, setError] = useState<string | null>(null)
+  const messageIsComplete = message.status?.type === 'complete'
+  const r = useRadius()
+  const d = useDensity()
+
+  // Parse and validate JSON in useMemo - only recomputes when code changes
+  const parsedSpec = useMemo(() => {
+    const trimmedCode = code.trim()
+    if (!trimmedCode) return null
+
+    try {
+      return JSON.parse(trimmedCode) as Record<string, unknown>
+    } catch {
+      return null
+    }
+  }, [code])
+
+  // Only render when we have valid JSON AND message is complete
+  const shouldRender = messageIsComplete && parsedSpec !== null
+
+  useEffect(() => {
+    if (!containerRef.current || !shouldRender) {
+      return
+    }
+
+    setError(null)
+
+    const runChart = async () => {
+      try {
+        // Clean up any existing view
+        if (viewRef.current) {
+          viewRef.current.finalize()
+          viewRef.current = null
+        }
+
+        const chart = parse(parsedSpec)
+        const view = new View(chart, {
+          container: containerRef.current ?? undefined,
+          renderer: 'svg',
+          hover: true,
+          logLevel: Warn,
+        })
+        viewRef.current = view
+
+        await view.runAsync()
+      } catch (err) {
+        console.error('Failed to render chart:', err)
+        setError(err instanceof Error ? err.message : 'Failed to render chart')
+      }
+    }
+
+    runChart()
+
+    return () => {
+      if (viewRef.current) {
+        viewRef.current.finalize()
+        viewRef.current = null
+      }
+    }
+  }, [shouldRender, parsedSpec])
+
+  return (
+    <div
+      className={cn(
+        // the after:hidden is to prevent assistant-ui from showing its default code block loading indicator
+        'relative flex min-h-[400px] w-fit max-w-full min-w-[400px] items-center justify-center border p-6 after:hidden',
+        r('lg'),
+        d('p-lg')
+      )}
+    >
+      {!shouldRender && !error && (
+        <div className="shimmer text-muted-foreground bg-background/80 absolute inset-0 z-10 flex items-center justify-center">
+          Rendering chart...
+        </div>
+      )}
+
+      {error && (
+        <div className="bg-background absolute inset-0 z-10 flex items-center justify-center gap-2 text-rose-500">
+          <AlertCircleIcon name="alert-circle" className="h-4 w-4" />
+          {error}
+        </div>
+      )}
+
+      <div ref={containerRef} className={!shouldRender ? 'hidden' : 'block'} />
+    </div>
+  )
+}

package/src/plugins/chart/index.ts

@@ -0,0 +1,27 @@
+import { Plugin } from '@/types/plugins'
+import { ChartRenderer } from './component'
+
+/**
+ * This plugin renders Vega charts.
+ */
+export const chart: Plugin = {
+  language: 'vega',
+  prompt: `When a user requests a chart or visualization, respond with a valid Vega specification (https://vega.github.io/vega/) in a code block annotated with the language identifier 'vega'.
+
+CRITICAL JSON REQUIREMENTS:
+- The code block MUST contain ONLY valid, parseable JSON
+- NO comments (no // or /* */ anywhere)
+- NO trailing commas
+- Use double quotes for all strings and keys
+- NO text before or after the JSON object
+- The JSON must start with { and end with }
+
+CONTENT GUIDELINES:
+- Outside the code block, describe trends and insights found in the data
+- Do not describe visual properties or technical implementation details
+- Do not mention "Vega" or other technical terms - this is user-facing
+
+The Vega spec will be parsed with JSON.parse() - if it fails, the chart will error. Ensure strict JSON validity.`,
+  Component: ChartRenderer,
+  Header: undefined,
+}

package/src/plugins/index.ts

@@ -0,0 +1,7 @@
+import type { Plugin } from '@/types/plugins'
+import { chart } from './chart'
+
+export const recommended: Plugin[] = [chart]
+export { chart } from './chart'
+
+export type { Plugin } from '@/types/plugins'

package/src/server.ts

@@ -0,0 +1,89 @@
+import { IncomingMessage, ServerResponse } from 'node:http'
+
+type ServerHandler<T> = (
+  req: IncomingMessage,
+  res: ServerResponse,
+  options?: T
+) => Promise<void>
+
+interface ServerHandlers {
+  /**
+   * Handler to create a new chat session token.
+   *
+   * @example
+   * ```typescript
+   * import { createElementsServerHandlers } from '@gram-ai/elements/server'
+   * import express from 'express'
+   * const app = express()
+   * const handlers = createElementsServerHandlers()
+   * app.post('/chat/session', handlers.session)
+   * app.listen(3000)
+   * ```
+   */
+  session: ServerHandler<SessionHandlerOptions>
+}
+
+export const createElementsServerHandlers = (): ServerHandlers => {
+  return {
+    session: sessionHandler,
+  }
+}
+
+interface SessionHandlerOptions {
+  /**
+   * The origin from which the token will be used
+   */
+  embedOrigin: string
+
+  /**
+   * Free-form user identifier
+   */
+  userIdentifier: string
+
+  /**
+   * Token expiration in seconds (max / default 3600)
+   * @default 3600
+   */
+  expiresAfter?: number
+}
+
+const sessionHandler: ServerHandler<SessionHandlerOptions> = async (
+  req,
+  res,
+  options
+) => {
+  const base = process.env.GRAM_API_URL ?? 'https://app.getgram.ai'
+  if (req.method === 'POST') {
+    const projectSlug = Array.isArray(req.headers['gram-project'])
+      ? req.headers['gram-project'][0]
+      : req.headers['gram-project']
+
+    fetch(base + '/rpc/chatSessions.create', {
+      method: 'POST',
+      body: JSON.stringify({
+        embed_origin: options?.embedOrigin,
+        user_identifier: options?.userIdentifier,
+        expires_after: options?.expiresAfter,
+      }),
+      headers: {
+        'Content-Type': 'application/json',
+        'Gram-Project': typeof projectSlug === 'string' ? projectSlug : '',
+        'Gram-Key': process.env.GRAM_API_KEY ?? '',
+      },
+    })
+      .then(async (response) => {
+        const body = await response.text()
+        res.writeHead(response.status, { 'Content-Type': 'application/json' })
+        res.end(body)
+      })
+      .catch((error) => {
+        console.error('Failed to create chat session:', error)
+        res.writeHead(500, { 'Content-Type': 'application/json' })
+        res.end(
+          JSON.stringify({
+            error: 'Failed to create chat session: ' + error.message,
+          })
+        )
+      })
+  }
+}