databricks-mason 0.1.3.dev0__tar.gz → 0.1.5.dev0__tar.gz
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_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/PKG-INFO +136 -36
- databricks_mason-0.1.5.dev0/README.md +313 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/pyproject.toml +15 -5
- databricks_mason-0.1.5.dev0/src/databricks_mason/__init__.py +52 -0
- databricks_mason-0.1.3.dev0/src/databricks_mason/client.py → databricks_mason-0.1.5.dev0/src/databricks_mason/_api_client.py +63 -41
- databricks_mason-0.1.5.dev0/src/databricks_mason/_pagination.py +15 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/agent_project.py +96 -9
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/auth.py +4 -4
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/cli.py +5 -5
- databricks_mason-0.1.5.dev0/src/databricks_mason/client.py +28 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/deploy.py +250 -123
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/dev.py +58 -46
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/help.py +5 -2
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/init.py +201 -31
- databricks_mason-0.1.5.dev0/src/databricks_mason/lakebase_durability_store.py +78 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/memory.py +35 -15
- databricks_mason-0.1.5.dev0/src/databricks_mason/memory_store.py +356 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/render.py +29 -1
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/runtime/__init__.py +8 -2
- databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/__init__.py +1 -0
- databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/app.py +227 -0
- databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/attempt.py +140 -0
- databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/recovery.py +112 -0
- databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/runtime.py +167 -0
- databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/store.py +783 -0
- databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/types.py +189 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/runtime/tool_manifest.py +25 -12
- databricks_mason-0.1.5.dev0/src/databricks_mason/session_store.py +358 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/sessions.py +22 -7
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/store_access.py +6 -1
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/tracing.py +10 -4
- databricks_mason-0.1.3.dev0/README.md +0 -215
- databricks_mason-0.1.3.dev0/src/databricks_mason/__init__.py +0 -89
- databricks_mason-0.1.3.dev0/src/databricks_mason/runtime/background.py +0 -37
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/.gitignore +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/NOTICE +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/errors.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/langgraph/__init__.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/langgraph/mcp.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/langgraph/memory.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/langgraph/session_store.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/mcp.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/memory_store_access.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/models.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/openai/__init__.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/openai/mcp.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/openai/memory.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/openai/sessions.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/project_config.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/py.typed +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/runtime/session_store_client.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/runtime/tracing.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/runtime/workspace.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/sandbox.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/session_store_access.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/templates/python_tool_langgraph.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/templates/python_tool_test.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/templates/sandbox_mcp.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/templates/sandbox_mcp_langgraph.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/timefmt.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/tools.py +0 -0
|
@@ -1,12 +1,12 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: databricks-mason
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.5.dev0
|
|
4
4
|
Summary: Databricks integration for Mason
|
|
5
5
|
Author-email: Databricks <agent-feedback@databricks.com>
|
|
6
6
|
License-File: NOTICE
|
|
7
7
|
Requires-Python: >=3.10
|
|
8
8
|
Requires-Dist: click>=8.1
|
|
9
|
-
Requires-Dist: databricks-sdk>=0.
|
|
9
|
+
Requires-Dist: databricks-sdk>=0.94.0
|
|
10
10
|
Requires-Dist: psycopg[binary]>=3.1
|
|
11
11
|
Requires-Dist: pyyaml>=6.0
|
|
12
12
|
Requires-Dist: rich>=13.7
|
|
@@ -14,6 +14,7 @@ Requires-Dist: tomli>=2.0
|
|
|
14
14
|
Requires-Dist: tomlkit>=0.13
|
|
15
15
|
Provides-Extra: runtime
|
|
16
16
|
Requires-Dist: databricks-agents>=1.9.3; extra == 'runtime'
|
|
17
|
+
Requires-Dist: databricks-ai-bridge[memory]>=0.21.0; extra == 'runtime'
|
|
17
18
|
Requires-Dist: databricks-langchain>=0.17.0; extra == 'runtime'
|
|
18
19
|
Requires-Dist: fastapi>=0.129.0; extra == 'runtime'
|
|
19
20
|
Requires-Dist: langchain-mcp-adapters>=0.2.1; extra == 'runtime'
|
|
@@ -21,17 +22,18 @@ Requires-Dist: langchain>=1.0.0; extra == 'runtime'
|
|
|
21
22
|
Requires-Dist: langgraph>=1.1.0; extra == 'runtime'
|
|
22
23
|
Requires-Dist: mlflow>=3.10.1; extra == 'runtime'
|
|
23
24
|
Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime'
|
|
24
|
-
Requires-Dist:
|
|
25
|
+
Requires-Dist: uvicorn>=0.20.0; extra == 'runtime'
|
|
25
26
|
Provides-Extra: runtime-openai
|
|
26
27
|
Requires-Dist: databricks-agents>=1.9.3; extra == 'runtime-openai'
|
|
28
|
+
Requires-Dist: databricks-ai-bridge[memory]>=0.21.0; extra == 'runtime-openai'
|
|
27
29
|
Requires-Dist: databricks-openai>=0.13.0; extra == 'runtime-openai'
|
|
28
30
|
Requires-Dist: fastapi>=0.129.0; extra == 'runtime-openai'
|
|
29
31
|
Requires-Dist: mlflow>=3.10.1; extra == 'runtime-openai'
|
|
30
32
|
Requires-Dist: openai-agents>=0.7.0; extra == 'runtime-openai'
|
|
31
33
|
Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime-openai'
|
|
32
|
-
Requires-Dist:
|
|
34
|
+
Requires-Dist: uvicorn>=0.20.0; extra == 'runtime-openai'
|
|
33
35
|
Provides-Extra: tracing
|
|
34
|
-
Requires-Dist: mlflow[databricks]>=3.
|
|
36
|
+
Requires-Dist: mlflow[databricks]>=3.10.1; extra == 'tracing'
|
|
35
37
|
Description-Content-Type: text/markdown
|
|
36
38
|
|
|
37
39
|
# `databricks-mason`
|
|
@@ -62,6 +64,12 @@ For tracing commands, install Mason with tracing extras:
|
|
|
62
64
|
pip install 'databricks-mason[tracing]'
|
|
63
65
|
```
|
|
64
66
|
|
|
67
|
+
For the SDK-hosted durable agent application, install the runtime extra:
|
|
68
|
+
|
|
69
|
+
```sh
|
|
70
|
+
pip install 'databricks-mason[runtime]'
|
|
71
|
+
```
|
|
72
|
+
|
|
65
73
|
## Shell completion
|
|
66
74
|
Add this to `~/.zshrc`:
|
|
67
75
|
```sh
|
|
@@ -90,28 +98,117 @@ You can also pass the global `--profile/-p` option before an individual command,
|
|
|
90
98
|
|
|
91
99
|
## Python SDK
|
|
92
100
|
|
|
93
|
-
|
|
94
|
-
`
|
|
95
|
-
|
|
101
|
+
`MasonClient` adds a small resource-oriented layer over the Mason API. Pass it an
|
|
102
|
+
authenticated Databricks `WorkspaceClient`, or omit the argument to use the
|
|
103
|
+
Databricks SDK's default authentication resolution:
|
|
96
104
|
|
|
97
105
|
```python
|
|
106
|
+
from databricks.sdk import WorkspaceClient
|
|
98
107
|
from databricks_mason import MasonClient
|
|
99
108
|
|
|
100
|
-
|
|
109
|
+
mason = MasonClient(WorkspaceClient(profile="my-workspace"))
|
|
110
|
+
|
|
111
|
+
session_store = mason.session_stores.create("support-agent-sessions")
|
|
112
|
+
session = session_store.add(actor_id="customer-123", session_id="case-456")
|
|
113
|
+
session.append_items(
|
|
114
|
+
[
|
|
115
|
+
{"type": "message", "role": "user", "content": "I need help with my cluster."},
|
|
116
|
+
{"type": "message", "role": "assistant", "content": "Let's take a look."},
|
|
117
|
+
]
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
memory_store = mason.memory_stores.create("coding-agent-memory")
|
|
121
|
+
memory = memory_store.add(
|
|
122
|
+
actor_id="alice",
|
|
123
|
+
path="/preferences/style.md",
|
|
124
|
+
content="The user prefers concise answers.",
|
|
125
|
+
)
|
|
126
|
+
results = memory_store.search(
|
|
127
|
+
actor_id="alice",
|
|
128
|
+
query="response preferences",
|
|
129
|
+
limit=10,
|
|
130
|
+
)
|
|
131
|
+
memory = memory.update(content="The user prefers very concise answers.")
|
|
132
|
+
memory.delete()
|
|
133
|
+
```
|
|
134
|
+
|
|
135
|
+
The root collections manage stores: `mason.memory_stores.create/get/list` and
|
|
136
|
+
`mason.session_stores.create/get/list`. A returned store owns operations on its
|
|
137
|
+
contents, such as `memory_store.add()`, `memory_store.get("memory-id")`,
|
|
138
|
+
`memory_store.list()`, and `memory_store.search()`, or `session_store.add()`,
|
|
139
|
+
`session_store.get("session-id")`, and `session_store.list()`. Returned memories,
|
|
140
|
+
sessions, and stores own their `update()` and `delete()` operations.
|
|
141
|
+
|
|
142
|
+
All `list()` methods return iterators that automatically consume server pages. List
|
|
143
|
+
`page_size` and search `limit` values must be between 1 and 100. `session.list_items()`
|
|
144
|
+
also auto-pages. `session.fork(...)` creates an independent copy, optionally through
|
|
145
|
+
a specific item. Deleting a session with descendants requires
|
|
146
|
+
`session.delete(force=True)` to cascade the deletion.
|
|
147
|
+
|
|
148
|
+
The resource layer intentionally does not mirror every API method. Its private
|
|
149
|
+
transport will be replaced by the generated `WorkspaceClient.mason` service when that
|
|
150
|
+
is released, without changing this public surface. Deployment, sandbox, tracing, and
|
|
151
|
+
the existing CLI commands remain separate.
|
|
101
152
|
|
|
102
|
-
|
|
103
|
-
print(store.name, store.display_name) # typed attribute access
|
|
153
|
+
## Agent application
|
|
104
154
|
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
155
|
+
`AgentApp` provides Mason's invocation HTTP contract, including foreground, streaming, background,
|
|
156
|
+
polling, and event endpoints. By default its state is process-local. Set `durable_runtime=True` to
|
|
157
|
+
use Lakebase persistence, heartbeats, and crash recovery after deployment:
|
|
158
|
+
|
|
159
|
+
```python
|
|
160
|
+
from databricks_mason import AgentApp, DurableAgentContext
|
|
161
|
+
|
|
162
|
+
app = AgentApp(durable_runtime=True)
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
@app.invoke
|
|
166
|
+
async def invoke(input: object, context: DurableAgentContext) -> object:
|
|
167
|
+
return await run_agent(input, session_id=context.session_id)
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
@app.on_recovery
|
|
171
|
+
async def recover(input: object, context: DurableAgentContext) -> object:
|
|
172
|
+
return await recover_agent(input, session_id=context.session_id)
|
|
108
173
|
```
|
|
109
174
|
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
`
|
|
175
|
+
The Mason server exposes `POST /api/invocations`, `GET /api/invocations/{invocation_id}`, and
|
|
176
|
+
`GET /api/invocations/{invocation_id}/events?after={cursor}`. Databricks Apps bearer-token requests
|
|
177
|
+
must use `/api/` routes
|
|
178
|
+
([Apps documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/connect-local)).
|
|
179
|
+
The client supplies a UUID `id`, which is also the idempotency key for every invocation mode:
|
|
180
|
+
|
|
181
|
+
- foreground sync returns `200` with the result under `output`;
|
|
182
|
+
- background sync returns `202` with a status URL;
|
|
183
|
+
- foreground streaming returns `200` server-sent events; and
|
|
184
|
+
- background streaming returns `202` with status and event URLs.
|
|
185
|
+
|
|
186
|
+
`input` and `output` may be any JSON value. Transport fields are not passed to the callback. A
|
|
187
|
+
top-level `session_id` is rejected, but a framework template may carry its own stable application
|
|
188
|
+
session inside `input`. Polling uses only the invocation ID and relies on Databricks Apps
|
|
189
|
+
authentication. Without the durable runtime, request state and events exist only in the serving
|
|
190
|
+
process and horizontally scaled clients need sticky routing. With the durable runtime, Mason
|
|
191
|
+
persists the input, attempt status, heartbeats, lifecycle events, application events, and output.
|
|
192
|
+
|
|
193
|
+
Durability is enabled by default for both framework templates. Mason writes the durability binding
|
|
194
|
+
to `agent.toml`, and `mason deploy` then attaches one Lakebase database for runtime durability,
|
|
195
|
+
chosen in this order:
|
|
196
|
+
|
|
197
|
+
1. Reuse the configured Session Store's Lakebase database.
|
|
198
|
+
2. Otherwise reuse or provision a dedicated `<app>-durability` Lakebase project.
|
|
199
|
+
|
|
200
|
+
Mason adds only its `databricks_mason_runtime_<app-hash>` schema and tables to the selected database,
|
|
201
|
+
giving each app one owned schema. A replacement worker claims a stale heartbeat and calls the
|
|
202
|
+
`@app.on_recovery` handler. If that handler is omitted, startup warns that automatic crash recovery
|
|
203
|
+
is disabled; register the same function for both decorators when replaying the initial invocation is
|
|
204
|
+
safe. Agent checkpoint restoration and idempotent external side effects remain the developer's
|
|
205
|
+
responsibility.
|
|
206
|
+
|
|
207
|
+
Bare `mason init`, `--framework langgraph`, and `--framework openai` scaffold `AgentApp` with its
|
|
208
|
+
durable runtime enabled. Pass `--no-durable-runtime` for the same Mason HTTP contract with
|
|
209
|
+
process-local state and no Lakebase provisioning. Pass `--server custom` for a minimal FastAPI
|
|
210
|
+
server with one foreground `/invocations` route and no Mason `AgentApp`. Use `--disable-chat-app`
|
|
211
|
+
independently for API-only Mason server output.
|
|
115
212
|
|
|
116
213
|
## Commands
|
|
117
214
|
|
|
@@ -119,15 +216,19 @@ accessors (`store.name`) while remaining plain dicts underneath — so `store["n
|
|
|
119
216
|
mason [-p <profile>] [-o text|json]
|
|
120
217
|
login [--profile P]
|
|
121
218
|
logout
|
|
122
|
-
init [--framework openai|langgraph] [--
|
|
219
|
+
init [--framework openai|langgraph] [--server mason|custom]
|
|
220
|
+
[--no-durable-runtime] [--disable-chat-app]
|
|
123
221
|
[--profile P] [--repo URL] [--ref REF] [directory]
|
|
124
222
|
dev [--source PATH] [--prepare-environment] [--app-port PORT]
|
|
125
|
-
[--
|
|
126
|
-
[--with-traces C.S] [--no-create-stores]
|
|
223
|
+
[--with-traces C.S]
|
|
127
224
|
memory
|
|
225
|
+
bind STORE [--source PATH] [--no-create-stores]
|
|
226
|
+
unbind [--source PATH]
|
|
128
227
|
stores create | list | get | update | delete
|
|
129
228
|
entries create | get | list | search | update | delete
|
|
130
229
|
sessions create | list | get | update | delete | fork
|
|
230
|
+
bind STORE [--source PATH] [--no-create-stores]
|
|
231
|
+
unbind [--source PATH]
|
|
131
232
|
stores create | list | get | update | delete
|
|
132
233
|
items list | append | pop | clear
|
|
133
234
|
tracing
|
|
@@ -141,9 +242,7 @@ mason [-p <profile>] [-o text|json]
|
|
|
141
242
|
add uc-function FUNCTION [--name NAME] [--source PATH]
|
|
142
243
|
add python NAME [--source PATH]
|
|
143
244
|
list [--source PATH]
|
|
144
|
-
deploy <name> --source PATH [--
|
|
145
|
-
[--session/-s N] [--actor-id ID]
|
|
146
|
-
[--with-traces C.S] [--no-create-stores]
|
|
245
|
+
deploy <name> --source PATH [--with-traces C.S] [--instances N]
|
|
147
246
|
deployments list | get | logs | start | stop | delete
|
|
148
247
|
```
|
|
149
248
|
|
|
@@ -222,28 +321,29 @@ mason init --framework langgraph \
|
|
|
222
321
|
--profile <profile> \
|
|
223
322
|
./my-agent
|
|
224
323
|
cd ./my-agent
|
|
225
|
-
|
|
324
|
+
mason dev
|
|
226
325
|
```
|
|
227
326
|
|
|
228
327
|
The chat app includes synchronous, SSE streaming, background polling, Session Store, Memory Store,
|
|
229
328
|
and HITL resume UI. The framework-specific overlay adds `ui/`, `runtime/ui.py`, the UI-enabled
|
|
230
329
|
`runtime/main.py`, and UI tests.
|
|
231
330
|
|
|
232
|
-
For the full deployed demo,
|
|
331
|
+
For the full deployed demo, bind both managed stores, then deploy:
|
|
233
332
|
|
|
234
333
|
```sh
|
|
235
|
-
mason
|
|
236
|
-
|
|
237
|
-
|
|
238
|
-
--actor-id alice
|
|
334
|
+
mason sessions bind mason-demo-sessions
|
|
335
|
+
mason memory bind mason-demo-memory
|
|
336
|
+
mason --profile <profile> deploy mason-agent-demo --source .
|
|
239
337
|
```
|
|
240
338
|
|
|
241
|
-
(
|
|
339
|
+
(Binding creates a missing store automatically; pass `--no-create-stores` to require it already
|
|
340
|
+
exists. The agent reads the bound stores from `agent.toml` at runtime; `deploy` grants the app's
|
|
341
|
+
service principal access to them.)
|
|
242
342
|
|
|
243
|
-
The
|
|
244
|
-
|
|
245
|
-
|
|
246
|
-
|
|
343
|
+
The chat UI generates a stable application session UUID in browser local storage, places it inside
|
|
344
|
+
the durable invocation's opaque `input`, and creates a fresh invocation UUID per turn. The
|
|
345
|
+
`__Host-databricks-app-router` cookie remains independent: API clients may reuse it for sticky
|
|
346
|
+
replica routing, but it is neither authentication nor the template's application session state.
|
|
247
347
|
|
|
248
348
|
The generated `README.md` documents every request the client makes: config discovery, sync and SSE
|
|
249
349
|
invocations, background submission and polling, session transcript loading, HITL resume, and memory
|
|
@@ -0,0 +1,313 @@
|
|
|
1
|
+
# `databricks-mason`
|
|
2
|
+
|
|
3
|
+
Mason is an experimental CLI for Databricks custom agent preview APIs and
|
|
4
|
+
deployments. It manages memory, sessions, tracing, and deployments from one
|
|
5
|
+
authenticated command.
|
|
6
|
+
|
|
7
|
+
> The underlying APIs are in preview and may need workspace enablement.
|
|
8
|
+
|
|
9
|
+
## Installation
|
|
10
|
+
|
|
11
|
+
From PyPI:
|
|
12
|
+
|
|
13
|
+
```sh
|
|
14
|
+
pip install databricks-mason
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
From source:
|
|
18
|
+
|
|
19
|
+
```sh
|
|
20
|
+
pip install 'git+https://github.com/databricks/databricks-ai-bridge.git#subdirectory=integrations/mason'
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
For tracing commands, install Mason with tracing extras:
|
|
24
|
+
|
|
25
|
+
```sh
|
|
26
|
+
pip install 'databricks-mason[tracing]'
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
For the SDK-hosted durable agent application, install the runtime extra:
|
|
30
|
+
|
|
31
|
+
```sh
|
|
32
|
+
pip install 'databricks-mason[runtime]'
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
## Shell completion
|
|
36
|
+
Add this to `~/.zshrc`:
|
|
37
|
+
```sh
|
|
38
|
+
eval "$(_MASON_COMPLETE=zsh_source mason)"
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
## Authentication
|
|
42
|
+
|
|
43
|
+
Mason uses [Databricks authentication](https://docs.databricks.com/aws/en/dev-tools/cli/authentication).
|
|
44
|
+
Ask Mason to authenticate and remember a named profile:
|
|
45
|
+
|
|
46
|
+
```sh
|
|
47
|
+
mason login --profile <profile>
|
|
48
|
+
mason sessions stores list
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
`mason login` validates existing credentials first. If credentials are missing or rejected in
|
|
52
|
+
an interactive terminal, Mason runs `databricks auth login --profile <profile>`, revalidates the
|
|
53
|
+
profile, and stores the selection in `~/.mason/config.json`. This browser-based setup requires
|
|
54
|
+
the Databricks CLI. In non-interactive environments, authenticate the profile before running
|
|
55
|
+
Mason. `mason logout` forgets the saved selection without revoking the underlying credentials.
|
|
56
|
+
|
|
57
|
+
If Databricks SDK default authentication is already configured, you can skip `mason login`.
|
|
58
|
+
You can also pass the global `--profile/-p` option before an individual command, for example
|
|
59
|
+
`mason --profile <profile> mcp list`. Use `--output json` for scripting.
|
|
60
|
+
|
|
61
|
+
## Python SDK
|
|
62
|
+
|
|
63
|
+
`MasonClient` adds a small resource-oriented layer over the Mason API. Pass it an
|
|
64
|
+
authenticated Databricks `WorkspaceClient`, or omit the argument to use the
|
|
65
|
+
Databricks SDK's default authentication resolution:
|
|
66
|
+
|
|
67
|
+
```python
|
|
68
|
+
from databricks.sdk import WorkspaceClient
|
|
69
|
+
from databricks_mason import MasonClient
|
|
70
|
+
|
|
71
|
+
mason = MasonClient(WorkspaceClient(profile="my-workspace"))
|
|
72
|
+
|
|
73
|
+
session_store = mason.session_stores.create("support-agent-sessions")
|
|
74
|
+
session = session_store.add(actor_id="customer-123", session_id="case-456")
|
|
75
|
+
session.append_items(
|
|
76
|
+
[
|
|
77
|
+
{"type": "message", "role": "user", "content": "I need help with my cluster."},
|
|
78
|
+
{"type": "message", "role": "assistant", "content": "Let's take a look."},
|
|
79
|
+
]
|
|
80
|
+
)
|
|
81
|
+
|
|
82
|
+
memory_store = mason.memory_stores.create("coding-agent-memory")
|
|
83
|
+
memory = memory_store.add(
|
|
84
|
+
actor_id="alice",
|
|
85
|
+
path="/preferences/style.md",
|
|
86
|
+
content="The user prefers concise answers.",
|
|
87
|
+
)
|
|
88
|
+
results = memory_store.search(
|
|
89
|
+
actor_id="alice",
|
|
90
|
+
query="response preferences",
|
|
91
|
+
limit=10,
|
|
92
|
+
)
|
|
93
|
+
memory = memory.update(content="The user prefers very concise answers.")
|
|
94
|
+
memory.delete()
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
The root collections manage stores: `mason.memory_stores.create/get/list` and
|
|
98
|
+
`mason.session_stores.create/get/list`. A returned store owns operations on its
|
|
99
|
+
contents, such as `memory_store.add()`, `memory_store.get("memory-id")`,
|
|
100
|
+
`memory_store.list()`, and `memory_store.search()`, or `session_store.add()`,
|
|
101
|
+
`session_store.get("session-id")`, and `session_store.list()`. Returned memories,
|
|
102
|
+
sessions, and stores own their `update()` and `delete()` operations.
|
|
103
|
+
|
|
104
|
+
All `list()` methods return iterators that automatically consume server pages. List
|
|
105
|
+
`page_size` and search `limit` values must be between 1 and 100. `session.list_items()`
|
|
106
|
+
also auto-pages. `session.fork(...)` creates an independent copy, optionally through
|
|
107
|
+
a specific item. Deleting a session with descendants requires
|
|
108
|
+
`session.delete(force=True)` to cascade the deletion.
|
|
109
|
+
|
|
110
|
+
The resource layer intentionally does not mirror every API method. Its private
|
|
111
|
+
transport will be replaced by the generated `WorkspaceClient.mason` service when that
|
|
112
|
+
is released, without changing this public surface. Deployment, sandbox, tracing, and
|
|
113
|
+
the existing CLI commands remain separate.
|
|
114
|
+
|
|
115
|
+
## Agent application
|
|
116
|
+
|
|
117
|
+
`AgentApp` provides Mason's invocation HTTP contract, including foreground, streaming, background,
|
|
118
|
+
polling, and event endpoints. By default its state is process-local. Set `durable_runtime=True` to
|
|
119
|
+
use Lakebase persistence, heartbeats, and crash recovery after deployment:
|
|
120
|
+
|
|
121
|
+
```python
|
|
122
|
+
from databricks_mason import AgentApp, DurableAgentContext
|
|
123
|
+
|
|
124
|
+
app = AgentApp(durable_runtime=True)
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
@app.invoke
|
|
128
|
+
async def invoke(input: object, context: DurableAgentContext) -> object:
|
|
129
|
+
return await run_agent(input, session_id=context.session_id)
|
|
130
|
+
|
|
131
|
+
|
|
132
|
+
@app.on_recovery
|
|
133
|
+
async def recover(input: object, context: DurableAgentContext) -> object:
|
|
134
|
+
return await recover_agent(input, session_id=context.session_id)
|
|
135
|
+
```
|
|
136
|
+
|
|
137
|
+
The Mason server exposes `POST /api/invocations`, `GET /api/invocations/{invocation_id}`, and
|
|
138
|
+
`GET /api/invocations/{invocation_id}/events?after={cursor}`. Databricks Apps bearer-token requests
|
|
139
|
+
must use `/api/` routes
|
|
140
|
+
([Apps documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/connect-local)).
|
|
141
|
+
The client supplies a UUID `id`, which is also the idempotency key for every invocation mode:
|
|
142
|
+
|
|
143
|
+
- foreground sync returns `200` with the result under `output`;
|
|
144
|
+
- background sync returns `202` with a status URL;
|
|
145
|
+
- foreground streaming returns `200` server-sent events; and
|
|
146
|
+
- background streaming returns `202` with status and event URLs.
|
|
147
|
+
|
|
148
|
+
`input` and `output` may be any JSON value. Transport fields are not passed to the callback. A
|
|
149
|
+
top-level `session_id` is rejected, but a framework template may carry its own stable application
|
|
150
|
+
session inside `input`. Polling uses only the invocation ID and relies on Databricks Apps
|
|
151
|
+
authentication. Without the durable runtime, request state and events exist only in the serving
|
|
152
|
+
process and horizontally scaled clients need sticky routing. With the durable runtime, Mason
|
|
153
|
+
persists the input, attempt status, heartbeats, lifecycle events, application events, and output.
|
|
154
|
+
|
|
155
|
+
Durability is enabled by default for both framework templates. Mason writes the durability binding
|
|
156
|
+
to `agent.toml`, and `mason deploy` then attaches one Lakebase database for runtime durability,
|
|
157
|
+
chosen in this order:
|
|
158
|
+
|
|
159
|
+
1. Reuse the configured Session Store's Lakebase database.
|
|
160
|
+
2. Otherwise reuse or provision a dedicated `<app>-durability` Lakebase project.
|
|
161
|
+
|
|
162
|
+
Mason adds only its `databricks_mason_runtime_<app-hash>` schema and tables to the selected database,
|
|
163
|
+
giving each app one owned schema. A replacement worker claims a stale heartbeat and calls the
|
|
164
|
+
`@app.on_recovery` handler. If that handler is omitted, startup warns that automatic crash recovery
|
|
165
|
+
is disabled; register the same function for both decorators when replaying the initial invocation is
|
|
166
|
+
safe. Agent checkpoint restoration and idempotent external side effects remain the developer's
|
|
167
|
+
responsibility.
|
|
168
|
+
|
|
169
|
+
Bare `mason init`, `--framework langgraph`, and `--framework openai` scaffold `AgentApp` with its
|
|
170
|
+
durable runtime enabled. Pass `--no-durable-runtime` for the same Mason HTTP contract with
|
|
171
|
+
process-local state and no Lakebase provisioning. Pass `--server custom` for a minimal FastAPI
|
|
172
|
+
server with one foreground `/invocations` route and no Mason `AgentApp`. Use `--disable-chat-app`
|
|
173
|
+
independently for API-only Mason server output.
|
|
174
|
+
|
|
175
|
+
## Commands
|
|
176
|
+
|
|
177
|
+
```text
|
|
178
|
+
mason [-p <profile>] [-o text|json]
|
|
179
|
+
login [--profile P]
|
|
180
|
+
logout
|
|
181
|
+
init [--framework openai|langgraph] [--server mason|custom]
|
|
182
|
+
[--no-durable-runtime] [--disable-chat-app]
|
|
183
|
+
[--profile P] [--repo URL] [--ref REF] [directory]
|
|
184
|
+
dev [--source PATH] [--prepare-environment] [--app-port PORT]
|
|
185
|
+
[--with-traces C.S]
|
|
186
|
+
memory
|
|
187
|
+
bind STORE [--source PATH] [--no-create-stores]
|
|
188
|
+
unbind [--source PATH]
|
|
189
|
+
stores create | list | get | update | delete
|
|
190
|
+
entries create | get | list | search | update | delete
|
|
191
|
+
sessions create | list | get | update | delete | fork
|
|
192
|
+
bind STORE [--source PATH] [--no-create-stores]
|
|
193
|
+
unbind [--source PATH]
|
|
194
|
+
stores create | list | get | update | delete
|
|
195
|
+
items list | append | pop | clear
|
|
196
|
+
tracing
|
|
197
|
+
setup --catalog C --schema S [--experiment E]
|
|
198
|
+
list | get | instrument
|
|
199
|
+
mcp
|
|
200
|
+
list [--schema CATALOG.SCHEMA]
|
|
201
|
+
tools
|
|
202
|
+
add sandbox --scope SCOPE [--scope SCOPE ...] [--source PATH]
|
|
203
|
+
add mcp SERVICE [--name NAME] [--source PATH]
|
|
204
|
+
add uc-function FUNCTION [--name NAME] [--source PATH]
|
|
205
|
+
add python NAME [--source PATH]
|
|
206
|
+
list [--source PATH]
|
|
207
|
+
deploy <name> --source PATH [--with-traces C.S] [--instances N]
|
|
208
|
+
deployments list | get | logs | start | stop | delete
|
|
209
|
+
```
|
|
210
|
+
|
|
211
|
+
## Command help
|
|
212
|
+
|
|
213
|
+
Use the conventional help flag at any command level. Every command's help includes runnable
|
|
214
|
+
examples:
|
|
215
|
+
|
|
216
|
+
```sh
|
|
217
|
+
mason --help
|
|
218
|
+
mason deploy --help
|
|
219
|
+
mason sessions items append --help
|
|
220
|
+
```
|
|
221
|
+
|
|
222
|
+
For the shortest path from a blank directory to a running and deployed agent:
|
|
223
|
+
|
|
224
|
+
```sh
|
|
225
|
+
mason login --profile <profile>
|
|
226
|
+
mason init my-agent
|
|
227
|
+
cd my-agent
|
|
228
|
+
mason dev
|
|
229
|
+
mason deploy my-agent
|
|
230
|
+
```
|
|
231
|
+
|
|
232
|
+
## Agent tools
|
|
233
|
+
|
|
234
|
+
`mason init` writes portable tool intent to `agent.toml` and template provenance to
|
|
235
|
+
`.mason/project.toml`. The manifest runtime is currently implemented only by the in-repository
|
|
236
|
+
`agent-langgraph` template; `mason tools add` fails explicitly for other frameworks until they
|
|
237
|
+
provide an adapter at the same runtime seam.
|
|
238
|
+
|
|
239
|
+
Remote tools update only `agent.toml`; they do not generate framework source. The LangGraph runtime
|
|
240
|
+
loads the manifest and materializes its native MCP tools when the agent runs, so a direct manifest
|
|
241
|
+
edit and a CLI edit have the same behavior:
|
|
242
|
+
|
|
243
|
+
```sh
|
|
244
|
+
mason tools add sandbox --scope table:samples.nyctaxi.trips
|
|
245
|
+
mason tools add mcp system.ai.web_search
|
|
246
|
+
mason tools add uc-function catalog.schema.lookup_ticket
|
|
247
|
+
mason tools add python lookup-ticket
|
|
248
|
+
mason tools remove mcp system.ai.web_search
|
|
249
|
+
mason tools list
|
|
250
|
+
```
|
|
251
|
+
|
|
252
|
+
For MCP services, the remove command accepts the same service name as the add command. You can also
|
|
253
|
+
remove any binding by the ID shown in `mason tools list`, for example `mason tools remove
|
|
254
|
+
web_search`. Removal updates only `agent.toml`; Python source and test files remain user-owned.
|
|
255
|
+
|
|
256
|
+
Discover the MCP Services available to your user before adding one. By default Mason lists the
|
|
257
|
+
Databricks-managed services in `system.ai`; pass `--schema catalog.schema` for another Unity Catalog
|
|
258
|
+
schema. Text output includes a copyable add command, while `--output json` returns normalized service
|
|
259
|
+
records for scripts:
|
|
260
|
+
|
|
261
|
+
```sh
|
|
262
|
+
mason mcp list
|
|
263
|
+
mason mcp list --schema main.tools
|
|
264
|
+
```
|
|
265
|
+
|
|
266
|
+
The Python command additionally creates user-owned `agent/tools/<name>.py` and
|
|
267
|
+
`tests/tools/test_<name>.py` files using the LangGraph-native `@tool` decorator. `mason dev` and
|
|
268
|
+
`mason deploy` preserve `agent.toml`; they do not generate or patch agent source.
|
|
269
|
+
|
|
270
|
+
Sandbox scopes default to read-only access. Repeat `--scope` to allow more than one resource, use
|
|
271
|
+
`volume:` or `workspace:` for those resource types, and use `--permission read_write` only when the
|
|
272
|
+
agent needs writes. Every sandbox call carries this fixed downscope in MCP `_meta`, outside the tool
|
|
273
|
+
arguments controlled by the model.
|
|
274
|
+
|
|
275
|
+
## Initialize the chat app demo
|
|
276
|
+
|
|
277
|
+
The chat app is a LangGraph-specific init overlay, not a command that mutates an existing project.
|
|
278
|
+
It is included by default for `--framework langgraph`; pass `--disable-chat-app` to scaffold the
|
|
279
|
+
API-only backend instead.
|
|
280
|
+
|
|
281
|
+
```sh
|
|
282
|
+
mason init --framework langgraph \
|
|
283
|
+
--profile <profile> \
|
|
284
|
+
./my-agent
|
|
285
|
+
cd ./my-agent
|
|
286
|
+
mason dev
|
|
287
|
+
```
|
|
288
|
+
|
|
289
|
+
The chat app includes synchronous, SSE streaming, background polling, Session Store, Memory Store,
|
|
290
|
+
and HITL resume UI. The framework-specific overlay adds `ui/`, `runtime/ui.py`, the UI-enabled
|
|
291
|
+
`runtime/main.py`, and UI tests.
|
|
292
|
+
|
|
293
|
+
For the full deployed demo, bind both managed stores, then deploy:
|
|
294
|
+
|
|
295
|
+
```sh
|
|
296
|
+
mason sessions bind mason-demo-sessions
|
|
297
|
+
mason memory bind mason-demo-memory
|
|
298
|
+
mason --profile <profile> deploy mason-agent-demo --source .
|
|
299
|
+
```
|
|
300
|
+
|
|
301
|
+
(Binding creates a missing store automatically; pass `--no-create-stores` to require it already
|
|
302
|
+
exists. The agent reads the bound stores from `agent.toml` at runtime; `deploy` grants the app's
|
|
303
|
+
service principal access to them.)
|
|
304
|
+
|
|
305
|
+
The chat UI generates a stable application session UUID in browser local storage, places it inside
|
|
306
|
+
the durable invocation's opaque `input`, and creates a fresh invocation UUID per turn. The
|
|
307
|
+
`__Host-databricks-app-router` cookie remains independent: API clients may reuse it for sticky
|
|
308
|
+
replica routing, but it is neither authentication nor the template's application session state.
|
|
309
|
+
|
|
310
|
+
The generated `README.md` documents every request the client makes: config discovery, sync and SSE
|
|
311
|
+
invocations, background submission and polling, session transcript loading, HITL resume, and memory
|
|
312
|
+
entry operations. Capability colors are automatic from `/api/demo/config`; only the
|
|
313
|
+
sync/streaming/background transport selector is manual.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "databricks-mason"
|
|
3
|
-
version = "0.1.
|
|
3
|
+
version = "0.1.5.dev0"
|
|
4
4
|
description = "Databricks integration for Mason"
|
|
5
5
|
authors = [
|
|
6
6
|
{ name="Databricks", email="agent-feedback@databricks.com" },
|
|
@@ -9,7 +9,7 @@ readme = "README.md"
|
|
|
9
9
|
requires-python = ">=3.10"
|
|
10
10
|
dependencies = [
|
|
11
11
|
"click>=8.1",
|
|
12
|
-
"databricks-sdk>=0.
|
|
12
|
+
"databricks-sdk>=0.94.0",
|
|
13
13
|
"psycopg[binary]>=3.1",
|
|
14
14
|
"PyYAML>=6.0",
|
|
15
15
|
"rich>=13.7",
|
|
@@ -19,7 +19,7 @@ dependencies = [
|
|
|
19
19
|
|
|
20
20
|
[project.optional-dependencies]
|
|
21
21
|
tracing = [
|
|
22
|
-
"mlflow[databricks]>=3.
|
|
22
|
+
"mlflow[databricks]>=3.10.1",
|
|
23
23
|
]
|
|
24
24
|
# The agent-side runtime helpers a deployed agent imports. Kept as extras so a plain
|
|
25
25
|
# `pip install databricks-mason` (the CLI) stays light; each template depends on the extra for its
|
|
@@ -27,24 +27,26 @@ tracing = [
|
|
|
27
27
|
# OpenAI Agents SDK adapter (databricks_mason.openai). Both carry the shared framework-neutral stack
|
|
28
28
|
# (databricks_mason.runtime); the framework SDKs differ, so an agent installs only the one it uses.
|
|
29
29
|
runtime = [
|
|
30
|
+
"databricks-ai-bridge[memory]>=0.21.0",
|
|
30
31
|
"databricks-langchain>=0.17.0",
|
|
31
32
|
"langgraph>=1.1.0",
|
|
32
33
|
"langchain>=1.0.0",
|
|
33
34
|
"langchain-mcp-adapters>=0.2.1",
|
|
34
35
|
"fastapi>=0.129.0",
|
|
35
36
|
"mlflow>=3.10.1",
|
|
36
|
-
"uuid-utils>=0.10.0",
|
|
37
37
|
"opentelemetry-exporter-otlp-proto-grpc>=1.25.0",
|
|
38
38
|
"databricks-agents>=1.9.3",
|
|
39
|
+
"uvicorn>=0.20.0",
|
|
39
40
|
]
|
|
40
41
|
runtime-openai = [
|
|
42
|
+
"databricks-ai-bridge[memory]>=0.21.0",
|
|
41
43
|
"openai-agents>=0.7.0",
|
|
42
44
|
"databricks-openai>=0.13.0",
|
|
43
45
|
"fastapi>=0.129.0",
|
|
44
46
|
"mlflow>=3.10.1",
|
|
45
|
-
"uuid-utils>=0.10.0",
|
|
46
47
|
"opentelemetry-exporter-otlp-proto-grpc>=1.25.0",
|
|
47
48
|
"databricks-agents>=1.9.3",
|
|
49
|
+
"uvicorn>=0.20.0",
|
|
48
50
|
]
|
|
49
51
|
|
|
50
52
|
[project.scripts]
|
|
@@ -57,7 +59,12 @@ dev = [
|
|
|
57
59
|
{ include-group = "tests" },
|
|
58
60
|
]
|
|
59
61
|
tests = [
|
|
62
|
+
"databricks-ai-bridge[memory]>=0.21.0",
|
|
63
|
+
"fastapi>=0.129.0",
|
|
64
|
+
"httpx>=0.28.1",
|
|
60
65
|
"pytest==9.0.2",
|
|
66
|
+
"pytest-asyncio==1.3.0",
|
|
67
|
+
"uvicorn>=0.20.0",
|
|
61
68
|
]
|
|
62
69
|
|
|
63
70
|
[build-system]
|
|
@@ -75,6 +82,9 @@ include = [
|
|
|
75
82
|
[tool.hatch.build.targets.wheel]
|
|
76
83
|
packages = ["src/databricks_mason"]
|
|
77
84
|
|
|
85
|
+
[tool.uv.sources]
|
|
86
|
+
databricks-ai-bridge = { path = "../..", editable = true }
|
|
87
|
+
|
|
78
88
|
[tool.ruff]
|
|
79
89
|
include = ["pyproject.toml", "src/**/*.py", "tests/**/*.py"]
|
|
80
90
|
extend = "../../pyproject.toml"
|