databricks-agentbricks 0.2.0__py3-none-any.whl

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

@@ -0,0 +1,144 @@
+"""Translate between Agent Bricks invocations and the framework-native agent entrypoint."""
+
+from collections.abc import AsyncGenerator, AsyncIterator
+from typing import Any
+
+from agent.agent import recovery_input, run_agent
+from langchain.messages import AIMessageChunk
+from langgraph.types import Command
+
+from databricks_agentkit import InvocationContext
+from databricks_agentkit.runtime.auth import AuthError
+
+
+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]) -> Any:
+    if (resume := payload.get("resume")) is not None:
+        return Command(resume=resume)
+    messages = payload.get("messages") or []
+    if not isinstance(messages, list):
+        raise ValueError("messages must be a list")
+    return {"messages": messages}
+
+
+async def invoke(value: Any, context: InvocationContext) -> dict:
+    payload = _payload(value)
+    _check_user_input(payload, context)
+    return await _invoke_agent(_agent_input(payload), payload, context)
+
+
+async def recover(value: Any, context: InvocationContext) -> dict:
+    payload = _payload(value)
+    _check_user_input(payload, context)
+    if getattr(context, "request_auth", None) is not None:
+        raise AuthError(
+            "MCP_USER_AUTH_BACKGROUND_UNSUPPORTED",
+            "Request-user invocations cannot be recovered in the background.",
+            400,
+        )
+    session_id = _session_id(payload, context)
+    actor = _actor(payload, session_id)
+    agent_input = await recovery_input(
+        _agent_input(payload),
+        session_id=session_id,
+        actor=actor,
+        invocation_id=context.invocation_id,
+    )
+    return await _invoke_agent(agent_input, payload, context)
+
+
+def _check_user_input(payload: dict[str, Any], context: InvocationContext) -> None:
+    auth = getattr(context, "request_auth", None)
+    if auth is not None 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,
+        )
+
+
+async def _invoke_agent(
+    agent_input: Any,
+    payload: dict[str, Any],
+    context: InvocationContext,
+) -> dict:
+    session_id = _session_id(payload, context)
+    actor = _actor(payload, session_id)
+    auth = getattr(context, "request_auth", None)
+    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
+    user_auth = auth is not None
+    auth_kwargs = {"workspace_client_for": auth.client_for} if user_auth else {}
+    model = payload.get("model")
+    outputs = []
+    async for event in _serialize_events(
+        run_agent(
+            agent_input,
+            session_id=internal_session_id,
+            actor=actor,
+            model=model if isinstance(model, str) else None,
+            invocation_id=context.invocation_id,
+            **auth_kwargs,
+        )
+    ):
+        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(async_stream: AsyncIterator[Any]) -> AsyncGenerator[dict, None]:
+    async for mode, payload in async_stream:
+        if mode == "updates":
+            if interrupts := payload.get("__interrupt__"):
+                for item in interrupts:
+                    yield {"type": "interrupt", "id": item.id, "value": item.value}
+                continue
+            for node_data in payload.values():
+                messages = node_data.get("messages", []) if isinstance(node_data, dict) else []
+                for message in messages:
+                    yield {"type": "message", "message": message.model_dump()}
+        elif mode == "messages":
+            try:
+                chunk = payload[0]
+                if isinstance(chunk, AIMessageChunk) and (content := chunk.content):
+                    yield {"type": "delta", "content": content, "id": chunk.id}
+            except (KeyError, IndexError, TypeError):
+                continue

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

