databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentbricks/templates/custom-agent-langgraph/agent/agent.py

@@ -0,0 +1,27 @@
+from typing import Any
+
+from databricks_langchain import ChatDatabricks
+from langchain.agents import create_agent
+
+from databricks_agentkit import workspace_client, workspace_headers
+
+MODEL = "databricks-gpt-5-2"
+
+
+class _RoutedChatDatabricks(ChatDatabricks):
+    def _get_client_kwargs(self) -> dict[str, Any]:
+        kwargs = super()._get_client_kwargs()
+        if headers := workspace_headers():
+            kwargs["default_headers"] = headers
+        return kwargs
+
+
+async def invoke(messages: list[dict[str, Any]]) -> dict[str, Any]:
+    agent = create_agent(
+        model=_RoutedChatDatabricks(endpoint=MODEL, workspace_client=workspace_client()),
+        tools=[],
+    )
+    result = await agent.ainvoke({"messages": messages})
+    message = result["messages"][-1]
+    output = message.model_dump() if hasattr(message, "model_dump") else message
+    return {"output": [output]}

databricks_agentbricks/templates/custom-agent-langgraph/app.yaml

@@ -0,0 +1 @@
+command: ["uv", "run", "start-server"]

databricks_agentbricks/templates/custom-agent-langgraph/pyproject.toml

