Databricks Free Edition + Lakebase で「記憶を持つ」LangGraphエージェントを動かす

Posted at 2026-06-18

2026年6月17日に、Databricks Free Editionへ5つの新機能 (Genie Code、サーバーレスGPU、Lakebase、Agent Bricks、Lakeflow Designer) が一気に追加されました。これでFree Editionでもデータエンジニアリングから分析、機械学習、アプリ開発、AIエージェントまでひと通り触れるようになっています。

本記事では、このうち Lakebase を題材に、LangGraphエージェントの会話メモリ (短期メモリ) をLakebase上のPostgresに永続化する例を、Free Edition上で動かすところまでをまとめます。実は以前Free Editionで同じことを試したときは、接続前の import psycopg の段階でカーネルごとクラッシュして到達できませんでした。今回はその原因と回避策込みで、最後まで動かせたので再現手順として残しておきます。

Free Editionに追加された5つの新機能

機能	概要
Genie Code	自然言語の指示からコードを生成・実行・反復してくれる。データ分析やパイプライン整備、可視化をエージェントが自律的に進める
サーバーレスGPU	提供状況に応じてGPUが利用可能。ニューラルネットの学習、ファインチューニング、推論、大規模データ処理などのディープラーニング系ワークロード向け
Lakebase	フルマネージドのPostgres互換データベース。データアプリやAIエージェントの状態ストアとして、レイクハウス上で直接トランザクション処理を扱える
Agent Bricks	プロダクション級のAIエージェントを構築するためのフレームワーク。ツール、メモリ、オーケストレーション、評価といったコンポーネントが組み込み済み
Lakeflow Designer	データパイプラインをビジュアルに構築。データフローの設計やステップの接続を白紙からではなく組み立てられる

今回はこの中の Lakebase を、AIエージェントの状態ストアとして使います。

今回作るもの

「名前を覚えるエージェント」というシンプルなシナリオで、メモリがLakebaseに永続化されることを確認します。やることは次の3つです。

LangGraphで、チャットモデル1ノードだけの最小エージェントを作る
LangGraphの PostgresSaver (チェックポインタ) の保存先をLakebaseにする
同じスレッドでは名前を覚え、別スレッドでは覚えていないことを確認し、実際にLakebaseの中身をSQLで覗く

LangGraphのチェックポインタは、エージェントの会話状態をスレッド単位 (thread_id) でデータベースに保存する仕組みです。これをLakebaseに向けることで、プロセスを再起動しても会話を再開できる「記憶を持つエージェント」になります。

事前準備その1: ノートブックの環境バージョンをv2にする

ここが今回いちばんの勘所です。先に書いておきます。

psycopgはネイティブライブラリ (libpq) を同梱しており、サーバーレスの新しい環境バージョンでは import psycopg の時点でカーネルがクラッシュすることがあります。実際にFree Editionで Fatal error: The Python kernel is unresponsive. が出たのはこれが原因でした。コードでもLakebaseでもなく、ライブラリのロード段階で落ちていたわけです。

対処はシンプルで、ノートブック右側の Environment (環境) パネルを開き、Environment versionをv2に設定 してから実行します。これで import psycopg が通るようになります。本記事は環境バージョンv2を前提としています。

事前準備その2: Lakebaseプロジェクトを作成する

画面右上のアプリスイッチャーから Lakebase Postgres を開きます。
Autoscaling を選び、New project をクリックします。プロジェクト名 (例: agent-memory-fe) を入力し、Postgresバージョンはデフォルトの17のままにします。
作成されると production ブランチと、デフォルトの databricks_postgres データベースが用意されます。コンピュートが起動するまで少し待ちます。

事前準備その3: Connectダイアログから接続情報を取得する

プロジェクトの Connect ダイアログを開き、ホスト名 と トークン をコピーします。トークンは60分有効です。プロジェクトの作成者には自動的にPostgresロールが払い出されるため、自分のユーザー (メールアドレス) でそのまま接続できます。

トークン発行はSDKで自動化することもできますが、ノートブックでの検証では Connect ダイアログのトークンをそのまま使うのが最短です。本番のアプリや常駐プロセスでは、60分ごとのトークンローテーションを実装することになります。

ライブラリのインストール

環境バージョンv2で動作確認できた構成です。databricks-langchain をインストールすると langgraph が巻き戻ることがあるため、最後に langgraph を明示的に再アップグレードします。

%pip install --quiet \
  "psycopg[binary]==3.3.3" \
  "langgraph-checkpoint-postgres==3.0.5" \
  "databricks-langchain==0.19.0"
%pip install --quiet --upgrade "langgraph>=1.1.10"