@@ -0,0 +1,25 @@
+"""Run the LangGraph 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-langgraph/tests/test_agent.py

@@ -0,0 +1,211 @@
+"""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 types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+from agent.tools import all_tools
+from langchain_core.tools import BaseTool
+from runtime.adapter import _serialize_events
+
+from databricks_agentkit.langgraph.session_store import (
+    checkpointer,
+    invocation_metadata,
+    thread_config,
+)
+
+
+def test_tools_autoregister():
+    tools = all_tools()
+    assert tools, "expected the sample tool to auto-register"
+    assert all(isinstance(t, BaseTool) for t in tools)
+    assert {"get_current_time", "send_message"} <= {t.name for t in tools}
+
+
+def test_gated_tool_is_in_require_approval():
+    # The gated demo tool must exist and be listed for approval, or the HITL demo does nothing.
+    from agent.agent import REQUIRE_APPROVAL
+
+    assert REQUIRE_APPROVAL.get("send_message")
+    assert "send_message" in {t.name for t in all_tools()}
+
+
+class _FakeInterrupt:
+    def __init__(self, value, id):  # mirrors langgraph.types.Interrupt's `.value` / `.id`
+        self.value, self.id = value, id
+
+
+async def _aiter(events):
+    for e in events:
+        yield e
+
+
+@pytest.mark.asyncio
+async def test_serialize_events_relays_interrupt_as_native_event():
+    hitl = {"action_requests": [{"name": "send_message", "args": {"recipient": "x", "body": "y"}}]}
+    stream = _aiter([("updates", {"__interrupt__": (_FakeInterrupt(hitl, "int-1"),)})])
+    events = [e async for e in _serialize_events(stream)]
+    assert events == [{"type": "interrupt", "id": "int-1", "value": hitl}]
+
+
+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_chat_model_forwards_account_routing_header(monkeypatch):
+    from agent.agent import _RoutedChatDatabricks
+
+    monkeypatch.setenv("DATABRICKS_WORKSPACE_ID", "123456")
+    model = _RoutedChatDatabricks(endpoint="test-endpoint")
+
+    assert model._get_client_kwargs()["default_headers"] == {"X-Databricks-Org-Id": "123456"}
+
+
+def test_thread_config_from_session_id():
+    # actor_id rides alongside thread_id — the durable saver maps it onto the Session's actor.
+    assert thread_config("abc-123") == {
+        "configurable": {"thread_id": "abc-123", "actor_id": "abc-123"}
+    }
+
+
+def test_thread_config_uses_supplied_actor():
+    # A caller-supplied actor (e.g. the signed-in user) partitions the durable store per user.
+    assert thread_config("abc-123", "alice") == {
+        "configurable": {"thread_id": "abc-123", "actor_id": "alice"}
+    }
+
+
+def test_checkpointer_is_shared(monkeypatch):
+    # In-memory by default (no AGENT_SESSION_STORE); built once and shared so multi-turn works.
+    import databricks_agentkit.langgraph.session_store as ss
+
+    monkeypatch.setattr(ss, "_saver", None)  # reset the process-wide saver
+    assert checkpointer() is checkpointer()
+
+
+def test_session_store_selects_durable_saver(monkeypatch):
+    # AGENT_SESSION_STORE must route to the durable Session Store saver, not stay in-memory. Stub the
+    # REST client so it stays hermetic (no network); the saver builds without touching the API.
+    import databricks_agentkit.langgraph.session_store as ss
+
+    monkeypatch.setattr(ss, "_saver", None)
+    monkeypatch.setenv("AGENT_SESSION_STORE", "my-store")
+    monkeypatch.setattr(ss, "SessionStoreClient", lambda *a, **k: _FakeStoreClient())
+    saver = checkpointer()
+    assert isinstance(saver, ss.DatabricksSessionStoreSaver)
+
+
+class _FakeStoreClient:
+    def set_session_store(self, name):
+        return self
+
+
+@pytest.mark.asyncio
+async def test_recovery_input_resumes_current_checkpoint(monkeypatch):
+    import agent.agent as agent_module
+
+    class Saver:
+        async def aget_tuple(self, config):
+            return SimpleNamespace(metadata=invocation_metadata("inv-1"))
+
+    monkeypatch.setattr(agent_module, "checkpointer", lambda: Saver())
+    original = {"messages": [{"role": "user", "content": "hi"}]}
+    recovered = await agent_module.recovery_input(
+        original,
+        session_id="session-1",
+        actor="actor-1",
+        invocation_id="inv-1",
+    )
+
+    assert recovered is None
+
+
+@pytest.mark.asyncio
+async def test_recovery_replays_input_without_current_checkpoint(monkeypatch):
+    import agent.agent as agent_module
+
+    class Saver:
+        async def aget_tuple(self, config):
+            return None
+
+    monkeypatch.setattr(agent_module, "checkpointer", lambda: Saver())
+    original = {"messages": [{"role": "user", "content": "hi"}]}
+    recovered = await agent_module.recovery_input(
+        original,
+        session_id="session-1",
+        actor="actor-1",
+        invocation_id="inv-1",
+    )
+
+    assert recovered is original
+
+
+@pytest.mark.asyncio
+async def test_adapter_calls_same_run_agent_for_invoke_and_recovery(monkeypatch):
+    import runtime.adapter as adapter
+
+    calls = []
+
+    async def fake_run_agent(agent_input, **kwargs):
+        calls.append((agent_input, kwargs))
+        if False:
+            yield None
+
+    async def fake_recovery_input(agent_input, **_kwargs):
+        return None
+
+    monkeypatch.setattr(adapter, "run_agent", fake_run_agent)
+    monkeypatch.setattr(adapter, "recovery_input", fake_recovery_input)
+    context = SimpleNamespace(
+        invocation_id="inv-1",
+        session_id="runtime-session",
+        emit=AsyncMock(),
+    )
+    payload = {"session_id": "session-1", "messages": [{"role": "user", "content": "hi"}]}
+
+    await adapter.invoke(payload, context)
+    await adapter.recover(payload, context)
+
+    assert calls[0][0] == {"messages": payload["messages"]}
+    assert calls[1][0] is None
+    assert calls[0][1] == calls[1][1]
+
+
+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()
+    events = [
+        event
+        async for event in run_agent(
+            {"messages": [{"role": "user", "content": "Reply with the single word: pong"}]},
+            session_id="test-e2e",
+            actor="test-actor",
+        )
+    ]
+    assert events

