databricks-mason 0.1.3.dev0__tar.gz → 0.1.4.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.4.dev0}/PKG-INFO +60 -29
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/README.md +57 -26
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/pyproject.toml +3 -3
- databricks_mason-0.1.4.dev0/src/databricks_mason/__init__.py +41 -0
- databricks_mason-0.1.3.dev0/src/databricks_mason/client.py → databricks_mason-0.1.4.dev0/src/databricks_mason/_api_client.py +63 -41
- databricks_mason-0.1.4.dev0/src/databricks_mason/_pagination.py +15 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/agent_project.py +42 -7
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/auth.py +4 -4
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/cli.py +5 -5
- databricks_mason-0.1.4.dev0/src/databricks_mason/client.py +28 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/deploy.py +169 -102
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/dev.py +29 -41
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/help.py +5 -2
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/memory.py +28 -13
- databricks_mason-0.1.4.dev0/src/databricks_mason/memory_store.py +356 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/render.py +29 -1
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/tool_manifest.py +25 -12
- databricks_mason-0.1.4.dev0/src/databricks_mason/session_store.py +358 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/sessions.py +10 -4
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/store_access.py +6 -1
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/tracing.py +10 -4
- databricks_mason-0.1.3.dev0/src/databricks_mason/__init__.py +0 -89
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/.gitignore +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/NOTICE +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/errors.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/init.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/langgraph/__init__.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/langgraph/mcp.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/langgraph/memory.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/langgraph/session_store.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/mcp.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/memory_store_access.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/models.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/openai/__init__.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/openai/mcp.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/openai/memory.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/openai/sessions.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/project_config.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/py.typed +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/__init__.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/background.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/session_store_client.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/tracing.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/workspace.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/sandbox.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/session_store_access.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/templates/python_tool_langgraph.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/templates/python_tool_test.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/templates/sandbox_mcp.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/templates/sandbox_mcp_langgraph.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/timefmt.py +0 -0
- {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.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.4.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
|
|
@@ -31,7 +31,7 @@ Requires-Dist: openai-agents>=0.7.0; extra == 'runtime-openai'
|
|
|
31
31
|
Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime-openai'
|
|
32
32
|
Requires-Dist: uuid-utils>=0.10.0; extra == 'runtime-openai'
|
|
33
33
|
Provides-Extra: tracing
|
|
34
|
-
Requires-Dist: mlflow[databricks]>=3.
|
|
34
|
+
Requires-Dist: mlflow[databricks]>=3.10.1; extra == 'tracing'
|
|
35
35
|
Description-Content-Type: text/markdown
|
|
36
36
|
|
|
37
37
|
# `databricks-mason`
|
|
@@ -90,28 +90,57 @@ You can also pass the global `--profile/-p` option before an individual command,
|
|
|
90
90
|
|
|
91
91
|
## Python SDK
|
|
92
92
|
|
|
93
|
-
|
|
94
|
-
`
|
|
95
|
-
|
|
93
|
+
`MasonClient` adds a small resource-oriented layer over the Mason API. Pass it an
|
|
94
|
+
authenticated Databricks `WorkspaceClient`, or omit the argument to use the
|
|
95
|
+
Databricks SDK's default authentication resolution:
|
|
96
96
|
|
|
97
97
|
```python
|
|
98
|
+
from databricks.sdk import WorkspaceClient
|
|
98
99
|
from databricks_mason import MasonClient
|
|
99
100
|
|
|
100
|
-
|
|
101
|
+
mason = MasonClient(WorkspaceClient(profile="my-workspace"))
|
|
102
|
+
|
|
103
|
+
session_store = mason.session_stores.create("support-agent-sessions")
|
|
104
|
+
session = session_store.add(actor_id="customer-123", session_id="case-456")
|
|
105
|
+
session.append_items(
|
|
106
|
+
[
|
|
107
|
+
{"type": "message", "role": "user", "content": "I need help with my cluster."},
|
|
108
|
+
{"type": "message", "role": "assistant", "content": "Let's take a look."},
|
|
109
|
+
]
|
|
110
|
+
)
|
|
111
|
+
|
|
112
|
+
memory_store = mason.memory_stores.create("coding-agent-memory")
|
|
113
|
+
memory = memory_store.add(
|
|
114
|
+
actor_id="alice",
|
|
115
|
+
path="/preferences/style.md",
|
|
116
|
+
content="The user prefers concise answers.",
|
|
117
|
+
)
|
|
118
|
+
results = memory_store.search(
|
|
119
|
+
actor_id="alice",
|
|
120
|
+
query="response preferences",
|
|
121
|
+
limit=10,
|
|
122
|
+
)
|
|
123
|
+
memory = memory.update(content="The user prefers very concise answers.")
|
|
124
|
+
memory.delete()
|
|
125
|
+
```
|
|
101
126
|
|
|
102
|
-
|
|
103
|
-
|
|
127
|
+
The root collections manage stores: `mason.memory_stores.create/get/list` and
|
|
128
|
+
`mason.session_stores.create/get/list`. A returned store owns operations on its
|
|
129
|
+
contents, such as `memory_store.add()`, `memory_store.get("memory-id")`,
|
|
130
|
+
`memory_store.list()`, and `memory_store.search()`, or `session_store.add()`,
|
|
131
|
+
`session_store.get("session-id")`, and `session_store.list()`. Returned memories,
|
|
132
|
+
sessions, and stores own their `update()` and `delete()` operations.
|
|
104
133
|
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
134
|
+
All `list()` methods return iterators that automatically consume server pages. List
|
|
135
|
+
`page_size` and search `limit` values must be between 1 and 100. `session.list_items()`
|
|
136
|
+
also auto-pages. `session.fork(...)` creates an independent copy, optionally through
|
|
137
|
+
a specific item. Deleting a session with descendants requires
|
|
138
|
+
`session.delete(force=True)` to cascade the deletion.
|
|
109
139
|
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
`databricks_mason.AgentCliError`. Deployment, sandbox, and tracing remain CLI-only.
|
|
140
|
+
The resource layer intentionally does not mirror every API method. Its private
|
|
141
|
+
transport will be replaced by the generated `WorkspaceClient.mason` service when that
|
|
142
|
+
is released, without changing this public surface. Deployment, sandbox, tracing, and
|
|
143
|
+
the existing CLI commands remain separate.
|
|
115
144
|
|
|
116
145
|
## Commands
|
|
117
146
|
|
|
@@ -122,12 +151,15 @@ mason [-p <profile>] [-o text|json]
|
|
|
122
151
|
init [--framework openai|langgraph] [--disable-chat-app]
|
|
123
152
|
[--profile P] [--repo URL] [--ref REF] [directory]
|
|
124
153
|
dev [--source PATH] [--prepare-environment] [--app-port PORT]
|
|
125
|
-
[--
|
|
126
|
-
[--with-traces C.S] [--no-create-stores]
|
|
154
|
+
[--with-traces C.S]
|
|
127
155
|
memory
|
|
156
|
+
bind STORE [--source PATH] [--no-create-stores]
|
|
157
|
+
unbind [--source PATH]
|
|
128
158
|
stores create | list | get | update | delete
|
|
129
159
|
entries create | get | list | search | update | delete
|
|
130
160
|
sessions create | list | get | update | delete | fork
|
|
161
|
+
bind STORE [--source PATH] [--no-create-stores]
|
|
162
|
+
unbind [--source PATH]
|
|
131
163
|
stores create | list | get | update | delete
|
|
132
164
|
items list | append | pop | clear
|
|
133
165
|
tracing
|
|
@@ -141,9 +173,7 @@ mason [-p <profile>] [-o text|json]
|
|
|
141
173
|
add uc-function FUNCTION [--name NAME] [--source PATH]
|
|
142
174
|
add python NAME [--source PATH]
|
|
143
175
|
list [--source PATH]
|
|
144
|
-
deploy <name> --source PATH [--
|
|
145
|
-
[--session/-s N] [--actor-id ID]
|
|
146
|
-
[--with-traces C.S] [--no-create-stores]
|
|
176
|
+
deploy <name> --source PATH [--with-traces C.S]
|
|
147
177
|
deployments list | get | logs | start | stop | delete
|
|
148
178
|
```
|
|
149
179
|
|
|
@@ -229,16 +259,17 @@ The chat app includes synchronous, SSE streaming, background polling, Session St
|
|
|
229
259
|
and HITL resume UI. The framework-specific overlay adds `ui/`, `runtime/ui.py`, the UI-enabled
|
|
230
260
|
`runtime/main.py`, and UI tests.
|
|
231
261
|
|
|
232
|
-
For the full deployed demo,
|
|
262
|
+
For the full deployed demo, bind both managed stores, then deploy:
|
|
233
263
|
|
|
234
264
|
```sh
|
|
235
|
-
mason
|
|
236
|
-
|
|
237
|
-
|
|
238
|
-
--actor-id alice
|
|
265
|
+
mason sessions bind mason-demo-sessions
|
|
266
|
+
mason memory bind mason-demo-memory
|
|
267
|
+
mason --profile <profile> deploy mason-agent-demo --source .
|
|
239
268
|
```
|
|
240
269
|
|
|
241
|
-
(
|
|
270
|
+
(Binding creates a missing store automatically; pass `--no-create-stores` to require it already
|
|
271
|
+
exists. The agent reads the bound stores from `agent.toml` at runtime; `deploy` grants the app's
|
|
272
|
+
service principal access to them.)
|
|
242
273
|
|
|
243
274
|
The Databricks Apps `__Host-databricks-app-router` cookie is both the sticky routing key and the
|
|
244
275
|
application session id. The browser sends it automatically; API clients must reuse it as a cookie.
|
|
@@ -54,28 +54,57 @@ You can also pass the global `--profile/-p` option before an individual command,
|
|
|
54
54
|
|
|
55
55
|
## Python SDK
|
|
56
56
|
|
|
57
|
-
|
|
58
|
-
`
|
|
59
|
-
|
|
57
|
+
`MasonClient` adds a small resource-oriented layer over the Mason API. Pass it an
|
|
58
|
+
authenticated Databricks `WorkspaceClient`, or omit the argument to use the
|
|
59
|
+
Databricks SDK's default authentication resolution:
|
|
60
60
|
|
|
61
61
|
```python
|
|
62
|
+
from databricks.sdk import WorkspaceClient
|
|
62
63
|
from databricks_mason import MasonClient
|
|
63
64
|
|
|
64
|
-
|
|
65
|
+
mason = MasonClient(WorkspaceClient(profile="my-workspace"))
|
|
66
|
+
|
|
67
|
+
session_store = mason.session_stores.create("support-agent-sessions")
|
|
68
|
+
session = session_store.add(actor_id="customer-123", session_id="case-456")
|
|
69
|
+
session.append_items(
|
|
70
|
+
[
|
|
71
|
+
{"type": "message", "role": "user", "content": "I need help with my cluster."},
|
|
72
|
+
{"type": "message", "role": "assistant", "content": "Let's take a look."},
|
|
73
|
+
]
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
memory_store = mason.memory_stores.create("coding-agent-memory")
|
|
77
|
+
memory = memory_store.add(
|
|
78
|
+
actor_id="alice",
|
|
79
|
+
path="/preferences/style.md",
|
|
80
|
+
content="The user prefers concise answers.",
|
|
81
|
+
)
|
|
82
|
+
results = memory_store.search(
|
|
83
|
+
actor_id="alice",
|
|
84
|
+
query="response preferences",
|
|
85
|
+
limit=10,
|
|
86
|
+
)
|
|
87
|
+
memory = memory.update(content="The user prefers very concise answers.")
|
|
88
|
+
memory.delete()
|
|
89
|
+
```
|
|
65
90
|
|
|
66
|
-
|
|
67
|
-
|
|
91
|
+
The root collections manage stores: `mason.memory_stores.create/get/list` and
|
|
92
|
+
`mason.session_stores.create/get/list`. A returned store owns operations on its
|
|
93
|
+
contents, such as `memory_store.add()`, `memory_store.get("memory-id")`,
|
|
94
|
+
`memory_store.list()`, and `memory_store.search()`, or `session_store.add()`,
|
|
95
|
+
`session_store.get("session-id")`, and `session_store.list()`. Returned memories,
|
|
96
|
+
sessions, and stores own their `update()` and `delete()` operations.
|
|
68
97
|
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
98
|
+
All `list()` methods return iterators that automatically consume server pages. List
|
|
99
|
+
`page_size` and search `limit` values must be between 1 and 100. `session.list_items()`
|
|
100
|
+
also auto-pages. `session.fork(...)` creates an independent copy, optionally through
|
|
101
|
+
a specific item. Deleting a session with descendants requires
|
|
102
|
+
`session.delete(force=True)` to cascade the deletion.
|
|
73
103
|
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
`databricks_mason.AgentCliError`. Deployment, sandbox, and tracing remain CLI-only.
|
|
104
|
+
The resource layer intentionally does not mirror every API method. Its private
|
|
105
|
+
transport will be replaced by the generated `WorkspaceClient.mason` service when that
|
|
106
|
+
is released, without changing this public surface. Deployment, sandbox, tracing, and
|
|
107
|
+
the existing CLI commands remain separate.
|
|
79
108
|
|
|
80
109
|
## Commands
|
|
81
110
|
|
|
@@ -86,12 +115,15 @@ mason [-p <profile>] [-o text|json]
|
|
|
86
115
|
init [--framework openai|langgraph] [--disable-chat-app]
|
|
87
116
|
[--profile P] [--repo URL] [--ref REF] [directory]
|
|
88
117
|
dev [--source PATH] [--prepare-environment] [--app-port PORT]
|
|
89
|
-
[--
|
|
90
|
-
[--with-traces C.S] [--no-create-stores]
|
|
118
|
+
[--with-traces C.S]
|
|
91
119
|
memory
|
|
120
|
+
bind STORE [--source PATH] [--no-create-stores]
|
|
121
|
+
unbind [--source PATH]
|
|
92
122
|
stores create | list | get | update | delete
|
|
93
123
|
entries create | get | list | search | update | delete
|
|
94
124
|
sessions create | list | get | update | delete | fork
|
|
125
|
+
bind STORE [--source PATH] [--no-create-stores]
|
|
126
|
+
unbind [--source PATH]
|
|
95
127
|
stores create | list | get | update | delete
|
|
96
128
|
items list | append | pop | clear
|
|
97
129
|
tracing
|
|
@@ -105,9 +137,7 @@ mason [-p <profile>] [-o text|json]
|
|
|
105
137
|
add uc-function FUNCTION [--name NAME] [--source PATH]
|
|
106
138
|
add python NAME [--source PATH]
|
|
107
139
|
list [--source PATH]
|
|
108
|
-
deploy <name> --source PATH [--
|
|
109
|
-
[--session/-s N] [--actor-id ID]
|
|
110
|
-
[--with-traces C.S] [--no-create-stores]
|
|
140
|
+
deploy <name> --source PATH [--with-traces C.S]
|
|
111
141
|
deployments list | get | logs | start | stop | delete
|
|
112
142
|
```
|
|
113
143
|
|
|
@@ -193,16 +223,17 @@ The chat app includes synchronous, SSE streaming, background polling, Session St
|
|
|
193
223
|
and HITL resume UI. The framework-specific overlay adds `ui/`, `runtime/ui.py`, the UI-enabled
|
|
194
224
|
`runtime/main.py`, and UI tests.
|
|
195
225
|
|
|
196
|
-
For the full deployed demo,
|
|
226
|
+
For the full deployed demo, bind both managed stores, then deploy:
|
|
197
227
|
|
|
198
228
|
```sh
|
|
199
|
-
mason
|
|
200
|
-
|
|
201
|
-
|
|
202
|
-
--actor-id alice
|
|
229
|
+
mason sessions bind mason-demo-sessions
|
|
230
|
+
mason memory bind mason-demo-memory
|
|
231
|
+
mason --profile <profile> deploy mason-agent-demo --source .
|
|
203
232
|
```
|
|
204
233
|
|
|
205
|
-
(
|
|
234
|
+
(Binding creates a missing store automatically; pass `--no-create-stores` to require it already
|
|
235
|
+
exists. The agent reads the bound stores from `agent.toml` at runtime; `deploy` grants the app's
|
|
236
|
+
service principal access to them.)
|
|
206
237
|
|
|
207
238
|
The Databricks Apps `__Host-databricks-app-router` cookie is both the sticky routing key and the
|
|
208
239
|
application session id. The browser sends it automatically; API clients must reuse it as a cookie.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "databricks-mason"
|
|
3
|
-
version = "0.1.
|
|
3
|
+
version = "0.1.4.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
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
"""High-level Python client and framework-neutral runtime helpers for Mason."""
|
|
2
|
+
|
|
3
|
+
from typing import TYPE_CHECKING
|
|
4
|
+
|
|
5
|
+
from databricks_mason.client import MasonClient
|
|
6
|
+
from databricks_mason.memory_store import Memory, MemorySearchResult, MemoryStore
|
|
7
|
+
from databricks_mason.session_store import Session, SessionItem, SessionStore
|
|
8
|
+
|
|
9
|
+
if TYPE_CHECKING:
|
|
10
|
+
from databricks_mason.runtime import (
|
|
11
|
+
configure_tracing,
|
|
12
|
+
tag_session,
|
|
13
|
+
workspace_client,
|
|
14
|
+
workspace_headers,
|
|
15
|
+
)
|
|
16
|
+
|
|
17
|
+
__all__ = [
|
|
18
|
+
"MasonClient",
|
|
19
|
+
"Memory",
|
|
20
|
+
"MemorySearchResult",
|
|
21
|
+
"MemoryStore",
|
|
22
|
+
"Session",
|
|
23
|
+
"SessionItem",
|
|
24
|
+
"SessionStore",
|
|
25
|
+
"configure_tracing",
|
|
26
|
+
"tag_session",
|
|
27
|
+
"workspace_client",
|
|
28
|
+
"workspace_headers",
|
|
29
|
+
]
|
|
30
|
+
|
|
31
|
+
_RUNTIME_REEXPORTS = frozenset(
|
|
32
|
+
{"configure_tracing", "tag_session", "workspace_client", "workspace_headers"}
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def __getattr__(name: str) -> object:
|
|
37
|
+
if name in _RUNTIME_REEXPORTS:
|
|
38
|
+
import importlib
|
|
39
|
+
|
|
40
|
+
return getattr(importlib.import_module("databricks_mason.runtime"), name)
|
|
41
|
+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
|
@@ -1,12 +1,9 @@
|
|
|
1
|
-
"""
|
|
2
|
-
|
|
3
|
-
|
|
4
|
-
|
|
5
|
-
|
|
6
|
-
|
|
7
|
-
configured vanity host and add the resolved workspace id as a routing header, matching
|
|
8
|
-
the authentication behavior of the Databricks CLI.
|
|
9
|
-
Deployment is handled separately (deploy.py) since it wraps the `databricks apps` CLI.
|
|
1
|
+
"""Private transport for the agents/v1 memory and session APIs.
|
|
2
|
+
|
|
3
|
+
The public SDK is the resource-oriented :class:`databricks_mason.MasonClient`.
|
|
4
|
+
This module temporarily owns the one-method-per-endpoint transport used by that
|
|
5
|
+
wrapper and the CLI. It can be replaced by the generated ``WorkspaceClient.mason``
|
|
6
|
+
service without changing the public resource surface.
|
|
10
7
|
"""
|
|
11
8
|
|
|
12
9
|
from __future__ import annotations
|
|
@@ -37,6 +34,11 @@ def _query(**kwargs: Any) -> dict[str, Any]:
|
|
|
37
34
|
return {k: v for k, v in kwargs.items() if v is not None and v != ""}
|
|
38
35
|
|
|
39
36
|
|
|
37
|
+
def _body(**kwargs: Any) -> dict[str, Any]:
|
|
38
|
+
"""Build a request body while retaining meaningful empty values."""
|
|
39
|
+
return {k: v for k, v in kwargs.items() if v is not None}
|
|
40
|
+
|
|
41
|
+
|
|
40
42
|
def _as(cls: type, resp: Any) -> Any:
|
|
41
43
|
"""Wrap a JSON response in a typed model, passing non-dicts through unchanged."""
|
|
42
44
|
return cls(resp) if isinstance(resp, dict) else resp
|
|
@@ -45,8 +47,7 @@ def _as(cls: type, resp: Any) -> Any:
|
|
|
45
47
|
def memory_store_path(name: str) -> str:
|
|
46
48
|
"""Normalize a store id or name into the `memory-stores/{id}` resource segment.
|
|
47
49
|
|
|
48
|
-
|
|
49
|
-
instead of building a URL like `memory-stores/` and leaking a raw ENDPOINT_NOT_FOUND.
|
|
50
|
+
Validate locally so malformed resource names do not produce misleading endpoint errors.
|
|
50
51
|
"""
|
|
51
52
|
raw = (name or "").strip()
|
|
52
53
|
if raw.startswith("memory-stores/"):
|
|
@@ -60,7 +61,7 @@ def memory_store_path(name: str) -> str:
|
|
|
60
61
|
|
|
61
62
|
|
|
62
63
|
def session_store_path(name: str) -> str:
|
|
63
|
-
"""Normalize
|
|
64
|
+
"""Normalize a session store name into the `session-stores/{name}` resource segment."""
|
|
64
65
|
raw = (name or "").strip()
|
|
65
66
|
if raw.startswith("session-stores/"):
|
|
66
67
|
raw = raw[len("session-stores/") :]
|
|
@@ -109,19 +110,19 @@ def _workspace_client(profile: Optional[str]) -> WorkspaceClient:
|
|
|
109
110
|
)
|
|
110
111
|
|
|
111
112
|
|
|
112
|
-
class
|
|
113
|
-
"""
|
|
113
|
+
class _MasonApiClient:
|
|
114
|
+
"""Private transport for the agents/v1 API until the generated SDK is available."""
|
|
114
115
|
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
116
|
+
def __init__(
|
|
117
|
+
self,
|
|
118
|
+
profile: Optional[str] = None,
|
|
119
|
+
*,
|
|
120
|
+
workspace_client: Optional[WorkspaceClient] = None,
|
|
121
|
+
) -> None:
|
|
122
|
+
if profile is not None and workspace_client is not None:
|
|
123
|
+
raise ValueError("profile and workspace_client are mutually exclusive")
|
|
123
124
|
try:
|
|
124
|
-
self._w = _workspace_client(profile)
|
|
125
|
+
self._w = workspace_client or _workspace_client(profile)
|
|
125
126
|
except Exception as exc: # noqa: BLE001 - surfaced as a clean CLI error
|
|
126
127
|
raise AgentCliError(
|
|
127
128
|
f"Could not initialize Databricks auth: {exc}",
|
|
@@ -190,7 +191,7 @@ class MasonClient:
|
|
|
190
191
|
*,
|
|
191
192
|
retry_transient: bool = False,
|
|
192
193
|
) -> models.MemoryStore:
|
|
193
|
-
body =
|
|
194
|
+
body = _body(display_name=display_name, description=description)
|
|
194
195
|
return _as(
|
|
195
196
|
models.MemoryStore,
|
|
196
197
|
self._do(
|
|
@@ -219,7 +220,7 @@ class MasonClient:
|
|
|
219
220
|
def update_memory_store(
|
|
220
221
|
self, name: str, display_name: Optional[str] = None, description: Optional[str] = None
|
|
221
222
|
) -> models.MemoryStore:
|
|
222
|
-
body =
|
|
223
|
+
body = _body(display_name=display_name, description=description)
|
|
223
224
|
if not body:
|
|
224
225
|
raise AgentCliError("No fields to update. Provide a display name and/or description.")
|
|
225
226
|
mask = ",".join(body.keys())
|
|
@@ -248,7 +249,7 @@ class MasonClient:
|
|
|
248
249
|
session_id: Optional[str] = None,
|
|
249
250
|
source_type: Optional[str] = None,
|
|
250
251
|
) -> models.MemoryEntry:
|
|
251
|
-
body =
|
|
252
|
+
body = _body(
|
|
252
253
|
actor_id=actor_id,
|
|
253
254
|
path=path,
|
|
254
255
|
content=content,
|
|
@@ -261,9 +262,16 @@ class MasonClient:
|
|
|
261
262
|
self._do("POST", f"{_BASE}/{memory_store_path(store)}/entries", body=body),
|
|
262
263
|
)
|
|
263
264
|
|
|
264
|
-
def get_memory_entry(
|
|
265
|
+
def get_memory_entry(
|
|
266
|
+
self, store: str, entry: str, read_mask: Optional[str] = None
|
|
267
|
+
) -> models.MemoryEntry:
|
|
265
268
|
return _as(
|
|
266
|
-
models.MemoryEntry,
|
|
269
|
+
models.MemoryEntry,
|
|
270
|
+
self._do(
|
|
271
|
+
"GET",
|
|
272
|
+
f"{_BASE}/{memory_entry_path(store, entry)}",
|
|
273
|
+
query=_query(read_mask=read_mask),
|
|
274
|
+
),
|
|
267
275
|
)
|
|
268
276
|
|
|
269
277
|
def list_memory_entries(
|
|
@@ -274,6 +282,7 @@ class MasonClient:
|
|
|
274
282
|
session_id: Optional[str] = None,
|
|
275
283
|
page_size: Optional[int] = None,
|
|
276
284
|
page_token: Optional[str] = None,
|
|
285
|
+
read_mask: Optional[str] = None,
|
|
277
286
|
) -> models.MemoryEntryList:
|
|
278
287
|
return _as(
|
|
279
288
|
models.MemoryEntryList,
|
|
@@ -286,14 +295,31 @@ class MasonClient:
|
|
|
286
295
|
session_id=session_id,
|
|
287
296
|
page_size=page_size,
|
|
288
297
|
page_token=page_token,
|
|
298
|
+
read_mask=read_mask,
|
|
289
299
|
),
|
|
290
300
|
),
|
|
291
301
|
)
|
|
292
302
|
|
|
293
303
|
def search_memory_entries(
|
|
294
|
-
self,
|
|
304
|
+
self,
|
|
305
|
+
store: str,
|
|
306
|
+
actor_id: str,
|
|
307
|
+
query: str,
|
|
308
|
+
limit: Optional[int] = None,
|
|
309
|
+
page_size: Optional[int] = None,
|
|
310
|
+
path_prefix: Optional[str] = None,
|
|
311
|
+
session_id: Optional[str] = None,
|
|
312
|
+
read_mask: Optional[str] = None,
|
|
295
313
|
) -> models.MemorySearchResult:
|
|
296
|
-
body =
|
|
314
|
+
body = _body(
|
|
315
|
+
actor_id=actor_id,
|
|
316
|
+
query=query,
|
|
317
|
+
limit=limit,
|
|
318
|
+
page_size=page_size,
|
|
319
|
+
path_prefix=path_prefix,
|
|
320
|
+
session_id=session_id,
|
|
321
|
+
read_mask=read_mask,
|
|
322
|
+
)
|
|
297
323
|
return _as(
|
|
298
324
|
models.MemorySearchResult,
|
|
299
325
|
self._do(
|
|
@@ -311,7 +337,7 @@ class MasonClient:
|
|
|
311
337
|
content: Optional[str] = None,
|
|
312
338
|
description: Optional[str] = None,
|
|
313
339
|
) -> models.MemoryEntry:
|
|
314
|
-
body =
|
|
340
|
+
body = _body(content=content, description=description)
|
|
315
341
|
if not body:
|
|
316
342
|
raise AgentCliError("No fields to update. Provide content and/or a description.")
|
|
317
343
|
return _as(
|
|
@@ -332,7 +358,7 @@ class MasonClient:
|
|
|
332
358
|
*,
|
|
333
359
|
retry_transient: bool = False,
|
|
334
360
|
) -> models.SessionStore:
|
|
335
|
-
body =
|
|
361
|
+
body = _body(description=description, metadata=metadata)
|
|
336
362
|
return _as(
|
|
337
363
|
models.SessionStore,
|
|
338
364
|
self._do(
|
|
@@ -362,7 +388,7 @@ class MasonClient:
|
|
|
362
388
|
def update_session_store(
|
|
363
389
|
self, name: str, description: Optional[str] = None, metadata: Optional[dict] = None
|
|
364
390
|
) -> models.SessionStore:
|
|
365
|
-
body =
|
|
391
|
+
body = _body(description=description, metadata=metadata)
|
|
366
392
|
if not body:
|
|
367
393
|
raise AgentCliError("No fields to update. Provide a description and/or metadata.")
|
|
368
394
|
mask = ",".join(body.keys())
|
|
@@ -389,7 +415,7 @@ class MasonClient:
|
|
|
389
415
|
parent_session_id: Optional[str] = None,
|
|
390
416
|
metadata: Optional[dict] = None,
|
|
391
417
|
) -> models.Session:
|
|
392
|
-
body =
|
|
418
|
+
body = _body(actor_id=actor_id, parent_session_id=parent_session_id, metadata=metadata)
|
|
393
419
|
return _as(
|
|
394
420
|
models.Session,
|
|
395
421
|
self._do(
|
|
@@ -433,7 +459,7 @@ class MasonClient:
|
|
|
433
459
|
"PATCH",
|
|
434
460
|
f"{_BASE}/session-stores/{store}/sessions/{session_id}",
|
|
435
461
|
query={"update_mask": "metadata"},
|
|
436
|
-
body=
|
|
462
|
+
body=_body(metadata=metadata),
|
|
437
463
|
),
|
|
438
464
|
)
|
|
439
465
|
|
|
@@ -441,7 +467,7 @@ class MasonClient:
|
|
|
441
467
|
return self._do(
|
|
442
468
|
"DELETE",
|
|
443
469
|
f"{_BASE}/session-stores/{store}/sessions/{session_id}",
|
|
444
|
-
query=
|
|
470
|
+
query={"force": True} if force else None,
|
|
445
471
|
)
|
|
446
472
|
|
|
447
473
|
def fork_session(
|
|
@@ -453,7 +479,7 @@ class MasonClient:
|
|
|
453
479
|
session_id: Optional[str] = None,
|
|
454
480
|
metadata: Optional[dict] = None,
|
|
455
481
|
) -> models.Session:
|
|
456
|
-
body =
|
|
482
|
+
body = _body(
|
|
457
483
|
source_session_id=source_session_id,
|
|
458
484
|
actor_id=actor_id,
|
|
459
485
|
up_to_item_id=up_to_item_id,
|
|
@@ -509,7 +535,3 @@ class MasonClient:
|
|
|
509
535
|
return self._do(
|
|
510
536
|
"POST", f"{_BASE}/session-stores/{store}/sessions/{session_id}/items:clear", body={}
|
|
511
537
|
)
|
|
512
|
-
|
|
513
|
-
|
|
514
|
-
# Backwards-compatible alias for the pre-1.0 internal name.
|
|
515
|
-
AgentApiClient = MasonClient
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
"""Shared pagination validation for typed Mason resources."""
|
|
2
|
+
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
_MAX_PAGE_SIZE = 100
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def validate_page_size(page_size: Optional[int]) -> None:
|
|
9
|
+
if page_size is not None and not 1 <= page_size <= _MAX_PAGE_SIZE:
|
|
10
|
+
raise ValueError(f"page_size must be between 1 and {_MAX_PAGE_SIZE}")
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def validate_limit(limit: Optional[int]) -> None:
|
|
14
|
+
if limit is not None and not 1 <= limit <= _MAX_PAGE_SIZE:
|
|
15
|
+
raise ValueError(f"limit must be between 1 and {_MAX_PAGE_SIZE}")
|