From 4c9e2c472e183fac313c170b19e23b9eb8bd8e20 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Thu, 12 Nov 2026 17:22:47 +0000 Subject: [PATCH 1/4] feat(appkit): support Genie agent mode with reasoning and visualizations The genie plugin only used the Genie chat API, which returns no agentic reasoning and no rendered charts. Genie agent mode does, through a separate API. Spaces registered under the new `agents` config map use agent mode on the same routes, exports and agent tools as chat spaces: - Sending a message calls POST /api/2.0/genie/agents/{id}/responses with visualizations enabled and maps its stream onto the existing events: message_start, status (per tool call), a new `reasoning` event per reasoning step, then message_result and query_result. The completed message is read back for its query and chart attachments; the answer text, which that message doesn't carry, is added as a text attachment. - History and pending-message polling add the answer text and reasoning from the conversation's items endpoint. - A new visualization route streams the PNG Genie rendered for a `viz` attachment, so the UI can show it in an . GenieChat shows the reasoning while the agent works and shows Genie's chart in place of the inferred chart for queries that have one. Co-authored-by: Isaac Signed-off-by: Fabian Jakobs --- .../client/src/routes/genie.route.tsx | 43 +- apps/dev-playground/server/index.ts | 2 + docs/docs/plugins/genie.md | 32 ++ .../__tests__/genie-chat-message.test.tsx | 85 ++++ .../genie/__tests__/use-genie-chat.test.ts | 85 ++++ .../react/genie/genie-chat-message-list.tsx | 2 + .../src/react/genie/genie-chat-message.tsx | 67 ++- packages/appkit-ui/src/react/genie/types.ts | 4 + .../src/react/genie/use-genie-chat.ts | 57 ++- packages/appkit/src/connectors/genie/agent.ts | 94 ++++ .../appkit/src/connectors/genie/client.ts | 284 ++++++++++- packages/appkit/src/connectors/genie/index.ts | 2 +- .../src/connectors/genie/tests/agent.test.ts | 455 ++++++++++++++++++ packages/appkit/src/plugins/genie/genie.ts | 138 +++++- .../appkit/src/plugins/genie/manifest.json | 7 + .../src/plugins/genie/tests/genie.test.ts | 270 ++++++++++- packages/appkit/src/plugins/genie/types.ts | 7 + packages/shared/src/genie.ts | 13 + 18 files changed, 1604 insertions(+), 43 deletions(-) create mode 100644 packages/appkit-ui/src/react/genie/__tests__/genie-chat-message.test.tsx create mode 100644 packages/appkit-ui/src/react/genie/__tests__/use-genie-chat.test.ts create mode 100644 packages/appkit/src/connectors/genie/agent.ts create mode 100644 packages/appkit/src/connectors/genie/tests/agent.test.ts 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/plugins/genie.md b/docs/docs/plugins/genie.md index 386ae7e2f..63552e307 100644 --- a/docs/docs/plugins/genie.md +++ b/docs/docs/plugins/genie.md @@ -11,6 +11,7 @@ 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 +- Agent mode: the agent's reasoning, streamed as it works, and the charts it draws - On-behalf-of (OBO) user execution ## Basic usage @@ -31,6 +32,7 @@ await createApp({ | Option | Type | Default | Description | |--------|------|---------|-------------| | `spaces` | `Record` | `{ default: DATABRICKS_GENIE_SPACE_ID }` | Map of alias names to Genie Space IDs | +| `agents` | `Record` | `{}` | Map of alias names to Genie Space IDs used in [agent mode](#agent-mode) | | `timeout` | `number` | `120000` | Polling timeout in ms. Set to `0` for indefinite | ### Space aliases @@ -48,6 +50,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`. +- Charts arrive as `viz` attachments. The [visualization endpoint](#get-a-visualization) serves each one as a PNG, and `GenieChat` shows it in place of the chart it would otherwise infer from the query result. + +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 +92,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 +112,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 +125,14 @@ 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. + ## 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..d25855fb3 --- /dev/null +++ b/packages/appkit-ui/src/react/genie/__tests__/genie-chat-message.test.tsx @@ -0,0 +1,85 @@ +import { cleanup, render, screen } from "@testing-library/react"; +import { afterEach, describe, expect, test, vi } from "vitest"; + +vi.mock("../genie-query-visualization", () => ({ + GenieQueryVisualization: () =>
, +})); + +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("shows Genie's chart instead of the inferred one for a query with a visualization", () => { + render( + , + ); + + const img = screen.getByRole("img", { name: "Franchises by Country" }); + expect(img.getAttribute("src")).toBe("/api/genie/a/viz-1.png"); + expect(screen.queryByTestId("inferred-chart")).toBeNull(); + }); + + test("keeps the inferred chart for a query without a visualization", () => { + render( + , + ); + + expect(screen.getByTestId("inferred-chart")).toBeTruthy(); + 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__/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..c7f4bb0e7 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, + ); return (
@@ -122,16 +171,24 @@ export function GenieChatMessage({
- {queryResult != null && ( - - - - )} + {queryVisualizations.length > 0 + ? queryVisualizations.map((viz) => ( + + )) + : queryResult != null && ( + + + + )}
); })}
)} + + {unmatchedVisualizations.map((viz) => ( + + ))} ); diff --git a/packages/appkit-ui/src/react/genie/types.ts b/packages/appkit-ui/src/react/genie/types.ts index a4ac3c09d..033a34476 100644 --- a/packages/appkit-ui/src/react/genie/types.ts +++ b/packages/appkit-ui/src/react/genie/types.ts @@ -24,6 +24,10 @@ export interface GenieMessageItem { attachments: GenieAttachmentResponse[]; queryResults: Map; error?: string; + /** Agent mode: the agent's reasoning steps, in order. */ + reasoning?: string[]; + /** Image URL of each visualization attachment, keyed by attachment ID. */ + visualizationUrls?: Record; } export interface UseGenieChatOptions { diff --git a/packages/appkit-ui/src/react/genie/use-genie-chat.ts b/packages/appkit-ui/src/react/genie/use-genie-chat.ts index 7688e2aa9..a58de3d5c 100644 --- a/packages/appkit-ui/src/react/genie/use-genie-chat.ts +++ b/packages/appkit-ui/src/react/genie/use-genie-chat.ts @@ -54,7 +54,31 @@ function makeUserItem( }; } -function makeAssistantItem(msg: GenieMessageResponse): GenieMessageItem { +/** URLs of the plugin route that serves each visualization attachment as a PNG. */ +function visualizationUrls( + basePath: string, + alias: string, + msg: GenieMessageResponse, +): Record | undefined { + const entries = (msg.attachments ?? []).flatMap((att) => + att.viz && att.attachmentId + ? [ + [ + att.attachmentId, + `${basePath}/${encodeURIComponent(alias)}/conversations/${encodeURIComponent(msg.conversationId)}` + + `/messages/${encodeURIComponent(msg.messageId)}/attachments/${encodeURIComponent(att.attachmentId)}/visualization`, + ], + ] + : [], + ); + return entries.length > 0 ? Object.fromEntries(entries) : undefined; +} + +function makeAssistantItem( + msg: GenieMessageResponse, + basePath: string, + alias: string, +): GenieMessageItem { return { id: msg.messageId, role: "assistant", @@ -63,6 +87,8 @@ function makeAssistantItem(msg: GenieMessageResponse): GenieMessageItem { attachments: msg.attachments ?? [], queryResults: new Map(), error: msg.error, + reasoning: msg.reasoning, + visualizationUrls: visualizationUrls(basePath, alias, msg), }; } @@ -74,7 +100,11 @@ function makeAssistantItem(msg: GenieMessageResponse): GenieMessageItem { * attachments yet, we emit an empty assistant placeholder so the UI can * show a loading indicator and later poll for the completed response. */ -function messageResultToItems(msg: GenieMessageResponse): GenieMessageItem[] { +function messageResultToItems( + msg: GenieMessageResponse, + basePath: string, + alias: string, +): GenieMessageItem[] { const hasAttachments = (msg.attachments?.length ?? 0) > 0; if (!hasAttachments && TERMINAL_STATUSES.has(msg.status)) { @@ -93,7 +123,7 @@ function messageResultToItems(msg: GenieMessageResponse): GenieMessageItem[] { }, ]; } - return [makeUserItem(msg, "-user"), makeAssistantItem(msg)]; + return [makeUserItem(msg, "-user"), makeAssistantItem(msg, basePath, alias)]; } /** @@ -128,7 +158,7 @@ function fetchConversationPage( const event = JSON.parse(message.data) as GenieStreamEvent; switch (event.type) { case "message_result": - items.push(...messageResultToItems(event.message)); + items.push(...messageResultToItems(event.message, basePath, alias)); break; case "query_result": for (let i = items.length - 1; i >= 0; i--) { @@ -224,8 +254,23 @@ export function useGenieChat(options: UseGenieChatOptions): UseGenieChatReturn { break; } + case "reasoning": { + setMessages((prev) => { + const last = prev[prev.length - 1]; + if (last?.role !== "assistant") return prev; + return [ + ...prev.slice(0, -1), + { + ...last, + reasoning: [...(last.reasoning ?? []), event.content], + }, + ]; + }); + break; + } + case "message_result": { - const item = makeAssistantItem(event.message); + const item = makeAssistantItem(event.message, basePath, alias); setMessages((prev) => { const last = prev[prev.length - 1]; if (!last || last.role !== "assistant") return prev; @@ -272,7 +317,7 @@ export function useGenieChat(options: UseGenieChatOptions): UseGenieChatReturn { } } }, - [persistInUrl, urlParamName], + [persistInUrl, urlParamName, basePath, alias], ); processStreamEventRef.current = processStreamEvent; diff --git a/packages/appkit/src/connectors/genie/agent.ts b/packages/appkit/src/connectors/genie/agent.ts new file mode 100644 index 000000000..52ec2a746 --- /dev/null +++ b/packages/appkit/src/connectors/genie/agent.ts @@ -0,0 +1,94 @@ +/** + * Genie agent mode answers with a Responses-style protocol: each turn is a + * list of output items (reasoning, tool calls, assistant messages). These + * helpers pick out the answer text and the reasoning. Queries and charts are + * read from the message projection instead, as in chat mode. + */ +import type { GenieMessageResponse } from "./types"; + +/** The fields of an agent output item that AppKit reads. */ +export interface GenieAgentItem { + type: string; + role?: string; + name?: string; + content?: Array<{ text?: string }>; + metadata?: { message_id?: string }; +} + +/** One SSE event from `POST /api/2.0/genie/agents/{id}/responses`. */ +export interface GenieAgentStreamEvent { + type: string; + response?: { + id: string; + conversation_id: string; + error?: { message?: string }; + }; + item?: GenieAgentItem; + message?: string; +} + +/** The answer text and reasoning an agent produced for one message. */ +export interface GenieAgentTurn { + answer: string[]; + reasoning: string[]; +} + +// Agent tools mapped onto status names the chat UI already labels. +const TOOL_STATUS: Record = { + execute_sql: "EXECUTING_QUERY", + generate_visualization: "GENERATING_VISUALIZATION", +}; + +export function newAgentTurn(): GenieAgentTurn { + return { answer: [], reasoning: [] }; +} + +/** The chat-mode status for a tool call item, if AppKit knows the tool. */ +export function agentToolStatus(item: GenieAgentItem): string | undefined { + return item.type === "function_call" && item.name + ? TOOL_STATUS[item.name] + : undefined; +} + +/** + * Adds an item's text to the turn. Returns the text when the item is a + * reasoning step, so a live stream can forward it. + */ +export function addAgentItem( + turn: GenieAgentTurn, + item: GenieAgentItem, +): string | undefined { + const text = (item.content ?? []) + .map((part) => part.text) + .filter(Boolean) + .join("\n\n"); + if (!text) return undefined; + if (item.type === "reasoning") { + turn.reasoning.push(text); + return text; + } + if (item.type === "message" && item.role === "assistant") { + turn.answer.push(text); + } + return undefined; +} + +/** + * The message projection of an agent turn has its queries and charts but no + * answer text. Add the answer as a text attachment, which is where chat mode + * carries it, and the reasoning. + */ +export function withAgentTurn( + message: GenieMessageResponse, + turn: GenieAgentTurn | undefined, +): GenieMessageResponse { + if (!turn) return message; + const answer = turn.answer.join("\n\n"); + return { + ...message, + attachments: answer + ? [...(message.attachments ?? []), { text: { content: answer } }] + : message.attachments, + reasoning: turn.reasoning, + }; +} diff --git a/packages/appkit/src/connectors/genie/client.ts b/packages/appkit/src/connectors/genie/client.ts index 73d56d5f6..d88c7ede1 100644 --- a/packages/appkit/src/connectors/genie/client.ts +++ b/packages/appkit/src/connectors/genie/client.ts @@ -1,5 +1,15 @@ import { createLogger } from "../../logging"; +import { readSseEvents } from "../../stream/sse-reader"; import type { GenieMessage, WorkspaceClient } from "../../workspace-client"; +import { + addAgentItem, + agentToolStatus, + type GenieAgentItem, + type GenieAgentStreamEvent, + type GenieAgentTurn, + newAgentTurn, + withAgentTurn, +} from "./agent"; import { genieConnectorDefaults } from "./defaults"; import { type Pollable, pollWaiter } from "./poll-waiter"; import type { @@ -26,12 +36,22 @@ type CreateMessageWaiter = Pollable; const DEFAULT_WAIT_TIMEOUT_MS = 10 * 60_000; const MAX_POLL_INTERVAL_MS = 10_000; +// The agent items endpoint returns at most 100 items per page. +const AGENT_ITEMS_PAGE_SIZE = 100; +const MAX_AGENT_ITEM_PAGES = 20; + +// IDs Genie hands out (hex ids, `pub_toolu_…_output` attachment ids). +const GENIE_ID_PATTERN = /^[\w-]+$/; + interface GenieConnectorConfig { timeout?: number; maxMessages?: number; } -function mapAttachments(message: GenieMessage): GenieAttachmentResponse[] { +function mapAttachments( + message: GenieMessage, + includeViz: boolean, +): GenieAttachmentResponse[] { return ( message.attachments?.map(({ attachmentId, attachment }) => ({ attachmentId, @@ -52,18 +72,28 @@ function mapAttachments(message: GenieMessage): GenieAttachmentResponse[] { attachment?.$case === "suggestedQuestions" ? attachment.suggestedQuestions.questions : undefined, + viz: + includeViz && attachment?.$case === "viz" + ? { + title: attachment.viz.title, + queryAttachmentId: attachment.viz.queryAttachmentId, + } + : undefined, })) ?? [] ); } -function toMessageResponse(message: GenieMessage): GenieMessageResponse { +function toMessageResponse( + message: GenieMessage, + includeViz = false, +): GenieMessageResponse { return { messageId: message.messageId ?? "", conversationId: message.conversationId ?? "", spaceId: message.spaceId ?? "", status: message.status ?? "COMPLETED", content: message.content ?? "", - attachments: mapAttachments(message), + attachments: mapAttachments(message, includeViz), error: message.error?.error, }; } @@ -94,6 +124,18 @@ function classifyGenieError(error: unknown): string { return message || GenieErrors.REQUEST_FAILED; } +/** Agent endpoints answer `{"error": {"type", "message"}}`, a shape ApiError doesn't parse. */ +function classifyGenieAgentError(error: unknown): string { + const nested = (error as { response?: { error?: { message?: unknown } } }) + ?.response?.error?.message; + return typeof nested === "string" ? nested : classifyGenieError(error); +} + +/** Whether `id` is safe to put into a Genie resource path. */ +export function isGenieId(id: string): boolean { + return GENIE_ID_PATTERN.test(id); +} + export class GenieConnector { private readonly config: Required; @@ -197,7 +239,7 @@ export class GenieConnector { workspaceClient: WorkspaceClient, spaceId: string, conversationId: string, - options?: { pageSize?: number; pageToken?: string }, + options?: { pageSize?: number; pageToken?: string; includeViz?: boolean }, ): Promise<{ messages: GenieMessageResponse[]; nextPageToken: string | null; @@ -212,7 +254,9 @@ export class GenieConnector { ...(options?.pageToken ? { pageToken: options.pageToken } : {}), }); - const messages = (response.messages ?? []).reverse().map(toMessageResponse); + const messages = (response.messages ?? []) + .reverse() + .map((message) => toMessageResponse(message, options?.includeViz)); return { messages, @@ -298,6 +342,176 @@ export class GenieConnector { } } + /** + * Agent mode: runs one agent turn (`POST /genie/agents/{id}/responses`, which + * the SDK doesn't cover) and maps its stream onto the chat-mode events, plus + * a `reasoning` event per reasoning step. Once the turn completes, the + * message is read back for its query and chart attachments. + */ + async *streamAgentResponse( + workspaceClient: WorkspaceClient, + agentId: string, + content: string, + conversationId: string | undefined, + options?: { signal?: AbortSignal }, + ): AsyncGenerator { + const signal = options?.signal; + try { + const response = await workspaceClient.request({ + method: "POST", + path: `/api/2.0/genie/agents/${encodeURIComponent(agentId)}/responses`, + headers: new Headers({ + "Content-Type": "application/json", + Accept: "text/event-stream", + }), + body: JSON.stringify({ + enable_viz: true, + conversation_id: conversationId, + input: [ + { + type: "message", + role: "user", + content: [{ type: "input_text", text: content }], + }, + ], + }), + signal, + }); + if (!response.body) throw new Error(GenieErrors.REQUEST_FAILED); + + const turn = newAgentTurn(); + let started: { conversationId: string; messageId: string } | undefined; + let completed = false; + for await (const { data } of readSseEvents(response.body, signal)) { + const event = JSON.parse(data) as GenieAgentStreamEvent; + if (event.type === "response.created" && event.response) { + started = { + conversationId: event.response.conversation_id, + messageId: event.response.id, + }; + yield { type: "message_start", ...started, spaceId: agentId }; + } else if (event.type === "response.output_item.done" && event.item) { + const status = agentToolStatus(event.item); + if (status) yield { type: "status", status }; + const reasoning = addAgentItem(turn, event.item); + if (reasoning) yield { type: "reasoning", content: reasoning }; + } else if (event.type === "response.completed") { + completed = true; + } else if (event.type === "response.failed" || event.type === "error") { + throw new Error( + event.response?.error?.message ?? + event.message ?? + GenieErrors.REQUEST_FAILED, + ); + } + } + if (signal?.aborted) return; + if (!started || !completed) { + throw new Error("Genie agent response ended before it completed"); + } + + const message = await workspaceClient.genie.genieGetConversationMessage( + { spaceId: agentId, ...started }, + { signal }, + ); + const messageResponse = withAgentTurn( + toMessageResponse(message, true), + turn, + ); + yield { type: "message_result", message: messageResponse }; + yield* this.emitQueryResults( + workspaceClient, + agentId, + started.conversationId, + started.messageId, + messageResponse, + ); + } catch (error) { + if (signal?.aborted) return; + logger.error( + "Genie agent response error (agentId=%s, conversationId=%s): %O", + agentId, + conversationId ?? "new", + error, + ); + yield { type: "error", error: classifyGenieAgentError(error) }; + } + } + + /** + * Agent mode: adds the answer text and reasoning to messages read from the + * message projection, which has neither. They come from the conversation's + * items (`GET /genie/agents/{id}/conversations/{cid}/items`). + */ + private async withAgentTurns( + workspaceClient: WorkspaceClient, + agentId: string, + conversationId: string, + messages: GenieMessageResponse[], + signal?: AbortSignal, + ): Promise { + const turns = new Map(); + let after: string | undefined; + // shortcut: reads every item of the conversation for each history page; align item pages with message pages if long agent conversations matter. + for (let page = 0; page < MAX_AGENT_ITEM_PAGES; page++) { + const response = await workspaceClient.request({ + method: "GET", + path: `/api/2.0/genie/agents/${encodeURIComponent(agentId)}/conversations/${encodeURIComponent(conversationId)}/items`, + query: { + limit: String(AGENT_ITEMS_PAGE_SIZE), + ...(after ? { after } : {}), + }, + signal, + }); + const body = (await response.json()) as { + data?: GenieAgentItem[]; + has_more?: boolean; + last_id?: string; + }; + for (const item of body.data ?? []) { + const messageId = item.metadata?.message_id; + if (!messageId) continue; + let turn = turns.get(messageId); + if (!turn) { + turn = newAgentTurn(); + turns.set(messageId, turn); + } + addAgentItem(turn, item); + } + if (!body.has_more || !body.last_id) break; + after = body.last_id; + } + return messages.map((message) => + withAgentTurn(message, turns.get(message.messageId)), + ); + } + + /** + * Downloads the PNG Genie rendered for a visualization attachment. The + * message must be COMPLETED. + */ + async downloadVisualization( + workspaceClient: WorkspaceClient, + spaceId: string, + conversationId: string, + messageId: string, + attachmentId: string, + signal?: AbortSignal, + ): Promise { + // The SDK puts `name` into the URL unencoded, so only plain ids may reach it. + if (![spaceId, conversationId, messageId, attachmentId].every(isGenieId)) { + throw new Error("Invalid Genie id"); + } + const { contents } = + await workspaceClient.genie.downloadMessageAttachmentVisualization( + { + name: `spaces/${spaceId}/conversations/${conversationId}/messages/${messageId}/attachments/${attachmentId}`, + }, + { signal }, + ); + return contents; + } + private async *emitQueryResults( workspaceClient: WorkspaceClient, spaceId: string, @@ -347,18 +561,33 @@ export class GenieConnector { pageSize?: number; pageToken?: string; signal?: AbortSignal; + /** The space is used in agent mode. */ + agent?: boolean; }, ): AsyncGenerator { const includeQueryResults = options?.includeQueryResults !== false; try { - const { messages: messageResponses, nextPageToken } = - await this.listConversationMessages( - workspaceClient, - spaceId, - conversationId, - { pageSize: options?.pageSize, pageToken: options?.pageToken }, - ); + const page = await this.listConversationMessages( + workspaceClient, + spaceId, + conversationId, + { + pageSize: options?.pageSize, + pageToken: options?.pageToken, + includeViz: options?.agent, + }, + ); + const { nextPageToken } = page; + const messageResponses = options?.agent + ? await this.withAgentTurns( + workspaceClient, + spaceId, + conversationId, + page.messages, + options.signal, + ) + : page.messages; for (const messageResponse of messageResponses) { yield { type: "message_result", message: messageResponse }; @@ -449,7 +678,13 @@ export class GenieConnector { spaceId: string, conversationId: string, messageId: string, - options?: { timeout?: number; pollInterval?: number; signal?: AbortSignal }, + options?: { + timeout?: number; + pollInterval?: number; + signal?: AbortSignal; + /** The space is used in agent mode. */ + agent?: boolean; + }, ): AsyncGenerator { const pollInterval = options?.pollInterval ?? 3_000; const signal = options?.signal; @@ -472,7 +707,16 @@ export class GenieConnector { const isTerminal = message.status === "COMPLETED" || message.status === "FAILED"; if (isTerminal) { - const messageResponse = toMessageResponse(message); + let messageResponse = toMessageResponse(message, options?.agent); + if (options?.agent) { + [messageResponse] = await this.withAgentTurns( + workspaceClient, + spaceId, + conversationId, + [messageResponse], + signal, + ); + } yield { type: "message_result", message: messageResponse }; yield* this.emitQueryResults( workspaceClient, @@ -534,6 +778,7 @@ export class GenieConnector { workspaceClient: WorkspaceClient, spaceId: string, conversationId: string, + options?: { agent?: boolean }, ): Promise { const allMessages: GenieMessageResponse[] = []; let pageToken: string | undefined; @@ -546,16 +791,25 @@ export class GenieConnector { { pageSize: genieConnectorDefaults.pageSize, pageToken, + includeViz: options?.agent, }, ); allMessages.push(...messages); pageToken = nextPageToken ?? undefined; } while (pageToken && allMessages.length < this.config.maxMessages); + const kept = allMessages.slice(0, this.config.maxMessages); return { conversationId, spaceId, - messages: allMessages.slice(0, this.config.maxMessages), + messages: options?.agent + ? await this.withAgentTurns( + workspaceClient, + spaceId, + conversationId, + kept, + ) + : kept, }; } } diff --git a/packages/appkit/src/connectors/genie/index.ts b/packages/appkit/src/connectors/genie/index.ts index 7f23e1546..7063c1acb 100644 --- a/packages/appkit/src/connectors/genie/index.ts +++ b/packages/appkit/src/connectors/genie/index.ts @@ -1,2 +1,2 @@ -export { GenieConnector } from "./client"; +export { GenieConnector, isGenieId } from "./client"; export type { GenieConversationHistoryResponse } from "./types"; diff --git a/packages/appkit/src/connectors/genie/tests/agent.test.ts b/packages/appkit/src/connectors/genie/tests/agent.test.ts new file mode 100644 index 000000000..c96aff015 --- /dev/null +++ b/packages/appkit/src/connectors/genie/tests/agent.test.ts @@ -0,0 +1,455 @@ +import { describe, expect, test, vi } from "vitest"; + +import { GenieConnector } from "../client"; +import type { GenieStreamEvent } from "../types"; + +// Shapes trimmed from a real agent-mode turn ("Show total sales by franchise country as a bar chart"). +const AGENT_ID = "agent-1"; +const CONV = "conv-1"; +const MSG = "msg-1"; + +const created = { + type: "response.created", + response: { id: MSG, conversation_id: CONV, status: "in_progress" }, +}; +const reasoningItem = (text: string) => ({ + type: "reasoning", + status: "completed", + content: [{ type: "reasoning_text", text }], + metadata: { message_id: MSG }, +}); +const toolCall = (name: string) => ({ + type: "function_call", + status: "completed", + name, + arguments: "{}", + metadata: { message_id: MSG }, +}); +const answerItem = { + type: "message", + role: "assistant", + status: "completed", + content: [ + { type: "output_text", text: "## Franchises by country" }, + { + type: "output_text", + text: "", + metadata: { viz: { attachment_id: "viz-1" } }, + }, + { type: "output_text", text: "Japan has the most." }, + ], + metadata: { message_id: MSG }, +}; +const done = (item: object) => ({ type: "response.output_item.done", item }); +const completed = { + type: "response.completed", + response: { id: MSG, conversation_id: CONV, status: "completed" }, +}; + +/** The message projection of the same turn: queries and charts, no answer text. */ +const projection = { + messageId: MSG, + conversationId: CONV, + spaceId: AGENT_ID, + status: "COMPLETED", + content: "Show franchises by country", + attachments: [ + { + attachmentId: "query-1", + attachment: { + $case: "query", + query: { + title: "Franchises by country", + query: "SELECT country, count(*) FROM franchises GROUP BY 1", + statementId: "stmt-1", + }, + }, + }, + { + attachmentId: "viz-1", + attachment: { + $case: "viz", + viz: { title: "Franchises by Country", queryAttachmentId: "query-1" }, + }, + }, + ], +}; + +const expectedMessage = { + messageId: MSG, + conversationId: CONV, + spaceId: AGENT_ID, + status: "COMPLETED", + content: "Show franchises by country", + attachments: [ + { + attachmentId: "query-1", + query: { + title: "Franchises by country", + query: "SELECT country, count(*) FROM franchises GROUP BY 1", + statementId: "stmt-1", + }, + }, + { + attachmentId: "viz-1", + viz: { title: "Franchises by Country", queryAttachmentId: "query-1" }, + }, + { text: { content: "## Franchises by country\n\nJapan has the most." } }, + ], + reasoning: [ + "I'll chart franchises by country.", + "There are no sales figures.", + ], +}; + +function sse(events: object[]): Response { + return new Response( + events + .map( + (e) => + `event:${(e as { type: string }).type}\ndata:${JSON.stringify(e)}\n\n`, + ) + .join(""), + ); +} + +function createMockWorkspaceClient() { + return { + request: vi.fn(), + genie: { + genieGetConversationMessage: vi.fn().mockResolvedValue(projection), + genieListConversationMessages: vi.fn(), + genieGetMessageAttachmentQueryResult: vi.fn().mockResolvedValue({ + statementResponse: { + manifest: { schema: { columns: [] } }, + result: {}, + }, + }), + downloadMessageAttachmentVisualization: vi.fn(), + }, + } as any; +} + +async function collect( + gen: AsyncGenerator, +): Promise { + const events: GenieStreamEvent[] = []; + for await (const event of gen) events.push(event); + return events; +} + +const fullTurn = [ + created, + { + type: "response.output_item.added", + item: reasoningItem("I'll chart franchises by country."), + }, + done(reasoningItem("I'll chart franchises by country.")), + done(reasoningItem("There are no sales figures.")), + done(toolCall("execute_sql")), + done(toolCall("generate_visualization")), + done(answerItem), + completed, +]; + +describe("GenieConnector agent mode", () => { + describe("streamAgentResponse", () => { + test("maps an agent turn onto the chat-mode events plus reasoning", async () => { + const ws = createMockWorkspaceClient(); + ws.request.mockResolvedValue(sse(fullTurn)); + + const events = await collect( + new GenieConnector().streamAgentResponse( + ws, + AGENT_ID, + "Show franchises", + undefined, + ), + ); + + expect(events).toEqual([ + { + type: "message_start", + conversationId: CONV, + messageId: MSG, + spaceId: AGENT_ID, + }, + { type: "reasoning", content: "I'll chart franchises by country." }, + { type: "reasoning", content: "There are no sales figures." }, + { type: "status", status: "EXECUTING_QUERY" }, + { type: "status", status: "GENERATING_VISUALIZATION" }, + { type: "message_result", message: expectedMessage }, + expect.objectContaining({ + type: "query_result", + attachmentId: "query-1", + statementId: "stmt-1", + }), + ]); + expect(ws.genie.genieGetConversationMessage).toHaveBeenCalledWith( + { spaceId: AGENT_ID, conversationId: CONV, messageId: MSG }, + expect.anything(), + ); + }); + + test("starts a new conversation with viz enabled", async () => { + const ws = createMockWorkspaceClient(); + ws.request.mockResolvedValue(sse(fullTurn)); + + await collect( + new GenieConnector().streamAgentResponse(ws, AGENT_ID, "Hi", undefined), + ); + + const req = ws.request.mock.calls[0][0]; + expect(req).toMatchObject({ + method: "POST", + path: "/api/2.0/genie/agents/agent-1/responses", + }); + expect(req.headers.get("Accept")).toBe("text/event-stream"); + expect(JSON.parse(req.body)).toEqual({ + enable_viz: true, + input: [ + { + type: "message", + role: "user", + content: [{ type: "input_text", text: "Hi" }], + }, + ], + }); + }); + + test("continues an existing conversation", async () => { + const ws = createMockWorkspaceClient(); + ws.request.mockResolvedValue(sse(fullTurn)); + + await collect( + new GenieConnector().streamAgentResponse( + ws, + AGENT_ID, + "And now?", + CONV, + ), + ); + + expect(JSON.parse(ws.request.mock.calls[0][0].body)).toMatchObject({ + conversation_id: CONV, + }); + }); + + test("yields an error when the agent fails", async () => { + const ws = createMockWorkspaceClient(); + ws.request.mockResolvedValue( + sse([ + created, + { + type: "response.failed", + response: { + id: MSG, + conversation_id: CONV, + error: { message: "boom" }, + }, + }, + ]), + ); + + const events = await collect( + new GenieConnector().streamAgentResponse(ws, AGENT_ID, "Hi", undefined), + ); + + expect(events.map((e) => e.type)).toEqual(["message_start", "error"]); + expect(events[1]).toEqual({ type: "error", error: "boom" }); + }); + + test("yields an error when the stream ends before the turn completes", async () => { + const ws = createMockWorkspaceClient(); + ws.request.mockResolvedValue( + sse([created, done(reasoningItem("thinking"))]), + ); + + const events = await collect( + new GenieConnector().streamAgentResponse(ws, AGENT_ID, "Hi", undefined), + ); + + expect(events.at(-1)).toEqual({ + type: "error", + error: "Genie agent response ended before it completed", + }); + expect(ws.genie.genieGetConversationMessage).not.toHaveBeenCalled(); + }); + + test("surfaces the message of an agent API error body", async () => { + const ws = createMockWorkspaceClient(); + ws.request.mockRejectedValue( + Object.assign(new Error("Response from server (Not Found)"), { + response: { + error: { + type: "NOT_FOUND", + message: "NOT_FOUND: Space with id agent-1 not found", + }, + }, + }), + ); + + const events = await collect( + new GenieConnector().streamAgentResponse(ws, AGENT_ID, "Hi", undefined), + ); + + expect(events).toEqual([ + { type: "error", error: "NOT_FOUND: Space with id agent-1 not found" }, + ]); + }); + }); + + describe("history", () => { + // The items endpoint pages oldest first; a turn can span pages. + function mockItemPages(ws: ReturnType) { + ws.request + .mockResolvedValueOnce( + Response.json({ + data: [ + { + type: "message", + role: "user", + content: [ + { type: "input_text", text: "Show franchises by country" }, + ], + metadata: { message_id: MSG }, + }, + reasoningItem("I'll chart franchises by country."), + reasoningItem("There are no sales figures."), + ], + has_more: true, + last_id: "item-3", + }), + ) + .mockResolvedValueOnce( + Response.json({ + data: [toolCall("execute_sql"), answerItem], + has_more: false, + }), + ); + } + + test("streamConversation adds answer text and reasoning from the items", async () => { + const ws = createMockWorkspaceClient(); + ws.genie.genieListConversationMessages.mockResolvedValue({ + messages: [projection], + }); + mockItemPages(ws); + + const events = await collect( + new GenieConnector().streamConversation(ws, AGENT_ID, CONV, { + agent: true, + includeQueryResults: false, + }), + ); + + expect(events[0]).toEqual({ + type: "message_result", + message: expectedMessage, + }); + expect(ws.request).toHaveBeenCalledTimes(2); + expect(ws.request.mock.calls[0][0]).toMatchObject({ + method: "GET", + path: "/api/2.0/genie/agents/agent-1/conversations/conv-1/items", + query: { limit: "100" }, + }); + expect(ws.request.mock.calls[1][0].query).toEqual({ + limit: "100", + after: "item-3", + }); + }); + + test("streamGetMessage adds answer text and reasoning once the message is done", async () => { + const ws = createMockWorkspaceClient(); + mockItemPages(ws); + + const events = await collect( + new GenieConnector().streamGetMessage(ws, AGENT_ID, CONV, MSG, { + agent: true, + }), + ); + + expect(events.find((e) => e.type === "message_result")).toEqual({ + type: "message_result", + message: expectedMessage, + }); + }); + + test("getConversation adds answer text and reasoning", async () => { + const ws = createMockWorkspaceClient(); + ws.genie.genieListConversationMessages.mockResolvedValue({ + messages: [projection], + }); + mockItemPages(ws); + + const history = await new GenieConnector().getConversation( + ws, + AGENT_ID, + CONV, + { agent: true }, + ); + + expect(history.messages).toEqual([expectedMessage]); + }); + + test("chat-mode history doesn't read agent items", async () => { + const ws = createMockWorkspaceClient(); + ws.genie.genieListConversationMessages.mockResolvedValue({ + messages: [projection], + }); + + await collect( + new GenieConnector().streamConversation(ws, AGENT_ID, CONV, { + includeQueryResults: false, + }), + ); + + expect(ws.request).not.toHaveBeenCalled(); + }); + }); + + describe("downloadVisualization", () => { + test("requests the attachment's rendered PNG", async () => { + const ws = createMockWorkspaceClient(); + const png = new ReadableStream(); + ws.genie.downloadMessageAttachmentVisualization.mockResolvedValue({ + contents: png, + }); + + const result = await new GenieConnector().downloadVisualization( + ws, + AGENT_ID, + CONV, + MSG, + "viz-1_output", + ); + + expect(result).toBe(png); + expect( + ws.genie.downloadMessageAttachmentVisualization, + ).toHaveBeenCalledWith( + { + name: "spaces/agent-1/conversations/conv-1/messages/msg-1/attachments/viz-1_output", + }, + expect.anything(), + ); + }); + + test("rejects ids that could change the resource path", async () => { + const ws = createMockWorkspaceClient(); + + await expect( + new GenieConnector().downloadVisualization( + ws, + AGENT_ID, + CONV, + MSG, + "../../x", + ), + ).rejects.toThrow("Invalid Genie id"); + expect( + ws.genie.downloadMessageAttachmentVisualization, + ).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/packages/appkit/src/plugins/genie/genie.ts b/packages/appkit/src/plugins/genie/genie.ts index b90113473..960ca32ce 100644 --- a/packages/appkit/src/plugins/genie/genie.ts +++ b/packages/appkit/src/plugins/genie/genie.ts @@ -1,4 +1,6 @@ import { randomUUID } from "node:crypto"; +import { Readable } from "node:stream"; +import type { ReadableStream as NodeReadableStream } from "node:stream/web"; import type express from "express"; import type { @@ -9,7 +11,7 @@ import type { } from "shared"; import { z } from "zod"; -import { GenieConnector } from "../../connectors"; +import { GenieConnector, isGenieId } from "../../connectors"; import { getWorkspaceClient } from "../../context"; import { buildToolkitEntries } from "../../core/agent/build-toolkit"; import { @@ -54,7 +56,11 @@ export class GeniePlugin extends Plugin implements ToolProvider { }); const spaces = this.config.spaces ?? {}; - const missingAliases = Object.entries(spaces) + const agents = this.config.agents ?? {}; + const missingAliases = [ + ...Object.entries(spaces), + ...Object.entries(agents), + ] .filter(([, id]) => !id) .map(([alias]) => alias); if (missingAliases.length > 0) { @@ -70,7 +76,18 @@ export class GeniePlugin extends Plugin implements ToolProvider { ); } - for (const alias of Object.keys(spaces)) { + const sharedAliases = Object.keys(agents).filter((alias) => + Object.hasOwn(spaces, alias), + ); + if (sharedAliases.length > 0) { + throw new Error( + `GeniePlugin: ${sharedAliases.map((a) => `"${a}"`).join(", ")} ` + + "configured in both `spaces` and `agents`. An alias is either a chat space or an agent; " + + "use a different alias for the agent.", + ); + } + + for (const alias of [...Object.keys(spaces), ...Object.keys(agents)]) { Object.assign(this.tools, this._defineSpaceTools(alias)); } } @@ -127,7 +144,12 @@ export class GeniePlugin extends Plugin implements ToolProvider { } private resolveSpaceId(alias: string): string | null { - return this.config.spaces?.[alias] ?? null; + return this.config.spaces?.[alias] ?? this.config.agents?.[alias] ?? null; + } + + /** Whether the alias uses Genie agent mode (an `agents` entry). */ + private isAgent(alias: string): boolean { + return Object.hasOwn(this.config.agents ?? {}, alias); } /** @@ -166,6 +188,15 @@ export class GeniePlugin extends Plugin implements ToolProvider { await this._asUserScoped(req)._handleGetMessage(req, res); }, }); + + this.route(router, { + name: "getVisualization", + method: "get", + path: "/:alias/conversations/:conversationId/messages/:messageId/attachments/:attachmentId/visualization", + handler: async (req: express.Request, res: express.Response) => { + await this._asUserScoped(req)._handleGetVisualization(req, res); + }, + }); } async _handleSendMessage( @@ -216,13 +247,21 @@ export class GeniePlugin extends Plugin implements ToolProvider { await this.executeStream( res, (signal) => - this.genieConnector.streamSendMessage( - workspaceClient, - spaceId, - content, - conversationId, - { timeout, signal }, - ), + this.isAgent(alias) + ? this.genieConnector.streamAgentResponse( + workspaceClient, + spaceId, + content, + conversationId, + { signal }, + ) + : this.genieConnector.streamSendMessage( + workspaceClient, + spaceId, + content, + conversationId, + { timeout, signal }, + ), streamSettings, ); } @@ -272,7 +311,12 @@ export class GeniePlugin extends Plugin implements ToolProvider { workspaceClient, spaceId, conversationId, - { includeQueryResults, pageToken, signal }, + { + includeQueryResults, + pageToken, + signal, + agent: this.isAgent(alias), + }, ), streamSettings, ); @@ -325,12 +369,69 @@ export class GeniePlugin extends Plugin implements ToolProvider { spaceId, conversationId, messageId, - { timeout, signal }, + { timeout, signal, agent: this.isAgent(alias) }, ), streamSettings, ); } + /** Streams the PNG Genie rendered for a visualization attachment, so the UI can show it in an ``. */ + async _handleGetVisualization( + req: express.Request, + res: express.Response, + ): Promise { + const { alias, conversationId, messageId, attachmentId } = req.params; + const spaceId = this.resolveSpaceId(alias); + + if (!spaceId) { + res.status(404).json({ error: `Unknown space alias: ${alias}` }); + return; + } + if (![conversationId, messageId, attachmentId].every(isGenieId)) { + res + .status(400) + .json({ error: "Invalid conversation, message or attachment id" }); + return; + } + + const workspaceClient = getWorkspaceClient(); + const result = await this.execute( + (signal) => + this.genieConnector.downloadVisualization( + workspaceClient, + spaceId, + conversationId, + messageId, + attachmentId, + signal, + ), + { default: genieStreamDefaults.default }, + ); + + if (!result.ok) { + // The SDK leaves the message empty for Genie's `{"error": {...}}` bodies. + res + .status(result.status) + .json({ error: result.message || "Visualization not available" }); + return; + } + + res.setHeader("Content-Type", "image/png"); + res.setHeader("X-Content-Type-Options", "nosniff"); + // The chart of a completed message never changes. + res.setHeader("Cache-Control", "private, max-age=3600"); + if (!result.data) { + res.end(); + return; + } + const image = Readable.fromWeb(result.data as NodeReadableStream); + image.on("error", (error) => { + logger.error("Visualization stream error: %O", error); + res.destroy(); + }); + image.pipe(res); + } + async getConversation( alias: string, conversationId: string, @@ -353,6 +454,7 @@ export class GeniePlugin extends Plugin implements ToolProvider { workspaceClient, spaceId, conversationId, + { agent: this.isAgent(alias) }, ); } @@ -372,6 +474,16 @@ export class GeniePlugin extends Plugin implements ToolProvider { throw new Error(`Unknown space alias: ${alias}`); } const workspaceClient = getWorkspaceClient(); + if (this.isAgent(alias)) { + yield* this.genieConnector.streamAgentResponse( + workspaceClient, + spaceId, + content, + conversationId, + { signal: options?.signal }, + ); + return; + } const timeout = options?.timeout ?? this.config.timeout ?? 120_000; yield* this.genieConnector.streamSendMessage( workspaceClient, diff --git a/packages/appkit/src/plugins/genie/manifest.json b/packages/appkit/src/plugins/genie/manifest.json index bc42be21c..9442c2b9a 100644 --- a/packages/appkit/src/plugins/genie/manifest.json +++ b/packages/appkit/src/plugins/genie/manifest.json @@ -47,6 +47,13 @@ "type": "string" } }, + "agents": { + "type": "object", + "description": "Map of alias names to Genie Space IDs used in agent mode (streamed reasoning and PNG visualizations)", + "additionalProperties": { + "type": "string" + } + }, "timeout": { "type": "number", "default": 120000, diff --git a/packages/appkit/src/plugins/genie/tests/genie.test.ts b/packages/appkit/src/plugins/genie/tests/genie.test.ts index 09b26d4c1..c01a21a10 100644 --- a/packages/appkit/src/plugins/genie/tests/genie.test.ts +++ b/packages/appkit/src/plugins/genie/tests/genie.test.ts @@ -1,3 +1,5 @@ +import { PassThrough } from "node:stream"; + import { createMockRequest, createMockResponse, @@ -78,6 +80,7 @@ function createMockGenieService() { ); const genieListConversationMessages = vi.fn(); + const downloadMessageAttachmentVisualization = vi.fn(); return { genieStartConversation, @@ -85,6 +88,7 @@ function createMockGenieService() { genieGetConversationMessage, genieGetMessageAttachmentQueryResult, genieListConversationMessages, + downloadMessageAttachmentVisualization, register, }; } @@ -93,6 +97,8 @@ describe("Genie Plugin", () => { let config: IGenieConfig; let serviceContextMock: Awaited>; let mockGenieService: ReturnType; + // Raw workspace requests (the agent-mode endpoints the SDK doesn't cover). + let mockRequest: ReturnType; beforeEach(async () => { config = { @@ -127,9 +133,11 @@ describe("Genie Plugin", () => { }, }); + mockRequest = vi.fn(); serviceContextMock = await mockServiceContext({ userDatabricksClient: { genie: mockGenieService, + request: mockRequest, }, }); }); @@ -161,7 +169,7 @@ describe("Genie Plugin", () => { expect.any(Function), ); - expect(router.get).toHaveBeenCalledTimes(2); + expect(router.get).toHaveBeenCalledTimes(3); expect(router.get).toHaveBeenCalledWith( "/:alias/conversations/:conversationId", expect.any(Function), @@ -170,6 +178,10 @@ describe("Genie Plugin", () => { "/:alias/conversations/:conversationId/messages/:messageId", expect.any(Function), ); + expect(router.get).toHaveBeenCalledWith( + "/:alias/conversations/:conversationId/messages/:messageId/attachments/:attachmentId/visualization", + expect.any(Function), + ); }); }); @@ -1390,4 +1402,260 @@ describe("Genie Plugin", () => { ); }); }); + describe("agent mode", () => { + const userHeaders = { + "x-forwarded-access-token": "user-token", + "x-forwarded-user": "user-1", + }; + + function agentTurn(): Response { + const events = [ + { + type: "response.created", + response: { id: "agent-msg", conversation_id: "agent-conv" }, + }, + { + type: "response.output_item.done", + item: { + type: "reasoning", + content: [{ type: "reasoning_text", text: "Counting franchises." }], + }, + }, + { + type: "response.output_item.done", + item: { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "Japan has the most." }], + }, + }, + { type: "response.completed", response: { id: "agent-msg" } }, + ]; + return new Response( + events.map((e) => `data:${JSON.stringify(e)}\n\n`).join(""), + ); + } + + test("defines tools for agent aliases", () => { + const plugin = new GeniePlugin({ + ...config, + agents: { analyst: "agent-space-id" }, + }); + const names = plugin.getAgentTools().map((t) => t.name); + + expect(names).toContain("analyst.sendMessage"); + expect(names).toContain("analyst.getConversation"); + expect(names).toHaveLength(6); + }); + + test("throws at construction when an alias is both a space and an agent", () => { + expect( + () => + new GeniePlugin({ ...config, agents: { myspace: "agent-space-id" } }), + ).toThrow(/"myspace" configured in both `spaces` and `agents`/); + }); + + test("throws at construction when an agent id is missing", () => { + expect( + () => new GeniePlugin({ ...config, agents: { analyst: undefined } }), + ).toThrow(/alias "analyst" was configured with a missing Genie Space ID/); + }); + + test("sends messages to the agent responses endpoint", async () => { + const plugin = new GeniePlugin({ + ...config, + agents: { analyst: "agent-space-id" }, + }); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + mockRequest.mockResolvedValue(agentTurn()); + mockGenieService.register("agent-conv", "agent-msg"); + + const mockRes = createMockResponse(); + await getHandler("POST", "/:alias/messages")( + createMockRequest({ + params: { alias: "analyst" }, + body: { content: "Franchises by country?" }, + headers: userHeaders, + }), + mockRes, + ); + + expect(mockGenieService.genieStartConversation).not.toHaveBeenCalled(); + expect(mockRequest).toHaveBeenCalledWith( + expect.objectContaining({ + method: "POST", + path: "/api/2.0/genie/agents/agent-space-id/responses", + }), + ); + const stream = expectStream(mockRes); + await stream.toEmit("message_start", "reasoning", "message_result"); + const events = await stream.collect(); + expect(events.find((e) => e.type === "message_result")).toMatchObject({ + message: { + reasoning: ["Counting franchises."], + attachments: [{ text: { content: "Japan has the most." } }], + }, + }); + }); + + test("reads agent conversation history with the items endpoint", async () => { + const plugin = new GeniePlugin({ + ...config, + agents: { analyst: "agent-space-id" }, + }); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + mockGenieService.genieListConversationMessages.mockResolvedValue({ + messages: [], + }); + mockRequest.mockResolvedValue(Response.json({ data: [] })); + + await getHandler("GET", "/:alias/conversations/:conversationId")( + createMockRequest({ + params: { alias: "analyst", conversationId: "agent-conv" }, + headers: userHeaders, + }), + createMockResponse(), + ); + + expect(mockRequest).toHaveBeenCalledWith( + expect.objectContaining({ + path: "/api/2.0/genie/agents/agent-space-id/conversations/agent-conv/items", + }), + ); + }); + + test("chat aliases keep using the chat API", async () => { + const plugin = new GeniePlugin({ + ...config, + agents: { analyst: "agent-space-id" }, + }); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + await getHandler("POST", "/:alias/messages")( + createMockRequest({ + params: { alias: "myspace" }, + body: { content: "Top customers?" }, + headers: userHeaders, + }), + createMockResponse(), + ); + + expect(mockGenieService.genieStartConversation).toHaveBeenCalled(); + expect(mockRequest).not.toHaveBeenCalled(); + }); + }); + + describe("getVisualization", () => { + const vizPath = + "/:alias/conversations/:conversationId/messages/:messageId/attachments/:attachmentId/visualization"; + + function vizRequest(params: Record) { + return createMockRequest({ + params: { + alias: "myspace", + conversationId: "conv-1", + messageId: "msg-1", + attachmentId: "viz-1_output", + ...params, + }, + headers: { + "x-forwarded-access-token": "user-token", + "x-forwarded-user": "user-1", + }, + }); + } + + // A real Writable, so the handler can pipe the PNG into it. + function streamResponse() { + const res: any = new PassThrough(); + res.status = vi.fn().mockReturnValue(res); + res.json = vi.fn().mockReturnValue(res); + res.setHeader = vi.fn().mockReturnValue(res); + return res; + } + + async function readAll(res: PassThrough): Promise { + const chunks: Buffer[] = []; + for await (const chunk of res) chunks.push(chunk as Buffer); + return Buffer.concat(chunks); + } + + test("streams the PNG Genie rendered for the attachment", async () => { + const plugin = new GeniePlugin(config); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + const png = Buffer.from([0x89, 0x50, 0x4e, 0x47]); + mockGenieService.downloadMessageAttachmentVisualization.mockResolvedValue( + { + contents: new Response(png).body, + }, + ); + const res = streamResponse(); + + await getHandler("GET", vizPath)(vizRequest({}), res); + + expect(await readAll(res)).toEqual(png); + expect(res.setHeader).toHaveBeenCalledWith("Content-Type", "image/png"); + expect(res.setHeader).toHaveBeenCalledWith( + "X-Content-Type-Options", + "nosniff", + ); + expect( + mockGenieService.downloadMessageAttachmentVisualization, + ).toHaveBeenCalledWith( + { + name: "spaces/test-space-id/conversations/conv-1/messages/msg-1/attachments/viz-1_output", + }, + expect.anything(), + ); + }); + + test("returns 400 for ids that could change the resource path", async () => { + const plugin = new GeniePlugin(config); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + const res = streamResponse(); + + await getHandler("GET", vizPath)( + vizRequest({ attachmentId: "../../x" }), + res, + ); + + expect(res.status).toHaveBeenCalledWith(400); + expect( + mockGenieService.downloadMessageAttachmentVisualization, + ).not.toHaveBeenCalled(); + }); + + test("returns 404 for an unknown alias", async () => { + const plugin = new GeniePlugin(config); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + const res = streamResponse(); + + await getHandler("GET", vizPath)(vizRequest({ alias: "unknown" }), res); + + expect(res.status).toHaveBeenCalledWith(404); + }); + + test("passes on Genie's status when the attachment isn't a visualization", async () => { + const plugin = new GeniePlugin(config); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + mockGenieService.downloadMessageAttachmentVisualization.mockRejectedValue( + Object.assign(new Error(""), { statusCode: 404 }), + ); + const res = streamResponse(); + + await getHandler("GET", vizPath)(vizRequest({}), res); + + expect(res.status).toHaveBeenCalledWith(404); + expect(res.json).toHaveBeenCalledWith({ + error: "Visualization not available", + }); + }); + }); }); diff --git a/packages/appkit/src/plugins/genie/types.ts b/packages/appkit/src/plugins/genie/types.ts index ef1ed4c03..2e78aa962 100644 --- a/packages/appkit/src/plugins/genie/types.ts +++ b/packages/appkit/src/plugins/genie/types.ts @@ -15,6 +15,13 @@ export interface IGenieConfig extends BasePluginConfig { * fails fast rather than masquerading as an unknown-alias 404 at request time. */ spaces?: Record; + /** + * Map of alias → Genie Space ID for spaces used in agent mode. Agent mode + * streams the agent's reasoning and returns its charts as PNG + * visualizations. Any Genie space can be used as an agent. Aliases share one + * namespace with `spaces`, so an alias can't be in both maps. + */ + agents?: Record; /** Genie polling timeout in ms. Set to 0 for indefinite. Default: 120000 (2 min) */ timeout?: number; } diff --git a/packages/shared/src/genie.ts b/packages/shared/src/genie.ts index 9ea830a55..885908392 100644 --- a/packages/shared/src/genie.ts +++ b/packages/shared/src/genie.ts @@ -7,6 +7,8 @@ export type GenieStreamEvent = spaceId: string; } | { type: "status"; status: string } + /** Agent mode only: one step of the agent's reasoning, sent as it happens. */ + | { type: "reasoning"; content: string } | { type: "message_result"; message: GenieMessageResponse } | { type: "query_result"; @@ -57,6 +59,8 @@ export interface GenieMessageResponse { content: string; attachments?: GenieAttachmentResponse[]; error?: string; + /** Agent mode only: the agent's reasoning steps, in order. */ + reasoning?: string[]; } export interface GenieAttachmentResponse { @@ -69,4 +73,13 @@ export interface GenieAttachmentResponse { }; text?: { content?: string }; suggestedQuestions?: string[]; + /** + * A chart Genie rendered for a query. Fetch it as a PNG from the plugin's + * `visualization` route with this attachment's `attachmentId`. + */ + viz?: { + title?: string; + /** The query attachment the chart was generated from. */ + queryAttachmentId?: string; + }; } From 439a47651ef9d1e0e23b550e24cf98b34b788160 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Thu, 12 Nov 2026 17:33:59 +0000 Subject: [PATCH 2/4] feat(appkit): show Genie's charts for chat spaces too The Genie chat API also returns `viz` attachments when Genie draws a chart, and the genie plugin dropped them. Map them in chat mode as well, so chat spaces and agents behave the same: the UI loads Genie's PNG through the visualization route and shows it in place of the inferred chart for queries that have one. Answers without a chart are unchanged. Co-authored-by: Isaac Signed-off-by: Fabian Jakobs --- docs/docs/plugins/genie.md | 5 +- .../appkit/src/connectors/genie/client.ts | 29 +++------- .../src/connectors/genie/tests/client.test.ts | 55 +++++++++++++++++++ 3 files changed, 67 insertions(+), 22 deletions(-) diff --git a/docs/docs/plugins/genie.md b/docs/docs/plugins/genie.md index 63552e307..b7a3b4e45 100644 --- a/docs/docs/plugins/genie.md +++ b/docs/docs/plugins/genie.md @@ -11,6 +11,7 @@ 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 @@ -64,7 +65,7 @@ genie({ 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`. -- Charts arrive as `viz` attachments. The [visualization endpoint](#get-a-visualization) serves each one as a PNG, and `GenieChat` shows it in place of the chart it would otherwise infer from the query result. +- 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. @@ -133,6 +134,8 @@ GET /api/genie/:alias/conversations/:conversationId/messages/:messageId/attachme 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/src/connectors/genie/client.ts b/packages/appkit/src/connectors/genie/client.ts index d88c7ede1..60695d02e 100644 --- a/packages/appkit/src/connectors/genie/client.ts +++ b/packages/appkit/src/connectors/genie/client.ts @@ -48,10 +48,7 @@ interface GenieConnectorConfig { maxMessages?: number; } -function mapAttachments( - message: GenieMessage, - includeViz: boolean, -): GenieAttachmentResponse[] { +function mapAttachments(message: GenieMessage): GenieAttachmentResponse[] { return ( message.attachments?.map(({ attachmentId, attachment }) => ({ attachmentId, @@ -73,7 +70,7 @@ function mapAttachments( ? attachment.suggestedQuestions.questions : undefined, viz: - includeViz && attachment?.$case === "viz" + attachment?.$case === "viz" ? { title: attachment.viz.title, queryAttachmentId: attachment.viz.queryAttachmentId, @@ -83,17 +80,14 @@ function mapAttachments( ); } -function toMessageResponse( - message: GenieMessage, - includeViz = false, -): GenieMessageResponse { +function toMessageResponse(message: GenieMessage): GenieMessageResponse { return { messageId: message.messageId ?? "", conversationId: message.conversationId ?? "", spaceId: message.spaceId ?? "", status: message.status ?? "COMPLETED", content: message.content ?? "", - attachments: mapAttachments(message, includeViz), + attachments: mapAttachments(message), error: message.error?.error, }; } @@ -239,7 +233,7 @@ export class GenieConnector { workspaceClient: WorkspaceClient, spaceId: string, conversationId: string, - options?: { pageSize?: number; pageToken?: string; includeViz?: boolean }, + options?: { pageSize?: number; pageToken?: string }, ): Promise<{ messages: GenieMessageResponse[]; nextPageToken: string | null; @@ -254,9 +248,7 @@ export class GenieConnector { ...(options?.pageToken ? { pageToken: options.pageToken } : {}), }); - const messages = (response.messages ?? []) - .reverse() - .map((message) => toMessageResponse(message, options?.includeViz)); + const messages = (response.messages ?? []).reverse().map(toMessageResponse); return { messages, @@ -414,10 +406,7 @@ export class GenieConnector { { spaceId: agentId, ...started }, { signal }, ); - const messageResponse = withAgentTurn( - toMessageResponse(message, true), - turn, - ); + const messageResponse = withAgentTurn(toMessageResponse(message), turn); yield { type: "message_result", message: messageResponse }; yield* this.emitQueryResults( workspaceClient, @@ -575,7 +564,6 @@ export class GenieConnector { { pageSize: options?.pageSize, pageToken: options?.pageToken, - includeViz: options?.agent, }, ); const { nextPageToken } = page; @@ -707,7 +695,7 @@ export class GenieConnector { const isTerminal = message.status === "COMPLETED" || message.status === "FAILED"; if (isTerminal) { - let messageResponse = toMessageResponse(message, options?.agent); + let messageResponse = toMessageResponse(message); if (options?.agent) { [messageResponse] = await this.withAgentTurns( workspaceClient, @@ -791,7 +779,6 @@ export class GenieConnector { { pageSize: genieConnectorDefaults.pageSize, pageToken, - includeViz: options?.agent, }, ); allMessages.push(...messages); diff --git a/packages/appkit/src/connectors/genie/tests/client.test.ts b/packages/appkit/src/connectors/genie/tests/client.test.ts index 00deb2998..f4b685446 100644 --- a/packages/appkit/src/connectors/genie/tests/client.test.ts +++ b/packages/appkit/src/connectors/genie/tests/client.test.ts @@ -732,6 +732,61 @@ describe("GenieConnector", () => { }); }); + test("handles viz attachments", async () => { + const msg = makeGenieMessage({ + attachments: [ + { + attachmentId: "att-viz", + attachment: { + $case: "viz", + viz: { title: "Sales by region", queryAttachmentId: "att-1" }, + }, + }, + ], + }); + mockStart(ws, msg); + + const events = await collect( + connector.streamSendMessage(ws, "space-1", "q", undefined), + ); + + const msgResult = events.find((e) => e.type === "message_result") as any; + expect(msgResult.message.attachments[0]).toEqual({ + attachmentId: "att-viz", + viz: { title: "Sales by region", queryAttachmentId: "att-1" }, + }); + }); + + test("maps viz attachments in conversation history", async () => { + ws.genie.genieListConversationMessages.mockResolvedValue({ + messages: [ + makeGenieMessage({ + attachments: [ + { + attachmentId: "att-viz", + attachment: { + $case: "viz", + viz: { title: "Sales by region", queryAttachmentId: "att-1" }, + }, + }, + ], + }), + ], + }); + + const events = await collect( + connector.streamConversation(ws, "space-1", "conv-1", { + includeQueryResults: false, + }), + ); + + const msgResult = events.find((e) => e.type === "message_result") as any; + expect(msgResult.message.attachments[0].viz).toEqual({ + title: "Sales by region", + queryAttachmentId: "att-1", + }); + }); + test("returns empty array when message has no attachments", async () => { const msg = makeGenieMessage({ attachments: undefined }); mockStart(ws, msg); From bd169e00959bc47443fc1ad4599d8bdeeac988e6 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Fri, 13 Nov 2026 09:11:35 +0000 Subject: [PATCH 3/4] fix(appkit-ui): keep the Genie result table when Genie draws a chart For a query with a Genie chart, GenieChatMessage replaced the whole GenieQueryVisualization with the PNG, so the Table tab and its rows were gone for every chart answer, in chat mode too. GenieQueryVisualization now takes an optional chartImageUrl: the Chart tab shows that image and the Table tab is unchanged. If the image fails to load, it falls back to the inferred chart. Charts whose query result didn't load are still shown as plain images. Co-authored-by: Isaac Signed-off-by: Fabian Jakobs --- .../genie/GenieQueryVisualization.mdx | 5 +- .../__tests__/genie-chat-message.test.tsx | 39 +++++++++-- .../genie-query-visualization.test.tsx | 66 +++++++++++++++++++ .../src/react/genie/genie-chat-message.tsx | 26 +++++--- .../react/genie/genie-query-visualization.tsx | 59 ++++++++++++----- 5 files changed, 165 insertions(+), 30 deletions(-) create mode 100644 packages/appkit-ui/src/react/genie/__tests__/genie-query-visualization.test.tsx 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/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 index d25855fb3..862d24d57 100644 --- 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 @@ -2,7 +2,9 @@ import { cleanup, render, screen } from "@testing-library/react"; import { afterEach, describe, expect, test, vi } from "vitest"; vi.mock("../genie-query-visualization", () => ({ - GenieQueryVisualization: () =>
, + GenieQueryVisualization: (props: { chartImageUrl?: string }) => ( +
+ ), })); import { GenieChatMessage } from "../genie-chat-message"; @@ -28,7 +30,7 @@ function assistant(overrides: Partial): GenieMessageItem { describe("GenieChatMessage", () => { afterEach(cleanup); - test("shows Genie's chart instead of the inferred one for a query with a visualization", () => { + 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("inferred-chart")).toBeNull(); + expect(screen.queryByTestId("query-result")).toBeNull(); }); - test("keeps the inferred chart for a query without a visualization", () => { + test("leaves the query result alone for a query without a visualization", () => { render( { />, ); - expect(screen.getByTestId("inferred-chart")).toBeTruthy(); + const result = screen.getByTestId("query-result"); + expect(result.getAttribute("data-chart-image")).toBeNull(); expect(screen.queryByRole("img")).toBeNull(); }); 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/genie-chat-message.tsx b/packages/appkit-ui/src/react/genie/genie-chat-message.tsx index c7f4bb0e7..17f445449 100644 --- a/packages/appkit-ui/src/react/genie/genie-chat-message.tsx +++ b/packages/appkit-ui/src/react/genie/genie-chat-message.tsx @@ -149,6 +149,11 @@ export function GenieChatMessage({ const queryVisualizations = visualizations.filter( (viz) => 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 (
@@ -171,15 +176,18 @@ export function GenieChatMessage({
- {queryVisualizations.length > 0 - ? queryVisualizations.map((viz) => ( - - )) - : queryResult != null && ( - - - - )} + {queryResult != null && ( + + + + )} + {standaloneVisualizations.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 && (