diff --git a/apps/dev-playground/client/src/routes/genie.route.tsx b/apps/dev-playground/client/src/routes/genie.route.tsx index 3b47ce372..871c2d001 100644 --- a/apps/dev-playground/client/src/routes/genie.route.tsx +++ b/apps/dev-playground/client/src/routes/genie.route.tsx @@ -1,11 +1,32 @@ -import { GenieChat } from "@databricks/appkit-ui/react"; +import { Button, GenieChat } from "@databricks/appkit-ui/react"; import { createFileRoute } from "@tanstack/react-router"; +import { useState } from "react"; export const Route = createFileRoute("/genie")({ component: GenieRoute, }); +type Mode = "chat" | "agent"; + +// Genie plugin aliases configured in server/index.ts. +const ALIASES: Record = { chat: "demo", agent: "agent" }; + function GenieRoute() { + const [mode, setMode] = useState(() => + new URLSearchParams(window.location.search).get("mode") === "agent" + ? "agent" + : "chat", + ); + + const switchMode = (next: Mode) => { + // Chat and agent mode don't share conversations. + const url = new URL(window.location.href); + url.searchParams.set("mode", next); + url.searchParams.delete("conversationId"); + window.history.replaceState({}, "", url.toString()); + setMode(next); + }; + return (
@@ -16,11 +37,29 @@ function GenieRoute() {

Ask natural language questions about your data using AI/BI Genie. + Agent mode shows the agent's reasoning and its charts.

+
+ + +
+
- +
diff --git a/apps/dev-playground/server/index.ts b/apps/dev-playground/server/index.ts index f2c22a7ca..e876cb315 100644 --- a/apps/dev-playground/server/index.ts +++ b/apps/dev-playground/server/index.ts @@ -67,6 +67,8 @@ createApp({ analytics({}), genie({ spaces: { demo: process.env.DATABRICKS_GENIE_SPACE_ID ?? "placeholder" }, + // Any Genie space also works as an agent. + agents: { agent: process.env.DATABRICKS_GENIE_SPACE_ID ?? "placeholder" }, }), ...(process.env.LAKEBASE_ENDPOINT ? [lakebase()] : []), lakebaseExamples(), diff --git a/docs/docs/api/appkit-ui/genie/GenieQueryVisualization.mdx b/docs/docs/api/appkit-ui/genie/GenieQueryVisualization.mdx index 8f7645214..44dd3cd33 100644 --- a/docs/docs/api/appkit-ui/genie/GenieQueryVisualization.mdx +++ b/docs/docs/api/appkit-ui/genie/GenieQueryVisualization.mdx @@ -7,9 +7,10 @@ Renders a chart + data table for a Genie query result. Renders a chart + data table for a Genie query result. +- When `chartImageUrl` is set: the Chart tab shows that image; if it fails to load, the inferred chart - When a chart type can be inferred: shows Tabs with "Chart" (default) and "Table" - When no chart fits: shows only the data table -- When data is empty/malformed: renders nothing +- When data is empty/malformed: renders only the chart image, if any **Source:** [`packages/appkit-ui/src/react/genie/genie-query-visualization.tsx`](https://github.com/databricks/appkit/blob/main/packages/appkit-ui/src/react/genie/genie-query-visualization.tsx) @@ -20,6 +21,8 @@ Renders a chart + data table for a Genie query result. | Prop | Type | Required | Default | Description | |------|------|----------|---------|-------------| | `data` | `GenieStatementResponse` | ✓ | - | Raw statement_response from the Genie API | +| `chartImageUrl` | `string` | | - | URL of a chart Genie rendered for this query, shown in place of the inferred chart | +| `chartImageAlt` | `string` | | - | Alt text for `chartImageUrl` | | `className` | `string` | | - | Additional CSS classes | diff --git a/docs/docs/plugins/genie.md b/docs/docs/plugins/genie.md index 386ae7e2f..d4d0dcb6c 100644 --- a/docs/docs/plugins/genie.md +++ b/docs/docs/plugins/genie.md @@ -11,6 +11,8 @@ Integrates [Databricks AI/BI Genie](https://docs.databricks.com/en/genie/index.h - SSE streaming with real-time status updates - Conversation history replay with automatic reconnection - Query result attachment fetching +- Genie's charts, served as PNG images +- Agent mode: the agent's reasoning, streamed as it works, and the charts it draws - On-behalf-of (OBO) user execution ## Basic usage @@ -31,7 +33,8 @@ await createApp({ | Option | Type | Default | Description | |--------|------|---------|-------------| | `spaces` | `Record` | `{ default: DATABRICKS_GENIE_SPACE_ID }` | Map of alias names to Genie Space IDs | -| `timeout` | `number` | `120000` | Polling timeout in ms. Set to `0` for indefinite | +| `agents` | `Record` | `{}` | Map of alias names to Genie Space IDs used in [agent mode](#agent-mode) | +| `timeout` | `number` | `120000` | Timeout in ms for a chat message or an agent turn. Set to `0` for indefinite | ### Space aliases @@ -48,6 +51,26 @@ genie({ If you omit `spaces`, the plugin reads `DATABRICKS_GENIE_SPACE_ID` from the environment and registers it under the `default` alias. +### Agent mode + +Genie agent mode answers like an agent: it reasons about the question, runs queries, and draws charts. Register a space under `agents` instead of `spaces` to use it in agent mode. Any Genie space can be used as an agent: + +```ts +genie({ + spaces: { sales: "01ABCDEF12345678" }, + agents: { analyst: "01ABCDEF12345678" }, +}); +``` + +An agent alias works with the same routes, the same `sendMessage` and `getConversation` exports, and the same `` component as a chat space. The difference is in what it returns: + +- `reasoning` events stream the agent's reasoning while it works. The completed message carries the steps in `reasoning`. +- Visualizations are always turned on. The charts the agent draws arrive as `viz` attachments, like a chat space's charts (see [Get a visualization](#get-a-visualization)). + +An alias can be either a chat space or an agent, not both. To use one space both ways, give it two aliases, as in the example above. Conversations don't carry over between the two modes. + +Agent mode uses the Genie agent API (`/api/2.0/genie/agents/{id}/responses` and `/api/2.0/genie/agents/{id}/conversations/{id}/items`), which the Databricks SDK doesn't cover yet. + ### Finding your Genie Space ID You can find the Space ID from the **About** tab on your Genie space page in Databricks: @@ -70,6 +93,7 @@ The genie plugin exposes these endpoints (mounted under `/api/genie`): - `POST /api/genie/:alias/messages` — Send a message to a Genie space (SSE stream) - `GET /api/genie/:alias/conversations/:conversationId` — Replay conversation history (SSE stream) +- `GET /api/genie/:alias/conversations/:conversationId/messages/:messageId/attachments/:attachmentId/visualization` — Download a chart as a PNG ### Send a message @@ -89,6 +113,7 @@ The response is an SSE stream that emits these event types: |------------|-------------| | `message_start` | Conversation and message IDs assigned | | `status` | Processing status updates (e.g. `ASKING_AI`, `EXECUTING_QUERY`) | +| `reasoning` | Agent mode only: one step of the agent's reasoning | | `message_result` | Final message with text and query attachments | | `query_result` | Tabular data for a query attachment | | `error` | Error details | @@ -101,6 +126,16 @@ GET /api/genie/:alias/conversations/:conversationId Returns an SSE stream of `message_result` and `query_result` events for all messages in the conversation. +### Get a visualization + +``` +GET /api/genie/:alias/conversations/:conversationId/messages/:messageId/attachments/:attachmentId/visualization +``` + +Returns the PNG Genie rendered for a `viz` attachment, so it can be used directly as an `` source. Use the `attachmentId` of the `viz` attachment, not its `queryAttachmentId`. The message must be completed. + +Chat spaces and agents both return `viz` attachments when Genie draws a chart. For a query with a `viz` attachment, `GenieChat` shows Genie's chart in place of the chart it would otherwise infer from the query result. + ## Programmatic access The plugin exports `sendMessage` and `getConversation` for server-side use: diff --git a/packages/appkit-ui/src/react/genie/__tests__/genie-chat-message.test.tsx b/packages/appkit-ui/src/react/genie/__tests__/genie-chat-message.test.tsx new file mode 100644 index 000000000..862d24d57 --- /dev/null +++ b/packages/appkit-ui/src/react/genie/__tests__/genie-chat-message.test.tsx @@ -0,0 +1,114 @@ +import { cleanup, render, screen } from "@testing-library/react"; +import { afterEach, describe, expect, test, vi } from "vitest"; + +vi.mock("../genie-query-visualization", () => ({ + GenieQueryVisualization: (props: { chartImageUrl?: string }) => ( +
+ ), +})); + +import { GenieChatMessage } from "../genie-chat-message"; +import type { GenieMessageItem, GenieStatementResponse } from "../types"; + +const queryResult: GenieStatementResponse = { + manifest: { schema: { columns: [{ name: "country" }] } }, + result: { dataArray: [["Japan"]] }, +}; + +function assistant(overrides: Partial): GenieMessageItem { + return { + id: "msg-1", + role: "assistant", + content: "Japan has the most.", + status: "COMPLETED", + attachments: [], + queryResults: new Map(), + ...overrides, + }; +} + +describe("GenieChatMessage", () => { + afterEach(cleanup); + + test("hands Genie's chart to the query result, which keeps its table", () => { + render( + , + ); + + const result = screen.getByTestId("query-result"); + expect(result.getAttribute("data-chart-image")).toBe( + "/api/genie/a/viz-1.png", + ); + expect(screen.queryByRole("img")).toBeNull(); + }); + + test("shows Genie's chart as an image while there's no query result", () => { + render( + , + ); + + const img = screen.getByRole("img", { name: "Franchises by Country" }); + expect(img.getAttribute("src")).toBe("/api/genie/a/viz-1.png"); + expect(screen.queryByTestId("query-result")).toBeNull(); + }); + + test("leaves the query result alone for a query without a visualization", () => { + render( + , + ); + + const result = screen.getByTestId("query-result"); + expect(result.getAttribute("data-chart-image")).toBeNull(); + expect(screen.queryByRole("img")).toBeNull(); + }); + + test("shows the agent's reasoning", () => { + render( + , + ); + + expect(screen.getByText("Thinking")).toBeTruthy(); + expect(screen.getByText("Counting franchises.")).toBeTruthy(); + }); +}); diff --git a/packages/appkit-ui/src/react/genie/__tests__/genie-query-visualization.test.tsx b/packages/appkit-ui/src/react/genie/__tests__/genie-query-visualization.test.tsx new file mode 100644 index 000000000..8f86337ec --- /dev/null +++ b/packages/appkit-ui/src/react/genie/__tests__/genie-query-visualization.test.tsx @@ -0,0 +1,66 @@ +import { cleanup, fireEvent, render, screen } from "@testing-library/react"; +import { afterEach, describe, expect, test, vi } from "vitest"; + +vi.mock("../../charts/base", () => ({ + BaseChart: () =>
, +})); + +import { GenieQueryVisualization } from "../genie-query-visualization"; +import type { GenieStatementResponse } from "../types"; + +const data: GenieStatementResponse = { + manifest: { + schema: { + columns: [ + { name: "country", typeName: "STRING" }, + { name: "franchises", typeName: "INT" }, + ], + }, + }, + result: { + dataArray: [ + ["Japan", "20"], + ["US", "16"], + ], + }, +}; + +describe("GenieQueryVisualization", () => { + afterEach(cleanup); + + test("shows the chart image in the Chart tab and keeps the Table tab", () => { + render( + , + ); + + const img = screen.getByRole("img", { name: "Franchises by Country" }); + expect(img.getAttribute("src")).toBe("/viz.png"); + expect(screen.queryByTestId("inferred-chart")).toBeNull(); + expect(screen.getByRole("tab", { name: "Table" })).toBeTruthy(); + expect(screen.getByRole("cell", { name: "Japan" })).toBeTruthy(); + expect( + screen.queryByRole("button", { name: "Change chart type" }), + ).toBeNull(); + }); + + test("falls back to the inferred chart when the image fails to load", () => { + render(); + + fireEvent.error(screen.getByRole("img")); + + expect(screen.queryByRole("img")).toBeNull(); + expect(screen.getByTestId("inferred-chart")).toBeTruthy(); + expect(screen.getByRole("cell", { name: "Japan" })).toBeTruthy(); + }); + + test("shows the inferred chart without a chart image", () => { + render(); + + expect(screen.getByTestId("inferred-chart")).toBeTruthy(); + expect(screen.queryByRole("img")).toBeNull(); + }); +}); diff --git a/packages/appkit-ui/src/react/genie/__tests__/use-genie-chat.test.ts b/packages/appkit-ui/src/react/genie/__tests__/use-genie-chat.test.ts new file mode 100644 index 000000000..45630f8f6 --- /dev/null +++ b/packages/appkit-ui/src/react/genie/__tests__/use-genie-chat.test.ts @@ -0,0 +1,85 @@ +import { act, renderHook } from "@testing-library/react"; +import { afterEach, describe, expect, test, vi } from "vitest"; + +import type { GenieStreamEvent } from "../types"; + +// Capture the SSE callbacks so the test can play server events. +let onMessage: ((msg: { data: string }) => Promise) | undefined; +const mockConnectSSE = vi.fn().mockImplementation((opts: any) => { + onMessage = opts.onMessage; + return new Promise(() => {}); +}); + +vi.mock("@/js", () => ({ + connectSSE: (...args: unknown[]) => mockConnectSSE(...args), +})); + +import { useGenieChat } from "../use-genie-chat"; + +async function emit(event: GenieStreamEvent) { + await act(async () => { + await onMessage?.({ data: JSON.stringify(event) }); + }); +} + +describe("useGenieChat", () => { + afterEach(() => { + mockConnectSSE.mockClear(); + }); + + test("shows reasoning while streaming and links visualizations to the plugin route", async () => { + const { result } = renderHook(() => + useGenieChat({ + alias: "my agent", + basePath: "/api/genie", + persistInUrl: false, + }), + ); + + act(() => result.current.sendMessage("Franchises by country?")); + await emit({ + type: "message_start", + conversationId: "conv-1", + messageId: "msg-1", + spaceId: "s", + }); + await emit({ type: "reasoning", content: "Counting franchises." }); + await emit({ type: "reasoning", content: "Charting them." }); + + expect(result.current.messages.at(-1)).toMatchObject({ + role: "assistant", + content: "", + reasoning: ["Counting franchises.", "Charting them."], + }); + + await emit({ + type: "message_result", + message: { + messageId: "msg-1", + conversationId: "conv-1", + spaceId: "s", + status: "COMPLETED", + content: "Franchises by country?", + attachments: [ + { attachmentId: "query-1", query: { title: "Franchises" } }, + { + attachmentId: "viz-1_output", + viz: { title: "Chart", queryAttachmentId: "query-1" }, + }, + { text: { content: "Japan has the most." } }, + ], + reasoning: ["Counting franchises.", "Charting them."], + }, + }); + + expect(result.current.messages.at(-1)).toMatchObject({ + id: "msg-1", + content: "Japan has the most.", + reasoning: ["Counting franchises.", "Charting them."], + visualizationUrls: { + "viz-1_output": + "/api/genie/my%20agent/conversations/conv-1/messages/msg-1/attachments/viz-1_output/visualization", + }, + }); + }); +}); diff --git a/packages/appkit-ui/src/react/genie/genie-chat-message-list.tsx b/packages/appkit-ui/src/react/genie/genie-chat-message-list.tsx index 016eae8a1..756ca1088 100644 --- a/packages/appkit-ui/src/react/genie/genie-chat-message-list.tsx +++ b/packages/appkit-ui/src/react/genie/genie-chat-message-list.tsx @@ -27,6 +27,7 @@ interface GenieChatMessageListProps { const STATUS_LABELS: Record = { ASKING_AI: "Asking AI...", EXECUTING_QUERY: "Executing query...", + GENERATING_VISUALIZATION: "Generating visualization...", FILTERING_RESULTS: "Filtering results...", COMPLETED: "Done", }; @@ -206,6 +207,7 @@ export function GenieChatMessageList({ (msg) => msg.role !== "assistant" || msg.content || + msg.reasoning?.length || (msg.id !== "" && TERMINAL_STATUSES.has(msg.status)), ) .map((msg) => ( diff --git a/packages/appkit-ui/src/react/genie/genie-chat-message.tsx b/packages/appkit-ui/src/react/genie/genie-chat-message.tsx index ca58669bc..17f445449 100644 --- a/packages/appkit-ui/src/react/genie/genie-chat-message.tsx +++ b/packages/appkit-ui/src/react/genie/genie-chat-message.tsx @@ -38,6 +38,26 @@ function isQueryAttachment(att: GenieAttachmentResponse): boolean { return !!(att.query?.title || att.query?.query); } +interface Visualization { + url: string; + title?: string; + queryAttachmentId?: string; +} + +/** A chart Genie rendered, shown as the PNG the genie plugin serves. */ +function GenieVisualizationImage({ viz }: { viz: Visualization }) { + return ( + + {viz.title + + ); +} + /** Renders a single Genie message bubble with optional expandable SQL query attachments. */ export function GenieChatMessage({ message, @@ -45,6 +65,16 @@ export function GenieChatMessage({ }: GenieChatMessageProps) { const isUser = message.role === "user"; const queryAttachments = message.attachments.filter(isQueryAttachment); + const visualizations: Visualization[] = message.attachments.flatMap((att) => { + const url = + att.attachmentId && message.visualizationUrls?.[att.attachmentId]; + return att.viz && url ? [{ url, ...att.viz }] : []; + }); + const queryIds = new Set(queryAttachments.map((att) => att.attachmentId)); + // Charts whose query isn't among the attachments still get shown, after the queries. + const unmatchedVisualizations = visualizations.filter( + (viz) => !queryIds.has(viz.queryAttachmentId), + ); const html = useMemo( () => message.content @@ -81,6 +111,22 @@ export function GenieChatMessage({ : "bg-muted", )} > + {message.reasoning && message.reasoning.length > 0 && ( +
+ + Thinking + +
    + {message.reasoning.map((step, i) => ( +
  • {step}
    • + ))} +
+
+ )} + {html && (
viz.queryAttachmentId === att.attachmentId, + ); + // The first chart goes into the result's Chart tab; any others, or all + // of them while there's no result, are shown as plain images. + const [chartViz, ...otherVisualizations] = queryVisualizations; + const standaloneVisualizations = + queryResult != null ? otherVisualizations : queryVisualizations; return (
@@ -124,14 +178,25 @@ export function GenieChatMessage({ {queryResult != null && ( - + )} + {standaloneVisualizations.map((viz) => ( + + ))}
); })}
)} + + {unmatchedVisualizations.map((viz) => ( + + ))}
); diff --git a/packages/appkit-ui/src/react/genie/genie-query-visualization.tsx b/packages/appkit-ui/src/react/genie/genie-query-visualization.tsx index 93c11f01f..60e356dfb 100644 --- a/packages/appkit-ui/src/react/genie/genie-query-visualization.tsx +++ b/packages/appkit-ui/src/react/genie/genie-query-visualization.tsx @@ -47,6 +47,10 @@ const CHART_TYPE_LABELS: Record = { interface GenieQueryVisualizationProps { /** Raw statement_response from the Genie API */ data: GenieStatementResponse; + /** URL of a chart Genie rendered for this query, shown in place of the inferred chart */ + chartImageUrl?: string; + /** Alt text for `chartImageUrl` */ + chartImageAlt?: string; /** Additional CSS classes */ className?: string; } @@ -54,12 +58,15 @@ interface GenieQueryVisualizationProps { /** * Renders a chart + data table for a Genie query result. * + * - When `chartImageUrl` is set: the Chart tab shows that image; if it fails to load, the inferred chart * - When a chart type can be inferred: shows Tabs with "Chart" (default) and "Table" * - When no chart fits: shows only the data table - * - When data is empty/malformed: renders nothing + * - When data is empty/malformed: renders only the chart image, if any */ export function GenieQueryVisualization({ data, + chartImageUrl, + chartImageAlt, className, }: GenieQueryVisualizationProps) { const transformed = useMemo(() => transformGenieData(data), [data]); @@ -76,8 +83,24 @@ export function GenieQueryVisualization({ const [chartTypeOverride, setChartTypeOverride] = useState( null, ); + const [imageFailed, setImageFailed] = useState(false); - if (!transformed || transformed.rows.length === 0) return null; + const chartImage = + chartImageUrl && !imageFailed ? ( + {chartImageAlt setImageFailed(true)} + /> + ) : null; + + if (!transformed || transformed.rows.length === 0) { + return ( + chartImage &&
{chartImage}
+ ); + } const { rows, columns } = transformed; const truncated = rows.length > TABLE_ROW_LIMIT; @@ -118,7 +141,22 @@ export function GenieQueryVisualization({ ); - if (!inference || !activeChartType) { + const chart = + chartImage ?? + (inference && activeChartType ? ( + + + + ) : null); + + if (!chart) { return
{dataTable}
; } @@ -129,7 +167,7 @@ export function GenieQueryVisualization({ Chart Table - {compatibleTypes.length > 1 && ( + {!chartImage && compatibleTypes.length > 1 && (