ClickHouse chat agent

import { prompts } from "@trigger.dev/sdk";
import { chat } from "@trigger.dev/sdk/ai";
import { anthropic } from "@ai-sdk/anthropic";
import { createProviderRegistry, stepCountIs, streamText } from "ai";
import { z } from "zod";
import { catalogPromptSection } from "../lib/catalog";

const registry = createProviderRegistry({ anthropic });

const systemPrompt = prompts.define({
  id: "clickhouse-analyst",
  model: "anthropic:claude-opus-4-8",
  variables: z.object({ componentReference: z.string() }),
  content: `You are a ClickHouse data analyst. ...

## renderVisualization spec reference

{{componentReference}}`,
});

export const clickhouseAgent = chat.agent({
  id: "clickhouse-agent",
  idleTimeoutInSeconds: 300,
  // Declared on the config so tool results survive history re-conversion across turns
  tools: { listTables, describeTable, runQuery, renderVisualization },

  onChatStart: async () => {
    // Latest prompt version (or an active dashboard override), with the
    // component reference generated from the catalog so it always matches
    // the deployed code.
    const resolved = await systemPrompt.resolve({
      componentReference: catalogPromptSection(),
    });
    chat.prompt.set(resolved);
  },

  run: async ({ messages, tools, signal }) => {
    return streamText({
      // Fallback model only — placed BEFORE the spread so the stored
      // prompt's model (including dashboard overrides) wins when set.
      model: anthropic("claude-opus-4-8"),
      // Wires up prepareStep (compaction, steering, background injection),
      // plus the system prompt + model + config + telemetry from chat.prompt().
      ...chat.toStreamTextOptions({ registry }),
      messages,
      tools,
      stopWhen: stepCountIs(15),
      abortSignal: signal,
    });
  },
});

import { defineCatalog } from "@json-render/core";
import { schema } from "@json-render/react/schema";
import { shadcnComponentDefinitions } from "@json-render/shadcn/catalog";

export const catalog = defineCatalog(schema, {
  components: {
    // Layout & text from the stock shadcn catalog
    Card: shadcnComponentDefinitions.Card,
    Grid: shadcnComponentDefinitions.Grid,
    Table: shadcnComponentDefinitions.Table,
    // ...plus custom BarChart, LineChart, AreaChart, PieChart, Stat, PointMap
  },
  actions: {},
});

// Generates a component reference (props as JSON schema, from the same zod
// definitions) for the system prompt — the prompt can't drift from the code.
export function catalogPromptSection(): string {
  /* ... */
}

// Validates a spec against the catalog; errors are phrased for the model
// to correct and retry.
export function validateSpec(spec: VisualizationSpec) {
  /* ... */
}

const renderVisualization = tool({
  description:
    "Render charts, tables and stat cards for the user, instead of describing data as text.",
  inputSchema: z.object({
    spec: z.object({
      root: z.string(),
      elements: z.record(
        z.string(),
        z.object({
          type: z.string(),
          props: z.record(z.string(), z.unknown()),
          children: z.array(z.string()).optional(),
        })
      ),
    }),
  }),
  execute: async ({ spec }) => {
    const result = validateSpec(spec);
    if (!result.ok) {
      // The model reads these, fixes the spec, and calls the tool again
      return { ok: false, errors: result.errors };
    }
    return { ok: true, note: "Rendered to the user. Add at most a one-sentence takeaway." };
  },
});

"use client";

import { useChat } from "@ai-sdk/react";
import { useTriggerChatTransport } from "@trigger.dev/sdk/chat/react";
import type { clickhouseAgent } from "@/trigger/clickhouse-agent";
import { mintChatAccessToken, startChatSession } from "@/app/actions";

export function Chat() {
  const transport = useTriggerChatTransport<typeof clickhouseAgent>({
    task: "clickhouse-agent",
    accessToken: ({ chatId }) => mintChatAccessToken(chatId),
    startSession: ({ chatId, clientData }) => startChatSession({ chatId, clientData }),
  });

  const { messages, sendMessage, stop, status } = useChat({ transport });
  // Render text parts as markdown; render tool-renderVisualization parts
  // with json-render's <Renderer spec={...} registry={registry} />
}

import { defineRegistry } from "@json-render/react";
import { shadcnComponents } from "@json-render/shadcn";
import { catalog } from "./catalog";

export const { registry } = defineRegistry(catalog, {
  components: {
    Card: shadcnComponents.Card,
    Table: shadcnComponents.Table,
    // ...
    BarChart: ({ props }) => <BarChartView {...props} />,
    PointMap: ({ props }) => <PointMapView {...props} />,
  },
});

const READ_ONLY_STATEMENTS = /^\s*(select|with|show|describe|desc|explain|exists)\b/i;

const runQuery = tool({
  description: "Run a read-only SQL query against ClickHouse and get the results as JSON rows.",
  inputSchema: z.object({
    query: z.string().describe("The ClickHouse SQL query to run"),
  }),
  execute: async ({ query }) => {
    if (!READ_ONLY_STATEMENTS.test(query)) {
      return { error: "Only read-only statements are allowed." };
    }
    try {
      const result = await getClickHouse().query({
        query,
        format: "JSONEachRow",
        clickhouse_settings: {
          // readonly=2: reads only (no writes/DDL), but per-query settings
          // like the limits below are still allowed.
          readonly: "2",
          max_result_rows: "1000",
          result_overflow_mode: "break",
          max_execution_time: 30,
        },
      });
      const rows = await result.json();
      return { rowCount: rows.length, rows };
    } catch (error) {
      // Return ClickHouse errors to the model so it can fix the query and retry.
      return { error: error instanceof Error ? error.message : String(error) };
    }
  },
});

TRIGGER_PROJECT_REF=proj_xxxxxxxxxxxxxxxxxxxxxxxx
TRIGGER_SECRET_KEY=tr_dev_xxxxxxxxxxxxxxxxxxxxxxxx