Skip to main content
Last updated on

Tutorial: Add sessions and memory to an existing agent

Beta

This feature is in Beta.

If you already built an agent in Python, you can give it durable conversation history and long-term memory on Databricks without rewriting it. In this tutorial, you connect the agent to a session store that persists each conversation, a memory store that holds durable facts across conversations, and a memory pipeline that distills the stored conversations into memory. The steps cover agents built with the OpenAI Agents SDK, LangGraph, or any other Python framework, wherever the agent runs.

tip

If you don't have an agent yet, start with the Agent Bricks quickstart instead. The Agent Bricks CLI scaffolds a new agent with sessions and memory already set up.

In this example, a support agent should resume a customer's conversation where it left off, and remember the customer's preferences in later, separate conversations.

For how sessions and memory relate, see Agent memory and sessions.

Requirements​

  • An existing agent written in Python 3.10 or above. The framework-specific steps cover the OpenAI Agents SDK and LangGraph. The steps for other frameworks call the AgentKit SDK directly.
  • Databricks authentication. The AgentKit SDK uses the default authentication of the Databricks SDK for Python. For local development, run databricks auth login --host <workspace-url>, or set the DATABRICKS_HOST and DATABRICKS_TOKEN environment variables.
  • The Databricks CLI, to check the status of memory pipeline runs.
  • To create a memory pipeline, permission to query the model that the pipeline uses, and the other access listed in the memory pipeline requirements.

Step 1: Install the AgentKit SDK and the Agent Bricks CLI​

The databricks-agentbricks package contains the following:

  • The AgentKit SDK (databricks_agentkit), the Python client for session and memory stores.
  • The Agent Bricks CLI (agentbricks).
  • Adapters for the OpenAI Agents SDK and LangGraph that connect the session and memory stores to each framework's existing session and tool interfaces. The adapters use the framework packages that your agent already has installed.

Install version 0.5.0 or above into your agent's environment:

Bash
pip install --upgrade "databricks-agentbricks>=0.5.0"

Step 2: Create a session store and a memory store​

Your agent uses two stores, which Databricks manages for you:

  • A session store holds conversation history. Each session in it is the transcript of one conversation, which the agent reloads to continue that conversation where it left off.
  • A memory store holds durable facts about each user, such as preferences and past decisions. The agent searches it to recall what it learned in earlier, separate conversations.

The store is the security boundary: any principal with access to a store can read and write all of its entries. To isolate tenants from one another, create a separate pair of stores for each tenant. See Partition and secure memory.

Create each store one time and reuse it for every conversation. Your agent refers to both stores by name.

Python
from databricks.sdk import WorkspaceClient
from databricks_agentkit import AgentKitClient

client = AgentKitClient(WorkspaceClient())
client.session_stores.create("support-agent-sessions", description="Support agent conversations")
client.memory_stores.create("support-agent-memory", description="Support agent long-term memory")

Store names must be unique in the workspace. A memory store name must be 3 to 56 characters, start with a lowercase letter, end with a letter or number, and contain only lowercase letters, numbers, and hyphens. A session store name must be 3 to 63 characters.

Step 3: Choose actor and session IDs​

Every call that your agent makes to the stores carries two identifiers:

  • actor_id identifies whose conversations and memories they are. Both stores use it, and a memory pipeline writes each actor's memories under the same actor_id as that actor's sessions, so the agent recalls distilled memories only when the two match.
  • session_id identifies one conversation, so that a follow-up request resumes it.

For this tutorial, use test-user as the actor_id and conversation-1 as the session_id. The handle_message functions in the following steps take them as user_id and conversation_id.

In production, set actor_id in application code from trusted request context, such as the verified identity of the signed-in user, and never let the model or a user's message choose it. To share memory across a team or isolate tenants, see Partition and secure memory. For session_id, use your app's existing conversation or thread ID, or generate a UUID when a conversation starts.

Step 4: Persist conversations in the session store​

Conversation history is the ordered record of one conversation: the user's messages, the agent's replies, and any tool calls and results. A session store saves that history for each conversation, so the agent can reload it on the next request, even after a restart. Replace your agent's in-process conversation history with the session store that you created in Step 2.

The framework adapters plug into the interface that each framework already uses for session state, so your agent's model, tools, and prompts stay the same.

The OpenAI Agents SDK reads and writes conversation history through a Session object. session_store() returns a Session backed by the managed session store. Pass it to Runner.run:

Python
from agents import Agent, Runner
from databricks_agentkit.openai import session_store

