databricks-mason 0.1.0.dev0__tar.gz

databricks_mason-0.1.0.dev0/.gitignore

@@ -0,0 +1,22 @@
+# test
+.pytest_cache/
+mlruns/
+
+# Byte-compiled files
+__pycache__
+*.pyc
+
+# Editor files
+*.vscode
+.idea
+
+# Mac
+.DS_Store
+
+# packaging
+dist/
+*.egg-info/
+Pipfile
+docs/build
+
+**/uv.lock

databricks_mason-0.1.0.dev0/PKG-INFO

@@ -0,0 +1,80 @@
+Metadata-Version: 2.5
+Name: databricks-mason
+Version: 0.1.0.dev0
+Summary: Databricks integration for Mason
+Author-email: Databricks <agent-feedback@databricks.com>
+License: Apache-2.0
+Requires-Python: >=3.10
+Requires-Dist: click>=8.1
+Requires-Dist: databricks-sdk>=0.49
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: rich>=13.7
+Provides-Extra: tracing
+Requires-Dist: mlflow[databricks]>=3.9.0; extra == 'tracing'
+Description-Content-Type: text/markdown
+
+# `databricks-mason`
+
+Mason is an experimental CLI for Databricks custom agent preview APIs and
+deployments. It manages memory, sessions, tracing, and deployments from one
+authenticated command.
+
+> The underlying APIs are in preview and may need workspace enablement.
+
+## Installation
+
+From PyPI:
+
+```sh
+pip install databricks-mason
+```
+
+From source:
+
+```sh
+pip install 'git+https://github.com/databricks/databricks-ai-bridge.git#subdirectory=integrations/mason'
+```
+
+For tracing commands, install Mason with tracing extras:
+
+```sh
+pip install 'databricks-mason[tracing]'
+```
+
+## Authentication
+
+Mason uses [Databricks authentication](https://docs.databricks.com/aws/en/dev-tools/cli/authentication).
+If you do not already have credentials, authenticate a named profile first. You can
+then ask Mason to validate and remember that profile:
+
+```sh
+databricks auth login --profile <profile>
+mason login --profile <profile>
+mason sessions stores list
+```
+
+`mason login` does not create credentials; it stores the selected profile in
+`~/.mason/config.json`. `mason logout` forgets that selection without revoking the
+underlying credentials. If Databricks SDK default authentication is already configured,
+you can skip `mason login`. You can also pass `--profile/-p` for an individual command.
+Use `--output json` for scripting.
+
+## Commands
+
+```text
+mason [-p <profile>] [-o text|json]
+  login        [--profile P]
+  logout
+  memory
+    stores     create | list | get | update | delete
+    entries    create | get | list | search | update | delete
+  sessions     create | list | get | update | delete | fork
+    stores     create | list | get | update | delete
+    items      list | append | pop | clear
+  tracing
+    setup      --catalog C --schema S [--experiment E]
+    list | get | instrument
+  deploy       <name> --source PATH [--with-memory-store N]
+               [--with-session-store N] [--with-traces C.S] [--create-stores]
+  deployments  list | get | logs | start | stop | delete
+```

databricks_mason-0.1.0.dev0/README.md

@@ -0,0 +1,65 @@
+# `databricks-mason`
+
+Mason is an experimental CLI for Databricks custom agent preview APIs and
+deployments. It manages memory, sessions, tracing, and deployments from one
+authenticated command.
+
+> The underlying APIs are in preview and may need workspace enablement.
+
+## Installation
+
+From PyPI:
+
+```sh
+pip install databricks-mason
+```
+
+From source:
+
+```sh
+pip install 'git+https://github.com/databricks/databricks-ai-bridge.git#subdirectory=integrations/mason'
+```
+
+For tracing commands, install Mason with tracing extras:
+
+```sh
+pip install 'databricks-mason[tracing]'
+```
+
+## Authentication
+
+Mason uses [Databricks authentication](https://docs.databricks.com/aws/en/dev-tools/cli/authentication).
+If you do not already have credentials, authenticate a named profile first. You can
+then ask Mason to validate and remember that profile:
+
+```sh
+databricks auth login --profile <profile>
+mason login --profile <profile>
+mason sessions stores list
+```
+
+`mason login` does not create credentials; it stores the selected profile in
+`~/.mason/config.json`. `mason logout` forgets that selection without revoking the
+underlying credentials. If Databricks SDK default authentication is already configured,
+you can skip `mason login`. You can also pass `--profile/-p` for an individual command.
+Use `--output json` for scripting.
+
+## Commands
+
+```text
+mason [-p <profile>] [-o text|json]
+  login        [--profile P]
+  logout
+  memory
+    stores     create | list | get | update | delete
+    entries    create | get | list | search | update | delete
+  sessions     create | list | get | update | delete | fork
+    stores     create | list | get | update | delete
+    items      list | append | pop | clear
+  tracing
+    setup      --catalog C --schema S [--experiment E]
+    list | get | instrument
+  deploy       <name> --source PATH [--with-memory-store N]
+               [--with-session-store N] [--with-traces C.S] [--create-stores]
+  deployments  list | get | logs | start | stop | delete
+```

databricks_mason-0.1.0.dev0/pyproject.toml

@@ -0,0 +1,59 @@
+[project]
+name = "databricks-mason"
+version = "0.1.0.dev0"
+description = "Databricks integration for Mason"
+authors = [
+    { name="Databricks", email="agent-feedback@databricks.com" },
+]
+readme = "README.md"
+license = { text="Apache-2.0" }
+requires-python = ">=3.10"
+dependencies = [
+    "click>=8.1",
+    "databricks-sdk>=0.49",
+    "PyYAML>=6.0",
+    "rich>=13.7",
+]
+
+[project.optional-dependencies]
+tracing = [
+    "mlflow[databricks]>=3.9.0",
+]
+
+[project.scripts]
+mason = "databricks_mason.cli:main"
+
+[dependency-groups]
+dev = [
+    "ruff==0.14.10",
+    "ty==0.0.23",
+    { include-group = "tests" },
+]
+tests = [
+    "pytest==9.0.2",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build]
+include = [
+    "src/databricks_mason/*",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/databricks_mason"]
+
+[tool.ruff]
+include = ["pyproject.toml", "src/**/*.py", "tests/**/*.py"]
+extend = "../../pyproject.toml"
+
+[tool.pytest.ini_options]
+testpaths = ["tests/unit_tests"]
+
+[tool.ty.environment]
+root = ["./src", "./tests"]
+
+[tool.ty.src]
+include = ["./src", "./tests"]

databricks_mason-0.1.0.dev0/src/databricks_mason/__init__.py

@@ -0,0 +1 @@
+"""Databricks integration for Mason."""

databricks_mason-0.1.0.dev0/src/databricks_mason/auth.py

@@ -0,0 +1,83 @@
+"""`mason login` / `logout` — remember an optional Databricks profile.
+
+`login` validates a named profile and persists the selection; the root group falls back
+to it whenever `-p` is omitted. Without a saved profile, the Databricks SDK performs its
+normal default authentication resolution. `logout` removes only Mason's saved selection,
+not the underlying credentials. State lives in a small JSON file under `~/.mason`
+(override the directory with `MASON_CONFIG_HOME`, mainly for tests).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import pathlib
+from typing import Optional
+
+import click
+
+from databricks_mason import render
+from databricks_mason.client import AgentApiClient
+from databricks_mason.errors import AgentCliError
+
+
+def _config_file() -> pathlib.Path:
+    base = os.environ.get("MASON_CONFIG_HOME")
+    root = pathlib.Path(base) if base else pathlib.Path.home() / ".mason"
+    return root / "config.json"
+
+
+def load_default_profile() -> Optional[str]:
+    """The profile saved by `mason login`, or None if the user never logged in."""
+    try:
+        return json.loads(_config_file().read_text()).get("profile")
+    except (OSError, json.JSONDecodeError):
+        return None
+
+
+def _save_default_profile(profile: str) -> None:
+    path = _config_file()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps({"profile": profile}, indent=2) + "\n")
+
+
+@click.command()
+@click.option(
+    "--profile",
+    "-p",
+    default=None,
+    help="Profile to authenticate with and remember as the default.",
+)
+@click.pass_obj
+def login(obj, profile) -> None:
+    """Validate a profile's credentials and save it as the default, so later commands can omit -p."""
+    profile = profile or obj.profile
+    if not profile:
+        raise AgentCliError(
+            "No profile to save.",
+            hint="Pass one to remember, e.g. `mason login --profile my-workspace`.",
+        )
+    client = AgentApiClient(profile)
+    user = client.current_user  # round-trips current_user.me(), so a bad profile fails here
+    _save_default_profile(profile)
+    if obj.output == "json":
+        render.emit_json({"profile": profile, "user": user, "host": client.host})
+        return
+    render.success(
+        f"Logged in as {user}",
+        fields={"Profile": profile, "Host": client.host},
+        next_steps=["mason sessions stores list", "mason memory stores list"],
+    )
+
+
+@click.command()
+@click.pass_obj
+def logout(obj) -> None:
+    """Forget the saved profile selection without deleting its credentials."""
+    path = _config_file()
+    existed = path.exists()
+    path.unlink(missing_ok=True)
+    if obj.output == "json":
+        render.emit_json({"logged_out": existed})
+        return
+    render.success("Logged out" if existed else "No saved login to clear")

databricks_mason-0.1.0.dev0/src/databricks_mason/cli.py

@@ -0,0 +1,72 @@
+"""`mason` — the Databricks CLI for agent deployment, memory, and sessions.
+
+Root Click group. Global `--profile` and `--output` flow to every subcommand via
+`CliContext` on `ctx.obj`; subcommands build an `AgentApiClient` from it on demand.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import click
+
+from databricks_mason.auth import load_default_profile, login, logout
+from databricks_mason.client import AgentApiClient
+from databricks_mason.deploy import deploy, deployments
+from databricks_mason.memory import memory
+from databricks_mason.sessions import sessions
+from databricks_mason.tracing import tracing
+
+
+class CliContext:
+    """Shared per-invocation state: selected profile, output mode, lazily-built client."""
+
+    def __init__(self, profile: Optional[str], output: str):
+        self.profile = profile
+        self.output = output
+        self._client: Optional[AgentApiClient] = None
+
+    def client(self) -> AgentApiClient:
+        if self._client is None:
+            self._client = AgentApiClient(self.profile)
+        return self._client
+
+
+@click.group(context_settings={"help_option_names": ["-h", "--help"]})
+@click.option(
+    "--profile", "-p", default=None, help="~/.databrickscfg profile to authenticate with."
+)
+@click.option(
+    "--output",
+    "-o",
+    type=click.Choice(["text", "json"]),
+    default="text",
+    help="Output format (default: text).",
+)
+@click.version_option(package_name="databricks-mason", prog_name="mason")
+@click.pass_context
+def mason(ctx: click.Context, profile: Optional[str], output: str) -> None:
+    """Mason: deploy agents and manage their memory and sessions.
+
+    Targets the agents/v1 preview APIs served on a workspace; auth comes from a
+    .databrickscfg profile (pass --profile / -p, run `mason login` to save a default,
+    or rely on the SDK's default resolution).
+    """
+    ctx.obj = CliContext(profile=profile or load_default_profile(), output=output)
+
+
+mason.add_command(login)
+mason.add_command(logout)
+mason.add_command(memory)
+mason.add_command(sessions)
+mason.add_command(tracing)
+mason.add_command(deploy)
+mason.add_command(deployments)
+
+
+def main() -> None:
+    mason(prog_name="mason")
+
+
+if __name__ == "__main__":
+    main()

databricks_mason-0.1.0.dev0/src/databricks_mason/client.py

@@ -0,0 +1,298 @@
+"""Authenticated REST client for the agents/v1 memory and session APIs.
+
+Wraps a databricks-sdk `WorkspaceClient` so auth/host come from a `.databrickscfg`
+profile. Each method maps to one API operation.
+Deployment is handled separately (deploy.py) since it wraps the `databricks apps` CLI.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from databricks.sdk import WorkspaceClient
+
+from databricks_mason.errors import AgentCliError, wrap_api_error
+
+_BASE = "/api/agents/v1"
+
+
+def _query(**kwargs: Any) -> dict[str, Any]:
+    """Build a query dict, dropping None and empty values."""
+    return {k: v for k, v in kwargs.items() if v is not None and v != ""}
+
+
+def memory_store_path(name: str) -> str:
+    """Normalize a store id or name into the `memory-stores/{id}` resource segment."""
+    name = name.strip().strip("/")
+    return name if name.startswith("memory-stores/") else f"memory-stores/{name}"
+
+
+def memory_entry_path(store: str, entry: str) -> str:
+    entry = entry.strip().strip("/")
+    if entry.startswith("memory-stores/"):
+        return entry
+    return f"{memory_store_path(store)}/entries/{entry}"
+
+
+class AgentApiClient:
+    """Thin, authenticated wrapper over the agents/v1 REST surface."""
+
+    def __init__(self, profile: Optional[str] = None):
+        try:
+            self._w = WorkspaceClient(profile=profile)
+        except Exception as exc:  # noqa: BLE001 - surfaced as a clean CLI error
+            raise AgentCliError(
+                f"Could not initialize Databricks auth: {exc}",
+                hint="Check your profile (`databricks auth login --profile <name>`) "
+                "or pass --profile.",
+            ) from exc
+
+    @property
+    def host(self) -> str:
+        return self._w.config.host or "unknown"
+
+    @property
+    def current_user(self) -> str:
+        """The authenticated user's name (used to derive the app source workspace path)."""
+        return str(self._w.current_user.me().user_name or "unknown")
+
+    def _do(
+        self, method: str, path: str, *, query: Optional[dict] = None, body: Optional[dict] = None
+    ) -> Any:
+        try:
+            return self._w.api_client.do(method, path, query=query, body=body)
+        except Exception as exc:  # noqa: BLE001 - normalized to AgentCliError
+            raise wrap_api_error(exc) from exc
+
+    # --- memory stores -------------------------------------------------------
+
+    def create_memory_store(self, display_name: str, description: Optional[str] = None) -> dict:
+        body = _query(display_name=display_name, description=description)
+        return self._do("POST", f"{_BASE}/memory-stores", body=body)
+
+    def get_memory_store(self, name: str) -> dict:
+        return self._do("GET", f"{_BASE}/{memory_store_path(name)}")
+
+    def list_memory_stores(
+        self, page_size: Optional[int] = None, page_token: Optional[str] = None
+    ) -> dict:
+        return self._do(
+            "GET",
+            f"{_BASE}/memory-stores",
+            query=_query(page_size=page_size, page_token=page_token),
+        )
+
+    def update_memory_store(
+        self, name: str, display_name: Optional[str] = None, description: Optional[str] = None
+    ) -> dict:
+        body = _query(display_name=display_name, description=description)
+        mask = ",".join(body.keys())
+        return self._do(
+            "PATCH", f"{_BASE}/{memory_store_path(name)}", query=_query(update_mask=mask), body=body
+        )
+
+    def delete_memory_store(self, name: str) -> dict:
+        return self._do("DELETE", f"{_BASE}/{memory_store_path(name)}")
+
+    # --- memory entries ------------------------------------------------------
+
+    def create_memory_entry(
+        self,
+        store: str,
+        actor_id: str,
+        path: str,
+        content: Optional[str] = None,
+        description: Optional[str] = None,
+        session_id: Optional[str] = None,
+        source_type: Optional[str] = None,
+    ) -> dict:
+        body = _query(
+            actor_id=actor_id,
+            path=path,
+            content=content,
+            description=description,
+            session_id=session_id,
+            source_type=source_type,
+        )
+        return self._do("POST", f"{_BASE}/{memory_store_path(store)}/entries", body=body)
+
+    def get_memory_entry(self, store: str, entry: str) -> dict:
+        return self._do("GET", f"{_BASE}/{memory_entry_path(store, entry)}")
+
+    def list_memory_entries(
+        self,
+        store: str,
+        actor_id: str,
+        path_prefix: Optional[str] = None,
+        session_id: Optional[str] = None,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> dict:
+        return self._do(
+            "GET",
+            f"{_BASE}/{memory_store_path(store)}/entries",
+            query=_query(
+                actor_id=actor_id,
+                path_prefix=path_prefix,
+                session_id=session_id,
+                page_size=page_size,
+                page_token=page_token,
+            ),
+        )
+
+    def search_memory_entries(
+        self, store: str, actor_id: str, query: str, limit: Optional[int] = None
+    ) -> dict:
+        body = _query(actor_id=actor_id, query=query, limit=limit)
+        return self._do("POST", f"{_BASE}/{memory_store_path(store)}/entries:search", body=body)
+
+    def update_memory_entry(
+        self,
+        store: str,
+        entry: str,
+        content: Optional[str] = None,
+        description: Optional[str] = None,
+    ) -> dict:
+        body = _query(content=content, description=description)
+        return self._do("PATCH", f"{_BASE}/{memory_entry_path(store, entry)}", body=body)
+
+    def delete_memory_entry(self, store: str, entry: str) -> dict:
+        return self._do("DELETE", f"{_BASE}/{memory_entry_path(store, entry)}")
+
+    # --- session stores ------------------------------------------------------
+
+    def create_session_store(
+        self, name: str, description: Optional[str] = None, metadata: Optional[dict] = None
+    ) -> dict:
+        body = _query(description=description, metadata=metadata)
+        return self._do(
+            "POST", f"{_BASE}/session-stores", query={"session_store_name": name}, body=body
+        )
+
+    def get_session_store(self, name: str) -> dict:
+        return self._do("GET", f"{_BASE}/session-stores/{name}")
+
+    def list_session_stores(
+        self, page_size: Optional[int] = None, page_token: Optional[str] = None
+    ) -> dict:
+        return self._do(
+            "GET",
+            f"{_BASE}/session-stores",
+            query=_query(page_size=page_size, page_token=page_token),
+        )
+
+    def update_session_store(
+        self, name: str, description: Optional[str] = None, metadata: Optional[dict] = None
+    ) -> dict:
+        body = _query(description=description, metadata=metadata)
+        mask = ",".join(body.keys())
+        return self._do(
+            "PATCH", f"{_BASE}/session-stores/{name}", query=_query(update_mask=mask), body=body
+        )
+
+    def delete_session_store(self, name: str) -> dict:
+        return self._do("DELETE", f"{_BASE}/session-stores/{name}")
+
+    # --- sessions ------------------------------------------------------------
+
+    def create_session(
+        self,
+        store: str,
+        actor_id: str,
+        session_id: Optional[str] = None,
+        parent_session_id: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> dict:
+        body = _query(actor_id=actor_id, parent_session_id=parent_session_id, metadata=metadata)
+        return self._do(
+            "POST",
+            f"{_BASE}/session-stores/{store}/sessions",
+            query=_query(session_id=session_id),
+            body=body,
+        )
+
+    def list_sessions(
+        self,
+        store: str,
+        filter: Optional[str] = None,
+        order_by: Optional[str] = None,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> dict:
+        return self._do(
+            "GET",
+            f"{_BASE}/session-stores/{store}/sessions",
+            query=_query(
+                filter=filter, order_by=order_by, page_size=page_size, page_token=page_token
+            ),
+        )
+
+    def get_session(self, session_id: str, store: Optional[str] = None) -> dict:
+        if store:
+            return self._do("GET", f"{_BASE}/session-stores/{store}/sessions/{session_id}")
+        return self._do("GET", f"{_BASE}/sessions/{session_id}")
+
+    def update_session(self, store: str, session_id: str, metadata: dict) -> dict:
+        return self._do(
+            "PATCH",
+            f"{_BASE}/session-stores/{store}/sessions/{session_id}",
+            query={"update_mask": "metadata"},
+            body=_query(metadata=metadata),
+        )
+
+    def delete_session(self, store: str, session_id: str, force: bool = False) -> dict:
+        return self._do(
+            "DELETE",
+            f"{_BASE}/session-stores/{store}/sessions/{session_id}",
+            query=_query(force=force or None),
+        )
+
+    def fork_session(
+        self,
+        store: str,
+        source_session_id: str,
+        actor_id: str,
+        up_to_item_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> dict:
+        body = _query(
+            source_session_id=source_session_id,
+            actor_id=actor_id,
+            up_to_item_id=up_to_item_id,
+            session_id=session_id,
+            metadata=metadata,
+        )
+        return self._do("POST", f"{_BASE}/session-stores/{store}/sessions:fork", body=body)
+
+    # --- session items -------------------------------------------------------
+
+    def list_session_items(
+        self,
+        store: str,
+        session_id: str,
+        order_by: Optional[str] = None,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> dict:
+        return self._do(
+            "GET",
+            f"{_BASE}/session-stores/{store}/sessions/{session_id}/items",
+            query=_query(order_by=order_by, page_size=page_size, page_token=page_token),
+        )
+
+    def append_session_items(self, store: str, session_id: str, items: list[dict]) -> dict:
+        body = {"items": [{"data": item} for item in items]}
+        return self._do(
+            "POST", f"{_BASE}/session-stores/{store}/sessions/{session_id}/items:append", body=body
+        )
+
+    def pop_session_item(self, store: str, session_id: str) -> dict:
+        return self._do(
+            "POST", f"{_BASE}/session-stores/{store}/sessions/{session_id}/items:pop", body={}
+        )
+
+    def clear_session_items(self, store: str, session_id: str) -> dict:
+        return self._do(
+            "POST", f"{_BASE}/session-stores/{store}/sessions/{session_id}/items:clear", body={}
+        )