databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentbricks/cli/memory.py

@@ -0,0 +1,520 @@
+"""`ab memory` — manage workspace-scoped managed memory stores and entries."""
+
+from __future__ import annotations
+
+import pathlib
+import sys
+from typing import Any
+
+import click
+
+from databricks_agentbricks import render
+from databricks_agentbricks.errors import AgentCliError
+from databricks_agentbricks.render import field
+from databricks_agentkit import timefmt
+
+_BREADCRUMB = "Agent Memory"
+
+# Friendly aliases for the memory-entry source type, mapped to the API enum values.
+_SOURCE_TYPES = {
+    "agent": "MANAGED_MEMORY_ENTRY_SOURCE_TYPE_AGENT",
+    "unspecified": "MANAGED_MEMORY_ENTRY_SOURCE_TYPE_UNSPECIFIED",
+}
+
+
+def _normalize_source_type(value):
+    """Accept a friendly alias ('agent'/'unspecified') or the full enum; None passes through."""
+    if value is None:
+        return None
+    key = value.strip().lower()
+    if key in _SOURCE_TYPES:
+        return _SOURCE_TYPES[key]
+    if value in _SOURCE_TYPES.values():
+        return value
+    raise AgentCliError(f"Invalid --source-type {value!r}. Choose one of: agent, unspecified.")
+
+
+def _require_entry_store(store, entry) -> None:
+    """A store is needed unless the entry is a full `memory-stores/.../entries/...` name."""
+    if not store and not str(entry).strip().startswith("memory-stores/"):
+        raise AgentCliError(
+            "Provide --store, or pass the full entry resource name "
+            "(memory-stores/<store>/entries/<id>)."
+        )
+
+
+def _store_id(store: dict) -> str:
+    name = field(store, "name") or ""
+    return name.split("/")[-1] if name else "—"
+
+
+def _truncate(value: Any, length: int = 60) -> str:
+    text = "" if value is None else str(value)
+    return text if len(text) <= length else text[: length - 1] + "…"
+
+
+# --- group ------------------------------------------------------------------
+
+
+@click.group()
+def memory() -> None:
+    """Manage an agent's long-term memory: memory stores and their entries.
+
+    Memory is what an agent remembers across separate conversations — durable facts and
+    preferences (for example "prefers concise answers", or a saved profile detail), as opposed to
+    the turn-by-turn history of a single conversation (that is `ab sessions`).
+
+    A memory store is the managed store that holds this memory; each entry is a small document (a
+    path plus its content) partitioned by actor, so one store keeps every user's memories separate.
+    """
+
+
+@memory.group()
+def stores() -> None:
+    """Workspace-scoped managed memory stores."""
+
+
+@memory.group()
+def entries() -> None:
+    """Memory entries within a store, partitioned by actor."""
+
+
+# --- bind a store to an agent project (agent.toml) --------------------------
+
+
+def _source_option(function):
+    return click.option(
+        "--source",
+        type=click.Path(exists=True, file_okay=False, path_type=pathlib.Path),
+        default=pathlib.Path("."),
+        show_default=True,
+        help="Agent Bricks project containing agent.toml.",
+    )(function)
+
+
+@memory.command("bind")
+@click.argument("store")
+@_source_option
+@click.pass_obj
+def memory_bind(obj, store: str, source: pathlib.Path) -> None:
+    """Bind memory STORE to the agent by declaring it in agent.toml.
+
+    This only edits agent.toml — it does not create the store. `ab deploy` creates any declared
+    store that doesn't exist yet and grants the deployed app's service principal access to it.
+    """
+    from databricks_agentbricks.agent_project import AgentProject
+
+    project = AgentProject.load(source)
+    project.bind_memory_store(store)
+    project.write()
+    if obj.output == "json":
+        render.emit_json({"memory_store": store, "manifest": str(project.path)})
+        return
+    render.success(
+        f"Bound memory store '{store}'",
+        fields={"agent.toml": str(project.path)},
+        next_steps=[
+            (
+                f"ab memory stores create --name {store}",
+                "Create the store now without deploying",
+            ),
+            ("ab dev", "Re-run to pick up the store locally"),
+            ("ab deploy <name>", "Create it if missing and grant the app access"),
+        ],
+    )
+
+
+@memory.command("unbind")
+@_source_option
+@click.pass_obj
+def memory_unbind(obj, source: pathlib.Path) -> None:
+    """Remove the memory store binding from the agent's agent.toml.
+
+    Only edits agent.toml; the managed store itself is untouched (delete it with
+    `ab memory stores delete`).
+    """
+    from databricks_agentbricks.agent_project import AgentProject
+
+    project = AgentProject.load(source)
+    if project.unbind_memory_store():
+        project.write()
+        render.success("Removed memory store binding", fields={"agent.toml": str(project.path)})
+    else:
+        click.echo(f"No memory store binding in {project.path}.")
+
+
+# --- stores -----------------------------------------------------------------
+
+
+def _store_starter_code(obj, store: dict) -> list[tuple[str, str, str]]:
+    store_id = _store_id(store)
+    name = field(store, "name") or f"memory-stores/{store_id}"
+    return [
+        (
+            "curl",
+            "bash",
+            f"""
+curl -X POST "{obj.client().host}/api/2.0/agents/{name}/entries" \\
+  -H "Authorization: Bearer $DATABRICKS_TOKEN" -H "Content-Type: application/json" \\
+  -d '{{"actor_id": "alice", "path": "/preferences/style.md", "content": "Terse, code first."}}'
+""",
+        ),
+        (
+            "agentbricks",
+            "bash",
+            f"""
+ab memory entries create --store {store_id} \\
+  --actor-id alice --path /preferences/style.md --content "Terse, code first."
+ab memory entries search --store {store_id} --actor-id alice --query "style"
+""",
+        ),
+    ]
+
+
+def _store_created(store: dict):
+    # The API returns RFC 3339 `create_time`; older responses used epoch-millis
+    # `created_at`. Read the current field, falling back to the legacy one.
+    return field(store, "create_time") or field(store, "created_at")
+
+
+def _store_updated(store: dict):
+    return field(store, "update_time") or field(store, "updated_at")
+
+
+def _render_store_detail(obj, store: dict) -> None:
+    render.detail(
+        _BREADCRUMB,
+        field(store, "display_name") or _store_id(store),
+        {
+            "Name": field(store, "display_name"),
+            "Resource name": field(store, "name"),
+            "Workspace": field(store, "workspace_id"),
+            "Creator": field(store, "owner_user_id"),
+            "Storage": render.field(field(store, "storage_backend") or {}, "backend_id"),
+            "Description": field(store, "description"),
+            "Created": timefmt.absolute(_store_created(store)),
+            "Updated": timefmt.absolute(_store_updated(store)),
+        },
+        status="ACTIVE",
+        snippets=_store_starter_code(obj, store),
+    )
+
+
+@stores.command("create")
+@click.option(
+    "--display-name",
+    "--name",
+    "display_name",
+    required=True,
+    help="Workspace-unique display name (--name is accepted as an alias).",
+)
+@click.option("--description", default=None, help="Optional human-readable description.")
+@click.pass_obj
+def stores_create(obj, display_name, description) -> None:
+    """Create a memory store."""
+    with render.status(f"Creating memory store '{display_name}'…"):
+        data = obj.client().create_memory_store(display_name, description)
+    if obj.output == "json":
+        render.emit_json(data)
+        return
+    store_id = _store_id(data)
+    render.success(
+        f"Created memory store '{display_name}'",
+        fields={"Store ID": store_id, "Name": field(data, "name")},
+        next_steps=[
+            (
+                f"ab memory entries create --store {store_id} --actor-id <id> --path </p>",
+                "Add a memory entry for an actor",
+            ),
+            (f"ab memory stores get {store_id}", "View this store's details"),
+            (
+                f"ab memory bind {display_name}",
+                "Bind this store to the agent (wired in on dev/deploy)",
+            ),
+        ],
+    )
+
+
+@stores.command("list")
+@click.option("--page-size", type=int, default=25, show_default=True)
+@click.option("--page-token", default=None)
+@click.pass_obj
+def stores_list(obj, page_size, page_token) -> None:
+    """List memory stores in the workspace (25 per page; paginates interactively on a terminal)."""
+    client = obj.client()
+    if obj.output == "json":
+        render.emit_json(client.list_memory_stores(page_size, page_token))
+        return
+    # On a terminal, offer to fetch the next page instead of only printing the --page-token hint.
+    interactive = sys.stdin.isatty() and sys.stdout.isatty()
+    while True:
+        data = client.list_memory_stores(page_size, page_token)
+        items = field(data, "managed_memory_stores") or []
+        rows = [
+            [
+                field(s, "display_name"),
+                _store_id(s),
+                timefmt.relative(_store_created(s)),
+                timefmt.relative(_store_updated(s)),
+                _truncate(field(s, "description"), 40),
+            ]
+            for s in items
+        ]
+        token = field(data, "next_page_token")
+        render.resource_table(
+            "Managed Memory Stores",
+            [
+                ("Name", "left"),
+                ("Resource name", "left"),
+                ("Created", "left"),
+                ("Updated", "left"),
+                ("Description", "left"),
+            ],
+            rows,
+            # In an interactive session the prompt below replaces the token hint.
+            subtitle=None if (interactive and token) else _page_note(data),
+            no_wrap=[1],  # keep the resource name full-width; other columns narrow to fit
+        )
+        if not token or not interactive:
+            break
+        if not click.confirm("Show next page?", default=False):
+            break
+        page_token = token
+
+
+@stores.command("get")
+@click.argument("name")
+@click.pass_obj
+def stores_get(obj, name) -> None:
+    """Get a memory store by id or resource name."""
+    data = obj.client().get_memory_store(name)
+    if obj.output == "json":
+        render.emit_json(data)
+        return
+    _render_store_detail(obj, data)
+
+
+@stores.command("update")
+@click.argument("name")
+@click.option("--display-name", default=None)
+@click.option("--description", default=None)
+@click.pass_obj
+def stores_update(obj, name, display_name, description) -> None:
+    """Update a store's display name and/or description."""
+    data = obj.client().update_memory_store(name, display_name, description)
+    if obj.output == "json":
+        render.emit_json(data)
+        return
+    _render_store_detail(obj, data)
+
+
+@stores.command("delete")
+@click.argument("name")
+@click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.")
+@click.pass_obj
+def stores_delete(obj, name, yes) -> None:
+    """Delete (soft-delete) a memory store."""
+    render.confirm_destroy(f"memory store '{name}'", assume_yes=yes)
+    obj.client().delete_memory_store(name)
+    if obj.output == "json":
+        render.emit_json({"deleted": name})
+        return
+    render.success(f"Deleted memory store '{name}'")
+
+
+# --- entries ----------------------------------------------------------------
+
+
+def _render_entry_detail(entry: dict) -> None:
+    render.detail(
+        f"{_BREADCRUMB} Entry",
+        field(entry, "path") or "—",
+        {
+            "Name": field(entry, "name"),
+            "Actor": field(entry, "actor_id"),
+            "Session": field(entry, "session_id"),
+            "Path": field(entry, "path"),
+            "Source": field(entry, "source_type"),
+            "Description": field(entry, "description"),
+            "Content": field(entry, "content"),
+            "Created": timefmt.absolute(field(entry, "create_time")),
+            "Updated": timefmt.absolute(field(entry, "update_time")),
+        },
+        status="ACTIVE",
+    )
+
+
+@entries.command("create")
+@click.option("--store", required=True, help="Store id or resource name.")
+@click.option("--actor-id", required=True, help="Actor (partition) this entry belongs to.")
+@click.option("--path", required=True, help="Absolute path, e.g. /preferences/style.md.")
+@click.option(
+    "--content", default=None, help="Entry content (inline). Use --content-file for large content."
+)
+@click.option(
+    "--content-file",
+    "content_file",
+    type=click.Path(exists=True, dir_okay=False),
+    default=None,
+    help="Read entry content from a file (avoids shell arg-length limits on large content).",
+)
+@click.option("--description", default=None, help="Optional human-readable description.")
+@click.option("--session-id", default=None, help="Optional session id to associate the entry with.")
+@click.option(
+    "--source-type",
+    default=None,
+    help="Origin of the entry: 'agent' or 'unspecified'.",
+)
+@click.pass_obj
+def entries_create(
+    obj, store, actor_id, path, content, content_file, description, session_id, source_type
+) -> None:
+    """Create a memory entry."""
+    if content is not None and content_file is not None:
+        raise AgentCliError("Pass either --content or --content-file, not both.")
+    if content_file is not None:
+        content = pathlib.Path(content_file).read_text()
+    data = obj.client().create_memory_entry(
+        store, actor_id, path, content, description, session_id, _normalize_source_type(source_type)
+    )
+    if obj.output == "json":
+        render.emit_json(data)
+        return
+    render.success(f"Created memory entry '{path}'", fields={"Name": field(data, "name")})
+
+
+@entries.command("get")
+@click.option(
+    "--store", default=None, help="Store id/name (optional if ENTRY is a full resource name)."
+)
+@click.argument("entry")
+@click.pass_obj
+def entries_get(obj, store, entry) -> None:
+    """Get an entry by id or resource name (includes content)."""
+    _require_entry_store(store, entry)
+    data = obj.client().get_memory_entry(store, entry)
+    if obj.output == "json":
+        render.emit_json(data)
+        return
+    _render_entry_detail(data)
+
+
+@entries.command("list")
+@click.option("--store", required=True)
+@click.option("--actor-id", required=True, help="Required partition key.")
+@click.option("--path-prefix", default=None)
+@click.option("--session-id", default=None)
+@click.option("--page-size", type=int, default=None)
+@click.option("--page-token", default=None)
+@click.pass_obj
+def entries_list(obj, store, actor_id, path_prefix, session_id, page_size, page_token) -> None:
+    """List entries for an actor. The text view omits content; `-o json` includes it."""
+    data = obj.client().list_memory_entries(
+        store, actor_id, path_prefix, session_id, page_size, page_token
+    )
+    if obj.output == "json":
+        render.emit_json(data)
+        return
+    items = field(data, "managed_memory_entries") or []
+    rows = [
+        [
+            field(e, "path"),
+            field(e, "actor_id"),
+            field(e, "session_id"),
+            _truncate(field(e, "description"), 40),
+            timefmt.relative(field(e, "update_time")),
+        ]
+        for e in items
+    ]
+    render.resource_table(
+        f"Memory Entries · actor {actor_id}",
+        [
+            ("Path", "left"),
+            ("Actor", "left"),
+            ("Session", "left"),
+            ("Description", "left"),
+            ("Updated", "left"),
+        ],
+        rows,
+        subtitle=_page_note(data),
+    )
+
+
+@entries.command("search")
+@click.option("--store", required=True)
+@click.option("--actor-id", required=True)
+@click.option("--query", required=True)
+@click.option("--page-size", type=int, default=None)
+@click.pass_obj
+def entries_search(obj, store, actor_id, query, page_size) -> None:
+    """Full-text search an actor's entries, ranked (includes content)."""
+    data = obj.client().search_memory_entries(
+        store,
+        actor_id,
+        query,
+        page_size=page_size,
+    )
+    if obj.output == "json":
+        render.emit_json(data)
+        return
+    results = field(data, "results")
+    items = (
+        [field(result, "managed_memory_entry") for result in results]
+        if results is not None
+        else field(data, "managed_memory_entries") or []
+    )
+    rows = [
+        [
+            field(e, "path"),
+            field(e, "actor_id"),
+            _truncate(field(e, "content"), 50),
+            timefmt.relative(field(e, "update_time")),
+        ]
+        for e in items
+    ]
+    render.resource_table(
+        f"Memory Search · '{query}'",
+        [("Path", "left"), ("Actor", "left"), ("Content", "left"), ("Updated", "left")],
+        rows,
+    )
+
+
+@entries.command("update")
+@click.option(
+    "--store", default=None, help="Store id/name (optional if ENTRY is a full resource name)."
+)
+@click.argument("entry")
+@click.option("--content", default=None, help="New entry content.")
+@click.option("--description", default=None, help="New description.")
+@click.pass_obj
+def entries_update(obj, store, entry, content, description) -> None:
+    """Update an entry's content and/or description."""
+    _require_entry_store(store, entry)
+    data = obj.client().update_memory_entry(store, entry, content, description)
+    if obj.output == "json":
+        render.emit_json(data)
+        return
+    _render_entry_detail(data)
+
+
+@entries.command("delete")
+@click.option(
+    "--store", default=None, help="Store id/name (optional if ENTRY is a full resource name)."
+)
+@click.argument("entry")
+@click.option("--yes", "-y", is_flag=True, help="Skip the confirmation prompt.")
+@click.pass_obj
+def entries_delete(obj, store, entry, yes) -> None:
+    """Delete a memory entry."""
+    _require_entry_store(store, entry)
+    render.confirm_destroy(f"memory entry '{entry}'", assume_yes=yes)
+    obj.client().delete_memory_entry(store, entry)
+    if obj.output == "json":
+        render.emit_json({"deleted": entry})
+        return
+    render.success(f"Deleted memory entry '{entry}'")
+
+
+def _page_note(data: dict) -> str | None:
+    token = field(data, "next_page_token")
+    return f"More results available — pass --page-token {token}" if token else None