databricks-agentbricks 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (159) hide show
  1. databricks_agentbricks/__init__.py +1 -0
  2. databricks_agentbricks/_group.py +217 -0
  3. databricks_agentbricks/agent_project.py +657 -0
  4. databricks_agentbricks/app_resources.py +139 -0
  5. databricks_agentbricks/cli/__init__.py +12 -0
  6. databricks_agentbricks/cli/app.py +108 -0
  7. databricks_agentbricks/cli/app_auth.py +198 -0
  8. databricks_agentbricks/cli/auth.py +142 -0
  9. databricks_agentbricks/cli/deploy.py +954 -0
  10. databricks_agentbricks/cli/dev.py +310 -0
  11. databricks_agentbricks/cli/endpoint.py +145 -0
  12. databricks_agentbricks/cli/endpoint_examples.py +33 -0
  13. databricks_agentbricks/cli/endpoint_output.py +53 -0
  14. databricks_agentbricks/cli/endpoint_request.py +71 -0
  15. databricks_agentbricks/cli/endpoint_transport.py +148 -0
  16. databricks_agentbricks/cli/help.py +467 -0
  17. databricks_agentbricks/cli/init.py +496 -0
  18. databricks_agentbricks/cli/mcp.py +75 -0
  19. databricks_agentbricks/cli/memory.py +520 -0
  20. databricks_agentbricks/cli/sandbox.py +823 -0
  21. databricks_agentbricks/cli/sessions.py +588 -0
  22. databricks_agentbricks/cli/tools.py +456 -0
  23. databricks_agentbricks/cli/tracing.py +724 -0
  24. databricks_agentbricks/databricks_cli.py +38 -0
  25. databricks_agentbricks/errors.py +114 -0
  26. databricks_agentbricks/lakebase_runtime_store.py +107 -0
  27. databricks_agentbricks/legacy_lakebase_runtime_store.py +74 -0
  28. databricks_agentbricks/project_config.py +157 -0
  29. databricks_agentbricks/project_types.py +55 -0
  30. databricks_agentbricks/py.typed +0 -0
  31. databricks_agentbricks/render.py +349 -0
  32. databricks_agentbricks/templates/agent-bricks-migrate/SKILL.md +101 -0
  33. databricks_agentbricks/templates/agent-langgraph/.env.example +36 -0
  34. databricks_agentbricks/templates/agent-langgraph/.gitignore +199 -0
  35. databricks_agentbricks/templates/agent-langgraph/AGENTKIT_CONTRACT.md +95 -0
  36. databricks_agentbricks/templates/agent-langgraph/AGENTS.md +31 -0
  37. databricks_agentbricks/templates/agent-langgraph/CLAUDE.md +1 -0
  38. databricks_agentbricks/templates/agent-langgraph/README.md +141 -0
  39. databricks_agentbricks/templates/agent-langgraph/agent/__init__.py +0 -0
  40. databricks_agentbricks/templates/agent-langgraph/agent/agent.py +177 -0
  41. databricks_agentbricks/templates/agent-langgraph/agent/mcps.py +27 -0
  42. databricks_agentbricks/templates/agent-langgraph/agent/tools/__init__.py +24 -0
  43. databricks_agentbricks/templates/agent-langgraph/agent/tools/sample_tool.py +15 -0
  44. databricks_agentbricks/templates/agent-langgraph/agent/tools/send_message.py +17 -0
  45. databricks_agentbricks/templates/agent-langgraph/app.yaml +15 -0
  46. databricks_agentbricks/templates/agent-langgraph/pyproject.toml +42 -0
  47. databricks_agentbricks/templates/agent-langgraph/runtime/__init__.py +0 -0
  48. databricks_agentbricks/templates/agent-langgraph/runtime/adapter.py +144 -0
  49. databricks_agentbricks/templates/agent-langgraph/runtime/main.py +25 -0
  50. databricks_agentbricks/templates/agent-langgraph/tests/__init__.py +0 -0
  51. databricks_agentbricks/templates/agent-langgraph/tests/test_agent.py +211 -0
  52. databricks_agentbricks/templates/agent-langgraph/tests/test_runtime.py +19 -0
  53. databricks_agentbricks/templates/agent-langgraph/tests/test_workspace.py +28 -0
  54. databricks_agentbricks/templates/agent-openai/.env.example +37 -0
  55. databricks_agentbricks/templates/agent-openai/.gitignore +199 -0
  56. databricks_agentbricks/templates/agent-openai/AGENTKIT_CONTRACT.md +107 -0
  57. databricks_agentbricks/templates/agent-openai/AGENTS.md +76 -0
  58. databricks_agentbricks/templates/agent-openai/CLAUDE.md +1 -0
  59. databricks_agentbricks/templates/agent-openai/README.md +143 -0
  60. databricks_agentbricks/templates/agent-openai/agent/__init__.py +0 -0
  61. databricks_agentbricks/templates/agent-openai/agent/agent.py +205 -0
  62. databricks_agentbricks/templates/agent-openai/agent/mcps.py +28 -0
  63. databricks_agentbricks/templates/agent-openai/agent/tools/__init__.py +24 -0
  64. databricks_agentbricks/templates/agent-openai/agent/tools/sample_tool.py +15 -0
  65. databricks_agentbricks/templates/agent-openai/agent/tools/send_message.py +18 -0
  66. databricks_agentbricks/templates/agent-openai/app.yaml +15 -0
  67. databricks_agentbricks/templates/agent-openai/pyproject.toml +42 -0
  68. databricks_agentbricks/templates/agent-openai/runtime/__init__.py +0 -0
  69. databricks_agentbricks/templates/agent-openai/runtime/adapter.py +188 -0
  70. databricks_agentbricks/templates/agent-openai/runtime/main.py +25 -0
  71. databricks_agentbricks/templates/agent-openai/tests/__init__.py +0 -0
  72. databricks_agentbricks/templates/agent-openai/tests/test_agent.py +326 -0
  73. databricks_agentbricks/templates/agent-openai/tests/test_runtime.py +19 -0
  74. databricks_agentbricks/templates/agent-openai/tests/test_workspace.py +28 -0
  75. databricks_agentbricks/templates/custom-agent-langgraph/.env.example +1 -0
  76. databricks_agentbricks/templates/custom-agent-langgraph/.gitignore +11 -0
  77. databricks_agentbricks/templates/custom-agent-langgraph/README.md +22 -0
  78. databricks_agentbricks/templates/custom-agent-langgraph/agent/__init__.py +0 -0
  79. databricks_agentbricks/templates/custom-agent-langgraph/agent/agent.py +27 -0
  80. databricks_agentbricks/templates/custom-agent-langgraph/app.yaml +1 -0
  81. databricks_agentbricks/templates/custom-agent-langgraph/pyproject.toml +38 -0
  82. databricks_agentbricks/templates/custom-agent-langgraph/runtime/__init__.py +0 -0
  83. databricks_agentbricks/templates/custom-agent-langgraph/runtime/main.py +31 -0
  84. databricks_agentbricks/templates/custom-agent-langgraph/tests/test_runtime.py +29 -0
  85. databricks_agentbricks/templates/custom-agent-openai/.env.example +1 -0
  86. databricks_agentbricks/templates/custom-agent-openai/.gitignore +11 -0
  87. databricks_agentbricks/templates/custom-agent-openai/README.md +22 -0
  88. databricks_agentbricks/templates/custom-agent-openai/agent/__init__.py +0 -0
  89. databricks_agentbricks/templates/custom-agent-openai/agent/agent.py +29 -0
  90. databricks_agentbricks/templates/custom-agent-openai/app.yaml +1 -0
  91. databricks_agentbricks/templates/custom-agent-openai/pyproject.toml +38 -0
  92. databricks_agentbricks/templates/custom-agent-openai/runtime/__init__.py +0 -0
  93. databricks_agentbricks/templates/custom-agent-openai/runtime/main.py +31 -0
  94. databricks_agentbricks/templates/custom-agent-openai/tests/test_runtime.py +29 -0
  95. databricks_agentbricks/templates/sandbox_mcp.py +39 -0
  96. databricks_agentbricks/templates/sandbox_mcp_langgraph.py +39 -0
  97. databricks_agentbricks/templates/ui/agent-langgraph/CHAT_APP.md +49 -0
  98. databricks_agentbricks/templates/ui/agent-langgraph/runtime/main.py +26 -0
  99. databricks_agentbricks/templates/ui/agent-langgraph/runtime/ui.py +528 -0
  100. databricks_agentbricks/templates/ui/agent-langgraph/tests/test_demo_ui.py +542 -0
  101. databricks_agentbricks/templates/ui/agent-langgraph/ui/agentbricks-mark.png +0 -0
  102. databricks_agentbricks/templates/ui/agent-langgraph/ui/app.js +1464 -0
  103. databricks_agentbricks/templates/ui/agent-langgraph/ui/index.html +281 -0
  104. databricks_agentbricks/templates/ui/agent-langgraph/ui/styles.css +1913 -0
  105. databricks_agentbricks/templates/ui/agent-openai/CHAT_APP.md +54 -0
  106. databricks_agentbricks/templates/ui/agent-openai/runtime/main.py +26 -0
  107. databricks_agentbricks/templates/ui/agent-openai/runtime/ui.py +521 -0
  108. databricks_agentbricks/templates/ui/agent-openai/tests/test_demo_ui.py +512 -0
  109. databricks_agentbricks/templates/ui/agent-openai/ui/agentbricks-mark.png +0 -0
  110. databricks_agentbricks/templates/ui/agent-openai/ui/app.js +1428 -0
  111. databricks_agentbricks/templates/ui/agent-openai/ui/index.html +281 -0
  112. databricks_agentbricks/templates/ui/agent-openai/ui/styles.css +1913 -0
  113. databricks_agentbricks/theme.py +72 -0
  114. databricks_agentbricks-0.2.0.dist-info/METADATA +897 -0
  115. databricks_agentbricks-0.2.0.dist-info/RECORD +159 -0
  116. databricks_agentbricks-0.2.0.dist-info/WHEEL +4 -0
  117. databricks_agentbricks-0.2.0.dist-info/entry_points.txt +2 -0
  118. databricks_agentbricks-0.2.0.dist-info/licenses/NOTICE +39 -0
  119. databricks_agentkit/__init__.py +67 -0
  120. databricks_agentkit/_api_client.py +673 -0
  121. databricks_agentkit/_pagination.py +25 -0
  122. databricks_agentkit/client.py +31 -0
  123. databricks_agentkit/langgraph/__init__.py +98 -0
  124. databricks_agentkit/langgraph/genie.py +50 -0
  125. databricks_agentkit/langgraph/mcp.py +216 -0
  126. databricks_agentkit/langgraph/memory.py +67 -0
  127. databricks_agentkit/langgraph/session_store.py +485 -0
  128. databricks_agentkit/memory_store.py +371 -0
  129. databricks_agentkit/models.py +330 -0
  130. databricks_agentkit/openai/__init__.py +93 -0
  131. databricks_agentkit/openai/genie.py +50 -0
  132. databricks_agentkit/openai/mcp.py +189 -0
  133. databricks_agentkit/openai/memory.py +67 -0
  134. databricks_agentkit/openai/sessions.py +157 -0
  135. databricks_agentkit/py.typed +0 -0
  136. databricks_agentkit/runtime/ARCHITECTURE.md +64 -0
  137. databricks_agentkit/runtime/README.md +184 -0
  138. databricks_agentkit/runtime/__init__.py +69 -0
  139. databricks_agentkit/runtime/app.py +316 -0
  140. databricks_agentkit/runtime/auth.py +149 -0
  141. databricks_agentkit/runtime/durability/__init__.py +1 -0
  142. databricks_agentkit/runtime/durability/execution.py +133 -0
  143. databricks_agentkit/runtime/durability/heartbeat.py +64 -0
  144. databricks_agentkit/runtime/durability/lakebase_runtime_store.py +648 -0
  145. databricks_agentkit/runtime/durability/recovery.py +70 -0
  146. databricks_agentkit/runtime/durability/store.py +33 -0
  147. databricks_agentkit/runtime/execution.py +160 -0
  148. databricks_agentkit/runtime/genie.py +224 -0
  149. databricks_agentkit/runtime/mcp_auth.py +114 -0
  150. databricks_agentkit/runtime/model_services.py +90 -0
  151. databricks_agentkit/runtime/runtime.py +273 -0
  152. databricks_agentkit/runtime/session_store_client.py +120 -0
  153. databricks_agentkit/runtime/store.py +287 -0
  154. databricks_agentkit/runtime/tool_manifest.py +223 -0
  155. databricks_agentkit/runtime/tracing.py +71 -0
  156. databricks_agentkit/runtime/types.py +118 -0
  157. databricks_agentkit/runtime/workspace.py +27 -0
  158. databricks_agentkit/session_store.py +434 -0
  159. databricks_agentkit/timefmt.py +85 -0
@@ -0,0 +1,310 @@
1
+ """`ab dev` — run a scaffolded agent locally, wrapping `databricks apps run-local`.
2
+
3
+ Runs the app from its ``app.yaml`` exactly as the Databricks Apps runtime would locally: reads the
4
+ manifest's command + env, and (with ``--prepare-environment``) builds the venv via uv. This is the
5
+ local counterpart to ``ab deploy`` — same source dir, same manifest — so what runs here matches
6
+ what ships. Delegating to ``apps run-local`` means agentbricks inherits the Apps team's local-run behavior
7
+ rather than re-implementing it.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import pathlib
13
+ from typing import Optional
14
+
15
+ import click
16
+ import yaml
17
+
18
+ from databricks_agentbricks import render
19
+ from databricks_agentbricks.cli.deploy import (
20
+ _load_project,
21
+ resource_bindings,
22
+ )
23
+ from databricks_agentbricks.cli.endpoint_examples import print_agent_invoke_command
24
+ from databricks_agentbricks.cli.tracing import start_local_tracing_server, stop_local_tracing_server
25
+ from databricks_agentbricks.databricks_cli import _databricks
26
+ from databricks_agentbricks.errors import AgentCliError
27
+ from databricks_agentbricks.project_config import require_managed_tool_support
28
+ from databricks_agentbricks.project_types import AgentServer
29
+ from databricks_agentkit.runtime.store import RUNTIME_STORE_LOCAL_ENV
30
+ from databricks_agentkit.runtime.tool_manifest import MEMORY_STORE_ENV, SESSION_STORE_ENV
31
+
32
+ # Default local port; `databricks apps run-local` listens here unless --app-port overrides it.
33
+ _DEFAULT_APP_PORT = 8000
34
+ _LOCAL_APP_YAML = "app.agentbricksdev.yaml"
35
+
36
+ # Env vars that pin a package index for the *deployed* Apps build (a cloud-only workaround, see
37
+ # `ab deploy`). They point at an index the deploying environment can reach, which is not
38
+ # necessarily reachable from the local dev machine — so `ab dev`'s local `uv` build must ignore
39
+ # them and use the machine's own configured index instead.
40
+ _BUILD_INDEX_ENVS = frozenset({"PIP_INDEX_URL", "UV_INDEX_URL", "UV_DEFAULT_INDEX"})
41
+
42
+ # Workspace managed-resource env that `ab deploy` writes into app.yaml but `ab dev` must NOT
43
+ # inherit - dev runs every resource locally. These are stripped from the dev manifest:
44
+ # - tracing: dev sets up its OWN local MLflow server, so a stale workspace MLFLOW_EXPERIMENT_ID
45
+ # (which MLflow resolves ahead of MLFLOW_EXPERIMENT_NAME) would point the local agent at an
46
+ # experiment that doesn't exist on the local sqlite server. dev re-adds the local
47
+ # MLFLOW_TRACKING_URI / MLFLOW_EXPERIMENT_NAME itself.
48
+ # - stores: dev never uses the workspace memory/session stores (memory off, sessions in-process),
49
+ # so a prior deploy's AGENT_MEMORY_STORE / AGENT_SESSION_STORE must not quietly pull dev onto
50
+ # them. dev re-adds nothing here - the runtime falls back to its local defaults.
51
+ # (The Runtime Store's lakebase env needs no strip: dev already forces RUNTIME_STORE_LOCAL=true.)
52
+ _DEPLOY_TRACING_ENVS = frozenset(
53
+ {"MLFLOW_TRACKING_URI", "MLFLOW_EXPERIMENT_ID", "MLFLOW_TRACING_DESTINATION"}
54
+ )
55
+ _DEPLOY_RESOURCE_ENVS = _DEPLOY_TRACING_ENVS | {MEMORY_STORE_ENV, SESSION_STORE_ENV}
56
+
57
+
58
+ @click.command()
59
+ @click.option(
60
+ "--source",
61
+ default=".",
62
+ type=click.Path(exists=True, file_okay=False),
63
+ help="Local source directory to run (containing app.yaml). Defaults to the current directory.",
64
+ )
65
+ @click.option(
66
+ "--prepare-environment/--no-prepare-environment",
67
+ default=None,
68
+ help="Build the app's environment with uv before running. Default: build only if no .venv "
69
+ "exists yet, and reuse it otherwise. Requires uv.",
70
+ )
71
+ @click.option("--app-port", type=int, default=None, help="Port to run the app on (default 8000).")
72
+ @click.pass_obj
73
+ def dev(
74
+ obj,
75
+ source: str,
76
+ prepare_environment: Optional[bool],
77
+ app_port: Optional[int],
78
+ ) -> None:
79
+ """Run your agent locally so you can try it before deploying.
80
+
81
+ Starts the agent on a local server — by default http://localhost:8000 — and prints where to
82
+ reach it: the chat UI if the project has one, otherwise a sample request against the agent's
83
+ API.
84
+
85
+ Auth uses your Databricks profile (`-p` / `ab login`), and the agent reaches Databricks model
86
+ serving through the AI Gateway on that profile — so there are no model keys to set up.
87
+
88
+ Under the hood this wraps `databricks apps run-local`: it reads the command + env from
89
+ `app.yaml` and runs the app the way the Apps runtime would, so local behavior matches a
90
+ deployment. The environment is built on the first run and reused after; pass
91
+ `--prepare-environment` to force a rebuild (e.g. after changing dependencies).
92
+
93
+ Everything runs locally: `ab dev` is a local deployment that does not depend on a Databricks
94
+ workspace for its resources. Tracing goes to a local MLflow tracking server (sqlite-backed, under `.agentbricks/`)
95
+ so traces are recorded on your machine with no workspace experiment or setup - open the printed
96
+ Traces URL to view them (`ab tracing unbind` doesn't affect dev; it only stops the deployed
97
+ agent's tracing). Long-term memory is off and conversation history is in-process (not durable):
98
+ the memory/session stores bound with `ab memory/sessions bind` are created and used only when you
99
+ `ab deploy`, not here. So there's nothing to provision and no service-principal grant to make;
100
+ that all happens at `ab deploy` time.
101
+ """
102
+ source_dir = pathlib.Path(source)
103
+ app_yaml = source_dir / "app.yaml"
104
+ if not app_yaml.exists():
105
+ raise AgentCliError(
106
+ f"No app.yaml found at {app_yaml}.",
107
+ hint="Run from a scaffolded project, or pass --source <dir> (see `ab init`).",
108
+ )
109
+
110
+ project = _load_project(source_dir)
111
+ if project is not None and project.tools:
112
+ require_managed_tool_support(source_dir)
113
+
114
+ # `ab dev` is a fully local sandbox: the Runtime Store and tracing already run locally, and
115
+ # memory/sessions follow suit here. Dev never reaches the workspace for stores (so it stays fast
116
+ # and works offline): long-term memory is off and conversation history is in-process (not
117
+ # durable), regardless of any binding. Stores are created and used only by `ab deploy`; the
118
+ # deploy-written store env is stripped from the dev manifest (see `_dev_entry_point`) so a prior
119
+ # deploy can't quietly pull dev onto the workspace stores. Read the bindings only to name them.
120
+ memory_store, session_store, trace_experiment = resource_bindings(source_dir)
121
+ if memory_store:
122
+ render.console().print(
123
+ f"[dim]Memory store '{memory_store}' is bound but `ab dev` runs with "
124
+ "long-term memory off. Run `ab deploy` to use bound store.[/]"
125
+ )
126
+ if session_store:
127
+ render.console().print(
128
+ f"[dim]Session store '{session_store}' is bound but `ab dev` keeps "
129
+ "conversation history in-process (not durable). Run `ab deploy` to use bound store.[/]"
130
+ )
131
+ if trace_experiment:
132
+ render.console().print(
133
+ f"[dim]Tracing experiment '{trace_experiment}' is bound but `ab dev` "
134
+ "traces to a local MLflow server. Run `ab deploy` to trace to the bound experiment.[/]"
135
+ )
136
+ local_env: dict[str, str] = {}
137
+ # Local tracing: start a local MLflow tracking server backed by sqlite under .agentbricks/ and point the
138
+ # agent at it via the dev-only manifest — for any project, regardless of framework/server. An agent
139
+ # that uses MLflow (autolog or `start_trace`) then traces to it; it's harmless for one that doesn't.
140
+ # Traces stay on the machine — no workspace experiment, no auth, no username needed — and the same
141
+ # server serves the trace UI. Launched via `uvx mlflow` so it needs neither the (skinny) CLI env nor
142
+ # the agent venv, and it owns the sqlite schema (so there's no client/server migration mismatch).
143
+ # Best-effort: any launch failure degrades to running without traces. `ab deploy` handles the
144
+ # managed workspace experiment instead.
145
+ tracing_server, tracing_env = start_local_tracing_server(source_dir)
146
+ # Everything after the server starts runs under try/finally, so any failure — e.g. a malformed
147
+ # app.yaml that `_dev_entry_point` rejects — still tears the local server down (and removes the
148
+ # dev-only manifest) instead of orphaning the process.
149
+ entry_point: Optional[pathlib.Path] = None
150
+ try:
151
+ trace_url: Optional[str] = None
152
+ if tracing_env:
153
+ local_env.update(tracing_env)
154
+ uri = tracing_env["MLFLOW_TRACKING_URI"]
155
+ name = tracing_env.get("MLFLOW_EXPERIMENT_NAME")
156
+ # Can't deep-link the experiment (created lazily on the first request; its local id isn't
157
+ # stable), so name it after the local MLflow UI URL so the user knows which one to open.
158
+ trace_url = f"{uri} (experiment name: {name})" if name else uri
159
+
160
+ # Default: prepare only when there's no venv yet, so repeat runs don't rebuild. Explicit
161
+ # --prepare-environment / --no-prepare-environment overrides the auto-detect.
162
+ if prepare_environment is None:
163
+ prepare_environment = not (source_dir / ".venv").exists()
164
+
165
+ args = ["apps", "run-local"]
166
+ if prepare_environment:
167
+ args.append("--prepare-environment")
168
+ if app_port is not None:
169
+ args += ["--app-port", str(app_port)]
170
+
171
+ # Run against a local-only manifest that forces the Runtime Store in-process, removes deploy-only
172
+ # package-index overrides, and injects any locally resolved store ids and the local tracing env.
173
+ entry_point = _dev_entry_point(app_yaml, local_env or None)
174
+ # run-local resolves this relative to cwd and rejects an absolute alternate-manifest path.
175
+ args += ["--entry-point", entry_point.name]
176
+
177
+ # `run-local` prints a generic "go to http://localhost:<port>" line that points at the chat UI —
178
+ # misleading for an API-only project, which serves no page there (404). Print an accurate line up
179
+ # front, keyed on whether this project actually carries the chat-app overlay.
180
+ _announce_local_url(
181
+ source_dir,
182
+ app_port or _DEFAULT_APP_PORT,
183
+ project.server if project else None,
184
+ trace_url,
185
+ )
186
+
187
+ # Run in the project dir so run-local finds the app; stream output (no capture).
188
+ _databricks(
189
+ args,
190
+ obj.profile,
191
+ cwd=str(source_dir),
192
+ action="Could not start the agent locally.",
193
+ )
194
+ finally:
195
+ # Remove the local-only manifest so a later `ab deploy` cannot sync it to the workspace, and
196
+ # stop the local tracing server — even if setup above raised before run-local.
197
+ if entry_point is not None:
198
+ entry_point.unlink(missing_ok=True)
199
+ if tracing_server is not None:
200
+ stop_local_tracing_server(tracing_server)
201
+
202
+
203
+ def _announce_local_url(
204
+ source_dir: pathlib.Path, port: int, server: AgentServer | None, trace_url: str | None = None
205
+ ) -> None:
206
+ """Print how to reach the running app: the chat UI if present, else a sample invoke request.
207
+
208
+ ``trace_url`` (when tracing is on) is shown alongside so a dev run surfaces where its traces land,
209
+ matching the ``Traces`` line ``ab deploy`` prints.
210
+ """
211
+ base = f"http://localhost:{port}"
212
+ deploy_name = source_dir.resolve().name
213
+ tool_step: str | tuple[str, str] = (
214
+ "Edit agent/agent.py to give the agent a tool"
215
+ if server == AgentServer.CUSTOM
216
+ else ("ab tools add mcp <service>", "Give the agent a tool")
217
+ )
218
+ if (source_dir / "runtime" / "ui.py").is_file():
219
+ fields = {"Chat UI": base}
220
+ if trace_url:
221
+ fields["Traces"] = trace_url
222
+ render.success(
223
+ "Starting agent",
224
+ fields=fields,
225
+ next_steps=[
226
+ f"Open {base} to chat with your agent",
227
+ tool_step,
228
+ ("ab memory bind <store>", "Attach a memory / session store"),
229
+ (f"ab deploy {deploy_name}", "Deploy it to Databricks"),
230
+ ],
231
+ )
232
+ else:
233
+ # No page is served at `/`, so give a copy-pasteable request instead of just the URL.
234
+ uses_runtime_api = server == AgentServer.AGENTBRICKS
235
+ endpoint = f"{base}/api/invocations" if uses_runtime_api else f"{base}/invocations"
236
+ body = (
237
+ '{"id": "00000000-0000-4000-8000-000000000000", '
238
+ '"input": [{"role": "user", "content": "hi"}]}'
239
+ if uses_runtime_api
240
+ else '{"input": [{"role": "user", "content": "hi"}]}'
241
+ )
242
+ sample = f"curl -X POST {endpoint} -H 'Content-Type: application/json' -d '{body}'"
243
+ fields = {"Invoke": f"POST {endpoint}"}
244
+ if trace_url:
245
+ fields["Traces"] = trace_url
246
+ render.success(
247
+ "Starting API-only agent (no chat UI — see `ab init --help`)",
248
+ fields=fields,
249
+ next_steps=[
250
+ (sample, "Send a test request"),
251
+ tool_step,
252
+ (f"ab deploy {deploy_name}", "Deploy it to Databricks"),
253
+ ],
254
+ )
255
+
256
+ print_agent_invoke_command(
257
+ f"--url {base}",
258
+ uses_runtime_api=(source_dir / "runtime" / "ui.py").is_file()
259
+ or server == AgentServer.AGENTBRICKS,
260
+ )
261
+
262
+
263
+ def _dev_entry_point(
264
+ app_yaml: pathlib.Path, extra_env: dict[str, str] | None = None
265
+ ) -> pathlib.Path:
266
+ """Write the local-only app manifest consumed by ``apps run-local``.
267
+
268
+ The manifest marks the process as local so Agent Bricks Runtime uses its in-memory store. Keeping this in
269
+ the entry point is more reliable than forwarding ``--env`` through the Databricks CLI and does
270
+ not mutate the deployable ``app.yaml``. Deploy-only package-index variables and the deploy-written
271
+ workspace resource env (see ``_DEPLOY_RESOURCE_ENVS`` - tracing + memory/session stores) are
272
+ removed so dev stays fully local. ``extra_env`` is merged in (overriding any same-named entries)
273
+ for dev-only overrides such as the local tracing config.
274
+ """
275
+ try:
276
+ doc = yaml.safe_load(app_yaml.read_text()) or {}
277
+ except yaml.YAMLError as exc:
278
+ raise AgentCliError(f"Could not parse {app_yaml}: {exc}") from exc
279
+ if not isinstance(doc, dict):
280
+ raise AgentCliError(f"Invalid {app_yaml}: top level must be an object.")
281
+ env = doc.get("env")
282
+ if env is not None and not isinstance(env, list):
283
+ raise AgentCliError(f"Invalid {app_yaml}: env must be a list.")
284
+ filtered = [
285
+ e for e in (env or []) if not (isinstance(e, dict) and e.get("name") in _BUILD_INDEX_ENVS)
286
+ ]
287
+ filtered = [
288
+ e
289
+ for e in filtered
290
+ if not (isinstance(e, dict) and e.get("name") == RUNTIME_STORE_LOCAL_ENV)
291
+ ]
292
+ # Drop the deploy-written workspace resource env (tracing + memory/session stores) so a prior
293
+ # `ab deploy` can't pull dev onto workspace resources: a stale MLFLOW_EXPERIMENT_ID would beat
294
+ # the local MLFLOW_EXPERIMENT_NAME dev re-adds below, and AGENT_MEMORY_STORE / AGENT_SESSION_STORE
295
+ # would silently point the local runtime at the workspace stores instead of its local defaults.
296
+ filtered = [
297
+ e for e in filtered if not (isinstance(e, dict) and e.get("name") in _DEPLOY_RESOURCE_ENVS)
298
+ ]
299
+ for name, value in (extra_env or {}).items():
300
+ filtered = [e for e in filtered if not (isinstance(e, dict) and e.get("name") == name)]
301
+ filtered.append({"name": name, "value": value})
302
+ filtered.append({"name": RUNTIME_STORE_LOCAL_ENV, "value": "true"})
303
+ doc["env"] = filtered
304
+ # The Apps CLI rejects hidden or hyphenated entry-point filenames.
305
+ dev_yaml = app_yaml.parent / _LOCAL_APP_YAML
306
+ try:
307
+ dev_yaml.write_text(yaml.safe_dump(doc, sort_keys=False))
308
+ except OSError as exc:
309
+ raise AgentCliError(f"Could not write {dev_yaml}: {exc}") from exc
310
+ return dev_yaml
@@ -0,0 +1,145 @@
1
+ """CLI entry point for invoking arbitrary HTTP endpoints."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ from dataclasses import replace
7
+ from typing import Optional
8
+ from uuid import uuid4
9
+
10
+ import click
11
+
12
+ from databricks_agentbricks.cli.deploy import _app_url, _prefixed_name
13
+ from databricks_agentbricks.cli.endpoint_output import SsePrinter, render_response
14
+ from databricks_agentbricks.cli.endpoint_request import build_request
15
+ from databricks_agentbricks.cli.endpoint_transport import HttpSession
16
+ from databricks_agentbricks.errors import AgentCliError
17
+ from databricks_agentkit._api_client import _workspace_client
18
+
19
+ _ROUTING_COOKIE = "__Host-databricks-app-router"
20
+
21
+
22
+ def _resolve_endpoint(
23
+ app: str | None,
24
+ url: str | None,
25
+ profile: Optional[str],
26
+ ) -> tuple[str, bool]:
27
+ if app and url:
28
+ raise AgentCliError("APP and --url are mutually exclusive.")
29
+ if url:
30
+ return url.rstrip("/"), False
31
+ if not app:
32
+ raise AgentCliError(
33
+ "Provide a Databricks App name or --url.",
34
+ hint="Use --url http://localhost:8000 when running the agent locally.",
35
+ )
36
+ app_name = _prefixed_name(app)
37
+ resolved_url = _app_url(app_name, profile)
38
+ if not resolved_url:
39
+ raise AgentCliError(f"Could not resolve a URL for Databricks App {app_name!r}.")
40
+ return resolved_url.rstrip("/"), True
41
+
42
+
43
+ def _authorization_header(profile: Optional[str]) -> str:
44
+ try:
45
+ client = _workspace_client(profile)
46
+ if client.config.auth_type == "pat":
47
+ raise AgentCliError(
48
+ "Databricks Apps API routes require OAuth; the selected profile uses a PAT.",
49
+ hint="Authenticate the same workspace with `databricks auth login`.",
50
+ )
51
+ authorization = client.config.authenticate().get("Authorization")
52
+ except AgentCliError:
53
+ raise
54
+ except Exception as exc: # noqa: BLE001 - render auth failures without a traceback
55
+ raise AgentCliError(f"Could not initialize endpoint authentication: {exc}.") from exc
56
+ if not authorization:
57
+ raise AgentCliError("Could not resolve an OAuth access token for the endpoint request.")
58
+ return authorization
59
+
60
+
61
+ def _platform_headers(
62
+ *,
63
+ authenticate: bool,
64
+ profile: Optional[str],
65
+ session_id: str | None,
66
+ ) -> dict[str, str]:
67
+ headers: dict[str, str] = {}
68
+ if authenticate:
69
+ headers["Authorization"] = _authorization_header(profile)
70
+ if session_id:
71
+ headers["Cookie"] = f"{_ROUTING_COOKIE}={session_id}"
72
+ return headers
73
+
74
+
75
+ @click.group()
76
+ def endpoint() -> None:
77
+ """Invoke arbitrary HTTP endpoints."""
78
+
79
+
80
+ @click.command("invoke")
81
+ @click.argument("app", required=False, metavar="[APP]")
82
+ @click.option("--url", default=None, help="Base URL for localhost or an arbitrary HTTP server.")
83
+ @click.option("--method", default="POST", show_default=True)
84
+ @click.option("--path", required=True, help="Request path, such as /api/invocations.")
85
+ @click.option("--query", "query", multiple=True, help="Query parameter as 'name=value'.")
86
+ @click.option("--json", "json_value", default=None, help="Complete JSON request body.")
87
+ @click.option("--sse", is_flag=True, help="Consume the response as Server-Sent Events.")
88
+ @click.option(
89
+ "--session-id",
90
+ default=None,
91
+ help="Application session id (default: generated for a Databricks App).",
92
+ )
93
+ @click.option("--timeout", type=click.FloatRange(min=0.1), default=300.0, show_default=True)
94
+ @click.option("--auth/--no-auth", default=None, help="Inject Databricks OAuth authentication.")
95
+ @click.pass_obj
96
+ def invoke(
97
+ obj,
98
+ app,
99
+ url,
100
+ method,
101
+ path,
102
+ query,
103
+ json_value,
104
+ sse,
105
+ session_id,
106
+ timeout,
107
+ auth,
108
+ ) -> None:
109
+ """Send one HTTP request to a Databricks App or arbitrary URL."""
110
+ base_url, is_app = _resolve_endpoint(app, url, obj.profile)
111
+ authenticate = is_app if auth is None else auth
112
+ session_id = session_id or (str(uuid4()) if is_app else None)
113
+ request = build_request(
114
+ base_url=base_url,
115
+ method=method,
116
+ path=path,
117
+ query=query,
118
+ json_value=json_value,
119
+ timeout=timeout,
120
+ sse=sse,
121
+ )
122
+ request = replace(
123
+ request,
124
+ headers={
125
+ **request.headers,
126
+ **_platform_headers(
127
+ authenticate=authenticate,
128
+ profile=obj.profile,
129
+ session_id=session_id,
130
+ ),
131
+ },
132
+ )
133
+ printer = SsePrinter(enabled=sse and obj.output == "text")
134
+ response = HttpSession().send(request, on_event=printer)
135
+ if not 200 <= response.status_code < 300:
136
+ raise AgentCliError(
137
+ f"Endpoint returned HTTP {response.status_code}.",
138
+ hint=json.dumps(response.body, default=str)[:1000]
139
+ if response.body is not None
140
+ else None,
141
+ )
142
+ render_response(response, output=obj.output, streamed=sse)
143
+
144
+
145
+ endpoint.add_command(invoke)
@@ -0,0 +1,33 @@
1
+ """Copy-pasteable examples for invoking Agent Bricks-generated agent endpoints."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import shlex
6
+
7
+ from rich.text import Text
8
+
9
+ from databricks_agentbricks import render
10
+
11
+
12
+ def agent_invoke_command(target: str, *, uses_runtime_api: bool) -> str:
13
+ """Build a single-line example, generating a fresh runtime invocation ID on every run."""
14
+ path = "/api/invocations" if uses_runtime_api else "/invocations"
15
+ body = '{"input":[{"role":"user","content":"hi"}]}'
16
+ if uses_runtime_api:
17
+ # Double quotes allow command substitution while preserving the JSON's literal quotes.
18
+ body = '{"id":"$(uuidgen)",' + body[1:]
19
+ json_arg = '"' + body.replace('"', r"\"") + '"'
20
+ else:
21
+ json_arg = shlex.quote(body)
22
+ return f"ab endpoint invoke {target} --path {path} --json {json_arg}"
23
+
24
+
25
+ def print_agent_invoke_command(target: str, *, uses_runtime_api: bool) -> None:
26
+ """Print only the new invocation example without changing the existing success panel."""
27
+ con = render.console()
28
+ con.print(Text("Invoke with Agent Bricks", style=render.SECONDARY))
29
+ # A separate, soft-wrapped line avoids copying panel borders or inserted line breaks.
30
+ con.print(
31
+ Text(agent_invoke_command(target, uses_runtime_api=uses_runtime_api), style=render.COMMAND),
32
+ soft_wrap=True,
33
+ )
@@ -0,0 +1,53 @@
1
+ """Output rendering for endpoint commands."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ from typing import Any
7
+
8
+ import click
9
+
10
+ from databricks_agentbricks import render
11
+ from databricks_agentbricks.cli.endpoint_transport import EndpointResponse
12
+
13
+
14
+ def render_response(response: EndpointResponse, *, output: str, streamed: bool) -> None:
15
+ """Render an endpoint response in JSON or human-readable form."""
16
+ if output == "json":
17
+ render.emit_json(_response_payload(response))
18
+ return
19
+ if streamed:
20
+ return
21
+ if isinstance(response.body, (dict, list)):
22
+ render.emit_json(response.body)
23
+ elif response.body is not None:
24
+ click.echo(str(response.body))
25
+
26
+
27
+ def _response_payload(response: EndpointResponse) -> dict[str, Any]:
28
+ payload: dict[str, Any] = {
29
+ "url": response.url,
30
+ "status_code": response.status_code,
31
+ "elapsed_seconds": round(response.elapsed_seconds, 6),
32
+ }
33
+ if response.events:
34
+ payload["events"] = list(response.events)
35
+ else:
36
+ payload["body"] = response.body
37
+ return payload
38
+
39
+
40
+ class SsePrinter:
41
+ """Print SSE data values as they arrive."""
42
+
43
+ def __init__(self, *, enabled: bool):
44
+ self.enabled = enabled
45
+
46
+ def __call__(self, event: dict[str, Any]) -> None:
47
+ if not self.enabled:
48
+ return
49
+ data = event.get("data")
50
+ if isinstance(data, str):
51
+ click.echo(data)
52
+ else:
53
+ click.echo(json.dumps(data, default=str))
@@ -0,0 +1,71 @@
1
+ """Generic HTTP request construction for Agent Bricks endpoint commands."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import urllib.parse
7
+ from typing import Any, Mapping
8
+
9
+ from databricks_agentbricks.cli.endpoint_transport import EndpointRequest
10
+ from databricks_agentbricks.errors import AgentCliError
11
+
12
+
13
+ def parse_pairs(values: tuple[str, ...], *, separator: str, description: str) -> dict[str, str]:
14
+ """Parse repeatable CLI key/value options."""
15
+ parsed: dict[str, str] = {}
16
+ for value in values:
17
+ key, found, item = value.partition(separator)
18
+ key = key.strip()
19
+ if not found or not key or "\n" in key or "\n" in item:
20
+ raise AgentCliError(f"Invalid {description} {value!r}.")
21
+ parsed[key] = item.strip()
22
+ return parsed
23
+
24
+
25
+ def load_json(value: str | None) -> Any:
26
+ """Validate and decode an optional JSON request body."""
27
+ if value is None:
28
+ return None
29
+ try:
30
+ return json.loads(value)
31
+ except json.JSONDecodeError as exc:
32
+ raise AgentCliError(f"Invalid JSON request body: {exc}.") from exc
33
+
34
+
35
+ def request_url(base_url: str, path: str, query: Mapping[str, str]) -> str:
36
+ """Join a request path to a base URL and merge query parameters."""
37
+ parsed_path = urllib.parse.urlsplit(path)
38
+ if parsed_path.scheme or parsed_path.netloc:
39
+ raise AgentCliError("--path must be relative to the selected endpoint URL.")
40
+ joined = urllib.parse.urljoin(f"{base_url}/", path)
41
+ if not query:
42
+ return joined
43
+ parsed = urllib.parse.urlsplit(joined)
44
+ merged_query = dict(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))
45
+ merged_query.update(query)
46
+ return urllib.parse.urlunsplit(parsed._replace(query=urllib.parse.urlencode(merged_query)))
47
+
48
+
49
+ def build_request(
50
+ *,
51
+ base_url: str,
52
+ method: str,
53
+ path: str,
54
+ query: tuple[str, ...],
55
+ json_value: str | None,
56
+ timeout: float,
57
+ sse: bool,
58
+ ) -> EndpointRequest:
59
+ """Materialize one generic HTTP request."""
60
+ body = load_json(json_value)
61
+ request_query = parse_pairs(query, separator="=", description="query parameter")
62
+ headers = {"Content-Type": "application/json"} if json_value is not None else {}
63
+ return EndpointRequest(
64
+ url=request_url(base_url, path, request_query),
65
+ method=method.upper(),
66
+ headers=headers,
67
+ body=body,
68
+ timeout=timeout,
69
+ sse=sse,
70
+ body_set=json_value is not None,
71
+ )