databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentbricks/databricks_cli.py

@@ -0,0 +1,38 @@
+"""Thin wrapper around the ``databricks`` CLI.
+
+Agent Bricks shells out to the ``databricks`` CLI for the operations it doesn't call over the REST client
+(``apps`` lifecycle, ``sync``, ``postgres`` endpoint/credential lookups). This is the single place
+that builds the argv, threads the ``--profile`` through, and normalizes a non-zero exit into an
+``AgentCliError``. Kept separate from any one caller so ``deploy``, ``dev``, and ``app_resources``
+share it without importing each other.
+"""
+
+from __future__ import annotations
+
+import subprocess
+from typing import Optional
+
+from databricks_agentbricks.errors import AgentCliError
+
+
+def _databricks(
+    args: list[str],
+    profile: Optional[str],
+    *,
+    capture: bool = False,
+    check: bool = True,
+    cwd: Optional[str] = None,
+    action: Optional[str] = None,
+) -> subprocess.CompletedProcess:
+    cmd = ["databricks", *args]
+    if profile:
+        cmd += ["--profile", profile]
+    result = subprocess.run(cmd, text=True, capture_output=capture, cwd=cwd)
+    if check and result.returncode != 0:
+        # Agent Bricks drives the `databricks apps` CLI as an implementation detail; surface a failure in
+        # Agent Bricks' own terms (`action`) rather than echoing the raw subcommand and --profile, which
+        # leaks the underlying tool at the customer. The captured stderr still rides along as the
+        # hint so debugging isn't lost.
+        detail = (result.stderr or result.stdout or "").strip() if capture else None
+        raise AgentCliError(action or "A Databricks CLI command failed.", hint=detail)
+    return result

databricks_agentbricks/errors.py

@@ -0,0 +1,114 @@
+"""Error types for the Agent Bricks CLI and mapping from Databricks REST errors.
+
+`AgentCliError` extends `click.ClickException` so a raised error prints as a clean
+one-liner (plus an optional hint) and exits non-zero, instead of dumping a traceback.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Optional
+
+import click
+from rich.console import Console
+
+# Error codes indicating that a preview API is unavailable in the workspace.
+_PREVIEW_ERROR_CODES = frozenset({"NOT_IMPLEMENTED", "UNIMPLEMENTED", "FEATURE_DISABLED"})
+
+# gRPC-style status codes the backend returns for transient failures (a cancelled or
+# deadline-exceeded RPC, a briefly unavailable service). Shared with the client, which
+# retries these before surfacing them. `client._do` retries; the hint here covers the
+# case where retries are still exhausted.
+TRANSIENT_ERROR_CODES = frozenset({"CANCELLED", "UNAVAILABLE", "DEADLINE_EXCEEDED", "ABORTED"})
+
+# Process-global output mode, set once by the root CLI group. When "json", errors are
+# emitted as a machine-readable JSON object instead of the styled text one-liner, so a
+# script driving `ab -o json` can parse failures instead of scraping human text.
+_OUTPUT_MODE = "text"
+
+
+def set_output_mode(mode: str) -> None:
+    """Record the CLI's --output mode so errors can render to match it."""
+    global _OUTPUT_MODE
+    _OUTPUT_MODE = mode
+
+
+_PREVIEW_HINT = (
+    "These 2.0 agents APIs are in preview and gated per workspace. This handler is "
+    "not enabled on the target workspace yet — try a different --profile or contact "
+    "your workspace administrator."
+)
+
+_TRANSIENT_HINT = (
+    "This is usually a transient backend issue — re-running the command often succeeds."
+)
+
+
+class AgentCliError(click.ClickException):
+    """A user-facing CLI error rendered without a Python traceback."""
+
+    def __init__(
+        self, message: str, *, error_code: Optional[str] = None, hint: Optional[str] = None
+    ):
+        super().__init__(message)
+        self.error_code = error_code
+        self.hint = hint
+
+    def show(self, file=None) -> None:
+        if _OUTPUT_MODE == "json":
+            payload: dict = {"message": self.message}
+            if self.error_code:
+                payload["code"] = self.error_code
+            if self.hint:
+                payload["hint"] = self.hint
+            click.echo(json.dumps({"error": payload}, indent=2), err=True)
+            return
+        # Render in the cargo/uv diagnostic grammar: a red `error:` (or cargo's
+        # `error[CODE]:`) keyword, the message, and an indented `help:` line carrying the fix.
+        from databricks_agentbricks import render
+
+        render.diagnostic(
+            "error",
+            self.message,
+            code=self.error_code,
+            help=self.hint,
+            con=Console(stderr=True),
+        )
+
+
+def wrap_api_error(exc: Exception) -> AgentCliError:
+    """Convert a databricks-sdk error (or any exception) into an `AgentCliError`.
+
+    The SDK raises `databricks.sdk.errors.DatabricksError` subclasses carrying an
+    `error_code` attribute; we stay duck-typed so we don't couple to the SDK's error
+    hierarchy or version.
+    """
+    # When the SDK exhausts its retry budget it raises a TimeoutError chaining the last underlying
+    # error as __cause__; surface that original service error, not the generic retry wrapper.
+    # __cause__ is typed BaseException, so use a separate BaseException-typed local rather than
+    # reassigning the Exception-typed parameter.
+    effective: BaseException = exc
+    if isinstance(exc, TimeoutError) and exc.__cause__ is not None:
+        effective = exc.__cause__
+    error_code = getattr(effective, "error_code", None)
+    # A message-less DatabricksError wraps `IOError(None)`, so `str(effective)` is the literal
+    # "None". Treat that (and an empty string) as "no detail" so we surface the error code
+    # instead of a bare `Error [CANCELLED]: None`.
+    detail = str(effective).strip()
+    if detail == "None":
+        detail = ""
+    message = detail or _no_detail_message(error_code) or effective.__class__.__name__
+    if error_code in _PREVIEW_ERROR_CODES:
+        hint = _PREVIEW_HINT
+    elif error_code in TRANSIENT_ERROR_CODES:
+        hint = _TRANSIENT_HINT
+    else:
+        hint = None
+    return AgentCliError(message, error_code=error_code, hint=hint)
+
+
+def _no_detail_message(error_code: Optional[str]) -> Optional[str]:
+    """Fallback message when the server returns an error code but no human-readable detail."""
+    if not error_code:
+        return None
+    return f"The server returned {error_code} with no further detail."

