Skill v1.0.1

currentAutomated scan100/100

databricks/app-templates/agent-memory

+1 new

──Details

PublishedAugust 23, 2026 at 11:47 AM

Content Hashsha256:129c93de48075244...

Git SHAfdc1b49d6504

Bump Typepatch

Compare with v1.0.0

──Files

Files (1 file, 11.8 KB)

SKILL.md11.8 KBactive

SKILL.md · 387 lines · 11.8 KB

version: "1.0.1" name: agent-memory description: "Add memory capabilities to your agent. Use when: (1) User asks about 'memory', 'state', 'remember', 'conversation history', (2) Want to persist conversations or user preferences, (3) Adding checkpointing or long-term storage."

Adding Memory to Your Agent

Note: This template does not include memory by default. Use this skill to add memory capabilities. For a pre-configured memory template, see:
- agent-langgraph-advanced - Short-term and long-term memory with long-running background tasks

Memory Types

Type	Use Case	Storage	Identifier
Short-term	Conversation history within a session	`AsyncCheckpointSaver`	`thread_id`
Long-term	User facts that persist across sessions	`AsyncDatabricksStore`	`user_id`

Prerequisites

Add memory dependency to pyproject.toml:

``toml dependencies = [ "databricks-langchain[memory]", ] ``

Then run uv sync

Configure Lakebase - See lakebase-setup skill for:

Creating/configuring Lakebase instance
Initializing tables (CRITICAL first-time step)

Quick Setup Summary

Adding memory requires changes to 4 files:

File	What to Add
`pyproject.toml`	Memory dependency
`.env`	Lakebase env vars (for local dev)
`databricks.yml`	Lakebase database resource + env vars in config block
`agent_server/agent.py`	Memory tools and AsyncDatabricksStore

Key Principles

Before implementing memory, understand these patterns from the production implementation.

1. Factory Function Pattern

Memory tools should be returned from a factory function, not defined as standalone functions:

python

def memory_tools():
    @tool
    async def get_user_memory(query: str, config: RunnableConfig) -> str:
        ...
    @tool
    async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str:
        ...
    @tool
    async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str:
        ...
    return [get_user_memory, save_user_memory, delete_user_memory]

2. User ID Extraction

Extract user_id from the request, checking custom_inputs first. Return None (not a default) to let the caller decide:

python

def get_user_id(request: ResponsesAgentRequest) -> Optional[str]:
    custom_inputs = dict(request.custom_inputs or {})
    if "user_id" in custom_inputs:
        return custom_inputs["user_id"]
    if request.context and getattr(request.context, "user_id", None):
        return request.context.user_id
    return None

3. Separate Error Handling

Check user_id and store separately with distinct error messages:

python

user_id = config.get("configurable", {}).get("user_id")
if not user_id:
    return "Memory not available - no user_id provided."
store: Optional[BaseStore] = config.get("configurable", {}).get("store")
if not store:
    return "Memory not available - store not configured."

4. JSON Validation for Save

Validate JSON input before storing - the LLM may pass invalid JSON:

python

try:
    memory_data = json.loads(memory_data_json)
    if not isinstance(memory_data, dict):
        return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}"
    await store.aput(namespace, memory_key, memory_data)
except json.JSONDecodeError as e:
    return f"Failed to save memory: Invalid JSON - {e}"

5. Pass Store via RunnableConfig

Pass the store through config, not as a function parameter:

python

config = {"configurable": {"user_id": user_id, "store": store}}
# Tools access via: config.get("configurable", {}).get("store")

Complete Example

A full implementation is available in this skill's examples folder:

bash

# Copy to your project
cp .claude/skills/agent-memory/examples/memory_tools.py agent_server/

See examples/memory_tools.py for production-ready code including all helper functions.

Production Reference

For implementations in the pre-built templates:

File	Description
`agent-langgraph-advanced/agent_server/utils_memory.py`	Memory tools factory, helpers, error handling
`agent-langgraph-advanced/agent_server/agent.py`	Integration with agent, store initialization

Key functions:

memory_tools() - Factory returning get/save/delete tools
get_user_id() - Extract user_id from request
init_lakebase_config() - Read autoscaling Lakebase env vars into a LakebaseConfig
get_lakebase_access_error_message() - Helpful error messages

Configuration Files

Step 1: databricks.yml (Lakebase Resource)

Add the Lakebase database resource to your app:

yaml

resources:
  apps:
    agent_langgraph:
      name: "your-app-name"
      source_code_path: ./
 
      resources:
        # ... other resources (experiment, UC functions, etc.) ...
 
        # Autoscaling Lakebase instance for long-term memory
        - name: 'postgres'
          postgres:
            branch: "projects/<project-name>/branches/<branch-name>"
            database: "projects/<project-name>/branches/<branch-name>/databases/<database-id>"
            permission: 'CAN_CONNECT_AND_CREATE'

Important: The name: 'postgres' must match the value_from reference in the databricks.yml config.env block.

Step 2: databricks.yml config block (Environment Variables)

Add the Lakebase environment variables to your app's config.env in databricks.yml:

yaml

      config:
        command: ["uv", "run", "start-app"]
        env:
          # ... other env vars ...
 
          # Autoscaling Lakebase endpoint (resolved from postgres resource)
          - name: LAKEBASE_AUTOSCALING_ENDPOINT
            value_from: "postgres"
 
          # Embedding configuration
          - name: EMBEDDING_ENDPOINT
            value: "databricks-gte-large-en"
          - name: EMBEDDING_DIMS
            value: "1024"

