managed-memory

# Long-Term Memory with Databricks Managed Memory (UC memory-store) Give your agent **durable, cross-session memory** about each user, exposed as five tools (`save_memory`, `get_memory`, `list_memories`, `update_memory`, `delete_memory`). The tools are thin REST calls to the Unity Catalog **memory-store** APIs. > **Beta.** The Databricks memory-store APIs are in beta — APIs and behavior may change. > ### This is Databricks *managed* memory — NOT the self-hosted Lakebase memory > A memory store is a governed **Unity Catalog securable** you read/write purely over REST: **no > database to provision, no tables to create, no embedding endpoint, and no extra Python dependency** > (it uses the `databricks-sdk` already in the template). This is **different from** the > `agent-openai-memory` / `agent-langgraph-memory` skills, which persist to a **Lakebase** instance you > run yourself. It's **additive to short-term/session memory** (the OpenAI `AsyncDatabricksSession` or the > LangGraph checkpointer) — keep that. But it **is** the agent's long-term memory, and there should be only > one: if the template already has a long-term memory system, remove it before adding these tools. This skill is framework-agnostic and flexible with both the OpenAI Agents SDK and LangGraph; each step notes the small per-SDK difference. **For a pre-existing agent (not built from a default template) — still on Databricks Apps.** The core (memory-store REST API, the five tools, the scope-as-isolation rule, the grant calls in Steps 1–2) is identical; only the template specifics differ. Map the `agent_server/...` paths to your own modules and reuse the Databricks Apps primitives you already have: the **forwarded OBO user token** for the signed-in user's id (what `resolve_scope()` reads), `config.env` for `DATABRICKS_MEMORY_STORE`, and `databricks apps` to deploy. Two invariants never change: the tools authenticate via `WorkspaceClient()` as the **app service principal you grant on the store**, and you pass the **end user's id** as `scope` — fail closed, never the SP. ## Prerequisites — this is an add-on This skill **adds long-term memory to an agent that's already set up** — it doesn't scaffold one. If there's **no `.env`** (auth not configured), run the **quickstart** skill first — it sets the Databricks profile + MLflow experiment, and on the advanced templates provisions the Lakebase used for short-term *session* memory (which this skill leaves intact). Then come back here. **Verify the app already has everything it needs to run** first — the **quickstart** skill tells you what each template needs set up. ## Concepts | Object | What it is | |---|---| | **Memory store** | A UC securable `catalog.schema.name` (type `MEMORY_STORE`) — the governance object you grant on and the container for memories. Read/written over REST, no SQL. | | **Memory entry** | One memory: a `path` (e.g. `/memories/preferences/coffee.md`), a one-line `description`, and optional `contents`. | | **Scope** | The partition key the caller assigns — decides whose memories you read/write. **Per-user** (a private partition, the default), a **shared constant** (org/team-wide), or **your own logic** (per project/tenant, user×project); see **Scope strategy** below. | **Access is two separate questions:** - *Can the caller use the store?* → make sure the caller has `READ_MEMORY_STORE` to retrieve memory entries and `WRITE_MEMORY_STORE` to write them. When testing locally the tools are called with the developer's credentials; when the agent is deployed on Apps they run with the app's credentials. - *Whose memories?* → the explicit **`scope`**, set by your code: the end user's id for private per-user memory, or a shared org/team constant for memory common to everyone (see **Scope strategy** below). The SP can see every scope, so **`scope` is your isolation boundary**: always set it in trusted code (to the end user, or a deliberate shared constant), and **never let the model choose it**. ## Step 1 — Create or choose the memory store Have your admin or agent developer create a memory store you can read/write memory entries to. First establish workspace creds: ```bash export DATABRICKS_HOST="https://<your-workspace-host>" export TOKEN="$(databricks auth token -p <profile> | jq -r .access_token)" ``` **Ask the user with `AskUserQuestion`** — two setup choices, in one call: *1. The store* — *"Do you have an existing memory store you can manage, or should I create one?"* - **Use an existing store** — you own it, or hold MANAGE / MANAGE_ACCESS_CONTROL on it. - **Create a new store** — under a catalog + schema you choose; you become the owner (needs `CREATE_MEMORY_STORE` on that schema). *2. The scope strategy* — *"How should memories be partitioned: private per end user, shared across a team/org/project, or by your own logic?"* (see **Scope strategy** below for the tradeoffs) - **Per-user (recommended)** — each user gets a private partition; the default wiring. - **Custom (shared)** — fixed scope multiple users can access; if chosen, collect the scope id as a free-text follow-up and set it as a constant in `resolve_scope`. - **Custom logic** — partition some other way (per project/tenant, or user×project). Ask the user to describe their isolation model, then write `resolve_scope` to it, honoring the contract under **Scope strategy → Your own logic**. The scope answer routes `resolve_scope` (Steps 3–4) and the `MEMORY_INSTRUCTIONS` framing (Step 5) — wire whichever the user picked. Then collect the store details as normal chat messages (free-text — not `AskUserQuestion`) and run the matching API call. Run these yourself so you can see exactly what each does: ```bash # CREATE a store you own. Returns the securable {full_name, owner, memory_store_id, ...}. curl -sS -X POST "$DATABRICKS_HOST/api/2.1/unity-catalog/memory-stores" \ -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \ -d '{"name":"<name>","catalog_name":"<catalog>","schema_name":"<schema>","description":"Long-term memory for my agent"}' # OR CONFIRM an existing store resolves (catches a typo'd name now instead of as NOT_FOUND at runtime): curl -sS "$DATABRICKS_HOST/api/2.1/unity-catalog/memory-stores/<catalog.schema.name>" -H "Authorization: Bearer $TOKEN" ``` **Record the full name** in the **same** env var, in two places — `.env` (read locally) **and** `databricks.yml` under the app's `config.env` (the deployed app doesn't read `.env`): ```yaml config: env: - name: DATABRICKS_MEMORY_STORE value: "<catalog.schema.name>" ``` ## Step 2 — Grant on the store (API calls) The tools call the API **as whatever principal the agent runs as**: the **app service principal** once deployed, and the **developer's own user** when running locally (the agent's `WorkspaceClient()` picks up the local profile). Grant **three** things to **both** principals: `READ_MEMORY_STORE` + `WRITE_MEMORY_STORE` on the store, **plus `USE_CATALOG` on its catalog and `USE_SCHEMA` on its schema**. The last two are easy to miss and non-obvious: without them *every* entry (and conversation) call fails with `User does not have USE CATALOG` — Unity Catalog hides a securable whose parent catalog/schema the caller can't traverse. DAB has **no `MEMORY_STORE` grant** yet, so these are direct permissions API calls (not `databricks.yml`) — run the `PATCH`es below yourself. `STORE` is the full name: ```bash export STORE="<catalog.schema.name>" PERM="$DATABRICKS_HOST/api/2.1/unity-catalog/permissions/memory_store/$STORE" # Grant the DEVELOPER's user (for local testing — the local agent runs as them): curl -sS -X PATCH "$PERM" -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \ -d '{"changes":[{"principal":"<developer@org.com>","add":["READ_MEMORY_STORE","WRITE_MEMORY_STORE"]}]}' # Grant the DEPLOYED app's service principal (run AFTER deploy, once the app + its SP exist): APP_SP=$(databricks apps get <your-app> -o json | jq -r .service_principal_client_id) curl -sS -X PATCH "$PERM" -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \ -d "{\"changes\":[{\"principal\":\"$APP_SP\",\"add\":[\"READ_MEMORY_STORE\",\"WRITE_MEMORY_STORE\"]}]}" # ALSO grant USE_CATALOG on the catalog and USE_SCHEMA on the schema to the SAME principals. # Required — without them the calls above fail "User does not have USE CATALOG". Replace <catalog> # and <catalog.schema>, and run once per principal (developer + $APP_SP): for PRIN in "<developer@org.com>" "$APP_SP"; do curl -sS -X PATCH "$DATABRICKS_HOST/api/2.1/unity-catalog/permissions/catalog/<catalog>" \ -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \ -d "{\"changes\":[{\"principal\":\"$PRIN\",\"add\":[\"USE_CATALOG\"]}]}" curl -sS -X PATCH "$DATABRICKS_HOST/api/2.1/unity-catalog/permissions/schema/<catalog.schema>" \ -H "Authorization: Bearer $TOKEN" -H "Content-Type: application/json" \ -d "{\"changes\":[{\"principal\":\"$PRIN\",\"add\":[\"USE_SCHEMA\"]}]}" done # Verify current store grants: curl -sS "$PERM" -H "Authorization: Bearer $TOKEN" ``` ## Step 3 — Add the memory tools Put these in `agent_server/utils_memory.py` — use **(a) the shared core + the block for your SDK** ((b) for the OpenAI Agents SDK *or* (c) for LangGraph; not both — they each define `_scope` their own way). **No new dependency** — it uses the `databricks-sdk` already in the template. **Most templates have no `utils_memory.py` — create it** (note the OpenAI advanced template keeps its *session* plumbing in `utils.py`, not here, so you still create a fresh `utils_memory.py`). **The one exception is `agent-langgraph-advanced`**: its existing `utils_memory.py` already holds the Lakebase plumbing — short-term checkpointer **and** a long-term `AsyncDatabricksStore` + `memory_tools()`. There, add these functions to that **same file** (don't create a second one), keep the checkpointer, and **replace** the long-term store — only one long-term system (see the intro and Step 4). **(a) Shared core** — the REST calls and scope resolution (SDK-agnostic): ```python import os from databricks.sdk import WorkspaceClient from databricks.sdk.errors import DatabricksError from mlflow.genai.agent_server import get_request_headers from agent_server.utils import get_user_workspace_client # API: BASE = /api/2.1/unity-catalog/memory-stores/{DATABRICKS_MEMORY_STORE} # create POST {BASE}/entries?scope=… {path,contents,description,creation_reason,creation_source} (flat body; scope is a query param) # get GET {BASE}/entries:get ?scope,path -> {contents, description, ...} # list GET {BASE}/entries ?scope -> {entries:[{path,description,has_contents}]} (key omitted entirely when empty) # update PATCH{BASE}/entries {scope, path, [description], [one contents edit op]} (>=1 of the two) # delete DELETE {BASE}/entries ?scope,path _client: WorkspaceClient | None = None def _ws() -> WorkspaceClient: """The memory caller — the app SP when deployed, the developer when local. Per-user isolation is via `scope`, NEVER this identity (the SP can see every scope).""" global _client if _client is None: _client = WorkspaceClient() return _client def _entries(suffix: str = "") -> str: store = os.getenv("DATABRICKS_MEMORY_STORE") if not store: raise RuntimeError("DATABRICKS_MEMORY_STORE is not set — it must be the full catalog.schema.name.") return f"/api/2.1/unity-catalog/memory-stores/{store}/entries{suffix}" def resolve_scope(request=None) -> str | None: """The end user's id used as `scope`, or None if it can't be determined (the handler MUST fail closed). Deployed: the OBO forwarded token -> current_user.me().id — the ONLY trusted source. Local: an X-Forwarded-User header, the request's custom_inputs.user_id (what the bundled chat UI / preflight send). NEVER the app's own identity, and NEVER a client-supplied value (X-Forwarded-User / custom_inputs) when deployed — those are spoofable.""" headers = get_request_headers() or {} if headers.get("x-forwarded-access-token"): obo = get_user_workspace_client() # None if the OBO client can't be built (see auth_type note) return obo.current_user.me().id if obo else None # fail closed rather than AttributeError # Deployed -> the verified OBO token above is the only trusted source. DATABRICKS_APP_NAME is set by # the Apps runtime when deployed, so the client-supplied fallbacks below are LOCAL-DEV ONLY: if os.getenv("DATABRICKS_APP_NAME"): return None ci = dict(getattr(request, "custom_inputs", None) or {}) return headers.get("x-forwarded-user") or ci.get("user_id") # The five operations. `scope` is passed in (never model-supplied). Each returns a short string. def _save(scope, path, description, contents=""): try: _ws().api_client.do("POST", _entries(), query={"scope": scope}, body={ "path": path, "contents": contents, "description": description, "creation_reason": "CREATION_REASON_AGENT_INFERRED", "creation_source": "CREATION_SOURCE_ONLINE_AGENT"}) except DatabricksError as e: if e.error_code == "ALREADY_EXISTS": return f"A memory already exists at {path}; use update_memory to revise it." return f"Could not save {path}: {getattr(e, 'message', str(e))}" return f"Saved memory at {path}." def _get(scope, path): try: entry = _ws().api_client.do("GET", _entries(":get"), query={"scope": scope, "path": path}) except DatabricksError as e: if e.error_code == "NOT_FOUND": return f"No memory at {path}." return f"Could not read {path}: {getattr(e, 'message', str(e))}" # A brief memory may have empty contents — its description is then the memory. return entry.get("contents") or entry.get("description") or f"(empty memory at {path})" def _list(scope): try: