databricks-agentbricks 0.2.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- databricks_agentbricks/__init__.py +1 -0
- databricks_agentbricks/_group.py +217 -0
- databricks_agentbricks/agent_project.py +657 -0
- databricks_agentbricks/app_resources.py +139 -0
- databricks_agentbricks/cli/__init__.py +12 -0
- databricks_agentbricks/cli/app.py +108 -0
- databricks_agentbricks/cli/app_auth.py +198 -0
- databricks_agentbricks/cli/auth.py +142 -0
- databricks_agentbricks/cli/deploy.py +954 -0
- databricks_agentbricks/cli/dev.py +310 -0
- databricks_agentbricks/cli/endpoint.py +145 -0
- databricks_agentbricks/cli/endpoint_examples.py +33 -0
- databricks_agentbricks/cli/endpoint_output.py +53 -0
- databricks_agentbricks/cli/endpoint_request.py +71 -0
- databricks_agentbricks/cli/endpoint_transport.py +148 -0
- databricks_agentbricks/cli/help.py +467 -0
- databricks_agentbricks/cli/init.py +496 -0
- databricks_agentbricks/cli/mcp.py +75 -0
- databricks_agentbricks/cli/memory.py +520 -0
- databricks_agentbricks/cli/sandbox.py +823 -0
- databricks_agentbricks/cli/sessions.py +588 -0
- databricks_agentbricks/cli/tools.py +456 -0
- databricks_agentbricks/cli/tracing.py +724 -0
- databricks_agentbricks/databricks_cli.py +38 -0
- databricks_agentbricks/errors.py +114 -0
- databricks_agentbricks/lakebase_runtime_store.py +107 -0
- databricks_agentbricks/legacy_lakebase_runtime_store.py +74 -0
- databricks_agentbricks/project_config.py +157 -0
- databricks_agentbricks/project_types.py +55 -0
- databricks_agentbricks/py.typed +0 -0
- databricks_agentbricks/render.py +349 -0
- databricks_agentbricks/templates/agent-bricks-migrate/SKILL.md +101 -0
- databricks_agentbricks/templates/agent-langgraph/.env.example +36 -0
- databricks_agentbricks/templates/agent-langgraph/.gitignore +199 -0
- databricks_agentbricks/templates/agent-langgraph/AGENTKIT_CONTRACT.md +95 -0
- databricks_agentbricks/templates/agent-langgraph/AGENTS.md +31 -0
- databricks_agentbricks/templates/agent-langgraph/CLAUDE.md +1 -0
- databricks_agentbricks/templates/agent-langgraph/README.md +141 -0
- databricks_agentbricks/templates/agent-langgraph/agent/__init__.py +0 -0
- databricks_agentbricks/templates/agent-langgraph/agent/agent.py +177 -0
- databricks_agentbricks/templates/agent-langgraph/agent/mcps.py +27 -0
- databricks_agentbricks/templates/agent-langgraph/agent/tools/__init__.py +24 -0
- databricks_agentbricks/templates/agent-langgraph/agent/tools/sample_tool.py +15 -0
- databricks_agentbricks/templates/agent-langgraph/agent/tools/send_message.py +17 -0
- databricks_agentbricks/templates/agent-langgraph/app.yaml +15 -0
- databricks_agentbricks/templates/agent-langgraph/pyproject.toml +42 -0
- databricks_agentbricks/templates/agent-langgraph/runtime/__init__.py +0 -0
- databricks_agentbricks/templates/agent-langgraph/runtime/adapter.py +144 -0
- databricks_agentbricks/templates/agent-langgraph/runtime/main.py +25 -0
- databricks_agentbricks/templates/agent-langgraph/tests/__init__.py +0 -0
- databricks_agentbricks/templates/agent-langgraph/tests/test_agent.py +211 -0
- databricks_agentbricks/templates/agent-langgraph/tests/test_runtime.py +19 -0
- databricks_agentbricks/templates/agent-langgraph/tests/test_workspace.py +28 -0
- databricks_agentbricks/templates/agent-openai/.env.example +37 -0
- databricks_agentbricks/templates/agent-openai/.gitignore +199 -0
- databricks_agentbricks/templates/agent-openai/AGENTKIT_CONTRACT.md +107 -0
- databricks_agentbricks/templates/agent-openai/AGENTS.md +76 -0
- databricks_agentbricks/templates/agent-openai/CLAUDE.md +1 -0
- databricks_agentbricks/templates/agent-openai/README.md +143 -0
- databricks_agentbricks/templates/agent-openai/agent/__init__.py +0 -0
- databricks_agentbricks/templates/agent-openai/agent/agent.py +205 -0
- databricks_agentbricks/templates/agent-openai/agent/mcps.py +28 -0
- databricks_agentbricks/templates/agent-openai/agent/tools/__init__.py +24 -0
- databricks_agentbricks/templates/agent-openai/agent/tools/sample_tool.py +15 -0
- databricks_agentbricks/templates/agent-openai/agent/tools/send_message.py +18 -0
- databricks_agentbricks/templates/agent-openai/app.yaml +15 -0
- databricks_agentbricks/templates/agent-openai/pyproject.toml +42 -0
- databricks_agentbricks/templates/agent-openai/runtime/__init__.py +0 -0
- databricks_agentbricks/templates/agent-openai/runtime/adapter.py +188 -0
- databricks_agentbricks/templates/agent-openai/runtime/main.py +25 -0
- databricks_agentbricks/templates/agent-openai/tests/__init__.py +0 -0
- databricks_agentbricks/templates/agent-openai/tests/test_agent.py +326 -0
- databricks_agentbricks/templates/agent-openai/tests/test_runtime.py +19 -0
- databricks_agentbricks/templates/agent-openai/tests/test_workspace.py +28 -0
- databricks_agentbricks/templates/custom-agent-langgraph/.env.example +1 -0
- databricks_agentbricks/templates/custom-agent-langgraph/.gitignore +11 -0
- databricks_agentbricks/templates/custom-agent-langgraph/README.md +22 -0
- databricks_agentbricks/templates/custom-agent-langgraph/agent/__init__.py +0 -0
- databricks_agentbricks/templates/custom-agent-langgraph/agent/agent.py +27 -0
- databricks_agentbricks/templates/custom-agent-langgraph/app.yaml +1 -0
- databricks_agentbricks/templates/custom-agent-langgraph/pyproject.toml +38 -0
- databricks_agentbricks/templates/custom-agent-langgraph/runtime/__init__.py +0 -0
- databricks_agentbricks/templates/custom-agent-langgraph/runtime/main.py +31 -0
- databricks_agentbricks/templates/custom-agent-langgraph/tests/test_runtime.py +29 -0
- databricks_agentbricks/templates/custom-agent-openai/.env.example +1 -0
- databricks_agentbricks/templates/custom-agent-openai/.gitignore +11 -0
- databricks_agentbricks/templates/custom-agent-openai/README.md +22 -0
- databricks_agentbricks/templates/custom-agent-openai/agent/__init__.py +0 -0
- databricks_agentbricks/templates/custom-agent-openai/agent/agent.py +29 -0
- databricks_agentbricks/templates/custom-agent-openai/app.yaml +1 -0
- databricks_agentbricks/templates/custom-agent-openai/pyproject.toml +38 -0
- databricks_agentbricks/templates/custom-agent-openai/runtime/__init__.py +0 -0
- databricks_agentbricks/templates/custom-agent-openai/runtime/main.py +31 -0
- databricks_agentbricks/templates/custom-agent-openai/tests/test_runtime.py +29 -0
- databricks_agentbricks/templates/sandbox_mcp.py +39 -0
- databricks_agentbricks/templates/sandbox_mcp_langgraph.py +39 -0
- databricks_agentbricks/templates/ui/agent-langgraph/CHAT_APP.md +49 -0
- databricks_agentbricks/templates/ui/agent-langgraph/runtime/main.py +26 -0
- databricks_agentbricks/templates/ui/agent-langgraph/runtime/ui.py +528 -0
- databricks_agentbricks/templates/ui/agent-langgraph/tests/test_demo_ui.py +542 -0
- databricks_agentbricks/templates/ui/agent-langgraph/ui/agentbricks-mark.png +0 -0
- databricks_agentbricks/templates/ui/agent-langgraph/ui/app.js +1464 -0
- databricks_agentbricks/templates/ui/agent-langgraph/ui/index.html +281 -0
- databricks_agentbricks/templates/ui/agent-langgraph/ui/styles.css +1913 -0
- databricks_agentbricks/templates/ui/agent-openai/CHAT_APP.md +54 -0
- databricks_agentbricks/templates/ui/agent-openai/runtime/main.py +26 -0
- databricks_agentbricks/templates/ui/agent-openai/runtime/ui.py +521 -0
- databricks_agentbricks/templates/ui/agent-openai/tests/test_demo_ui.py +512 -0
- databricks_agentbricks/templates/ui/agent-openai/ui/agentbricks-mark.png +0 -0
- databricks_agentbricks/templates/ui/agent-openai/ui/app.js +1428 -0
- databricks_agentbricks/templates/ui/agent-openai/ui/index.html +281 -0
- databricks_agentbricks/templates/ui/agent-openai/ui/styles.css +1913 -0
- databricks_agentbricks/theme.py +72 -0
- databricks_agentbricks-0.2.0.dist-info/METADATA +897 -0
- databricks_agentbricks-0.2.0.dist-info/RECORD +159 -0
- databricks_agentbricks-0.2.0.dist-info/WHEEL +4 -0
- databricks_agentbricks-0.2.0.dist-info/entry_points.txt +2 -0
- databricks_agentbricks-0.2.0.dist-info/licenses/NOTICE +39 -0
- databricks_agentkit/__init__.py +67 -0
- databricks_agentkit/_api_client.py +673 -0
- databricks_agentkit/_pagination.py +25 -0
- databricks_agentkit/client.py +31 -0
- databricks_agentkit/langgraph/__init__.py +98 -0
- databricks_agentkit/langgraph/genie.py +50 -0
- databricks_agentkit/langgraph/mcp.py +216 -0
- databricks_agentkit/langgraph/memory.py +67 -0
- databricks_agentkit/langgraph/session_store.py +485 -0
- databricks_agentkit/memory_store.py +371 -0
- databricks_agentkit/models.py +330 -0
- databricks_agentkit/openai/__init__.py +93 -0
- databricks_agentkit/openai/genie.py +50 -0
- databricks_agentkit/openai/mcp.py +189 -0
- databricks_agentkit/openai/memory.py +67 -0
- databricks_agentkit/openai/sessions.py +157 -0
- databricks_agentkit/py.typed +0 -0
- databricks_agentkit/runtime/ARCHITECTURE.md +64 -0
- databricks_agentkit/runtime/README.md +184 -0
- databricks_agentkit/runtime/__init__.py +69 -0
- databricks_agentkit/runtime/app.py +316 -0
- databricks_agentkit/runtime/auth.py +149 -0
- databricks_agentkit/runtime/durability/__init__.py +1 -0
- databricks_agentkit/runtime/durability/execution.py +133 -0
- databricks_agentkit/runtime/durability/heartbeat.py +64 -0
- databricks_agentkit/runtime/durability/lakebase_runtime_store.py +648 -0
- databricks_agentkit/runtime/durability/recovery.py +70 -0
- databricks_agentkit/runtime/durability/store.py +33 -0
- databricks_agentkit/runtime/execution.py +160 -0
- databricks_agentkit/runtime/genie.py +224 -0
- databricks_agentkit/runtime/mcp_auth.py +114 -0
- databricks_agentkit/runtime/model_services.py +90 -0
- databricks_agentkit/runtime/runtime.py +273 -0
- databricks_agentkit/runtime/session_store_client.py +120 -0
- databricks_agentkit/runtime/store.py +287 -0
- databricks_agentkit/runtime/tool_manifest.py +223 -0
- databricks_agentkit/runtime/tracing.py +71 -0
- databricks_agentkit/runtime/types.py +118 -0
- databricks_agentkit/runtime/workspace.py +27 -0
- databricks_agentkit/session_store.py +434 -0
- databricks_agentkit/timefmt.py +85 -0
|
@@ -0,0 +1,95 @@
|
|
|
1
|
+
# Managed runtime integration contract for LangGraph
|
|
2
|
+
|
|
3
|
+
This document owns the requirements shared by generated projects and existing agents adopting
|
|
4
|
+
the managed runtime. The [README](README.md) owns commands and HTTP examples; adapter docstrings own API details.
|
|
5
|
+
The adjacent template is an example, not a required graph architecture. Runtime internals live in
|
|
6
|
+
the installed package's runtime guide.
|
|
7
|
+
|
|
8
|
+
## Command requirements
|
|
9
|
+
|
|
10
|
+
Agent Bricks CLI commands operate on an application directory; they do not rewrite its graph. Each capability
|
|
11
|
+
needs runtime wiring as well as manifest configuration.
|
|
12
|
+
|
|
13
|
+
| Command | Required integration |
|
|
14
|
+
| --- | --- |
|
|
15
|
+
| `dev`, `deploy` | Installable dependencies and an `app.yaml` command starting the application |
|
|
16
|
+
| `tools add mcp`, `uc-function`, `genie`, `sandbox` | Load bindings into model definitions and tool execution |
|
|
17
|
+
| `sessions bind` | Select the bound checkpointer and supply session/actor identity |
|
|
18
|
+
| `memory bind` | Include bound, actor-scoped memory tools in definitions and execution |
|
|
19
|
+
| `tracing bind`, `tracing unbind` | Initialize tracing and open a root span from resolved config |
|
|
20
|
+
| `endpoint invoke` | An HTTP endpoint; the caller supplies its path and complete request body |
|
|
21
|
+
|
|
22
|
+
### Project and startup
|
|
23
|
+
|
|
24
|
+
Use `agent.toml` for framework, tool/store bindings, and tracing. The existing `.agentbricks/project.toml`
|
|
25
|
+
file records framework and template provenance; keep it when working with projects that have it.
|
|
26
|
+
Metadata alone does not integrate adapters. Install a compatible `databricks-agentbricks[langgraph]`
|
|
27
|
+
distribution, supply the real command in `app.yaml`, load configuration before adapters, and listen
|
|
28
|
+
on the app port. The Agent Bricks CLI and runtime find `agent.toml` from the working directory or
|
|
29
|
+
`AGENTBRICKS_PROJECT_ROOT`. Keep credentials out of `app.yaml`.
|
|
30
|
+
|
|
31
|
+
Store binding commands declare intent. `dev` and `deploy` resolve or provision declared stores and
|
|
32
|
+
supply runtime config; deploy grants app access. The resolved Memory Store ID reaches the runtime
|
|
33
|
+
as `AGENT_MEMORY_STORE`, rather than being stored as an ID in the manifest. Do not hardcode values
|
|
34
|
+
that make later bindings ineffective.
|
|
35
|
+
|
|
36
|
+
### Tools
|
|
37
|
+
|
|
38
|
+
Use `databricks_agentkit.langgraph.mcp_tools()` for MCP, UC function, and sandbox bindings, combining
|
|
39
|
+
them with existing servers without duplicates. It applies sandbox downscoping. Local Python tools
|
|
40
|
+
belong in source; the template discovers `BaseTool` objects through
|
|
41
|
+
[agent/tools/__init__.py](agent/tools/__init__.py).
|
|
42
|
+
|
|
43
|
+
Feed selected tools into both model definitions and graph execution while preserving approval
|
|
44
|
+
policy. New bindings must become effective on the next graph construction or restart without
|
|
45
|
+
another source edit.
|
|
46
|
+
|
|
47
|
+
### Sessions and memory
|
|
48
|
+
|
|
49
|
+
Use `checkpointer()` for managed state and merge `thread_config(session_id, actor)` into every
|
|
50
|
+
graph run. The helper selects the bound Session Store or an in-process saver; explicit arguments
|
|
51
|
+
and environment overrides take precedence. It caches the saver per process, so restart after a
|
|
52
|
+
binding change. The application state must be serializable by the selected saver.
|
|
53
|
+
|
|
54
|
+
Include `memory_tools(actor)` in model and execution wiring. It returns no tools when unconfigured.
|
|
55
|
+
Its closures capture actor identity, so never reuse them across users. The application owns trusted
|
|
56
|
+
actor/tenant authentication; a payload field alone does not establish identity.
|
|
57
|
+
|
|
58
|
+
### Tracing
|
|
59
|
+
|
|
60
|
+
Call `configure_tracing()` after loading config, then wrap agent execution, including streams, in
|
|
61
|
+
`start_trace(..., session_id=session_id)`. The root span lets framework autologging record children.
|
|
62
|
+
The helper enables tracing when destination and experiment settings are present and disables it
|
|
63
|
+
otherwise. Preserve existing tracing semantics when composing it.
|
|
64
|
+
|
|
65
|
+
## Managed server adapter
|
|
66
|
+
|
|
67
|
+
For the managed invocation protocol, construct `DurableAgentServer` in [runtime/main.py](runtime/main.py)
|
|
68
|
+
and register [runtime/adapter.py](runtime/adapter.py) hooks. Keep framework-native execution in
|
|
69
|
+
[agent/agent.py](agent/agent.py), independent of runtime request/context types.
|
|
70
|
+
|
|
71
|
+
The invoke hook translates opaque application input, runs the graph, translates native events,
|
|
72
|
+
calls `await context.emit(event)`, and returns JSON output. DurableAgentServer owns foreground/background
|
|
73
|
+
transport, polling, and replay. Invocation UUIDs differ from stable application session IDs; the
|
|
74
|
+
routing cookie is not the application session.
|
|
75
|
+
|
|
76
|
+
The example assumes `messages` state and message/update events. Custom state, outputs, and
|
|
77
|
+
interrupts require explicit mappings and must not be discarded to fit the example.
|
|
78
|
+
|
|
79
|
+
### Recovery and durability
|
|
80
|
+
|
|
81
|
+
Runtime Store persistence covers invocations and emitted events. Session Store persistence covers
|
|
82
|
+
graph checkpoints and paused interrupts. A durable Runtime Store alone does not preserve graph
|
|
83
|
+
state. `dev` uses process-local invocation storage; deployment attaches a Lakebase-backed
|
|
84
|
+
Runtime Store.
|
|
85
|
+
|
|
86
|
+
Register recovery when intended. Resume a checkpoint only when metadata associates it with the
|
|
87
|
+
current invocation; otherwise replay original input. Use synchronous checkpoint durability before
|
|
88
|
+
acknowledging progress. Recovery is at least once, so side effects must tolerate replay.
|
|
89
|
+
|
|
90
|
+
## Optional chat app
|
|
91
|
+
|
|
92
|
+
No CLI command requires the chat UI. The default overlay supplies it; `--disable-chat-app` omits
|
|
93
|
+
it. When adopting it, inspect `CHAT_APP.md`, `runtime/ui.py`, and browser code. Adapt model choice,
|
|
94
|
+
history, interrupts, and authentication to the actual graph. Its message-oriented assumptions must
|
|
95
|
+
not silently change custom application behavior.
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
# Agent Development Guide
|
|
2
|
+
|
|
3
|
+
This project is a LangGraph workload hosted by `databricks_agentkit.DurableAgentServer`.
|
|
4
|
+
|
|
5
|
+
Read [AGENTKIT_CONTRACT.md](AGENTKIT_CONTRACT.md) before changing integration points. It owns command
|
|
6
|
+
requirements, tool/state/tracing wiring, and recovery. [README.md](README.md) owns setup and client
|
|
7
|
+
examples. Keep this file as a development map rather than repeating those rules.
|
|
8
|
+
|
|
9
|
+
## Commands
|
|
10
|
+
|
|
11
|
+
```bash
|
|
12
|
+
ab dev
|
|
13
|
+
uv run pytest
|
|
14
|
+
ab --profile <profile> deploy <name> --source .
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
## Code map
|
|
18
|
+
|
|
19
|
+
| Change | File |
|
|
20
|
+
| --- | --- |
|
|
21
|
+
| Framework-native agent and `run_agent` | `agent/agent.py` |
|
|
22
|
+
| Local tools | `agent/tools/` |
|
|
23
|
+
| MCP servers | `agent/mcps.py` |
|
|
24
|
+
| Managed runtime `invoke`/`recover` hooks and input/output translation | `runtime/adapter.py` |
|
|
25
|
+
| `DurableAgentServer` construction and hook registration | `runtime/main.py` |
|
|
26
|
+
| Browser and managed-state routes | `runtime/ui.py` |
|
|
27
|
+
| Browser behavior | `ui/app.js` |
|
|
28
|
+
|
|
29
|
+
Shared adapters come from `databricks_agentkit.langgraph` and `databricks_agentkit.runtime`. Consult
|
|
30
|
+
their docstrings for API details. Update the shared contract and relevant tests when integration
|
|
31
|
+
requirements change.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
@AGENTS.md
|
|
@@ -0,0 +1,141 @@
|
|
|
1
|
+
# LangGraph agent template
|
|
2
|
+
|
|
3
|
+
A LangGraph agent served by `databricks_agentkit.DurableAgentServer`. The managed runtime keeps invocation state and events in
|
|
4
|
+
memory during `ab dev`. Deployment attaches a persistent Runtime Store, so invocation state and
|
|
5
|
+
events survive process loss and interrupted work can be recovered.
|
|
6
|
+
|
|
7
|
+
The generated project separates portable agent execution from the managed HTTP protocol:
|
|
8
|
+
|
|
9
|
+
```text
|
|
10
|
+
client -> runtime/main.py -> runtime/adapter.py -> agent/agent.py:run_agent
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
- `agent/agent.py` owns the framework-native agent, sessions, tools, MCP lifetime, and `run_agent`.
|
|
14
|
+
- `runtime/adapter.py` owns the agent-author integration hooks: runtime input/output translation plus
|
|
15
|
+
`invoke` and `recover`.
|
|
16
|
+
- `runtime/main.py` constructs the server and registers those hooks.
|
|
17
|
+
|
|
18
|
+
To bring an existing LangGraph agent, run `ab init --framework langgraph --existing .` in its
|
|
19
|
+
project and follow the generated prompt in your coding agent. The migration skill and this template share
|
|
20
|
+
[AGENTKIT_CONTRACT.md](AGENTKIT_CONTRACT.md), which owns integration requirements. This README owns
|
|
21
|
+
configuration and client examples; [AGENTS.md](AGENTS.md) provides the development map.
|
|
22
|
+
|
|
23
|
+
## Run locally
|
|
24
|
+
|
|
25
|
+
```bash
|
|
26
|
+
ab dev
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
The API is available at `http://localhost:8000/api/invocations`. Every request supplies a UUID `id`.
|
|
30
|
+
That ID is the invocation identifier and idempotency key. Agent-specific values live inside the
|
|
31
|
+
opaque `input` object:
|
|
32
|
+
|
|
33
|
+
```bash
|
|
34
|
+
SESSION_ID=$(uuidgen)
|
|
35
|
+
INVOCATION_ID=$(uuidgen)
|
|
36
|
+
|
|
37
|
+
curl -sS http://localhost:8000/api/invocations \
|
|
38
|
+
-H 'Content-Type: application/json' \
|
|
39
|
+
-d "{\"id\":\"$INVOCATION_ID\",\"input\":{\"session_id\":\"$SESSION_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"What time is it? Use your tool.\"}]}}"
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
Reuse `SESSION_ID` for multi-turn conversation state. Generate a new `INVOCATION_ID` for each turn.
|
|
43
|
+
Retrying the same request with the same invocation ID returns the persisted result; changing the
|
|
44
|
+
request while reusing the ID returns `409`.
|
|
45
|
+
|
|
46
|
+
## Invocation modes
|
|
47
|
+
|
|
48
|
+
- Foreground: omit `background` and `stream`; the response contains the agent result under `output`.
|
|
49
|
+
- Foreground streaming: set `stream: true`; the response is SSE backed by persisted events.
|
|
50
|
+
- Background: set `background: true`; poll the returned `status_url`.
|
|
51
|
+
- Background streaming: set both flags; the `202` response includes `status_url` and `events_url`.
|
|
52
|
+
|
|
53
|
+
```bash
|
|
54
|
+
INVOCATION_ID=$(uuidgen)
|
|
55
|
+
curl -sN http://localhost:8000/api/invocations \
|
|
56
|
+
-H 'Content-Type: application/json' \
|
|
57
|
+
-d "{\"id\":\"$INVOCATION_ID\",\"input\":{\"session_id\":\"$SESSION_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"Count to three.\"}]},\"stream\":true}"
|
|
58
|
+
|
|
59
|
+
INVOCATION_ID=$(uuidgen)
|
|
60
|
+
curl -sS http://localhost:8000/api/invocations \
|
|
61
|
+
-H 'Content-Type: application/json' \
|
|
62
|
+
-d "{\"id\":\"$INVOCATION_ID\",\"input\":{\"session_id\":\"$SESSION_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"Summarize durable agents.\"}]},\"background\":true}" | jq
|
|
63
|
+
curl -sS "http://localhost:8000/api/invocations/$INVOCATION_ID" | jq
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
SSE records contain events translated by `runtime/adapter.py`: token `delta`s, completed `message`s,
|
|
67
|
+
and HITL `interrupt`s. Replay from a cursor with
|
|
68
|
+
`GET /api/invocations/{id}/events?after={sequence}`.
|
|
69
|
+
|
|
70
|
+
## Human approval
|
|
71
|
+
|
|
72
|
+
`send_message` is gated by `HumanInTheLoopMiddleware`. Start a turn asking the agent to use that
|
|
73
|
+
tool. When the output or event stream contains an `interrupt`, submit a new invocation with the same
|
|
74
|
+
application session:
|
|
75
|
+
|
|
76
|
+
```json
|
|
77
|
+
{
|
|
78
|
+
"id": "<new-uuid>",
|
|
79
|
+
"input": {
|
|
80
|
+
"session_id": "<same-session-id>",
|
|
81
|
+
"resume": {"decisions": [{"type": "approve"}]}
|
|
82
|
+
}
|
|
83
|
+
}
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
The default checkpointer is process-local. Bind a managed Session Store to preserve multi-turn and
|
|
87
|
+
paused LangGraph state across restarts:
|
|
88
|
+
|
|
89
|
+
```bash
|
|
90
|
+
ab sessions bind my-agent-sessions
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
## Crash recovery
|
|
94
|
+
|
|
95
|
+
Deployment with a Runtime Store supports invocation recovery after worker loss. Graph checkpoint
|
|
96
|
+
persistence requires a Session Store separately. Recovery can replay work, so tools must tolerate
|
|
97
|
+
repeated side effects. See [Recovery and durability](AGENTKIT_CONTRACT.md#recovery-and-durability).
|
|
98
|
+
|
|
99
|
+
## Chat app
|
|
100
|
+
|
|
101
|
+
The browser UI is included by default. It generates a stable application session ID in local
|
|
102
|
+
storage, places it inside each invocation's `input`, and generates a fresh invocation UUID per turn.
|
|
103
|
+
Use `ab init --framework langgraph --disable-chat-app` for API-only output.
|
|
104
|
+
|
|
105
|
+
## Configure and deploy
|
|
106
|
+
|
|
107
|
+
- Change the model, instructions, tools, graph, and framework-native execution in `agent/agent.py`.
|
|
108
|
+
- Change `runtime/adapter.py` only to map a different application input/output contract.
|
|
109
|
+
- Add local tools under `agent/tools/`; modules are auto-discovered.
|
|
110
|
+
- Add MCP servers in `agent/mcps.py` or with `ab tools add mcp`.
|
|
111
|
+
- Bind long-term memory with `ab memory bind <store>`.
|
|
112
|
+
|
|
113
|
+
```bash
|
|
114
|
+
ab --profile <profile> deploy agent-langgraph --source .
|
|
115
|
+
```
|
|
116
|
+
|
|
117
|
+
When deployment provisions a dedicated Runtime Store, only the app-owned
|
|
118
|
+
`databricks_agentkit_runtime_<hash>` schema and runtime tables are added. Managed Runtime Store
|
|
119
|
+
deployments use their own default schema.
|
|
120
|
+
|
|
121
|
+
The `__Host-databricks-app-router` cookie may be supplied independently for sticky replica routing.
|
|
122
|
+
It is not authentication and is not used as the template's application session ID.
|
|
123
|
+
|
|
124
|
+
# Request-user authorization
|
|
125
|
+
|
|
126
|
+
Declared tools in `agent.toml` select `auth = "user"` or `auth = "app"`; legacy missing auth
|
|
127
|
+
continues to use the application/default identity. Deployed user tools require the request resolver
|
|
128
|
+
and never fall back to application credentials. Model, memory-service, session-service, and custom
|
|
129
|
+
MCP server credentials are unchanged.
|
|
130
|
+
|
|
131
|
+
`runtime/main.py` derives the invocation policy after `configure()`. The runtime adapter namespaces
|
|
132
|
+
public session IDs and actor values for the request owner, then passes only `workspace_client_for`
|
|
133
|
+
to the framework-native agent. Internal session keys are never returned to clients.
|
|
134
|
+
|
|
135
|
+
User-policy invocations use the existing Runtime for synchronous, streaming, and background calls,
|
|
136
|
+
including status polling, event replay, and invocation-ID idempotency. The Runtime Store persists no
|
|
137
|
+
credential; request-user authentication remains process-local for the active first attempt. A
|
|
138
|
+
replacement attempt fails with `MCP_USER_AUTH_RECOVERY_UNSUPPORTED` before agent code runs. Approval
|
|
139
|
+
interruptions remain unsupported and fail with `MCP_USER_AUTH_HITL_UNSUPPORTED`. The agent's
|
|
140
|
+
existing namespaced memory, conversation store, and checkpointer behavior is unchanged; OBO does not
|
|
141
|
+
add another saver.
|
|
File without changes
|
|
@@ -0,0 +1,177 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from collections.abc import AsyncGenerator, Callable
|
|
3
|
+
from typing import Any
|
|
4
|
+
|
|
5
|
+
from databricks.sdk import WorkspaceClient
|
|
6
|
+
from databricks_langchain import ChatDatabricks
|
|
7
|
+
from langchain.agents import create_agent
|
|
8
|
+
from langchain.agents.middleware import HumanInTheLoopMiddleware
|
|
9
|
+
|
|
10
|
+
from agent.mcps import build_mcp_servers
|
|
11
|
+
|
|
12
|
+
# Importing the tools package auto-registers every tool module.
|
|
13
|
+
from agent.tools import all_tools
|
|
14
|
+
from databricks_agentkit import workspace_client, workspace_headers
|
|
15
|
+
from databricks_agentkit.langgraph import (
|
|
16
|
+
checkpointer,
|
|
17
|
+
configure_tracing,
|
|
18
|
+
genie_tools,
|
|
19
|
+
mcp_tools,
|
|
20
|
+
memory_tools,
|
|
21
|
+
start_trace,
|
|
22
|
+
thread_config,
|
|
23
|
+
)
|
|
24
|
+
from databricks_agentkit.langgraph.session_store import (
|
|
25
|
+
invocation_id_from_metadata,
|
|
26
|
+
invocation_metadata,
|
|
27
|
+
)
|
|
28
|
+
from databricks_agentkit.runtime.auth import AuthError
|
|
29
|
+
|
|
30
|
+
# A Unity Catalog AI Gateway model service, served from the `system.ai` schema and queried through
|
|
31
|
+
# the gateway (see `use_ai_gateway=True` below). Swap for any `system.ai.*` model service your
|
|
32
|
+
# workspace exposes — the demo chat app's picker lists what's available.
|
|
33
|
+
MODEL = "system.ai.claude-sonnet-4-5"
|
|
34
|
+
|
|
35
|
+
# Tools that require human approval before they run. Map a tool name to True to allow every decision
|
|
36
|
+
# (approve / edit / reject / respond), or to a config dict to restrict them (see HumanInTheLoopMiddleware).
|
|
37
|
+
# When a listed tool is about to run, the agent pauses and emits an `interrupt` event; the client
|
|
38
|
+
# resumes by sending `resume` with the same session id. Empty this dict to disable approval gating.
|
|
39
|
+
REQUIRE_APPROVAL = {"send_message": True}
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
class _RoutedChatDatabricks(ChatDatabricks):
|
|
43
|
+
"""Forward account-host workspace routing to the underlying OpenAI clients."""
|
|
44
|
+
|
|
45
|
+
def _get_client_kwargs(self) -> dict[str, Any]:
|
|
46
|
+
kwargs = super()._get_client_kwargs()
|
|
47
|
+
if headers := workspace_headers():
|
|
48
|
+
kwargs["default_headers"] = headers
|
|
49
|
+
return kwargs
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def configure() -> None:
|
|
53
|
+
"""Wire up global state; call once at server startup (not at import)."""
|
|
54
|
+
_check_databricks_auth()
|
|
55
|
+
configure_tracing()
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def _check_databricks_auth() -> None:
|
|
59
|
+
"""Fail fast at startup with a clear message if Databricks auth isn't configured.
|
|
60
|
+
|
|
61
|
+
Without this, a missing/invalid profile only surfaces on the first model call — as a generic SDK
|
|
62
|
+
error buried in a request traceback. Resolving a WorkspaceClient here validates the same config
|
|
63
|
+
the model client uses, so the failure is immediate and actionable.
|
|
64
|
+
"""
|
|
65
|
+
try:
|
|
66
|
+
workspace_client()
|
|
67
|
+
except Exception as e:
|
|
68
|
+
profile = os.getenv("DATABRICKS_CONFIG_PROFILE")
|
|
69
|
+
target = (
|
|
70
|
+
f"profile {profile!r}" if profile else "the DEFAULT profile / DATABRICKS_HOST+TOKEN"
|
|
71
|
+
)
|
|
72
|
+
raise RuntimeError(
|
|
73
|
+
f"Databricks auth is not configured — the agent can't call the model. Tried {target}.\n"
|
|
74
|
+
"Fix one of:\n"
|
|
75
|
+
" • set DATABRICKS_CONFIG_PROFILE in .env to a profile from `databricks auth profiles`, or\n"
|
|
76
|
+
" • run `databricks auth login --profile <name>` to create one, or\n"
|
|
77
|
+
" • set DATABRICKS_HOST and DATABRICKS_TOKEN in .env.\n"
|
|
78
|
+
f"(underlying error: {e})"
|
|
79
|
+
) from e
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
async def create_agent_graph(
|
|
83
|
+
actor: str,
|
|
84
|
+
model: str | None = None,
|
|
85
|
+
*,
|
|
86
|
+
workspace_client_for: Callable[[str], WorkspaceClient] | None = None,
|
|
87
|
+
):
|
|
88
|
+
"""Build the LangGraph agent: local tools + long-term-memory tools + any MCP tools.
|
|
89
|
+
|
|
90
|
+
``actor`` is the identity whose long-term memory the agent reads/writes; it's captured in the
|
|
91
|
+
memory tools' closures (never exposed to the model).
|
|
92
|
+
|
|
93
|
+
``model`` selects the gateway model for this run; the chat UI passes the picker's choice and
|
|
94
|
+
everything else falls back to ``MODEL``. The agent is rebuilt per turn, so the model can vary
|
|
95
|
+
request to request.
|
|
96
|
+
"""
|
|
97
|
+
mcp = await mcp_tools(build_mcp_servers(), workspace_client_for=workspace_client_for)
|
|
98
|
+
tools = [
|
|
99
|
+
*all_tools(),
|
|
100
|
+
*memory_tools(actor),
|
|
101
|
+
*genie_tools(workspace_client_for=workspace_client_for),
|
|
102
|
+
*mcp,
|
|
103
|
+
]
|
|
104
|
+
middleware = (
|
|
105
|
+
[HumanInTheLoopMiddleware(interrupt_on=REQUIRE_APPROVAL)] if REQUIRE_APPROVAL else []
|
|
106
|
+
)
|
|
107
|
+
endpoint = model or MODEL
|
|
108
|
+
return create_agent(
|
|
109
|
+
# use_ai_gateway routes to the Unity Catalog AI Gateway (`<host>/ai-gateway/mlflow/v1`), so
|
|
110
|
+
# `endpoint` is a `system.ai.*` model name rather than a serving-endpoint name.
|
|
111
|
+
model=_RoutedChatDatabricks(
|
|
112
|
+
endpoint=endpoint, workspace_client=workspace_client(), use_ai_gateway=True
|
|
113
|
+
),
|
|
114
|
+
tools=tools,
|
|
115
|
+
middleware=middleware,
|
|
116
|
+
checkpointer=checkpointer(),
|
|
117
|
+
)
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
async def recovery_input(
|
|
121
|
+
agent_input: Any,
|
|
122
|
+
*,
|
|
123
|
+
session_id: str,
|
|
124
|
+
actor: str,
|
|
125
|
+
invocation_id: str,
|
|
126
|
+
) -> Any:
|
|
127
|
+
"""Return the input that should be passed to ``run_agent`` after worker loss.
|
|
128
|
+
|
|
129
|
+
A checkpoint tagged with this invocation means LangGraph can continue by receiving ``None``.
|
|
130
|
+
Otherwise recovery replays the original application input.
|
|
131
|
+
"""
|
|
132
|
+
checkpoint = await checkpointer().aget_tuple(thread_config(session_id, actor))
|
|
133
|
+
current_invocation_checkpointed = bool(
|
|
134
|
+
checkpoint and invocation_id_from_metadata(checkpoint.metadata) == invocation_id
|
|
135
|
+
)
|
|
136
|
+
return None if current_invocation_checkpointed else agent_input
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
async def run_agent(
|
|
140
|
+
agent_input: Any,
|
|
141
|
+
*,
|
|
142
|
+
session_id: str,
|
|
143
|
+
actor: str | None = None,
|
|
144
|
+
model: str | None = None,
|
|
145
|
+
invocation_id: str | None = None,
|
|
146
|
+
workspace_client_for: Callable[[str], WorkspaceClient] | None = None,
|
|
147
|
+
) -> AsyncGenerator[Any, None]:
|
|
148
|
+
"""Run the agent and yield native LangGraph stream events.
|
|
149
|
+
|
|
150
|
+
This is the framework-native entrypoint. It has no dependency on Agent Bricks request or context types,
|
|
151
|
+
so it can be called from another server, a notebook, or a test harness.
|
|
152
|
+
"""
|
|
153
|
+
actor = actor or session_id
|
|
154
|
+
graph = await create_agent_graph(actor, model, workspace_client_for=workspace_client_for)
|
|
155
|
+
config = thread_config(session_id, actor)
|
|
156
|
+
if invocation_id:
|
|
157
|
+
config["metadata"] = invocation_metadata(invocation_id)
|
|
158
|
+
|
|
159
|
+
last_update: Any = None
|
|
160
|
+
with start_trace(name="invoke", inputs=agent_input, session_id=session_id) as span:
|
|
161
|
+
async for event in graph.astream(
|
|
162
|
+
input=agent_input,
|
|
163
|
+
config=config,
|
|
164
|
+
stream_mode=["updates", "messages"],
|
|
165
|
+
durability="sync",
|
|
166
|
+
):
|
|
167
|
+
if event[0] == "updates":
|
|
168
|
+
last_update = event[1]
|
|
169
|
+
if workspace_client_for is not None and last_update.get("__interrupt__"):
|
|
170
|
+
raise AuthError(
|
|
171
|
+
"MCP_USER_AUTH_HITL_UNSUPPORTED",
|
|
172
|
+
"Request-user invocations do not support paused approvals.",
|
|
173
|
+
400,
|
|
174
|
+
)
|
|
175
|
+
yield event
|
|
176
|
+
if span is not None and last_update is not None:
|
|
177
|
+
span.set_outputs(last_update)
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
"""MCP servers to offer the agent — this is where you configure them.
|
|
2
|
+
|
|
3
|
+
Empty by default: the agent runs with only the MCP servers declared in ``agent.toml``. Add servers
|
|
4
|
+
to ``build_mcp_servers`` to offer more; ``agent.py`` passes them to ``mcp_tools``, which joins them
|
|
5
|
+
with the ``agent.toml`` servers and fetches their tools each request.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from databricks_langchain import DatabricksMCPServer
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def build_mcp_servers() -> list[DatabricksMCPServer]:
|
|
12
|
+
"""Return the MCP servers to offer the agent. Empty by default — add your own.
|
|
13
|
+
|
|
14
|
+
Example (a Databricks-managed MCP, authed as the app service principal)::
|
|
15
|
+
|
|
16
|
+
from databricks.sdk import WorkspaceClient
|
|
17
|
+
|
|
18
|
+
host = WorkspaceClient().config.host
|
|
19
|
+
return [
|
|
20
|
+
DatabricksMCPServer(
|
|
21
|
+
name="system-ai",
|
|
22
|
+
url=f"{host}/api/2.0/mcp/functions/system/ai",
|
|
23
|
+
workspace_client=WorkspaceClient(),
|
|
24
|
+
),
|
|
25
|
+
]
|
|
26
|
+
"""
|
|
27
|
+
return []
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
"""Agent tools package.
|
|
2
|
+
|
|
3
|
+
Every module here is auto-imported, and every LangChain ``BaseTool`` it defines (via the
|
|
4
|
+
``@tool`` decorator) is collected by ``all_tools()``. Drop a new ``*.py`` into this folder with a
|
|
5
|
+
``@tool``-decorated function and it's picked up automatically — no wiring to edit.
|
|
6
|
+
``create_agent_graph`` uses ``all_tools()``.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
import importlib
|
|
10
|
+
import inspect
|
|
11
|
+
import pkgutil
|
|
12
|
+
|
|
13
|
+
from langchain_core.tools import BaseTool
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def all_tools() -> list[BaseTool]:
|
|
17
|
+
"""Every BaseTool defined across the modules in this package."""
|
|
18
|
+
tools: list[BaseTool] = []
|
|
19
|
+
for module in pkgutil.iter_modules(__path__):
|
|
20
|
+
mod = importlib.import_module(f"{__name__}.{module.name}")
|
|
21
|
+
for _, obj in inspect.getmembers(mod, lambda o: isinstance(o, BaseTool)):
|
|
22
|
+
if obj not in tools: # a tool imported into several modules is collected once
|
|
23
|
+
tools.append(obj)
|
|
24
|
+
return tools
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
"""Sample tool. A working example — add your own tools as new files in this package.
|
|
2
|
+
|
|
3
|
+
Decorate a function with ``@tool`` (LangChain) and it becomes an agent tool; the package
|
|
4
|
+
auto-collects it via ``all_tools()``, which ``create_agent_graph`` uses.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from datetime import datetime
|
|
8
|
+
|
|
9
|
+
from langchain_core.tools import tool
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
@tool
|
|
13
|
+
def get_current_time() -> str:
|
|
14
|
+
"""Get the current date and time."""
|
|
15
|
+
return datetime.now().isoformat()
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""A side-effecting sample tool, gated by human approval.
|
|
2
|
+
|
|
3
|
+
Unlike ``get_current_time`` (a harmless read), this tool stands in for an action with real
|
|
4
|
+
consequences — sending a message. ``agent/agent.py`` lists it in ``REQUIRE_APPROVAL``, so before it
|
|
5
|
+
runs the agent pauses and asks a human to approve / edit / reject the call. Swap the body for a real
|
|
6
|
+
send; the approval gate is what the template is demonstrating.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from langchain_core.tools import tool
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
@tool
|
|
13
|
+
def send_message(recipient: str, body: str) -> str:
|
|
14
|
+
"""Send a message to a recipient. Use when the user asks to notify or message someone."""
|
|
15
|
+
# A real implementation would call an email/Slack/SMS API here. The point of this template is the
|
|
16
|
+
# human-approval gate in front of it (see REQUIRE_APPROVAL in agent/agent.py), not the delivery.
|
|
17
|
+
return f"Message sent to {recipient}: {body}"
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
command: ["uv", "run", "start-server"]
|
|
2
|
+
# databricks apps listen by default on port 8000
|
|
3
|
+
#
|
|
4
|
+
# Agent Bricks injects the deployed Runtime Store endpoint. The LangGraph checkpointer remains
|
|
5
|
+
# in-process unless a managed Session Store is bound. Tracing is bound with `ab tracing bind
|
|
6
|
+
# --experiment-name <path>` (ab deploy then wires the env below; `ab dev` traces to a local
|
|
7
|
+
# MLflow server). To set a deployment's tracing env by hand instead:
|
|
8
|
+
#
|
|
9
|
+
# env:
|
|
10
|
+
# # Tracing (set a destination + an experiment). Destination: MLFLOW_TRACKING_URI or
|
|
11
|
+
# # MLFLOW_TRACING_DESTINATION. Experiment: MLFLOW_EXPERIMENT_ID or MLFLOW_EXPERIMENT_NAME.
|
|
12
|
+
# - name: MLFLOW_TRACKING_URI
|
|
13
|
+
# value: "databricks"
|
|
14
|
+
# - name: MLFLOW_EXPERIMENT_ID
|
|
15
|
+
# value: "<experiment-id>"
|
|
@@ -0,0 +1,42 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "agent-langgraph"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "LangGraph agent on the Agent Bricks durable runtime for Databricks Apps"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
authors = [
|
|
7
|
+
{ name = "Agent Developer", email = "developer@example.com" }
|
|
8
|
+
]
|
|
9
|
+
requires-python = ">=3.11"
|
|
10
|
+
dependencies = [
|
|
11
|
+
# Agent Bricks includes the HTTP runtime; this extra adds the langgraph framework adapter.
|
|
12
|
+
"databricks-agentbricks[langgraph]>=0.2.0",
|
|
13
|
+
"uvicorn>=0.41.0",
|
|
14
|
+
"python-dotenv>=1.2.1",
|
|
15
|
+
"databricks-sdk>=0.124.0",
|
|
16
|
+
"databricks-agents>=1.9.3",
|
|
17
|
+
]
|
|
18
|
+
|
|
19
|
+
[build-system]
|
|
20
|
+
requires = ["hatchling"]
|
|
21
|
+
build-backend = "hatchling.build"
|
|
22
|
+
|
|
23
|
+
[tool.hatch.build.targets.wheel]
|
|
24
|
+
packages = ["agent", "runtime"]
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
[dependency-groups]
|
|
28
|
+
dev = [
|
|
29
|
+
"hatchling>=1.28.0",
|
|
30
|
+
"pytest>=9.0.2",
|
|
31
|
+
"pytest-asyncio>=1.3.0",
|
|
32
|
+
]
|
|
33
|
+
|
|
34
|
+
[tool.uv]
|
|
35
|
+
default-groups = ["dev"]
|
|
36
|
+
|
|
37
|
+
[tool.pytest.ini_options]
|
|
38
|
+
testpaths = ["tests"]
|
|
39
|
+
asyncio_mode = "auto"
|
|
40
|
+
|
|
41
|
+
[project.scripts]
|
|
42
|
+
start-server = "runtime.main:main"
|
|
File without changes
|