Important: LAKEBASE_AUTOSCALING_ENDPOINT uses value_from: "postgres" to resolve from the postgres resource at deploy time.

Step 3: .env (Local Development)

bash

# Lakebase configuration for long-term memory
LAKEBASE_AUTOSCALING_ENDPOINT=<your-endpoint>
EMBEDDING_ENDPOINT=databricks-gte-large-en
EMBEDDING_DIMS=1024

Integration Example

Minimal example showing how to integrate memory into your streaming function:

python

from agent_server.utils_memory import memory_tools, get_user_id
@stream()
async def streaming(request: ResponsesAgentRequest):
    user_id = get_user_id(request)
    async with AsyncDatabricksStore(
        autoscaling_endpoint=LAKEBASE_AUTOSCALING_ENDPOINT,
        embedding_endpoint=EMBEDDING_ENDPOINT,
        embedding_dims=EMBEDDING_DIMS,
    ) as store:
        await store.setup()  # Creates tables if needed
        tools = await mcp_client.get_tools() + memory_tools()
        config = {"configurable": {"user_id": user_id, "store": store}}
        agent = create_react_agent(model=model, tools=tools)
        async for event in agent.astream(messages, config):
            yield event

Initialize Tables and Deploy

Initialize Lakebase Tables (First Time Only)

Before deploying, initialize the tables locally:

bash

uv run python -c "$(cat <<'EOF'
import asyncio
from databricks_langchain import AsyncDatabricksStore
 
async def setup():
    async with AsyncDatabricksStore(
        autoscaling_endpoint="<your-endpoint>",
        embedding_endpoint="databricks-gte-large-en",
        embedding_dims=1024,
    ) as store:
        await store.setup()
        print("Tables created!")
 
asyncio.run(setup())
EOF
)"

Deploy

After initializing tables, deploy your agent. See deploy skill for full instructions.

Short-Term Memory

For conversation history within a session, use AsyncCheckpointSaver:

python

from databricks_langchain import AsyncCheckpointSaver
async with AsyncCheckpointSaver(autoscaling_endpoint=LAKEBASE_AUTOSCALING_ENDPOINT) as checkpointer:
    agent = create_react_agent(
        model=model,
        tools=tools,
        checkpointer=checkpointer,
    )
    config = {"configurable": {"thread_id": thread_id}}
    async for event in agent.astream(messages, config):
        yield event

See the agent-langgraph-advanced template for a complete implementation.

Testing Memory

Test Locally

bash

# Start the server
uv run start-app
 
# Save a memory
curl -X POST http://localhost:8000/invocations \
  -H "Content-Type: application/json" \
  -d '{
      "input": [{"role": "user", "content": "Remember that I am on the shipping team"}],
      "custom_inputs": {"user_id": "alice@example.com"}
  }'
 
# Recall the memory
curl -X POST http://localhost:8000/invocations \
  -H "Content-Type: application/json" \
  -d '{
      "input": [{"role": "user", "content": "What team am I on?"}],
      "custom_inputs": {"user_id": "alice@example.com"}
  }'
 
# Delete a memory
curl -X POST http://localhost:8000/invocations \
  -H "Content-Type: application/json" \
  -d '{
      "input": [{"role": "user", "content": "Forget what team I am on"}],
      "custom_inputs": {"user_id": "alice@example.com"}
  }'

Test Deployed App

bash

# Get OAuth token (PATs don't work for apps)
TOKEN=$(databricks auth token --host <workspace-url> | jq -r '.access_token')
 
# Test memory save
curl -X POST https://<app-url>/invocations \
  -H "Authorization: Bearer $TOKEN" \
  -H "Content-Type: application/json" \
  -d '{
      "input": [{"role": "user", "content": "Remember I prefer detailed explanations"}],
      "custom_inputs": {"user_id": "alice@example.com"}
  }'

First-Time Setup Checklist

[ ] Added databricks-langchain[memory] to pyproject.toml
[ ] Run uv sync to install dependencies
[ ] Created or identified Lakebase instance
[ ] Added Lakebase env vars to .env (for local dev)
[ ] Added postgres resource to databricks.yml
[ ] Added LAKEBASE_AUTOSCALING_ENDPOINT to databricks.yml config.env
[ ] Initialized tables locally by running await store.setup()
[ ] Deployed with databricks bundle deploy && databricks bundle run

Troubleshooting

Issue	Cause	Solution
"embedding_dims is required"	Missing parameter	Add `embedding_dims=1024` to AsyncDatabricksStore
"relation 'store' does not exist"	Tables not created	Run `await store.setup()` locally first
"Unable to resolve Lakebase endpoint 'None'"	Missing env var	Check `LAKEBASE_AUTOSCALING_ENDPOINT` in databricks.yml `config.env`
"permission denied for table store"	Missing grants	Add `postgres` resource to databricks.yml
"Memory not available - no user_id"	Missing user_id	Pass `custom_inputs.user_id` in request
Memory not persisting	Different user_ids	Use consistent user_id across requests
App not updated after deploy	Forgot to run bundle	Run `databricks bundle run agent_langgraph` after deploy

Pre-Built Memory Templates

For fully configured implementations without manual setup:

Template	Memory Type	Key Features
agent-langgraph-advanced	Short-term + Long-term	AsyncCheckpointSaver, AsyncDatabricksStore, memory tools

Next Steps

Configure Lakebase: see lakebase-setup skill
Test locally: see run-locally skill
Deploy: see deploy skill

← v1.0.0 All versions