mcpField
A
agent-memory

agent-memory

github.com/databricksエージェントワークフロー21737 回インストール1 週間前に更新
リポジトリ

agent-memory について

agent-memoryは、Databricks LangChainを使う開発者向けに、エージェントテンプレートへ短期・長期記憶を追加する実装ガイドです。依存関係、Lakebaseの設定、記憶ツールの統合、デプロイ時の確認事項を扱います。

主な機能

  1. 01セッション内の短期記憶を追加
  2. 02ユーザー情報をセッション間で保存
  3. 03記憶の読み書き・削除ツールの実装例
  4. 04Lakebaseの設定とデプロイ手順
  5. 05テストとトラブルシューティングの案内

活用シーン

  1. 01エージェントテンプレートに会話履歴を加える
  2. 02ユーザー情報をセッション間で保持する
  3. 03Lakebaseの記憶ストレージを設定・テストする

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/databricks/app-templates --skill agent-memory

インストール後、Claude Code で/agent-memoryを入力して読み込まれたか確認します。

同梱ファイルexamples/memory_tools.py

詳細

エージェントにメモリを追加する

注: このテンプレートには、デフォルトでメモリ機能は含まれていません。このスキルを使って、メモリ機能を追加してください。メモリが事前設定されたテンプレートについては、以下をご覧ください。

  • agent-langgraph-advanced - 長時間実行されるバックグラウンドタスクに対応した短期メモリと長期メモリ

メモリの種類

種類用途ストレージ識別子
短期セッション内の会話履歴AsyncCheckpointSaverthread_id
長期セッションをまたいで保持されるユーザー情報AsyncDatabricksStoreuser_id

前提条件

  1. pyproject.toml にメモリ依存関係を追加します。

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

    次に、uv sync を実行します。

  2. Lakebase を設定します。手順については lakebase-setup スキルをご覧ください。

    • Lakebase インスタンスの作成と設定
    • テーブルの初期化(初回に必ず実施する手順)

クイックセットアップの概要

メモリを追加するには、4 つのファイルを変更する必要があります。

ファイル追加する内容
pyproject.tomlメモリ依存関係
.envLakebase 環境変数(ローカル開発用)
databricks.ymlLakebase データベースリソースと 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 スキルを参照してください

コメント