databricks_agentbricks/lakebase_runtime_store.py

@@ -0,0 +1,107 @@
+"""Resolve service-managed Runtime Store backends from the API response."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from databricks_agentbricks.errors import AgentCliError
+from databricks_agentkit import models
+
+
+@dataclass(frozen=True)
+class RuntimeStoreBackend:
+    """Connection coordinates returned by the Runtime Store API."""
+
+    branch: str
+    database_id: str
+    username: str
+
+
+def get_or_create_backend(
+    client: Any, app: str, app_service_principal_id: str | None
+) -> RuntimeStoreBackend:
+    """Create or reuse the Runtime Store owned by a Databricks App."""
+    if not app_service_principal_id:
+        raise AgentCliError("Could not resolve the app's service principal for its Runtime Store.")
+    try:
+        store = client.create_runtime_store(
+            app,
+            app_service_principal_id,
+            app_name=app,
+            retry_transient=True,
+        )
+    except AgentCliError as exc:
+        if exc.error_code != "ALREADY_EXISTS":
+            raise
+        store = client.get_runtime_store(app)
+    return backend_from_api(app, app_service_principal_id, store)
+
+
+def delete(client: Any, app: str, app_service_principal_id: str) -> None:
+    """Delete the managed store while preserving the app when cleanup needs a retry."""
+    try:
+        store = client.get_runtime_store(app)
+        validate_owner(app, app_service_principal_id, store)
+        client.delete_runtime_store(app)
+    except AgentCliError as exc:
+        if exc.error_code == "NOT_FOUND":
+            return
+        raise AgentCliError(
+            f"Could not delete Runtime Store '{app}': {exc.message}",
+            error_code=exc.error_code,
+            hint=f"The deployment was retained. Retry `ab deployments delete {app}` "
+            "after resolving the Runtime Store error.",
+        ) from exc
+
+
+def _model(runtime_store: Any) -> models.RuntimeStore:
+    if not isinstance(runtime_store, dict):
+        raise AgentCliError("Runtime Store API returned an invalid resource.")
+    return models.RuntimeStore(runtime_store)
+
+
+def validate_owner(
+    app: str, app_service_principal_id: str, runtime_store: Any
+) -> models.RuntimeStore:
+    """Refuse to reuse or delete a Runtime Store owned by another app identity."""
+    store = _model(runtime_store)
+    if store.name != f"runtime-stores/{app}":
+        raise AgentCliError("Runtime Store API returned an unexpected resource name.")
+    owner = store.owner
+    app_owner = owner.app if owner is not None else None
+    if (
+        app_owner is None
+        or app_owner.name != app
+        or app_owner.service_principal_id != app_service_principal_id
+    ):
+        raise AgentCliError(f"Runtime Store '{app}' does not belong to this app identity.")
+    return store
+
+
+def backend_from_api(
+    app: str, app_service_principal_id: str, runtime_store: Any
+) -> RuntimeStoreBackend:
+    """Return only the connection coordinates supplied by the Runtime Store API."""
+    store = validate_owner(app, app_service_principal_id, runtime_store)
+    storage_backend = store.storage_backend
+    lakebase = storage_backend.lakebase if storage_backend is not None else None
+    branch = lakebase.branch if lakebase is not None else None
+    database_id = lakebase.database_id if lakebase is not None else None
+    owner = store.owner
+    app_owner = owner.app if owner is not None else None
+    username = app_owner.service_principal_id if app_owner is not None else None
+    if (
+        not isinstance(branch, str)
+        or not branch
+        or not isinstance(database_id, str)
+        or not database_id
+        or not isinstance(username, str)
+        or not username
+    ):
+        raise AgentCliError("Runtime Store API returned an incomplete Lakebase backend.")
+    return RuntimeStoreBackend(
+        branch=branch,
+        database_id=database_id,
+        username=username,
+    )

databricks_agentbricks/legacy_lakebase_runtime_store.py

@@ -0,0 +1,74 @@
+"""Legacy per-app Lakebase project provisioning for Agent Bricks Runtime Store."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from typing import Optional
+
+from databricks_agentbricks.app_resources import LakebaseBackend
+from databricks_agentbricks.databricks_cli import _databricks
+from databricks_agentbricks.errors import AgentCliError
+
+_BRANCH = "production"
+_DATABASE = "databricks-postgres"
+_ENDPOINT = "primary"
+_RESOURCE_NAME = "postgres-runtime-store"
+_SCHEMA_PREFIX = "databricks_agentkit_runtime_"
+
+
+def backend(app: str) -> LakebaseBackend:
+    """Return the dedicated per-app backend used before the managed API rollout."""
+    return LakebaseBackend(
+        project=_project_id(app),
+        branch=_BRANCH,
+        endpoint_id=_ENDPOINT,
+        database=_DATABASE,
+        schema=get_lakebase_schema(app),
+        tables=(),
+        resource_name=_RESOURCE_NAME,
+    )
+
+
+def get_or_create_backend(app: str, profile: Optional[str]) -> LakebaseBackend:
+    """Reuse or create the dedicated per-app project used by the legacy path."""
+    selected = backend(app)
+    project_path = f"projects/{selected.project}"
+    existing = _databricks(
+        ["postgres", "get-project", project_path], profile, capture=True, check=False
+    )
+    if existing.returncode == 0:
+        return selected
+
+    payload = {"spec": {"display_name": f"Agent Bricks Runtime Store for {app}"}}
+    created = _databricks(
+        ["postgres", "create-project", selected.project, "--json", json.dumps(payload)],
+        profile,
+        capture=True,
+        check=False,
+    )
+    if created.returncode == 0:
+        return selected
+
+    # A concurrent deploy can win the create race. Resolve the project before surfacing failure.
+    resolved = _databricks(
+        ["postgres", "get-project", project_path], profile, capture=True, check=False
+    )
+    if resolved.returncode == 0:
+        return selected
+    detail = (created.stderr or created.stdout or "").strip() or "unknown error"
+    raise AgentCliError(f"Could not create Runtime Store '{selected.project}'.", hint=detail)
+
+
+def _project_id(app: str) -> str:
+    normalized = re.sub(r"[^a-z0-9-]+", "-", app.lower()).strip("-") or "agentbricks-app"
+    if not normalized[0].isalpha():
+        normalized = f"agentbricks-{normalized}"
+    return f"{normalized}-runtime-store"[:63].rstrip("-")
+
+
+def get_lakebase_schema(app: str) -> str:
+    """Return the per-app Runtime Store schema name."""
+    digest = hashlib.sha256(app.encode("utf-8")).hexdigest()[:12]
+    return f"{_SCHEMA_PREFIX}{digest}"

databricks_agentbricks/project_config.py

@@ -0,0 +1,157 @@
+"""Persistent Agent Bricks project metadata and legacy framework detection."""
+
+from __future__ import annotations
+
+import json
+import pathlib
+import re
+from dataclasses import dataclass
+from typing import Any
+
+import tomli
+
+from databricks_agentbricks.errors import AgentCliError
+from databricks_agentbricks.project_types import AgentFramework, parse_framework
+
+_CONFIG_PATH = pathlib.Path(".agentbricks/project.toml")
+_SCHEMA_VERSION = 1
+_CUSTOM_SERVER_TEMPLATES = frozenset({"custom-agent-langgraph", "custom-agent-openai"})
+
+
+@dataclass(frozen=True)
+class ProjectMetadata:
+    """Template identity persisted by ``ab init``."""
+
+    framework: AgentFramework
+    template: str | None
+
+
+def write_project_metadata(
+    project: pathlib.Path,
+    *,
+    framework: str,
+    template: str,
+) -> pathlib.Path:
+    """Write the metadata consumed by template-aware Agent Bricks commands."""
+    selected_framework = parse_framework(framework)
+    target = project / _CONFIG_PATH
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(
+        f"schema_version = {_SCHEMA_VERSION}\n"
+        f"framework = {json.dumps(selected_framework.value)}\n"
+        f"template = {json.dumps(template)}\n",
+        encoding="utf-8",
+    )
+    return target
+
+
+def _read_toml(path: pathlib.Path, description: str) -> dict[str, Any]:
+    try:
+        with path.open("rb") as input_file:
+            value = tomli.load(input_file)
+    except (OSError, tomli.TOMLDecodeError) as exc:
+        raise AgentCliError(f"Could not read {description} at {path}: {exc}.") from exc
+    if not isinstance(value, dict):
+        raise AgentCliError(f"{description.capitalize()} at {path} must be a TOML table.")
+    return value
+
+
+def _load_persisted_metadata(project: pathlib.Path) -> ProjectMetadata | None:
+    path = project / _CONFIG_PATH
+    if not path.is_file():
+        return None
+    data = _read_toml(path, "Agent Bricks project config")
+    if data.get("schema_version") != _SCHEMA_VERSION:
+        raise AgentCliError(
+            f"Unsupported Agent Bricks project config schema in {path}.",
+            hint=f"Expected schema_version = {_SCHEMA_VERSION}.",
+        )
+    framework = parse_framework(data.get("framework"))
+    template = data.get("template")
+    if not isinstance(template, str) or not template:
+        raise AgentCliError(f"Agent Bricks project config at {path} must declare a template.")
+    return ProjectMetadata(framework=framework, template=template)
+
+
+def _dependency_name(requirement: str) -> str:
+    """Return a normalized distribution name from a PEP 508 requirement."""
+    name = re.split(r"[\s\[<>=!~;@]", requirement.strip(), maxsplit=1)[0]
+    return name.lower().replace("_", "-")
+
+
+def _infer_legacy_framework(project: pathlib.Path) -> ProjectMetadata:
+    pyproject = project / "pyproject.toml"
+    if not pyproject.is_file():
+        raise AgentCliError(
+            f"Could not determine the Agent Bricks framework for {project}.",
+            hint="Run `ab init` to create project metadata or pass `--framework`.",
+        )
+    data = _read_toml(pyproject, "pyproject")
+    project_table = data.get("project")
+    dependencies = project_table.get("dependencies", []) if isinstance(project_table, dict) else []
+    if not isinstance(dependencies, list) or not all(
+        isinstance(item, str) for item in dependencies
+    ):
+        dependencies = []
+    packages = {_dependency_name(item) for item in dependencies}
+    candidates = {
+        framework
+        for package, framework in (
+            ("databricks-openai", AgentFramework.OPENAI),
+            ("databricks-langchain", AgentFramework.LANGGRAPH),
+        )
+        if package in packages
+    }
+    if len(candidates) != 1:
+        detail = (
+            "both framework dependencies are present"
+            if candidates
+            else "no framework dependency was found"
+        )
+        raise AgentCliError(
+            f"Could not determine the Agent Bricks framework for {project}: {detail}.",
+            hint="Pass `--framework openai` or `--framework langgraph`.",
+        )
+    return ProjectMetadata(framework=candidates.pop(), template=None)
+
+
+def load_project_metadata(
+    project: pathlib.Path,
+    *,
+    framework_override: str | None = None,
+) -> ProjectMetadata:
+    """Load init metadata, with dependency inference for projects created before it existed."""
+    persisted = _load_persisted_metadata(project)
+    if framework_override is not None:
+        override = parse_framework(framework_override)
+        if persisted is not None and persisted.framework != override:
+            raise AgentCliError(
+                f"Framework override '{override.value}' conflicts with {persisted.framework.value!r} in "
+                f"{project / _CONFIG_PATH}."
+            )
+        return persisted or ProjectMetadata(framework=override, template=None)
+    return persisted or _infer_legacy_framework(project)
+
+
+def is_custom_server_template(template: str | None) -> bool:
+    """Whether ``template`` is one of the Agent Bricks custom HTTP server templates."""
+    return template in _CUSTOM_SERVER_TEMPLATES
+
+
+def uses_custom_server(project: pathlib.Path) -> bool:
+    """Whether persisted project metadata selects an Agent Bricks custom server template."""
+    metadata_path = project / _CONFIG_PATH
+    if not metadata_path.is_file():
+        return False
+    return is_custom_server_template(load_project_metadata(project).template)
+
+
+def require_managed_tool_support(project: pathlib.Path) -> None:
+    """Reject managed tool bindings for known templates that do not consume them."""
+    if not uses_custom_server(project):
+        return
+    raise AgentCliError(
+        "Managed tool bindings in agent.toml require an Agent Bricks server template.",
+        hint="Wire tools directly in agent/agent.py, or create a project with "
+        "`ab init --server agentbricks`.",
+    )

databricks_agentbricks/project_types.py

@@ -0,0 +1,55 @@
+"""Framework and server selections shared by the Agent Bricks CLI and project manifests."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+from databricks_agentbricks.errors import AgentCliError
+
+
+class AgentFramework(str, Enum):
+    """Frameworks supported by Agent Bricks project templates and adapters."""
+
+    LANGGRAPH = "langgraph"
+    OPENAI = "openai"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+class AgentServer(str, Enum):
+    """HTTP server ownership selected when an Agent Bricks project is created."""
+
+    AGENTBRICKS = "agentbricks"
+    CUSTOM = "custom"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+def parse_framework(value: object) -> AgentFramework:
+    """Validate a framework selection without exposing Enum errors to CLI users."""
+    if isinstance(value, str):
+        try:
+            return AgentFramework(value)
+        except ValueError:
+            pass
+    rendered = repr(str(value)) if isinstance(value, str) else "missing"
+    raise AgentCliError(
+        f"Unsupported Agent Bricks framework {rendered}.",
+        hint=f"Supported frameworks: {', '.join(framework.value for framework in AgentFramework)}.",
+    )
+
+
+def parse_server(value: object) -> AgentServer:
+    """Validate a server selection without exposing Enum errors to CLI users."""
+    if isinstance(value, str):
+        try:
+            return AgentServer(value)
+        except ValueError:
+            pass
+    rendered = repr(str(value)) if isinstance(value, str) else "missing"
+    raise AgentCliError(
+        f"Unsupported Agent Bricks server {rendered}.",
+        hint=f"Supported servers: {', '.join(server.value for server in AgentServer)}.",
+    )