SESSION_STORE = "support-agent-sessions"

# Your existing agent, unchanged.
agent = Agent(name="Support agent", instructions="You help customers with their Databricks workspaces.")

async def handle_message(user_id: str, conversation_id: str, message: str) -> str:
    session = session_store(conversation_id, actor=user_id, store=SESSION_STORE)
    result = await Runner.run(agent, message, session=session)
    return result.final_output

On the first message of a conversation, the adapter creates the session. On every run, the SDK loads the conversation's history from the session and appends the new user message, model output, and tool items.

Switching to the session store doesn't move conversations that your agent stored elsewhere. To carry an existing conversation over, append its items to a new session.

Step 5: Give the agent memory​

Session history covers one conversation. To use what the agent learned in earlier conversations, connect the agent to the memory store. There are two patterns, and you can combine them:

  • Memory tools: The model decides when to recall memory and when to save a new fact.
  • Recall before each run: Your application searches memory for the incoming message and adds the results to the agent's context, so the agent sees relevant memory without calling a tool.

If a memory pipeline writes your memory (see Step 6), the agent only needs to recall memory. You can skip tools that save memory.

The following example uses both patterns in a per-request copy of your agent:

  • memory_tools() returns a recall tool that searches the user's memory and a remember tool that saves a fact. The tools are bound to one actor, so the model can't read or write another user's memory.
  • recalled_facts() searches memory for the incoming message before the run and adds the results to the agent's instructions.
Python
from databricks.sdk import WorkspaceClient
from databricks_agentkit import AgentKitClient
from databricks_agentkit.openai import memory_tools, session_store

MEMORY_STORE = "support-agent-memory"
memory_store = AgentKitClient(WorkspaceClient()).memory_stores.get(MEMORY_STORE)

def recalled_facts(user_id: str, message: str) -> str:
    results = memory_store.search(actor_id=user_id, query=message, limit=5)
    return "\n".join(f"- {r.memory.content}" for r in results)

async def handle_message(user_id: str, conversation_id: str, message: str) -> str:
    instructions = agent.instructions  # Your agent's existing instructions, as a string.
    if facts := recalled_facts(user_id, message):  # Recall before the run.
        instructions += f"\n\nWhat you know about this user:\n{facts}"
    user_agent = agent.clone(
        instructions=instructions,
        tools=[*agent.tools, *memory_tools(user_id, store=MEMORY_STORE)],  # Memory tools.
    )
    session = session_store(conversation_id, actor=user_id, store=SESSION_STORE)
    result = await Runner.run(user_agent, message, session=session)
    return result.final_output

The recalled facts go into the instructions, not the input, so the session doesn't store them. To use only one pattern, drop the other: remove the memory_tools() line, or remove the recalled_facts() call. If you keep the tools, update your agent's instructions to tell it when to use them, for example: "Use recall to look up the user's preferences before you answer. Use remember to save durable facts that the user shares."

Step 6: Create a memory pipeline​

A memory pipeline reads the conversations in your session store and distills durable facts, preferences, and decisions for each actor into the memory store. It also merges overlapping entries over time, so your agent's memory stays consolidated without extra agent code.

  1. Create a pipeline that connects the two stores:

    Bash
    agentbricks memory pipeline create \
      --session-store support-agent-sessions \
      --memory-store support-agent-memory \
      --trigger scheduled \
      --model system.ai.claude-sonnet-5 \
      --instructions "Focus on durable customer preferences and account details. Ignore one-off troubleshooting steps."
    • --trigger scheduled runs the pipeline automatically about every 24 hours. Without it, the pipeline uses the default trigger, manual, and runs only when you start a run. See Run a pipeline on a schedule.
    • --model sets the model that distills memories. You can use any model that you have permission to query. If you omit --model, the pipeline uses a default model that Databricks chooses.
    • --instructions is optional and steers what the pipeline keeps. See Steer distillation with instructions.

    The command returns the pipeline's resource name, in the form memory-pipelines/<pipeline-id>. For every option, run agentbricks memory pipeline create --help or see the Agent Bricks CLI command reference.

  2. The scheduler chooses when the first scheduled run starts. To distill the conversations already in the session store, start a run now:

    Bash
    agentbricks memory pipeline run <pipeline-id>

    The command returns the run's resource name, in the form memory-pipelines/<pipeline-id>/runs/<run-id>.

  3. Check the run's result with the REST API. Set DATABRICKS_HOST to your workspace hostname, then run the following:

    Bash
    export DATABRICKS_TOKEN=$(databricks auth token --profile <profile> | jq -r .access_token)
    curl "https://${DATABRICKS_HOST}/api/2.0/agents/memory-pipelines/<pipeline-id>/runs/<run-id>" \
      -H "Authorization: Bearer ${DATABRICKS_TOKEN}"

    When the run's state is SUCCEEDED, the distilled memories are in the memory store.

