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.
Files changed (52) hide show
  1. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/PKG-INFO +60 -29
  2. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/README.md +57 -26
  3. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/pyproject.toml +3 -3
  4. databricks_mason-0.1.4.dev0/src/databricks_mason/__init__.py +41 -0
  5. 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
  6. databricks_mason-0.1.4.dev0/src/databricks_mason/_pagination.py +15 -0
  7. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/agent_project.py +42 -7
  8. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/auth.py +4 -4
  9. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/cli.py +5 -5
  10. databricks_mason-0.1.4.dev0/src/databricks_mason/client.py +28 -0
  11. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/deploy.py +169 -102
  12. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/dev.py +29 -41
  13. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/help.py +5 -2
  14. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/memory.py +28 -13
  15. databricks_mason-0.1.4.dev0/src/databricks_mason/memory_store.py +356 -0
  16. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/render.py +29 -1
  17. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/tool_manifest.py +25 -12
  18. databricks_mason-0.1.4.dev0/src/databricks_mason/session_store.py +358 -0
  19. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/sessions.py +10 -4
  20. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/store_access.py +6 -1
  21. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/tracing.py +10 -4
  22. databricks_mason-0.1.3.dev0/src/databricks_mason/__init__.py +0 -89
  23. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/.gitignore +0 -0
  24. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/NOTICE +0 -0
  25. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/errors.py +0 -0
  26. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/init.py +0 -0
  27. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/langgraph/__init__.py +0 -0
  28. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/langgraph/mcp.py +0 -0
  29. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/langgraph/memory.py +0 -0
  30. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/langgraph/session_store.py +0 -0
  31. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/mcp.py +0 -0
  32. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/memory_store_access.py +0 -0
  33. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/models.py +0 -0
  34. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/openai/__init__.py +0 -0
  35. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/openai/mcp.py +0 -0
  36. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/openai/memory.py +0 -0
  37. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/openai/sessions.py +0 -0
  38. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/project_config.py +0 -0
  39. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/py.typed +0 -0
  40. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/__init__.py +0 -0
  41. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/background.py +0 -0
  42. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/session_store_client.py +0 -0
  43. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/tracing.py +0 -0
  44. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/runtime/workspace.py +0 -0
  45. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/sandbox.py +0 -0
  46. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/session_store_access.py +0 -0
  47. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/templates/python_tool_langgraph.py +0 -0
  48. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/templates/python_tool_test.py +0 -0
  49. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/templates/sandbox_mcp.py +0 -0
  50. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/templates/sandbox_mcp_langgraph.py +0 -0
  51. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.4.dev0}/src/databricks_mason/timefmt.py +0 -0
  52. {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.dev0
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.49
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.9.0; extra == 'tracing'
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
- The same memory and session APIs are available programmatically through
94
- `MasonClient`, which authenticates exactly like the CLI (a `.databrickscfg` profile
95
- or the SDK's default resolution):
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
- client = MasonClient(profile="my-workspace") # or MasonClient() for default auth
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
- store = client.create_memory_store("my-store")
103
- print(store.name, store.display_name) # typed attribute access
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
- client.create_memory_entry("my-store", actor_id="alice", path="/notes/1.md", content="hi")
106
- for entry in client.list_memory_entries("my-store", actor_id="alice").entries:
107
- print(entry.path, entry.content)
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
- Each method maps to one `/api/agents/v1` operation. Responses come back as typed
111
- models (`MemoryStore`, `Session`, `SessionItemList`, ...) that expose attribute
112
- accessors (`store.name`) while remaining plain dicts underneath — so `store["name"]`,
113
- `json.dumps(store)`, and any new server-side fields keep working. API errors raise
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
- [--memory/-m N] [--session/-s N]
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 [--memory/-m N]
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, connect both managed stores:
262
+ For the full deployed demo, bind both managed stores, then deploy:
233
263
 
234
264
  ```sh
235
- mason --profile <profile> deploy mason-agent-demo --source . \
236
- --session mason-demo-sessions \
237
- --memory mason-demo-memory \
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
- (Missing stores are created automatically; pass `--no-create-stores` to require they already exist.)
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
- The same memory and session APIs are available programmatically through
58
- `MasonClient`, which authenticates exactly like the CLI (a `.databrickscfg` profile
59
- or the SDK's default resolution):
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
- client = MasonClient(profile="my-workspace") # or MasonClient() for default auth
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
- store = client.create_memory_store("my-store")
67
- print(store.name, store.display_name) # typed attribute access
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
- client.create_memory_entry("my-store", actor_id="alice", path="/notes/1.md", content="hi")
70
- for entry in client.list_memory_entries("my-store", actor_id="alice").entries:
71
- print(entry.path, entry.content)
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
- Each method maps to one `/api/agents/v1` operation. Responses come back as typed
75
- models (`MemoryStore`, `Session`, `SessionItemList`, ...) that expose attribute
76
- accessors (`store.name`) while remaining plain dicts underneath — so `store["name"]`,
77
- `json.dumps(store)`, and any new server-side fields keep working. API errors raise
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
- [--memory/-m N] [--session/-s N]
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 [--memory/-m N]
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, connect both managed stores:
226
+ For the full deployed demo, bind both managed stores, then deploy:
197
227
 
198
228
  ```sh
199
- mason --profile <profile> deploy mason-agent-demo --source . \
200
- --session mason-demo-sessions \
201
- --memory mason-demo-memory \
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
- (Missing stores are created automatically; pass `--no-create-stores` to require they already exist.)
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.dev0"
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.49",
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.9.0",
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
- """Authenticated client for the agents/v1 memory and session APIs.
2
-
3
- `MasonClient` is Mason's public Python entry point: construct it with a
4
- `.databrickscfg` profile (or rely on the SDK's default auth) and call one method per
5
- API operation. It wraps a databricks-sdk `WorkspaceClient` and returns typed models
6
- over the JSON responses from `/api/agents/v1`. Account-routed profiles retain their
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
- Validates client-side so an empty or wrong-typed argument raises a clear error
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/validate a session store name into the `session-stores/{name}` segment."""
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 MasonClient:
113
- """Thin, authenticated wrapper over the agents/v1 REST surface.
113
+ class _MasonApiClient:
114
+ """Private transport for the agents/v1 API until the generated SDK is available."""
114
115
 
115
- Example:
116
- >>> from databricks_mason import MasonClient
117
- >>> client = MasonClient(profile="my-workspace")
118
- >>> store = client.create_memory_store("my-store")
119
- >>> client.list_memory_stores()
120
- """
121
-
122
- def __init__(self, profile: Optional[str] = None):
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 = _query(display_name=display_name, description=description)
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 = _query(display_name=display_name, description=description)
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 = _query(
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(self, store: str, entry: str) -> models.MemoryEntry:
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, self._do("GET", f"{_BASE}/{memory_entry_path(store, entry)}")
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, store: str, actor_id: str, query: str, limit: Optional[int] = None
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 = _query(actor_id=actor_id, query=query, limit=limit)
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 = _query(content=content, description=description)
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 = _query(description=description, metadata=metadata)
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 = _query(description=description, metadata=metadata)
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 = _query(actor_id=actor_id, parent_session_id=parent_session_id, metadata=metadata)
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=_query(metadata=metadata),
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=_query(force=force or None),
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 = _query(
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}")