agent-memory について
agent-memoryは、Databricks LangChainを使う開発者向けに、エージェントテンプレートへ短期・長期記憶を追加する実装ガイドです。依存関係、Lakebaseの設定、記憶ツールの統合、デプロイ時の確認事項を扱います。
主な機能
- 01セッション内の短期記憶を追加
- 02ユーザー情報をセッション間で保存
- 03記憶の読み書き・削除ツールの実装例
- 04Lakebaseの設定とデプロイ手順
- 05テストとトラブルシューティングの案内
活用シーン
- 01エージェントテンプレートに会話履歴を加える
- 02ユーザー情報をセッション間で保持する
- 03Lakebaseの記憶ストレージを設定・テストする
インストール
コマンドをコピーしてターミナルで実行します。
npx skills add https:/ / github.com/ databricks/ app-templates --skill agent-memoryインストール後、Claude Code で/agent-memoryを入力して読み込まれたか確認します。
同梱ファイルexamples/memory_tools.py
詳細
エージェントにメモリを追加する
注: このテンプレートには、デフォルトでメモリ機能は含まれていません。このスキルを使って、メモリ機能を追加してください。メモリが事前設定されたテンプレートについては、以下をご覧ください。
- agent-langgraph-advanced - 長時間実行されるバックグラウンドタスクに対応した短期メモリと長期メモリ
メモリの種類
| 種類 | 用途 | ストレージ | 識別子 |
|---|---|---|---|
| 短期 | セッション内の会話履歴 | AsyncCheckpointSaver | thread_id |
| 長期 | セッションをまたいで保持されるユーザー情報 | AsyncDatabricksStore | user_id |
前提条件
-
pyproject.tomlにメモリ依存関係を追加します。dependencies = [ "databricks-langchain[memory]", ]次に、
uv syncを実行します。 -
Lakebase を設定します。手順については lakebase-setup スキルをご覧ください。
- Lakebase インスタンスの作成と設定
- テーブルの初期化(初回に必ず実施する手順)
クイックセットアップの概要
メモリを追加するには、4 つのファイルを変更する必要があります。
| ファイル | 追加する内容 |
|---|---|
pyproject.toml | メモリ依存関係 |
.env | Lakebase 環境変数(ローカル開発用) |
databricks.yml | Lakebase データベースリソースと config ブロック内の環境変数 |
agent_server/agent.py | メモリツールと AsyncDatabricksStore |
主な原則
メモリを実装する前に、本番環境の実装で使われている次のパターンを理解してください。
1. ファクトリー関数パターン
メモリツールは、単独の関数として定義するのではなく、ファクトリー関数から返すようにします。
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. ユーザー ID の抽出
リクエストから user_id を抽出します。まず custom_inputs を確認してください。呼び出し元が判断できるように、デフォルト値ではなく None を返します。
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. エラー処理の分離
user_id と store を個別に確認し、それぞれ異なるエラーメッセージを返します。
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 検証
保存する前に JSON 入力を検証します。LLM が無効な JSON を渡す場合があります。
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. RunnableConfig を介した Store の受け渡し
Store は関数のパラメーターとしてではなく、config を介して渡します。
config = {"configurable": {"user_id": user_id, "store": store}}
# Tools access via: config.get("configurable", {}).get("store")
完全な例
完全な実装は、このスキルの examples フォルダーにあります。
# Copy to your project
cp .claude/skills/agent-memory/examples/memory_tools.py agent_server/
本番環境で使用できるコードとすべてのヘルパー関数については、examples/memory_tools.py をご覧ください。
本番環境のリファレンス
事前構築済みテンプレートの実装については、以下をご覧ください。
| ファイル | 説明 |
|---|---|
agent-langgraph-advanced/agent_server/utils_memory.py | メモリツールのファクトリー、ヘルパー、エラー処理 |
agent-langgraph-advanced/agent_server/agent.py | エージェントとの統合、ストアの初期化 |
主な関数:
memory_tools()- get/save/delete ツールを返すファクトリーget_user_id()- リクエストから user_id を抽出init_lakebase_config()- オートスケーリング Lakebase 環境変数を読み込み、LakebaseConfigに格納get_lakebase_access_error_message()- 役立つエラーメッセージ
設定ファイル
ステップ 1:databricks.yml(Lakebase リソース)
アプリに Lakebase データベースリソースを追加します。
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'
重要: name: 'postgres' は、databricks.yml の config.env ブロック内にある value_from の参照先と一致させる必要があります。
ステップ 2:databricks.yml の config ブロック(環境変数)
databricks.yml のアプリの config.env に、Lakebase 環境変数を追加します。
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"
重要: LAKEBASE_AUTOSCALING_ENDPOINT は、デプロイ時に postgres リソースから値を解決するため、value_from: "postgres" を使用します。
ステップ 3:.env(ローカル開発)
# Lakebase configuration for long-term memory
LAKEBASE_AUTOSCALING_ENDPOINT=<your-endpoint>
EMBEDDING_ENDPOINT=databricks-gte-large-en
EMBEDDING_DIMS=1024
統合例
ストリーミング関数にメモリを統合する方法を示す最小限の例です。
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
テーブルの初期化とデプロイ
Lakebase テーブルの初期化(初回のみ)
デプロイする前に、ローカルでテーブルを初期化します。
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 スキルを参照してください。
短期記憶
セッション内の会話履歴には、AsyncCheckpointSaver を使用します。
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
完全な実装については、agent-langgraph-advanced テンプレートを参照してください。
メモリのテスト
ローカルでテストする
# 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": "[email protected]"}
}'
# 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": "[email protected]"}
}'
# 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": "[email protected]"}
}'
デプロイ済みのアプリをテストする
# 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": "[email protected]"}
}'
初回セットアップのチェックリスト
-
databricks-langchain[memory]をpyproject.tomlに追加した -
uv syncを実行して依存関係をインストールした - Lakebase インスタンスを作成、または特定した
- Lakebase の環境変数を
.envに追加した(ローカル開発用) -
postgresリソースをdatabricks.ymlに追加した -
LAKEBASE_AUTOSCALING_ENDPOINTをdatabricks.ymlのconfig.envに追加した - ローカルで
await store.setup()を実行してテーブルを初期化した -
databricks bundle deploy && databricks bundle runでデプロイした
トラブルシューティング
| 問題 | 原因 | 解決策 |
|---|---|---|
| 「embedding_dims is required」 | パラメーターが不足しています | AsyncDatabricksStore に embedding_dims=1024 を追加します |
| 「relation 'store' does not exist」 | テーブルが作成されていません | まずローカルで await store.setup() を実行します |
| 「Unable to resolve Lakebase endpoint 'None'」 | 環境変数が不足しています | databricks.yml の config.env にある LAKEBASE_AUTOSCALING_ENDPOINT を確認します |
| 「permission denied for table store」 | 権限の付与が不足しています | postgres リソースを databricks.yml に追加します |
| 「Memory not available - no user_id」 | user_id が不足しています | リクエストで custom_inputs.user_id を渡します |
| メモリが保持されない | user_id が異なっています | リクエスト間で一貫した user_id を使用します |
| デプロイ後にアプリが更新されない | bundle の実行を忘れています | デプロイ後に databricks bundle run agent_langgraph を実行します |
事前構築済みメモリテンプレート
手動セットアップなしですぐに使える実装については、以下を参照してください。
| テンプレート | メモリの種類 | 主な機能 |
|---|---|---|
| agent-langgraph-advanced | 短期+長期 | AsyncCheckpointSaver、AsyncDatabricksStore、メモリツール |
次のステップ
- Lakebase を設定する:lakebase-setup スキルを参照してください
- ローカルでテストする:run-locally スキルを参照してください
- デプロイする:deploy スキルを参照してください