%restart_python

なお、インストール時にpipの依存解決の警告 (langchainが要求する langgraph のバージョンと、再アップグレードで入る langgraph のバージョンの衝突など) が表示されますが、本検証の範囲では問題なく動作しました。

再起動後、import psycopg がクラッシュせずに通ることを確認します。ここで落ちる場合は、Environmentパネルの環境バージョンがv2になっているかを再確認してください。

import psycopg
print("psycopg:", psycopg.__version__)

psycopg: 3.3.3

接続情報の設定

Connect ダイアログでコピーした値を設定します。HOST は自分のプロジェクトのホスト名に、TOKEN はその場で発行された60分有効のトークンに置き換えてください。

HOST = "ep-xxxxx.database.<region>.cloud.databricks.com"
DBNAME = "databricks_postgres"
USER = "your.email@example.com"
TOKEN = "<Connectダイアログのトークンを貼る>"  # 60分有効

接続のスモークテスト

以前Free Editionで詰まったのがこの接続確立の段階でした。SELECT version() が返れば、Free EditionのサーバーレスコンピュートからLakebaseへ到達できたことになります。connect_timeout を入れておくと、コンピュートがゼロからの起動中で待たされているだけなのか、本当に接続できないのかを区別できます。

import psycopg

conn = psycopg.connect(
    host=HOST,
    dbname=DBNAME,
    user=USER,
    password=TOKEN,
    sslmode="require",
    connect_timeout=10,
)
with conn.cursor() as cur:
    cur.execute("SELECT version()")
    print(cur.fetchone()[0])
conn.close()