@@ -0,0 +1,38 @@
+[project]
+name = "custom-agent-langgraph"
+version = "0.1.0"
+description = "LangGraph agent with a custom FastAPI server"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "databricks-agentbricks>=0.2.0",
+    "databricks-langchain>=0.17.0",
+    "databricks-sdk>=0.124.0",
+    "fastapi>=0.129.0",
+    "langchain>=1.0.0",
+    "python-dotenv>=1.2.1",
+    "uvicorn>=0.41.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["agent", "runtime"]
+
+[dependency-groups]
+dev = [
+    "hatchling>=1.28.0",
+    "httpx>=0.28.1",
+    "pytest>=9.0.2",
+]
+
+[tool.uv]
+default-groups = ["dev"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[project.scripts]
+start-server = "runtime.main:main"

databricks_agentbricks/templates/custom-agent-langgraph/runtime/main.py

@@ -0,0 +1,31 @@
+import os
+from pathlib import Path
+from typing import Any
+
+import uvicorn
+from agent.agent import invoke
+from dotenv import load_dotenv
+from fastapi import FastAPI
+from pydantic import BaseModel, ConfigDict, Field
+
+# override=False so injected DATABRICKS_* (from `ab dev -p` or the deploy platform) win over a
+# checked-in .env.
+load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env", override=False)
+
+
+class InvocationRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    input: list[dict[str, Any]] = Field(default_factory=list)
+
+
+app = FastAPI(title="Custom LangGraph Agent")
+
+
+@app.post("/invocations")
+async def invocations(body: InvocationRequest) -> dict[str, Any]:
+    return await invoke(body.input)
+
+
+def main() -> None:
+    uvicorn.run("runtime.main:app", host="0.0.0.0", port=int(os.getenv("PORT", "8000")))

databricks_agentbricks/templates/custom-agent-langgraph/tests/test_runtime.py

@@ -0,0 +1,29 @@
+from unittest.mock import AsyncMock
+
+from fastapi.testclient import TestClient
+
+import runtime.main as runtime
+
+
+def test_foreground_invocation(monkeypatch) -> None:
+    invoke = AsyncMock(return_value={"output": [{"role": "assistant", "content": "hello"}]})
+    monkeypatch.setattr(runtime, "invoke", invoke)
+
+    response = TestClient(runtime.app).post(
+        "/invocations",
+        json={"input": [{"role": "user", "content": "hi"}]},
+    )
+
+    assert response.status_code == 200
+    assert response.json()["output"][0]["content"] == "hello"
+    invoke.assert_awaited_once_with([{"role": "user", "content": "hi"}])
+
+
+def test_custom_server_rejects_runtime_transport_fields() -> None:
+    response = TestClient(runtime.app).post(
+        "/invocations",
+        json={"input": [], "background": True},
+    )
+
+    assert response.status_code == 422
+    assert "/api/invocations" not in runtime.app.openapi()["paths"]

databricks_agentbricks/templates/custom-agent-openai/.env.example

@@ -0,0 +1 @@
+DATABRICKS_CONFIG_PROFILE=DEFAULT

databricks_agentbricks/templates/custom-agent-openai/.gitignore

@@ -0,0 +1,11 @@
+.env
+.venv/
+__pycache__/
+*.py[cod]
+.pytest_cache/
+dist/
+*.egg-info/
+
+# Agent Bricks local dev artifacts (`agentbricks dev`)
+.agentbricks/
+app.agentbricksdev.yaml

databricks_agentbricks/templates/custom-agent-openai/README.md

@@ -0,0 +1,22 @@
+# Custom OpenAI Agents server
+
+This template shows how to serve an OpenAI Agents SDK agent with an ordinary FastAPI application.
+It does not use the managed `DurableAgentServer` HTTP server or durable runtime.
+
+This template does not load managed tool bindings from `agent.toml`, so `ab tools add` is not
+supported. Wire framework-native Python tools and MCP servers directly in `agent/agent.py`.
+
+```bash
+ab dev
+```
+
+Call its single foreground endpoint:
+
+```bash
+curl -sS http://localhost:8000/invocations \
+  -H 'Content-Type: application/json' \
+  -d '{"input":[{"role":"user","content":"Hello"}]}'
+```
+
+Edit `runtime/main.py` to define your own HTTP contract. Edit `agent/agent.py` to change the model
+or agent behavior. Deploy with `ab --profile <profile> deploy custom-agent-openai`.

databricks_agentbricks/templates/custom-agent-openai/agent/agent.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+from agents import Agent, Runner, set_default_openai_api, set_default_openai_client
+from databricks_openai import AsyncDatabricksOpenAI
+
+from databricks_agentkit import workspace_client, workspace_headers
+
+MODEL = "databricks-gpt-5-2"
+
+
+def _configure_client() -> None:
+    headers = workspace_headers()
+    workspace = workspace_client()
+    client = (
+        AsyncDatabricksOpenAI(workspace_client=workspace, default_headers=headers)
+        if headers
+        else AsyncDatabricksOpenAI(workspace_client=workspace)
+    )
+    set_default_openai_client(client)
+    set_default_openai_api("chat_completions")
+
+
+async def invoke(messages: list[dict[str, Any]]) -> dict[str, Any]:
+    _configure_client()
+    result = await Runner.run(
+        Agent(name="Agent", instructions="You are a helpful assistant.", model=MODEL),
+        messages,
+    )
+    return {"output": [{"role": "assistant", "content": str(result.final_output)}]}

databricks_agentbricks/templates/custom-agent-openai/app.yaml

@@ -0,0 +1 @@
+command: ["uv", "run", "start-server"]

databricks_agentbricks/templates/custom-agent-openai/pyproject.toml

@@ -0,0 +1,38 @@
+[project]
+name = "custom-agent-openai"
+version = "0.1.0"
+description = "OpenAI Agents SDK agent with a custom FastAPI server"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "databricks-agentbricks>=0.2.0",
+    "databricks-openai>=0.13.0",
+    "databricks-sdk>=0.124.0",
+    "fastapi>=0.129.0",
+    "openai-agents>=0.8.0",
+    "python-dotenv>=1.2.1",
+    "uvicorn>=0.41.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["agent", "runtime"]
+
+[dependency-groups]
+dev = [
+    "hatchling>=1.28.0",
+    "httpx>=0.28.1",
+    "pytest>=9.0.2",
+]
+
+[tool.uv]
+default-groups = ["dev"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[project.scripts]
+start-server = "runtime.main:main"

databricks_agentbricks/templates/custom-agent-openai/runtime/main.py

@@ -0,0 +1,31 @@
+import os
+from pathlib import Path
+from typing import Any
+
+import uvicorn
+from agent.agent import invoke
+from dotenv import load_dotenv
+from fastapi import FastAPI
+from pydantic import BaseModel, ConfigDict, Field
+
+# override=False so injected DATABRICKS_* (from `ab dev -p` or the deploy platform) win over a
+# checked-in .env.
+load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env", override=False)
+
+
+class InvocationRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    input: list[dict[str, Any]] = Field(default_factory=list)
+
+
+app = FastAPI(title="Custom OpenAI Agent")
+
+
+@app.post("/invocations")
+async def invocations(body: InvocationRequest) -> dict[str, Any]:
+    return await invoke(body.input)
+
+
+def main() -> None:
+    uvicorn.run("runtime.main:app", host="0.0.0.0", port=int(os.getenv("PORT", "8000")))

databricks_agentbricks/templates/custom-agent-openai/tests/test_runtime.py

@@ -0,0 +1,29 @@
+from unittest.mock import AsyncMock
+
+from fastapi.testclient import TestClient
+
+import runtime.main as runtime
+
+
+def test_foreground_invocation(monkeypatch) -> None:
+    invoke = AsyncMock(return_value={"output": [{"role": "assistant", "content": "hello"}]})
+    monkeypatch.setattr(runtime, "invoke", invoke)
+
+    response = TestClient(runtime.app).post(
+        "/invocations",
+        json={"input": [{"role": "user", "content": "hi"}]},
+    )
+
+    assert response.status_code == 200
+    assert response.json()["output"][0]["content"] == "hello"
+    invoke.assert_awaited_once_with([{"role": "user", "content": "hi"}])
+
+
+def test_custom_server_rejects_runtime_transport_fields() -> None:
+    response = TestClient(runtime.app).post(
+        "/invocations",
+        json={"input": [], "stream": True},
+    )
+
+    assert response.status_code == 422
+    assert "/api/invocations" not in runtime.app.openapi()["paths"]

databricks_agentbricks/templates/sandbox_mcp.py

@@ -0,0 +1,39 @@
+from typing import Any
+
+from databricks.sdk import WorkspaceClient
+from databricks_openai.agents import McpServer  # ty: ignore[unresolved-import]
+
+# BEGIN: agentbricks add-sandbox
+# The downscope is fixed when this file is generated. It is sent as MCP metadata,
+# outside model-controlled tool arguments, on every sandbox call.
+_SANDBOX_DOWNSCOPE: dict[str, list[dict[str, str]]] = {}  # __AGENTBRICKS_SANDBOX_DOWNSCOPE__
+
+
+class _SandboxMcpServer(McpServer):
+    async def call_tool(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any] | None,
+        **kwargs: Any,
+    ) -> Any:
+        incoming_meta = kwargs.pop("meta", None)
+        meta = dict(incoming_meta) if isinstance(incoming_meta, dict) else {}
+        meta["downscope"] = _SANDBOX_DOWNSCOPE
+        return await super().call_tool(tool_name, arguments, meta=meta, **kwargs)
+
+
+def _build_sandbox_mcp_server() -> McpServer:
+    workspace_client = WorkspaceClient()
+    return _SandboxMcpServer(
+        url=(
+            f"{workspace_client.config.host.rstrip('/')}/ai-gateway/mcp-services/system.ai.sandbox"
+        ),
+        workspace_client=workspace_client,
+        timeout=120.0,
+        name="system.ai.sandbox",
+        # Production advertises ``sandbox``; older preview workspaces use ``run_code``.
+        tool_filter={"allowed_tool_names": ["sandbox", "run_code"]},
+    )
+
+
+# END: agentbricks add-sandbox

databricks_agentbricks/templates/sandbox_mcp_langgraph.py

@@ -0,0 +1,39 @@
+from typing import Any
+
+from databricks.sdk import WorkspaceClient
+from databricks_langchain import DatabricksMCPServer  # ty: ignore[unresolved-import]
+from langchain_mcp_adapters.sessions import create_session  # ty: ignore[unresolved-import]
+
+# BEGIN: agentbricks add-sandbox
+# The downscope is fixed when this file is generated. It is sent as MCP metadata,
+# outside model-controlled tool arguments, on every sandbox call.
+_SANDBOX_DOWNSCOPE: dict[str, list[dict[str, str]]] = {}  # __AGENTBRICKS_SANDBOX_DOWNSCOPE__
+
+
+def _build_sandbox_mcp_server() -> DatabricksMCPServer:
+    workspace_client = WorkspaceClient()
+    return DatabricksMCPServer(
+        name="system.ai.sandbox",
+        url=(
+            f"{workspace_client.config.host.rstrip('/')}/ai-gateway/mcp-services/system.ai.sandbox"
+        ),
+        workspace_client=workspace_client,
+        timeout=120.0,
+    )
+
+
+async def _sandbox_tool_interceptor(request: Any, handler: Any) -> Any:
+    if request.server_name != "system.ai.sandbox":
+        return await handler(request)
+
+    connection = _build_sandbox_mcp_server().to_connection_dict()
+    async with create_session(connection) as session:
+        await session.initialize()
+        return await session.call_tool(
+            request.name,
+            request.args,
+            meta={"downscope": _SANDBOX_DOWNSCOPE},
+        )
+
+
+# END: agentbricks add-sandbox

databricks_agentbricks/templates/ui/agent-langgraph/CHAT_APP.md

@@ -0,0 +1,49 @@
+# LangGraph Chat App Overlay
+
+`ab init --framework langgraph` copies this framework-specific overlay after the base
+`agent-langgraph` template (it is included by default; `--disable-chat-app` opts out). It is
+intentionally not a post-generation mutation command.
+
+## Installed files
+
+- `ui/` contains the zero-build chat client.
+- `runtime/ui.py` serves the assets and exposes demo APIs for memory and sessions.
+- `runtime/main.py` installs the chat routes on the base FastAPI runtime.
+- `tests/test_demo_ui.py` verifies the browser-facing routes.
+
+## Behavior
+
+The capability indicators are automatic. Streaming and background reflect the runtime contract;
+Session reflects checkpoint history; Memory requires `AGENT_MEMORY_STORE`; Traces reflects the MLflow
+tracing config (destination + experiment). Selecting Memory opens a pane for browsing the store's
+entries (searchable, filterable by actor, with a modal for each entry); the other indicators expand
+in place, and Traces links out to the MLflow experiment. The transport selector is the only manual
+capability choice.
+
+The composer's model picker lists the workspace's Unity Catalog AI Gateway chat model services — the
+`system.ai.*` model services from `GET /api/2.1/unity-catalog/model-services?parent=schemas/system.ai`
+(embeddings-only services filtered out), exposed as `GET /api/demo/models` and pinned to
+`agent.agent.MODEL` as the default. Each request sends the selected model as `model` in the
+invocation body; the agent is rebuilt per turn, so the picker changes the model for the next turn
+without a restart. Discovery is best-effort: if listing is unavailable (e.g. `system.ai` isn't
+readable), the picker falls back to just the default. Omitting `model` uses `MODEL`.
+
+The agent calls the chosen model through the gateway (`<host>/ai-gateway/mlflow/v1`) rather than
+`/serving-endpoints`, so `MODEL` is a `system.ai.*` model service name. The picker is capped
+(`_MODEL_LIMIT`, 20) with the default pinned first and the rest alphabetical, so truncation never
+drops the configured default. Transient list failures are retried in
+`databricks_agentkit.runtime.model_services` before the fallback applies.
+
+The UI reads local history from the LangGraph checkpoint and managed history from Session Store
+items. It keeps a stable application session UUID in browser local storage and includes it inside
+every durable invocation's `input`; each turn gets a separate invocation UUID. The router cookie is
+independent and may still provide sticky replica routing.
+
+The Sessions card creates new session UUIDs in the browser. With a managed Session Store,
+`GET /api/demo/sessions` lists the most recent sessions for the signed-in actor and each Open action
+calls `POST /api/demo/sessions/{session_id}/open`. Opening a session verifies ownership and reloads
+its transcript and pending LangGraph state. In local in-memory mode only the current browser session
+can be listed because there is no shared session index.
+
+Transcript responses include only user, assistant, tool, system, and human-decision message items;
+checkpoint fragments remain in Session Store but are never returned to the chat UI.

databricks_agentbricks/templates/ui/agent-langgraph/runtime/main.py

@@ -0,0 +1,26 @@
+"""Run the LangGraph agent with the optional Agent Bricks chat app."""
+
+import os
+from pathlib import Path
+
+import uvicorn
+from agent.agent import configure
+from dotenv import load_dotenv
+from runtime.adapter import invoke, recover
+from runtime.ui import install_ui
+
+from databricks_agentkit import DurableAgentServer
+
+# override=False so injected DATABRICKS_* (from `ab dev -p` or the deploy platform) win over a
+# checked-in .env.
+load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env", override=False)
+configure()
+
+app = DurableAgentServer()
+app.invoke(invoke)
+app.recover(recover)
+install_ui(app)
+
+
+def main() -> None:
+    uvicorn.run("runtime.main:app", host="0.0.0.0", port=int(os.getenv("PORT", "8000")))