Skip to content

Commit b074ca7

Browse files
hubertzub-dbhubertzub-db
committed
Chat plugin
Signed-off-by: Hubert Zub <hubert.zub@databricks.com>
1 parent b49eed7 commit b074ca7

33 files changed

+4121
-35
lines changed

‎integrations/appkit-agent/README.md‎

Lines changed: 122 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
# @databricks/appkit-agent
22

3-
Agent plugin for [Databricks AppKit](https://github.com/databricks/appkit). You can define an agent using one of the following approaches:
3+
Plugins for [Databricks AppKit](https://github.com/databricks/appkit):
44

5-
1. Declaratively define an agent by specifying `model`, `tools,` and `instructions`
6-
2. Implement a custom agent loop using **`AgentInterface`** — a contract for writing custom agent implementations that speak the OpenAI Responses API format (streaming + non-streaming).
5+
- **Agent plugin** — define an AI agent declaratively (model + tools + instructions) or bring a custom `AgentInterface` implementation. Speaks the OpenAI Responses API format with streaming.
6+
- **Chat plugin** — full-featured chat server with streaming, session management, optional PostgreSQL persistence, feedback via MLflow Traces, and stream resumption.
77

88
## Installation
99

@@ -23,26 +23,52 @@ If you use hosted tools (Genie, Vector Search, custom/external MCP servers):
2323
npm install @langchain/mcp-adapters
2424
```
2525

26+
For chat persistence (optional — without it the chat plugin runs in ephemeral mode):
27+
28+
```bash
29+
npm install pg
30+
```
31+
2632
## Quick Start
2733

34+
### Agent only
35+
2836
```typescript
2937
import { createApp, server } from "@databricks/appkit";
3038
import { agent } from "@databricks/appkit-agent";
3139

3240
const app = await createApp({
3341
plugins: [
34-
server(),
42+
server({ autoStart: true }),
3543
agent({
3644
model: "databricks-claude-sonnet-4-5",
3745
systemPrompt: "You are a helpful assistant.",
3846
}),
3947
],
4048
});
49+
```
4150

42-
app.server.start();
51+
### Agent + Chat
52+
53+
```typescript
54+
import { createApp, server } from "@databricks/appkit";
55+
import { agent, chat } from "@databricks/appkit-agent";
56+
57+
await createApp({
58+
plugins: [
59+
server({ autoStart: true }),
60+
agent({
61+
model: "databricks-claude-sonnet-4-5",
62+
systemPrompt: "You are a helpful assistant.",
63+
}),
64+
chat({
65+
backend: "agent",
66+
}),
67+
],
68+
});
4369
```
4470

45-
The plugin registers `POST /api/agent` which accepts the [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) request format with SSE streaming.
71+
The agent plugin registers `POST /api/agent` (OpenAI Responses API format with SSE streaming). The chat plugin registers routes under `/api/chat/` for streaming chat, history, feedback, and more.
4672

4773
## Environment Variables
4874

@@ -216,23 +242,105 @@ agent({ agentInstance: new MyAgent() });
216242

217243
The `StandardAgent` class (exported from this package) is the built-in implementation used when you pass `model` instead of `agentInstance`. It translates the underlying agent's stream events into Responses API format.
218244

245+
## Chat Plugin
246+
247+
The chat plugin provides a streaming chat server that can be wired to the agent plugin or to a remote Databricks serving endpoint.
248+
249+
### Configuration
250+
251+
```typescript
252+
import { chat } from "@databricks/appkit-agent";
253+
254+
chat({
255+
// Route chat through the local agent plugin
256+
backend: "agent",
257+
258+
// Or point to a remote proxy / serving endpoint:
259+
// backend: { proxy: "http://localhost:8000/invocations" },
260+
// backend: { endpoint: "databricks-claude-sonnet-4-5" },
261+
262+
// Optional: PostgreSQL pool for persistent chat history
263+
// pool: new Pool(),
264+
265+
// Auto-create the ai_chatbot schema and tables on startup (default: false)
266+
// autoMigrate: true,
267+
268+
// Enable thumbs up/down feedback (default: !!process.env.MLFLOW_EXPERIMENT_ID)
269+
// feedbackEnabled: true,
270+
271+
// Custom session resolver (default: x-forwarded-user headers → SCIM → env)
272+
// getSession: (req) => ({ user: { id: req.headers["x-user-id"] } }),
273+
});
274+
```
275+
276+
### Persistence
277+
278+
Without a `pool`, the chat plugin runs in **ephemeral mode** — conversations are not saved. To enable persistent chat history, pass a `pg.Pool`:
279+
280+
```typescript
281+
import { Pool } from "pg";
282+
283+
chat({
284+
backend: "agent",
285+
pool: new Pool({ connectionString: "postgres://..." }),
286+
autoMigrate: true,
287+
});
288+
```
289+
290+
When `autoMigrate: true` is set, the plugin creates the `ai_chatbot` schema and tables (`Chat`, `Message`, `Vote`) on startup using [Drizzle migrations](https://orm.drizzle.team/docs/migrations). Migration SQL is generated from the Drizzle schema definition, so it stays in sync automatically.
291+
292+
### Session Resolution
293+
294+
In production on Databricks Apps, sessions are resolved from headers set by the platform proxy (`x-forwarded-user`, `x-forwarded-email`). In local development, the plugin falls back to the SCIM `/Me` API (if Databricks auth is configured) or `process.env.USER`.
295+
296+
You can override this entirely with a custom `getSession` callback.
297+
298+
### Feedback
299+
300+
When feedback is enabled, thumbs up/down votes are submitted as assessments to the [MLflow Traces API](https://docs.databricks.com/en/mlflow/llm-tracing.html). Votes are also persisted in the database when a pool is configured.
301+
302+
### Chat API Routes
303+
304+
All routes are mounted under `/api/chat/`.
305+
306+
| Method | Path | Auth | Description |
307+
| -------- | ------------------------ | ------------ | ---------------------------------- |
308+
| `GET` | `/config` | — | Feature flags (history, feedback) |
309+
| `GET` | `/session` | — | Current user session |
310+
| `GET` | `/history` | required | Paginated chat list |
311+
| `GET` | `/messages/:id` | required+ACL | Messages for a chat |
312+
| `DELETE` | `/messages/:id/trailing` | required | Delete messages after a given one |
313+
| `POST` | `/feedback` | required | Submit thumbs up/down |
314+
| `GET` | `/feedback/chat/:chatId` | required+ACL | Votes for a chat |
315+
| `POST` | `/title` | required | Auto-generate title from message |
316+
| `PATCH` | `/:id/visibility` | required+ACL | Toggle public/private |
317+
| `GET` | `/:id/stream` | required | Resume an active stream |
318+
| `GET` | `/:id` | required+ACL | Get single chat |
319+
| `DELETE` | `/:id` | required+ACL | Delete a chat |
320+
| `POST` | `/` | required | Main chat handler (streaming) |
321+
219322
## API Reference
220323

221324
### Exports
222325

223-
| Export | Kind | Description |
224-
| --------------------- | -------------- | -------------------------------------------------------- |
225-
| `agent` | Plugin factory | Main entry point — call with config, pass to `createApp` |
226-
| `StandardAgent` | Class | Built-in `AgentInterface` implementation |
227-
| `createInvokeHandler` | Function | Express handler factory for the `/api/agent` endpoint |
228-
| `isFunctionTool` | Function | Type guard for `FunctionTool` |
229-
| `isHostedTool` | Function | Type guard for `HostedTool` |
326+
| Export | Kind | Description |
327+
| --------------------- | -------------- | ---------------------------------------------------------- |
328+
| `agent` | Plugin factory | Agent plugin — call with config, pass to `createApp` |
329+
| `chat` | Plugin factory | Chat plugin — call with config, pass to `createApp` |
330+
| `ChatPlugin` | Class | Chat plugin class (for `ChatPlugin.staticAssetsPath`, etc) |
331+
| `StandardAgent` | Class | Built-in `AgentInterface` implementation |
332+
| `createInvokeHandler` | Function | Express handler factory for the `/api/agent` endpoint |
333+
| `isFunctionTool` | Function | Type guard for `FunctionTool` |
334+
| `isHostedTool` | Function | Type guard for `HostedTool` |
230335

231336
### Types
232337

233338
| Type | Description |
234339
| --------------------- | ------------------------------------------------------------- |
235-
| `IAgentConfig` | Plugin configuration options |
340+
| `IAgentConfig` | Agent plugin configuration options |
341+
| `ChatConfig` | Chat plugin configuration options |
342+
| `ChatSession` | Session object with user info |
343+
| `GetSession` | Custom session resolver function type |
236344
| `AgentInterface` | Contract for custom agent implementations |
237345
| `AgentTool` | Union of `FunctionTool \| HostedTool` |
238346
| `FunctionTool` | Local tool with JSON Schema parameters and `execute` handler |

‎integrations/appkit-agent/drizzle.config.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
import { defineConfig } from "drizzle-kit";
2+
3+
export default defineConfig({
4+
schema: "./src/chat-plugin/schema.ts",
5+
out: "./drizzle",
6+
dialect: "postgresql",
7+
});

‎integrations/appkit-agent/drizzle/0000_lying_bill_hollister.sql‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
CREATE SCHEMA IF NOT EXISTS "ai_chatbot";
2+
--> statement-breakpoint
3+
CREATE TABLE "ai_chatbot"."Chat" (
4+
"id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL,
5+
"createdAt" timestamp NOT NULL,
6+
"title" text NOT NULL,
7+
"userId" text NOT NULL,
8+
"visibility" varchar DEFAULT 'private' NOT NULL,
9+
"lastContext" jsonb
10+
);
11+
--> statement-breakpoint
12+
CREATE TABLE "ai_chatbot"."Message" (
13+
"id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL,
14+
"chatId" uuid NOT NULL,
15+
"role" varchar NOT NULL,
16+
"parts" json NOT NULL,
17+
"attachments" json NOT NULL,
18+
"createdAt" timestamp NOT NULL,
19+
"traceId" text
20+
);
21+
--> statement-breakpoint
22+
CREATE TABLE "ai_chatbot"."Vote" (
23+
"chatId" uuid NOT NULL,
24+
"messageId" uuid NOT NULL,
25+
"isUpvoted" boolean NOT NULL,
26+
CONSTRAINT "Vote_chatId_messageId_pk" PRIMARY KEY("chatId","messageId")
27+
);
28+
--> statement-breakpoint
29+
ALTER TABLE "ai_chatbot"."Message" ADD CONSTRAINT "Message_chatId_Chat_id_fk" FOREIGN KEY ("chatId") REFERENCES "ai_chatbot"."Chat"("id") ON DELETE no action ON UPDATE no action;--> statement-breakpoint
30+
ALTER TABLE "ai_chatbot"."Vote" ADD CONSTRAINT "Vote_chatId_Chat_id_fk" FOREIGN KEY ("chatId") REFERENCES "ai_chatbot"."Chat"("id") ON DELETE no action ON UPDATE no action;--> statement-breakpoint
31+
ALTER TABLE "ai_chatbot"."Vote" ADD CONSTRAINT "Vote_messageId_Message_id_fk" FOREIGN KEY ("messageId") REFERENCES "ai_chatbot"."Message"("id") ON DELETE no action ON UPDATE no action;

0 commit comments

Comments
 (0)