PostgreSQL 17.10 (21f7c76) on x86_64-pc-linux-gnu, compiled by gcc (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0, 64-bit

PostgreSQL 17につながりました。これで接続基盤は確認できたので、エージェントのメモリを載せていきます。

チェックポインタ (短期メモリ) のセットアップ

LangGraphの PostgresSaver は、エージェントの会話状態をスレッド単位でチェックポイントとしてPostgresに保存します。PostgresSaver は autocommit=True と prepare_threshold=0 を要求するため、その指定で接続を開き、setup() で管理用テーブル (checkpoints など) を作成します。ここが「Lakebaseをエージェントの状態ストアにする」連携の中心部分です。

from psycopg import Connection
from psycopg.rows import dict_row
from langgraph.checkpoint.postgres import PostgresSaver

conn = Connection.connect(
    host=HOST,
    dbname=DBNAME,
    user=USER,
    password=TOKEN,
    sslmode="require",
    autocommit=True,
    prepare_threshold=0,
    row_factory=dict_row,
)
checkpointer = PostgresSaver(conn)
checkpointer.setup()
print("PostgresSaverのセットアップが完了しました")

PostgresSaverのセットアップが完了しました

エージェントの定義

ツールは使わず、チャットモデル1ノードのグラフにチェックポインタを差し込むだけです。同じ thread_id での会話履歴がLakebaseに永続化されるため、後続のターンでモデルが過去の発言を参照でき、結果として「名前を覚えている」ように振る舞います。

LLMにはFoundation Model APIsのエンドポイントを使います。Free Editionで利用できるエンドポイントは Serving タブで確認できます。databricks-meta-llama-3-3-70b-instruct が無い場合は、出ている別のモデル (例: databricks-gpt-oss-120b) に差し替えてください。

from databricks_langchain import ChatDatabricks
from langgraph.graph import StateGraph, START, MessagesState
from langchain_core.messages import SystemMessage, HumanMessage

LLM_ENDPOINT = "databricks-meta-llama-3-3-70b-instruct"
llm = ChatDatabricks(endpoint=LLM_ENDPOINT, temperature=0)

SYSTEM_PROMPT = "あなたは親切なアシスタントです。会話の中でユーザーが教えてくれた情報は覚えておき、日本語で簡潔に応答してください。"


def call_model(state: MessagesState):
    messages = [SystemMessage(SYSTEM_PROMPT)] + state["messages"]
    response = llm.invoke(messages)
    return {"messages": response}


builder = StateGraph(MessagesState)
builder.add_node("model", call_model)
builder.add_edge(START, "model")

graph = builder.compile(checkpointer=checkpointer)
print("エージェントをコンパイルしました")

エージェントをコンパイルしました

動作確認1: 同じスレッドで名前を覚えているか

thread_id を固定して2ターン会話します。1ターン目で名前を伝え、2ターン目で名前を尋ねます。チェックポインタ経由で1ターン目の履歴がLakebaseから読み出されるため、名前を答えられるはずです。

THREAD_1 = {"configurable": {"thread_id": "demo-thread-1"}}

# 1ターン目: 名前を伝える
res1 = graph.invoke(
    {"messages": [HumanMessage("私の名前はTakaです。覚えておいてください。")]},
    THREAD_1,
)
print("1ターン目:", res1["messages"][-1].content)

1ターン目: Takaさんでしたね。覚えておきます。どういたしまして。

# 2ターン目: 同じスレッドで名前を尋ねる
res2 = graph.invoke(
    {"messages": [HumanMessage("私の名前は何でしたっけ?")]},
    THREAD_1,
)
print("2ターン目:", res2["messages"][-1].content)

2ターン目: Takaさんです。

2ターン目のリクエストでは名前を渡していないのに、Takaさんです。 と答えています。1ターン目の履歴がLakebaseから読み戻されている証拠です。

動作確認2: 別スレッドでは覚えていないこと

メモリは thread_id でスコープされます。別のスレッドで同じ質問をすると、過去の履歴を参照できないため名前を答えられません。これにより、状態が確かにスレッド単位でLakebaseに分離保存されていることが確認できます。

THREAD_2 = {"configurable": {"thread_id": "demo-thread-2"}}

res3 = graph.invoke(
    {"messages": [HumanMessage("私の名前は何でしたっけ?")]},
    THREAD_2,
)
print("別スレッド:", res3["messages"][-1].content)

別スレッド: あなたの名前は私に教えていただいた記憶がありません。この会話が始まって間もないので、まだ名前を聞いていません。教えてくださる場合は、覚えておきます。

同じ質問でも、スレッドが違えば履歴が共有されないので名前を知りません。期待どおりの挙動です。

Lakebaseに実際に保存されているかを確認

チェックポイントが本当にPostgresに書き込まれているかを直接SQLで確認します。thread_id ごとに行が分かれていれば、状態ストアとして機能しています。

with conn.cursor() as cur:
    cur.execute(
        "SELECT thread_id, COUNT(*) AS checkpoints "
        "FROM checkpoints GROUP BY thread_id ORDER BY thread_id"
    )
    for row in cur.fetchall():
        print(f"thread_id={row['thread_id']}  checkpoints={row['checkpoints']}")

thread_id=demo-thread-1  checkpoints=6
thread_id=demo-thread-2  checkpoints=3

2ターン会話した demo-thread-1 の方がチェックポイント数が多く、スレッドごとに分かれて保存されているのが分かります。LakebaseのSQL Editorからも同じテーブルを覗けます。

保存済みの状態を読み戻す

get_state で、特定スレッドの最新状態に積み上がったメッセージ履歴を確認できます。プロセスを再起動しても、同じ thread_id を指定すればこの履歴から会話を再開できます。

state = graph.get_state(THREAD_1)
for m in state.values["messages"]:
    print(f"[{m.type}] {m.content}")

[human] 私の名前はTakaです。覚えておいてください。
[ai] Takaさんでしたね。覚えておきます。どういたしまして。
[human] 私の名前は何でしたっけ?
[ai] Takaさんです。

ハマりどころ: 環境バージョン

繰り返しになりますが、今回いちばんの落とし穴は環境バージョンでした。

import psycopg でカーネルが即死する場合、メモリ不足でもLakebase側の問題でもなく、psycopgのネイティブ層 (libpq) が新しいサーバーレス環境バージョンでロードに失敗しているのが原因のことがあります。import 単体で落ちるかどうかを確認すれば一発で切り分けられます。落ちる場合は、Environmentパネルで環境バージョンをv2に下げると通るようになります。

LangGraphの PostgresSaver はpsycopg専用なので、psycopgがロードできないとチェックポインタ自体が使えません。つまりこの環境バージョンの調整が、Lakebase連携全体の前提条件になります。

まとめと次のステップ

昨日Free Editionに追加されたLakebaseを、LangGraphエージェントの短期メモリ (会話状態) の永続化先として使えることを確認しました。
最大の勘所は ノートブックの環境バージョンをv2にする ことでした。新しい環境バージョンではpsycopgがネイティブ層のロードでカーネルごとクラッシュします。
接続は Connect ダイアログの60分トークンを使った単一接続で十分です。PostgresSaver の要件 (autocommit=True / prepare_threshold=0) を満たして setup() すれば、状態がスレッド単位でLakebaseに分離保存されることをSQLで確認できました。

次の発展として、PostgresStore とpgvectorを使ったセッション横断の長期メモリ (ユーザー単位で覚える) への拡張があります。短期メモリが通ったので、同じ接続基盤の上に積み上げられます。