databricks_agentbricks/templates/agent-langgraph/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-langgraph/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/agent-openai/.env.example

@@ -0,0 +1,37 @@
+# Copy to .env for local development:  cp .env.example .env
+#
+# The only required setting for local dev is a Databricks auth profile (used to call the model).
+# Everything else — MLflow tracing, long-term memory — is optional and off by default. Invocation
+# durability uses an in-memory store locally; conversation history uses an in-process SQLite session.
+
+# Databricks auth profile (from `databricks auth profiles`). Used to call the model endpoint.
+DATABRICKS_CONFIG_PROFILE=DEFAULT
+# Or use explicit host/token instead of a profile:
+# DATABRICKS_HOST=https://<your-workspace>.databricks.com
+# DATABRICKS_TOKEN=dapi...
+
+# --- Optional: MLflow tracing ---
+# Leave UNSET to skip tracing (local dev). To enable, set a destination AND an experiment (either
+# form of each works): destination = MLFLOW_TRACKING_URI or MLFLOW_TRACING_DESTINATION;
+# experiment = MLFLOW_EXPERIMENT_ID or MLFLOW_EXPERIMENT_NAME.
+# MLFLOW_TRACKING_URI="databricks"
+# MLFLOW_EXPERIMENT_ID=
+# MLFLOW_EXPERIMENT_NAME=
+# For fully-local tracing (no workspace), point MLFLOW_TRACKING_URI at a local OSS MLflow server
+# (`mlflow server`), e.g. MLFLOW_TRACKING_URI="http://localhost:5000".
+
+# --- Optional: long-term memory tools ---
+# Leave UNSET to skip. Set to a managed memory store ID to register the remember/recall tools
+# (persist/search facts across conversations); see databricks_agentkit/openai/memory.py. Memory is
+# partitioned per user: the agent uses the signed-in user as the actor (falls back to "agent"
+# locally); edit `_actor` in agent/agent.py to change that.
+# AGENT_MEMORY_STORE=<memory-store-id>
+
+# --- Optional: durable state (managed Session Store) ---
+# Leave UNSET to keep the in-process session (SQLiteSession in :memory:): multi-turn works within one
+# process but does not survive restarts/replicas. Set to an existing managed Session Store name to
+# persist the transcript through its REST API (DatabricksSessionStore) — durable conversation history
+# over RPCs (no DB connection). Access uses your normal Databricks auth (must have access to the
+# store). NOTE: this persists the transcript only; paused human-in-the-loop runs are in-process even
+# when this is set. See databricks_agentkit/openai/sessions.py.
+# AGENT_SESSION_STORE=my-agent-sessions

databricks_agentbricks/templates/agent-openai/.gitignore

@@ -0,0 +1,199 @@
+# Created by https://www.toptal.com/developers/gitignore/api/python
+# Edit at https://www.toptal.com/developers/gitignore?templates=python
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# VS Code
+.vscode/
+
+### Python Patch ###
+# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
+poetry.toml
+
+# ruff
+.ruff_cache/
+
+# LSP config files
+pyrightconfig.json
+
+# End of https://www.toptal.com/developers/gitignore/api/python
+
+.DS_*
+
+# Databricks / MLflow local artifacts
+**/mlruns/
+mlflow.db
+**/.databricks
+.claude/
+
+# Environment files
+**/.env
+**/.env.local
+
+# Local (Lakebase-free) SQLite session store
+local_agent_sessions.db
+local_agent_sessions.db-*
+*.db-journal
+
+# Agent Bricks local dev artifacts (`agentbricks dev`)
+.agentbricks/
+app.agentbricksdev.yaml