feat(appkit): add AgentPlugin for LangChain/LangGraph agents#166
feat(appkit): add AgentPlugin for LangChain/LangGraph agents#166hubertzub-db wants to merge 2 commits intodatabricks:mainfrom
Conversation
hubertzub-db
force-pushed
the
feat/agent-plugin
branch
12 times, most recently
from
51b8ebb to
e7b5ba3
Compare
| </p> | ||
| </div> | ||
|
|
||
| <AgentChat |
There was a problem hiding this comment.
N/B: we're keeping a minimal UI for now just to be able to chat, until we port the "full" chat UI as a separate plugin
| } | ||
|
|
||
| /** Renders a single chat message bubble (user or assistant with parts). */ | ||
| export function AgentChatMessage({ |
There was a problem hiding this comment.
N/B: initially I aimed to reuse UI bits from genie plugin, but it's probably not worth the effort since we want to move a fully featured chat UI as a separate plugin soon
hubertzub-db
force-pushed
the
feat/agent-plugin
branch
2 times, most recently
from
fd1bffc to
ba9ad8d
Compare
| case "user": | ||
| return new HumanMessage(content); | ||
| case "system": | ||
| return new SystemMessage(content); |
There was a problem hiding this comment.
Curious if there is a reason we dropped the "assistant" case here?
case "assistant":
return { role: "assistant", content } as any;
There was a problem hiding this comment.
oh that seems to be a migration artifact, fixing now
There was a problem hiding this comment.
New to this repo, so these might be obvious questions but curious about the following:
- Why are defining these types on our own rather than taking on the openai dependency?
- If we do it this way, are we worried about API drift? Between our definitions here and OpenAI resonses?
There was a problem hiding this comment.
So if I understand the current implementation of agent-langchain-ts in app-templates correctly, it's to avoid pulling a really large package just to get a handful of TS types. I think it makes sense to do the same here.
The types are standard, simple and expected to be stable in the long term so I think it's a justified risk compared to cost of pulling in a large dep.
|
|
||
| addTools: (tools: StructuredToolInterface[]) => this.addTools(tools), | ||
| addMcpServers: (servers: DatabricksMCPServer[]) => | ||
| this.addMcpServers(servers), |
There was a problem hiding this comment.
very nit: when we do sequential addTools and addMcpServers calls like this, aren't we re-building the agent twice? instead maybe we can introduce a wrapper/batch function which can looks something like below for a single rebuild:
await appkit.agent.configure({
tools: [...demoTools],
mcpServers: [ucServer],
}); // single rebuild
There was a problem hiding this comment.
good point! let me add another method
| const outputItem = { | ||
| id: `fco_${randomUUID()}`, | ||
| call_id: callId, | ||
| output: JSON.stringify(event.data?.output || ""), |
There was a problem hiding this comment.
UI nit: this returns a fill ToolMessage object with a lot of other metadata in it as you can see in the attached sc. Perhaps we can clean up this up a bit to show only the relevant sections to be more human readable?
There was a problem hiding this comment.
Hmm the approach was to pack as much data as possible, leaving it up to the UI to decide what's used for display (the current dev-playground chat UI just happens to display all for demo purposes). What do you think?
| Agent Chat | ||
| </h3> | ||
| <p className="text-muted-foreground mb-6 flex-grow"> | ||
| Chat with a LangChain/LangGraph AI agent powered by the AppKit |
There was a problem hiding this comment.
Minor: LangChain is an implementation detail, no need to mention it here and elsewhere
There was a problem hiding this comment.
In particular, we want the right to change our implementation to use another agent framework / SDK down the line if that's preferable
| </h3> | ||
| <p className="text-muted-foreground mb-6 flex-grow"> | ||
| Chat with a LangChain/LangGraph AI agent powered by the AppKit | ||
| Agent Plugin. Features Responses API SSE streaming and tool call |
There was a problem hiding this comment.
Nit: Responses API -> OpenResponses-compatible
There was a problem hiding this comment.
(To go from vendor-specific -> OSS, see https://www.openresponses.org/ for more context)
| import { toJSONSchema } from "zod/v4/core"; | ||
|
|
||
| /** | ||
| * Workaround: @databricks/langchainjs@0.1.0 calls `schema.toJSONSchema()` |
There was a problem hiding this comment.
Hm, curious why this is needed - do we need a bugfix in https://www.npmjs.com/package/@databricks/langchainjs?
There was a problem hiding this comment.
It's a workaround due to dev-playground using older zod. Agree, let me bump local zod instead
There was a problem hiding this comment.
update: actually it's not needed anymore after redoing tools to follow OpenResponses API
| agent({ | ||
| model: process.env.DATABRICKS_MODEL || "databricks-claude-sonnet-4-5", | ||
| systemPrompt: | ||
| "You are a helpful assistant. Use tools when appropriate — for example, use get_weather for weather questions, and get_current_time for time queries.", |
There was a problem hiding this comment.
Is it possible for users to pass tools here? It seems a little counterintuitive to have to call appkit.agent.addTools later. Or maybe here the user could have the option to pass either a list of tools, or a function that accepts an appkit and returns a list of tools, which would get called in the agent implementation to get the list of tools? In general, I think it'd be better to start with just a list of tools here for simplicity - we can see if it makes sense to support the (appkit) => listOftools type later
There was a problem hiding this comment.
Ah, I see it is possible to pass a list of tools directly - that makes sense, but I don't think we want to reuse the LangChain tool type (see my comment RE OpenResponses below). In particular I wonder if it makes sense to do something like support passing either typescript function objects (locally-defined tool implementations) or structured OpenResponses-style hosted tool definitions like {"type": "genie", "genie_space": {"id": ...}}. Happy to discuss more how the hosted tool piece should work - it should be possible to translate a spec of that format into a call to a Databricks managed MCP server
There was a problem hiding this comment.
- Correct, it's possible to add tools/MCPs either at initialization or later (anticipating for some plugin interop in the future).
- let me address in the comment below
| ...(process.env.APPKIT_E2E_TEST && { client: createMockClient() }), | ||
| }).then((appkit) => { | ||
| }).then(async (appkit) => { | ||
| // Add tools (and optionally MCP servers) after app creation |
There was a problem hiding this comment.
Ideally the interface/format for tools would match OpenResponses conventions, as outlined in this doc. AI can probably help make this change given https://github.com/openresponses/openresponses/blob/main/public/openapi/openapi.json#L1225 for reference :P
There was a problem hiding this comment.
redone tools API to use standard openresponses format
| @@ -0,0 +1,43 @@ | |||
| # Interface: AgentInterface | |||
|
|
|||
| Contract that agent implementations must fulfil. | |||
There was a problem hiding this comment.
Nice, AFAICT this allows basically plugging in our existing TypeScript agent implementation from https://github.com/databricks/app-templates/tree/main/agent-langchain-ts, which is great. Later we may want a way to get an agentInstance from an existing OpenResponses-compatible agent endpoint or Agent Brick on Databricks, e.g. for the chat UI. We can add support for that later though
There was a problem hiding this comment.
(Or maybe the chat UI plugin should just support running against such an endpoint, i.e. accept endpointName: agentBrickEndpointName as an option; that's probably cleaner)
There was a problem hiding this comment.
(Or maybe the chat UI plugin should just support running against such an endpoint, i.e. accept endpointName: agentBrickEndpointName as an option; that's probably cleaner)
yeah I think we should allow chatUI to be configured to target
- any hosted model on databricks
- hosted agent bricks
- agent plugin running in the same app
| ### mcpServers? | ||
|
|
||
| ```ts | ||
| optional mcpServers: DatabricksMCPServer[]; |
There was a problem hiding this comment.
Shouldn't be needed if we support OpenResponses/OpenAI Responses-style MCP tool references e.g.
{"type": "genie", "genie_space": {"id": "..."}},
{"type": "vector_search_index", "vector_search_index": {"name": "catalog.schema.index"}},
{"type": "custom_mcp_server", "custom_mcp_server": {"app_name": "..."}},
{"type": "external_mcp_server", "external_mcp_server": {"connection_name": "..."}}
Open to adding classes (GenieTool, VectorSearchIndexTool, etc) to simplify declaring these, but it should also be possible to just pass the OpenResponses-style tool definition here
There was a problem hiding this comment.
hmm I might not have enough context, why the DatabricksMCPServer abstraction was introduced then?
There was a problem hiding this comment.
It's useful as a building block for our AppKit implementation - it handles auth to MCP servers on Databricks. We can support e.g. a declarative Genie tool in AppKit by using the Genie managed MCP server via DatabricksMCPServer, but we shouldn't ask the user to pass in a DatabricksMCPServer reference to the Genie managed MCP server - it's better to just ask them for the Genie spaces they'd like to add as tools, etc.
Concretely this:
tools: [{"type": "genie", "genie_space": {"id": "..."}}, ...]
is better than this:
mcpServers: DatabricksMCPServer.fromGenieSpace(id)
since the former doesn't require learning which Databricks resources should be accessed via MCP servers, etc
Signed-off-by: Hubert Zub <hubert.zub@databricks.com>
Signed-off-by: Hubert Zub <hubert.zub@databricks.com>
hubertzub-db
force-pushed
the
feat/agent-plugin
branch
from
ba9ad8d to
f4e2280
Compare
|
|
||
| /** | ||
| * Tools to register with the agent. Accepts OpenResponses-aligned FunctionTool | ||
| * objects or LangChain StructuredToolInterface instances. |
There was a problem hiding this comment.
QQ: is LangChain StructuredToolInterface needed? My hope is to avoid exposing LangChain types in our user-facing AppKit interface, they should be an implementation detail. Instead we'd ideally support other OpenResponses-aligned types, e.g.
{"type": "genie", "genie_space": {"id": ...}}
In addition to the FunctionTool type
3 participants
Adds a new first-class plugin for building AI agents.
PR Reading Guide
This guide helps reviewers navigate the ~3,400-line commit. Changes are grouped by area with suggested reading order and file-by-file notes.
1. Agent plugin (core)
Location:
packages/appkit/src/plugins/agent/The plugin is a first-class AppKit plugin that exposes POST /api/agent and supports:
config.agentInstance(anyAgentInterfaceimplementation)Suggested order
types.tsIAgentConfig(agentInstance, model, systemPrompt, tools, mcpServers, etc.) andBasePluginConfigextension.agent-interface.tsInvokeParams,ResponseStreamEvent(SSE event union),ResponseOutputItem, andAgentInterface(invoke(),stream()). No OpenAI dependency; mirrors Responses API shape.manifest.jsonserving_endpointwithDATABRICKS_MODELenv.standard-agent.tsStandardAgentimplementsAgentInterface: wraps LangGraphcreateReactAgent, converts messages to/fromInvokeParamsand stream events to Responses API SSE. DefinesLangGraphAgentminimal interface.invoke-handler.tsinput,stream, optionalmodel), flattens history forInvokeParams, calls agentstream()orinvoke(), writes SSE (including[DONE]).agent.tsAgentPluginclass:setup()(useagentInstanceor build LangGraph agent with ChatDatabricks + tools/MCP),buildStandardAgent(),addTools()/addMcpServers(),injectRoutes()(POST/→ handler),exports()(invoke, stream, addTools, addMcpServers),abortActiveOperations()(MCP client cleanup). Factoryagent()viatoPlugin.index.tscreateInvokeHandler,StandardAgent,LangGraphAgent.Dependencies: Plugin uses optional peer deps:
@databricks/langchainjs,@langchain/core,@langchain/langgraph,@langchain/mcp-adapters(seepackages/appkit/package.json).2. Agent plugin – tests
Location:
packages/appkit/src/plugins/agent/tests/stub-agent.tsAgentInterfaceimplementation used by tests: echo-styleinvoke()andstream()that emit known SSE events.agent.test.tsAgentPlugin: factory/manifest/name,setup()withagentInstancevs config (model required),addTools/addMcpServers(only when not using agentInstance), route registration. UsesmockServiceContext,setupDatabricksEnv, mocksCacheManager.invoke-handler.test.tscreateInvokeHandler: streaming (SSE format,[DONE], output_item.added/done, text deltas), non-streaming (JSON array), error handling, request validation (Zod),flattenHistoryItemfor various message shapes. UsescreateMockRequest/createMockResponseand parses SSE fromres.writechunks.agent.integration.test.tscreateApp+ agent + server plugins (StubAgent),POST /api/agentstreaming and non-streaming, error cases, and (if present)addToolsaffecting behavior.How to run: From repo root,
pnpm testor filter:pnpm test -- packages/appkit/src/plugins/agent.3. appkit-ui changes
Location:
packages/appkit-ui/src/react/agent-chat/React components and hooks for a chat UI that talk to POST /invocations (Responses API SSE). The backend can expose that path by rewriting to
/api/agent(see dev-playground).Suggested order
types.tsAssistantPart(text, function_call, function_call_output),ChatMessage(user | assistant with parts),SSEEvent/SSEItem,UseAgentChatOptions/UseAgentChatReturn,AgentChatProps.utils.tsserializeForApi(msg): convertsChatMessageto request body (user content or assistant content array).tryFormatJson(s): pretty-print JSON for display.use-agent-chat.tshandleSubmit(POST withinput+stream: true, read SSE stream, parse events, update messages/parts),displayMessages(messages + current streaming message),isStreamingText. Handlesresponse.output_item.added,response.output_text.delta,response.completed,[DONE], errors.agent-chat-part.tsxtryFormatJson.agent-chat-message.tsxAgentChatPart.agent-chat.tsxAgentChat: scrollable message list, input form, usesuseAgentChat; auto-scroll on new content/streaming.index.ts@databricks/appkit-ui/react.Barrel:
packages/appkit-ui/src/react/index.ts– addsAgentChat(and any hook/type exports if present).4. apps/dev-playground updates
The playground wires the Agent Plugin into the app (server), adds an Agent Chat page (client), and documents the required env var.
Server
server/index.tsagent()is added to the plugins array with:model:process.env.DATABRICKS_MODELor fallback"databricks-claude-sonnet-4-5".systemPrompt: instructs the assistant to use tools (e.g.get_weather,get_current_time) when relevant..then()aftercreateApp,appkit.agent.addTools(demoTools)is called so the built-in LangGraph agent gets the demo tools without putting them in initial config (shows the “add tools after app creation” pattern)./invocationsrewrite: Insideserver.extend(), aPOST /invocationsroute is added that setsreq.url = "/api/agent"and re-dispatches into the Express app. That way the client can call the standard Databricks Apps path/invocationswhile the plugin still serves its canonical route/api/agent.server/agent-tools.ts(new file)StructuredTools used by the agent in the playground.get_weather: Takeslocation(string), returns a short fake weather summary (random condition and temp). Schema:z.object({ location: z.string().describe(...) }).get_current_time: Takes optionaltimezone(IANA, default UTC), returns current date/time in that zone. Schema:z.object({ timezone: z.string().optional().describe(...) }).patchZodSchema: Workaround because@databricks/langchainjs@0.1.0expectsschema.toJSONSchema()as a method, while zod v4 exposestoJSONSchema(schema)as a standalone. The helper attaches.toJSONSchemaon the schema so the LangChain/ChatDatabricks integration works. Both tools use it; they are exported asdemoTools.Client
client/src/routes/agent.route.tsx(new file)/agentroute viacreateFileRoute("/agent").POST /invocationsand Responses API SSE), then a single<AgentChat>that fills the remaining height.invokeUrl="/invocations"(hits the server rewrite above), plus placeholder, empty-message, andclassNamefor flex layout.client/src/routes/__root.tsx/agent, in the same style as Analytics, Genie, etc.client/src/routes/index.tsx/agent.client/src/routeTree.gen.ts/agentroute (no hand edits).Config
apps/dev-playground/.env.distDATABRICKS_MODEL=so developers know they must set the model serving endpoint name for the agent plugin (required when not usingagentInstance).5. Docs updates
Location:
docs/docs/api/appkit/Typedoc-generated API docs (and sidebar) for the new public surface.
Interface.AgentInterface.mdinvoke,stream).Interface.IAgentConfig.mdInterface.InvokeParams.mdInterface.StandardAgent.mdInterface.BasePluginConfig.mdTypeAlias.ResponseStreamEvent.mdVariable.agent.mdagentfactory.index.mdtypedoc-sidebar.ts6. Other changes
packages/appkit/package.jsonpackages/appkit/src/index.tsagent,AgentPlugin,AgentInterface).packages/appkit/src/plugins/index.tspnpm-lock.yamltemplate/appkit.plugins.jsonagentplugin entry (name, displayName, description, package, resources matching manifest).How to read this PR
Get the contract
Read
agent-interface.tsandtypes.tsso you knowAgentInterface,InvokeParams, and SSE event shapes.See how the plugin is used
Skim
apps/dev-playground/server/index.ts: plugin registration,addTools(demoTools), and the/invocations→/api/agentrewrite. Optionally openagent-tools.tsfor the demo tools.Trace the request path
Follow a request: invoke-handler (Zod + flatten history) → agent (get implementation) → standard-agent (LangGraph → SSE). Read
invoke-handler.tsthenstandard-agent.tsthen the streaming part ofagent.ts.Plugin lifecycle and API
Read
agent.ts:setup()(agentInstance vs build),buildStandardAgent(),addTools/addMcpServers,injectRoutes,exports(), cleanup inabortActiveOperations().UI and E2E flow
In appkit-ui read
use-agent-chat.ts(state + SSE parsing), thenagent-chat.tsxandagent-chat-message.tsx/agent-chat-part.tsx. In the app, openagent.route.tsxto see the page andinvokeUrl="/invocations".Tests
Run tests, then read
stub-agent.ts,invoke-handler.test.ts(SSE and validation),agent.test.ts(plugin behavior), andagent.integration.test.ts(HTTP).Manifest and config
Check
manifest.jsonandtemplate/appkit.plugins.jsonfor resources and template wiring.Docs
Use
docs/docs/api/appkit/for reference; no need to read every doc file unless you care about generated output.Review focus: Contract in
agent-interface.ts, correctness of SSE ininvoke-handlerandstandard-agent, plugin lifecycle and thread-safety ofaddTools/addMcpServers, and that the UI correctly parses the same SSE events the backend emits.