databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentbricks/templates/agent-openai/pyproject.toml

@@ -0,0 +1,42 @@
+[project]
+name = "agent-openai"
+version = "0.1.0"
+description = "OpenAI Agents SDK agent on the Agent Bricks durable runtime for Databricks Apps"
+readme = "README.md"
+authors = [
+    { name = "Agent Developer", email = "developer@example.com" }
+]
+requires-python = ">=3.11"
+dependencies = [
+    # Agent Bricks includes the HTTP runtime; this extra adds the openai framework adapter.
+    "databricks-agentbricks[openai]>=0.2.0",
+    "uvicorn>=0.41.0",
+    "python-dotenv>=1.2.1",
+    "databricks-sdk>=0.124.0",
+    "databricks-agents>=1.9.3",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["agent", "runtime"]
+
+
+[dependency-groups]
+dev = [
+    "hatchling>=1.28.0",
+    "pytest>=9.0.2",
+    "pytest-asyncio>=1.3.0",
+]
+
+[tool.uv]
+default-groups = ["dev"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+asyncio_mode = "auto"
+
+[project.scripts]
+start-server = "runtime.main:main"

databricks_agentbricks/templates/agent-openai/runtime/adapter.py

@@ -0,0 +1,188 @@
+"""Translate between Agent Bricks invocations and the framework-native agent entrypoint."""
+
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from agent.agent import resume_agent, run_agent
+from agents import RunResultStreaming
+from agents.items import ToolApprovalItem
+from openai.types.responses import ResponseTextDeltaEvent
+
+from databricks_agentkit import InvocationContext
+from databricks_agentkit.runtime.auth import AuthError
+
+_RECOVERY_INSTRUCTION = (
+    "This is a recovery attempt after a previous worker stopped before completing this invocation. "
+    "Continue the same task, but assume some tool calls or external side effects may already have "
+    "completed or may still be in progress. Check current state before repeating side effects and "
+    "avoid duplicate actions when possible."
+)
+
+
+def _payload(value: Any) -> dict[str, Any]:
+    if isinstance(value, list):
+        return {"messages": value}
+    if not isinstance(value, dict):
+        raise ValueError("input must be a message list or an object")
+    return value
+
+
+def _session_id(payload: dict[str, Any], context: InvocationContext) -> str:
+    value = payload.get("session_id") or context.session_id
+    if not isinstance(value, str) or not value:
+        raise ValueError("session_id must be a non-empty string")
+    return value
+
+
+def _actor(payload: dict[str, Any], session_id: str) -> str:
+    value = payload.get("actor") or session_id
+    if not isinstance(value, str) or not value:
+        raise ValueError("actor must be a non-empty string")
+    return value
+
+
+def _agent_input(
+    payload: dict[str, Any],
+    session_id: str,
+    *,
+    recovery: bool = False,
+) -> list[Any] | Any:
+    if (resume := payload.get("resume")) is not None:
+        if not isinstance(resume, dict):
+            raise ValueError("resume must be an object")
+        return resume_agent(session_id, resume)
+    messages = payload.get("messages") or []
+    if not isinstance(messages, list):
+        raise ValueError("messages must be a list")
+    if recovery:
+        return [{"role": "developer", "content": _RECOVERY_INSTRUCTION}, *messages]
+    return messages
+
+
+async def invoke(value: Any, context: InvocationContext) -> dict:
+    return await _invoke_agent(_payload(value), context)
+
+
+async def recover(value: Any, context: InvocationContext) -> dict:
+    return await _invoke_agent(_payload(value), context, recovery=True)
+
+
+async def _invoke_agent(
+    payload: dict[str, Any],
+    context: InvocationContext,
+    *,
+    recovery: bool = False,
+) -> dict:
+    session_id = _session_id(payload, context)
+    actor = _actor(payload, session_id)
+    auth = getattr(context, "request_auth", None)
+    user_auth = auth is not None
+    if user_auth and any(
+        payload.get(key) is not None for key in ("resume", "approval", "approvals")
+    ):
+        raise AuthError(
+            "MCP_USER_AUTH_HITL_UNSUPPORTED",
+            "Request-user invocations do not support approval or resume input.",
+            400,
+        )
+    if user_auth and recovery:
+        raise AuthError(
+            "MCP_USER_AUTH_BACKGROUND_UNSUPPORTED",
+            "Request-user invocations cannot be recovered in the background.",
+            400,
+        )
+    internal_session_id = (
+        auth.namespace("session", session_id) if auth and payload.get("session_id") else session_id
+    )
+    actor = auth.namespace("actor", actor) if auth else actor
+    auth_kwargs = {"workspace_client_for": auth.client_for} if user_auth else {}
+    model = payload.get("model")
+    outputs = []
+    async with run_agent(
+        _agent_input(payload, internal_session_id, recovery=recovery),
+        session_id=internal_session_id,
+        actor=actor,
+        model=model if isinstance(model, str) else None,
+        **auth_kwargs,
+    ) as result:
+        async for event in _serialize_events(result):
+            if user_auth and event.get("type") == "interrupt":
+                raise AuthError(
+                    "MCP_USER_AUTH_HITL_UNSUPPORTED",
+                    "Request-user invocations do not support paused approvals.",
+                    400,
+                )
+            await context.emit(event)
+            if event.get("type") in ("message", "interrupt"):
+                outputs.append(event)
+
+    interrupted = bool(outputs and outputs[-1].get("type") == "interrupt")
+    return {
+        "output": [event["message"] if event["type"] == "message" else event for event in outputs],
+        **({"session_id": session_id} if not user_auth or payload.get("session_id") else {}),
+        "status": "interrupted" if interrupted else "completed",
+    }
+
+
+async def _serialize_events(result: RunResultStreaming) -> AsyncGenerator[dict, None]:
+    async for event in result.stream_events():
+        if event.type == "raw_response_event":
+            if isinstance(event.data, ResponseTextDeltaEvent) and event.data.delta:
+                yield {"type": "delta", "content": event.data.delta, "id": event.data.item_id}
+        elif event.type == "run_item_stream_event":
+            if message := _normalize_item(event.item):
+                yield {"type": "message", "message": message}
+
+    for item in result.interruptions:
+        yield {"type": "interrupt", "id": item.call_id, "value": _approval_value(item)}
+
+
+def _approval_value(item: ToolApprovalItem) -> dict:
+    return {"action_requests": [{"name": item.tool_name, "args": _tool_args(item)}]}
+
+
+def _tool_args(item: ToolApprovalItem) -> Any:
+    import json
+
+    args = item.arguments
+    if isinstance(args, str):
+        try:
+            return json.loads(args)
+        except json.JSONDecodeError:
+            return {"arguments": args}
+    return args or {}
+
+
+def _normalize_item(item: Any) -> dict | None:
+    from agents import ItemHelpers
+    from agents.items import MessageOutputItem, ToolCallItem, ToolCallOutputItem
+
+    if isinstance(item, MessageOutputItem):
+        return {"role": "assistant", "content": ItemHelpers.text_message_output(item)}
+    if isinstance(item, ToolCallItem):
+        return {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [{"name": item.tool_name, "args": _tool_args_from_call(item)}],
+        }
+    if isinstance(item, ToolCallOutputItem):
+        return {"role": "tool", "name": _tool_call_name(item), "content": str(item.output)}
+    return None
+
+
+def _tool_args_from_call(item: Any) -> Any:
+    import json
+
+    raw = item.raw_item
+    args = raw.get("arguments") if isinstance(raw, dict) else getattr(raw, "arguments", None)
+    if isinstance(args, str):
+        try:
+            return json.loads(args)
+        except json.JSONDecodeError:
+            return {"arguments": args}
+    return args or {}
+
+
+def _tool_call_name(item: Any) -> str | None:
+    raw = item.raw_item
+    return raw.get("name") if isinstance(raw, dict) else getattr(raw, "name", None)

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

@@ -0,0 +1,25 @@
+"""Run the OpenAI Agents SDK agent through Agent Bricks Runtime."""
+
+import os
+from pathlib import Path
+
+# Importing the agent is side-effect-free (no env is read until configure()), so it sits up top.
+import uvicorn
+from agent.agent import configure
+from dotenv import load_dotenv
+
+from databricks_agentkit import DurableAgentServer
+from runtime.adapter import invoke, recover
+
+# 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)
+
+
+def main() -> None:
+    uvicorn.run("runtime.main:app", host="0.0.0.0", port=int(os.getenv("PORT", "8000")))

databricks_agentbricks/templates/agent-openai/tests/test_agent.py

@@ -0,0 +1,326 @@
+"""Smoke tests for the agent.
+
+Hermetic tests import only the leaf modules (tools, session store, event serialization) — no
+Databricks auth needed, so they run anywhere. The live test builds the full agent and calls the
+model; it is skipped unless a workspace profile is configured.
+"""
+
+import os
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from agent.agent import resume_agent
+from agent.tools import all_tools
+from agents import FunctionTool
+from runtime.adapter import _normalize_item, _serialize_events
+
+
+def test_tools_autoregister():
+    tools = all_tools()
+    assert tools, "expected the sample tool to auto-register"
+    assert all(isinstance(t, FunctionTool) for t in tools)
+    assert {"get_current_time", "send_message"} <= {t.name for t in tools}
+
+
+def test_gated_tool_needs_approval():
+    # The gated demo tool must exist, be listed for approval, and declare needs_approval, or the HITL
+    # demo does nothing.
+    from agent.agent import REQUIRE_APPROVAL
+
+    assert "send_message" in REQUIRE_APPROVAL
+    send = next(t for t in all_tools() if t.name == "send_message")
+    assert send.needs_approval is True
+
+
+class _FakeItem:
+    """Stand-in for an Agents SDK run item, matched by _normalize_item's isinstance checks."""
+
+
+def test_normalize_message_item():
+    from agents.items import MessageOutputItem
+
+    item = object.__new__(MessageOutputItem)
+    # ItemHelpers.text_message_output reads raw_item.content; give it a text part.
+    from openai.types.responses import ResponseOutputMessage, ResponseOutputText
+
+    item.raw_item = ResponseOutputMessage(
+        id="m1",
+        type="message",
+        role="assistant",
+        status="completed",
+        content=[ResponseOutputText(type="output_text", text="hello", annotations=[])],
+    )
+    assert _normalize_item(item) == {"role": "assistant", "content": "hello"}
+
+
+class _FakeToolApproval:
+    def __init__(self, name, args, call_id):
+        self.tool_name, self.arguments, self.call_id = name, args, call_id
+
+
+class _FakeStreamResult:
+    """Minimal RunResultStreaming stand-in: a delta, a message, then a pending interruption."""
+
+    def __init__(self, events, interruptions, state):
+        self._events, self.interruptions, self._state = events, interruptions, state
+        self.final_output = None
+
+    async def stream_events(self):
+        for event in self._events:
+            yield event
+
+    def to_state(self):
+        return self._state
+
+
+@pytest.mark.asyncio
+async def test_agent_events_omit_unavailable_mcp_servers(monkeypatch):
+    import agent.agent as agent_module
+
+    def server(name, *, connect_error=None, list_error=None, cleanup_error=None):
+        value = MagicMock(name=name)
+        value.name = name
+        value.tool_filter = lambda *_args: True
+        value.cache_tools_list = False
+        value.connect = AsyncMock(side_effect=connect_error)
+        value.cleanup = AsyncMock(side_effect=cleanup_error)
+        value.list_tools = AsyncMock(return_value=[], side_effect=list_error)
+        return value
+
+    healthy = server("healthy")
+    unavailable = [
+        server("connect-failure", connect_error=PermissionError("HTTP error 403")),
+        server(
+            "list-failure",
+            list_error=RuntimeError("tool discovery failed"),
+            cleanup_error=RuntimeError("cleanup failed"),
+        ),
+    ]
+    all_servers = [healthy, *unavailable]
+    tool_filters = [server.tool_filter for server in all_servers]
+
+    async def mcp_servers(_extra):
+        return [healthy, *unavailable]
+
+    create_agent = MagicMock(return_value=object())
+
+    monkeypatch.setattr(agent_module, "mcp_servers", mcp_servers)
+    monkeypatch.setattr(agent_module, "build_mcp_servers", lambda: [])
+    monkeypatch.setattr(agent_module, "create_agent", create_agent)
+    monkeypatch.setattr(agent_module, "session_store", lambda _session_id, _actor: None)
+    monkeypatch.setattr(
+        agent_module.Runner,
+        "run_streamed",
+        lambda *_args, **_kwargs: _FakeStreamResult([], [], None),
+    )
+
+    async with agent_module.run_agent([], session_id="s", actor="actor") as result:
+        assert [event async for event in result.stream_events()] == []
+    # create_agent(actor, mcp) — the healthy servers are the second positional arg.
+    assert create_agent.call_args.args[1] == [healthy]
+    assert healthy.cache_tools_list is True
+    assert all(
+        server.tool_filter is tool_filter
+        for server, tool_filter in zip(all_servers, tool_filters, strict=True)
+    )
+    for server in all_servers:
+        server.cleanup.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_agent_events_propagate_request_user_mcp_failure(monkeypatch):
+    import agent.agent as agent_module
+
+    server = MagicMock(name="request-user")
+    server.name = "request-user"
+    server._agentbricks_request_user = True
+    server.connect = AsyncMock(side_effect=PermissionError("request-user denied"))
+    server.cleanup = AsyncMock()
+
+    async def mcp_servers(_extra):
+        return [server]
+
+    monkeypatch.setattr(agent_module, "mcp_servers", mcp_servers)
+    monkeypatch.setattr(agent_module, "build_mcp_servers", lambda: [])
+    monkeypatch.setattr(agent_module, "create_agent", MagicMock(return_value=object()))
+    monkeypatch.setattr(agent_module, "session_store", lambda _session_id, _actor: None)
+    monkeypatch.setattr(
+        agent_module.Runner,
+        "run_streamed",
+        lambda *_args, **_kwargs: _FakeStreamResult([], [], None),
+    )
+
+    with pytest.raises(PermissionError, match="request-user denied"):
+        async with agent_module.run_agent([], session_id="s", actor="actor"):
+            pass
+
+
+@pytest.mark.asyncio
+async def test_serialize_events_relays_interrupt_as_native_event():
+    approval = _FakeToolApproval("send_message", '{"recipient": "x", "body": "y"}', "call-1")
+    sentinel_state = object()
+    result = _FakeStreamResult([], [approval], sentinel_state)
+
+    events = [e async for e in _serialize_events(result)]
+
+    assert events == [
+        {
+            "type": "interrupt",
+            "id": "call-1",
+            "value": {
+                "action_requests": [
+                    {"name": "send_message", "args": {"recipient": "x", "body": "y"}}
+                ]
+            },
+        }
+    ]
+
+
+def test_resume_agent_approves_pending_run(monkeypatch):
+    from agent.agent import _pending_runs
+
+    approved = []
+
+    class _State:
+        def get_interruptions(self):
+            return ["item-a"]
+
+        def approve(self, item):
+            approved.append(item)
+
+        def reject(self, item, rejection_message=None):
+            raise AssertionError("should not reject on approve")
+
+    _pending_runs["sess-2"] = _State()
+    resume_agent("sess-2", {"decisions": [{"type": "approve"}]})
+    assert approved == ["item-a"]
+    assert "sess-2" not in _pending_runs  # popped so it can't be resumed twice
+
+
+def test_resume_agent_without_pending_run_raises():
+    with pytest.raises(RuntimeError, match="No paused run"):
+        resume_agent("never-started", {"decisions": [{"type": "approve"}]})
+
+
+def test_configure_raises_clear_error_without_auth(monkeypatch):
+    from agent.agent import configure
+
+    monkeypatch.delenv("DATABRICKS_CONFIG_PROFILE", raising=False)
+    monkeypatch.delenv("DATABRICKS_HOST", raising=False)
+    monkeypatch.delenv("DATABRICKS_TOKEN", raising=False)
+    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", "/nonexistent-databrickscfg")
+    with pytest.raises(RuntimeError, match="Databricks auth is not configured"):
+        configure()
+
+
+def test_configure_routes_openai_client_to_workspace(monkeypatch):
+    import agent.agent as agent_module
+    import agents
+
+    workspace = object()
+    created = object()
+    client = MagicMock(return_value=created)
+    set_client = MagicMock()
+    set_api = MagicMock()
+
+    monkeypatch.setattr(agent_module, "workspace_client", lambda: workspace)
+    monkeypatch.setattr(
+        agent_module,
+        "workspace_headers",
+        lambda: {"X-Databricks-Org-Id": "123"},
+    )
+    monkeypatch.setattr(agent_module, "AsyncDatabricksOpenAI", client)
+    monkeypatch.setattr(agent_module, "configure_tracing", lambda: None)
+    monkeypatch.setattr(agents, "set_default_openai_client", set_client)
+    monkeypatch.setattr(agents, "set_default_openai_api", set_api)
+
+    agent_module.configure()
+
+    client.assert_called_once_with(
+        workspace_client=workspace,
+        default_headers={"X-Databricks-Org-Id": "123"},
+        use_ai_gateway=True,
+    )
+    set_client.assert_called_once_with(created)
+    set_api.assert_called_once_with("chat_completions")
+
+
+def test_session_store_defaults_to_in_process(monkeypatch):
+    import databricks_agentkit.openai.sessions as ss
+
+    monkeypatch.delenv("AGENT_SESSION_STORE", raising=False)
+    ss._local_sessions.clear()
+    # In-process default: same session id returns the same cached SQLiteSession (multi-turn works).
+    assert ss.session_store("abc-123") is ss.session_store("abc-123")
+
+
+def test_session_store_selects_durable_store(monkeypatch):
+    import databricks_agentkit.openai.sessions as ss
+
+    monkeypatch.setenv("AGENT_SESSION_STORE", "my-store")
+    monkeypatch.setattr(ss, "SessionStoreClient", lambda *a, **k: _FakeStoreClient())
+    store = ss.session_store("abc-123")
+    assert isinstance(store, ss.DatabricksSessionStore)
+
+
+class _FakeStoreClient:
+    def set_session_store(self, name):
+        return self
+
+
+@pytest.mark.asyncio
+async def test_adapter_recovery_marks_replayed_agent_input(monkeypatch):
+    import runtime.adapter as adapter
+
+    calls = []
+
+    @asynccontextmanager
+    async def fake_run_agent(agent_input, **kwargs):
+        calls.append((agent_input, kwargs))
+        yield _FakeStreamResult([], [], None)
+
+    monkeypatch.setattr(adapter, "run_agent", fake_run_agent)
+    payload = {"session_id": "session-1", "messages": [{"role": "user", "content": "hi"}]}
+    context = SimpleNamespace(session_id="runtime-session", emit=AsyncMock())
+
+    await adapter.invoke(payload, context)
+    await adapter.recover(payload, context)
+
+    assert calls == [
+        (payload["messages"], {"session_id": "session-1", "actor": "session-1", "model": None}),
+        (
+            [
+                {"role": "developer", "content": adapter._RECOVERY_INSTRUCTION},
+                *payload["messages"],
+            ],
+            {"session_id": "session-1", "actor": "session-1", "model": None},
+        ),
+    ]
+    assert payload["messages"] == [{"role": "user", "content": "hi"}]
+
+
+def _has_workspace_auth() -> bool:
+    return bool(
+        os.getenv("DATABRICKS_CONFIG_PROFILE")
+        or (os.getenv("DATABRICKS_HOST") and os.getenv("DATABRICKS_TOKEN"))
+    )
+
+
+@pytest.mark.skipif(
+    not _has_workspace_auth(),
+    reason="no Databricks profile configured; skipping live model call",
+)
+@pytest.mark.asyncio
+async def test_agent_responds_end_to_end():
+    from agent.agent import configure, run_agent
+
+    configure()
+    async with run_agent(
+        [{"role": "user", "content": "Reply with the single word: pong"}],
+        session_id="test-e2e",
+        actor="test-actor",
+    ) as result:
+        _ = [event async for event in result.stream_events()]
+        assert result.final_output

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

@@ -0,0 +1,19 @@
+from databricks_agentkit import DurableAgentServer
+from databricks_agentkit.runtime.store import InMemoryRuntimeStore
+
+
+async def _invoke(request, context):
+    return request
+
+
+def test_runtime_exposes_invocation_routes() -> None:
+    app = DurableAgentServer(runtime_store=InMemoryRuntimeStore())
+    app.invoke(_invoke)
+    app.recover(_invoke)
+    paths = app.openapi()["paths"]
+
+    assert paths["/api/invocations"]["post"]
+    assert paths["/api/invocations/{invocation_id}"]["get"]
+    assert paths["/api/invocations/{invocation_id}/events"]["get"]
+    assert "/invocations" not in paths
+    assert "/api/health" not in paths

databricks_agentbricks/templates/agent-openai/tests/test_workspace.py

@@ -0,0 +1,28 @@
+from databricks_agentkit.runtime import workspace
+
+
+class _FakeWorkspaceClient:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+
+def test_workspace_client_adds_account_routing_header(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_WORKSPACE_ID", " 123456 ")
+    monkeypatch.setattr(workspace, "WorkspaceClient", _FakeWorkspaceClient)
+
+    client = workspace.workspace_client()
+
+    assert client.kwargs == {
+        "custom_headers": {"X-Databricks-Org-Id": "123456"},
+    }
+    assert workspace.workspace_headers() == {"X-Databricks-Org-Id": "123456"}
+
+
+def test_workspace_client_uses_default_sdk_resolution_without_workspace_id(monkeypatch):
+    monkeypatch.delenv("DATABRICKS_WORKSPACE_ID", raising=False)
+    monkeypatch.setattr(workspace, "WorkspaceClient", _FakeWorkspaceClient)
+
+    client = workspace.workspace_client()
+
+    assert client.kwargs == {}
+    assert workspace.workspace_headers() == {}

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

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

databricks_agentbricks/templates/custom-agent-langgraph/.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-langgraph/README.md

@@ -0,0 +1,22 @@
+# Custom LangGraph server
+
+This template shows how to serve a LangGraph 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-langgraph`.