Pipeline runs use the identity that created the pipeline, not your agent's identity. A pipeline can rewrite, merge, or delete any entry under an actor_id that it processes, including entries that your agent saved with memory tools. For the REST API, run states, and limitations, see Agent memory pipelines.

Record a transcript for LangGraph agents​

A memory pipeline distills only user and assistant messages. The LangGraph checkpointer stores serialized checkpoints, so a pipeline that reads the checkpointer's session store finds no messages to distill. To use a pipeline with a LangGraph agent, also append each turn's messages to a separate session store, and create the pipeline on that store:

Python
transcripts = client.session_stores.create("support-agent-transcripts")  # Create once.

def start_conversation(user_id: str, conversation_id: str) -> None:
    transcripts.add(actor_id=user_id, session_id=conversation_id)

def record_turn(conversation_id: str, message: str, reply: str) -> None:
    transcripts.get(conversation_id).append_items(
        [
            {"type": "message", "role": "user", "content": message},
            {"type": "message", "role": "assistant", "content": reply},
        ]
    )

Call start_conversation when a conversation begins, before its first ainvoke, and call record_turn after each ainvoke. Pass --session-store support-agent-transcripts when you create the pipeline. Alternatively, skip the pipeline and let the agent save memory with memory tools.

Step 7: Grant your deployed agent access​

When you run the agent locally, it uses your own credentials. A deployed agent usually runs as a service principal, which needs access to every store that it reads or writes. Grant access with the AgentKit SDK, passing the service principal's application (client) ID:

Python
from databricks.sdk import WorkspaceClient
from databricks_agentkit import AgentKitClient

w = WorkspaceClient()
client = AgentKitClient(w)

# For an agent deployed as a Databricks app, look up the app's service principal.
sp_client_id = w.apps.get("support-agent").service_principal_client_id

client.session_stores.get("support-agent-sessions").grant_permission(sp_client_id)
client.memory_stores.get("support-agent-memory").grant_permission(sp_client_id)

grant_permission grants write access by default. To grant read-only access, pass permission="READ". If you created a transcript store for a LangGraph agent, grant access to it too. The memory pipeline doesn't need a grant for your agent's service principal, because pipeline runs use the identity that created the pipeline.

Step 8: Test the agent end to end​

  1. Call handle_message with user_id="test-user" and conversation_id="conversation-1", and share a durable preference, such as "I prefer email over phone, and I'm in the Pacific time zone." If you followed the steps for other frameworks, call start_conversation("test-user", "conversation-1") first. Send a follow-up message in the same conversation, and confirm that the agent has the earlier context.

  2. Confirm that the session store holds the conversation:

    Bash
    agentbricks sessions items list --store support-agent-sessions --session-id conversation-1

  3. Start a memory pipeline run and wait until its state is SUCCEEDED, as in Step 6.

  4. Confirm that the memory store holds the preference for the test user:

    Bash
    agentbricks memory entries search --store support-agent-memory --actor-id test-user --query "contact preferences"

  5. Start a new conversation for test-user with conversation_id="conversation-2", and ask "How should you contact me?" The agent answers from memory.

Troubleshooting​

  • PERMISSION_DENIED (HTTP 403) from the deployed agent. The agent's service principal can't access a store. Grant it access, as in Step 7.
  • A pipeline run succeeds but writes no entries. Confirm that the session store holds user and assistant messages in a shape that the pipeline reads. A pipeline skips tool calls, tool results, and LangGraph checkpoints. A run processes only an actor's sessions since that actor's last successful run.
  • The agent doesn't recall distilled memories. Confirm that the agent searches memory with the same actor_id that its sessions use.
  • Starting a run fails with ABORTED. Another run of the pipeline is in progress. Wait for it to finish, then start a new run.

Clean up​

To delete the resources that this tutorial created, run the following. Deleting a store deletes everything in it.

Bash
agentbricks memory pipeline delete <pipeline-id>
agentbricks sessions stores delete support-agent-sessions
agentbricks memory stores delete support-agent-memory

Next steps​

On this page