Lakebaseを使ったAIエージェントを実装してみる

これは何？

「Cross Data Platforms Meetup #1.5 ～DatabricksとSnowflakeのSummitについて熱く語り合おう～」での発表資料を記事に起こしたものです！
イベントページはこちら
発表資料はこちら

書いてあること/ないこと

• 書いてあること
  • Databricksの新機能、LakebaseをAIエージェントに組み込む際の実装手順
  • 実運用に載せるにあたって未だわかっていない部分
• 書いてないこと
  • Lakebaseの詳細について（他の方がたくさん触れていらっしゃるはずです！気になる方は参考資料をあたってください><）

Lakebaseとは

Databricks Data Intelligence Platform に統合されたフルマネージドの Postgres OLTP データベースエンジン（公式ドキュメントより）

要するに、Databricks 上で Postgres がネイティブに使えるようになりました！
（2025-07-07 時点では Public Preview、しかも日本リージョンにはまだ来ていません ><）

DatabricksでPostgres！？何に使うんだ？

公式やコミュニティで挙げられている主なユースケースは次のとおりです。

用途	目的
オンライン特徴量ストア	推論時にリアルタイムで特徴量を取得する
データ入力処理（受注・チャット等）	低レイテンシでトランザクション処理を実行
アプリ状態の保存	Databricks Apps でのフロントエンド状態を永続化
Model 推論結果の即時返却	パーソナライズ推薦やチャット応答など
AI エージェントのメモリ	対話履歴・思考を pgvector と連携して保存／検索
リアルタイム分析への同期	OLTP データを Delta 経由で即時分析ワークロードへ

→ 今回はAI エージェントのメモリとしての用途に着目します！

AI エージェントにおける短期記憶と長期記憶

LangGraph（AIエージェント実装用のパッケージ） には、会話セッション専用の短期記憶と、ユーザ横断で共有する長期記憶があります。どちらも Postgres による実装がパッケージ化されています。

項目	短期記憶 (Short-term)	長期記憶 (Long-term)
用途	進行中の会話コンテキストを保持し、同一スレッド内の応答精度・一貫性を向上	ユーザ／アプリ全体で共有し、学習内容やパーソナライズ情報を永続化
格納情報	メッセージ履歴、アップロードファイル、生成物などスレッド限定データ	セマンティック（事実）、エピソード（経験）、プロシージャル（ルール）など
代表的ストア	- InMemory + Checkpointer（標準） - SQLiteSaver langgraph-checkpoint-sqlite - PostgresSaver langgraph-checkpoint-postgres - RedisSaver langgraph-checkpoint-redis	- InMemoryStore（開発用） - PostgresStore（pgvector 検索対応） - RedisStore（高速＋ベクトル検索）

短期記憶をLakebaseに保存すると何が嬉しいの？

セッション情報を永続化できるので、「以前の会話を途中から再開」することができます

やってみよう

今回は Short-term Memory を持つエージェントを実装してみます！

1. 検証用アカウントの作成

Databricks の無料 Trial アカウントを登録
※ Free Edition では Lakebase が使えないので注意 ><

2. Lakebase を作成

1. [Compute] → [OLTP データベース] → [インスタンスの作成]
   
2. 名前とインスタンスサイズを入力し作成
   
3. ステータスが 利用可能 になれば完了！
   

3. 簡単なエージェントを実装

1. Lakebase インスタンスの OAuth トークン を発行 & コピー
   
2. Databricks Secret に登録
   
   • 手順詳細は公式ドキュメント参照
     1. https://docs.databricks.com/aws/en/dev-tools/cli/tutorial
     2. https://docs.databricks.com/aws/ja/security/secrets/
     3. https://docs.databricks.com/aws/ja/security/secrets/example-secret-workflow
3. Lakebaseへの接続文字列 をコピー
   
4. 必要ライブラリをインストール

   pip install langchain langgraph langgraph-checkpoint-postgres psycopg2-binary databricks-langchain
   

5. Lakebase へ接続

   from langgraph.chat_models import init_chat_model
   from langgraph.graph import StateGraph, MessagesState, START
   from langgraph.checkpoint.postgres import PostgresSaver
   from databricks_langchain import ChatDatabricks
   from psycopg import Connection
   
   # Lakebase 接続
   PGPASSWORD = dbutils.secrets.get(scope="kuropura-demo", key="pgpassword")
   DB_URI = (
       "postgresql://tomohiro.kamei%40supership.jp:"
       f"{PGPASSWORD}@instance-808f51e8-2f01-4c50-af54-c4c5cee0c9cf.database.cloud.databricks.com"
       ":5432/databricks_postgres?sslmode=require"
   )
   conn = Connection.connect(DB_URI, autocommit=True, prepare_threshold=0)
   

6. エージェント構築

   # Agent
   LLM_ENDPOINT_NAME = "databricks-meta-llama-3-1-405b-instruct"
   llm = ChatDatabricks(endpoint=LLM_ENDPOINT_NAME)
   
   def call_model(state: MessagesState):
       response = llm.invoke(state["messages"])
       return {"messages": response}
   
   builder = StateGraph(MessagesState)
   builder.add_node(call_model)
   builder.add_edge(START, "call_model")
   
   # Lakebase をメモリに
   checkpointer = PostgresSaver(conn)
   checkpointer.setup()
   
   # コンパイル時にチェックポイントを登録
   graph = builder.compile(checkpointer=checkpointer)
   

4. 動作確認

1. チャットしてみる

   config = {"configurable": {"thread_id": "123"}}
   
   for chunk in graph.stream(
       {"messages": [{"role": "user", "content": "こんにちは、私の名前は亀井です。覚えておいてね。"}]},
       config,
       stream_mode="values"
   ):
       chunk["messages"][-1].pretty_print()
   

   

2. Lakebase に session 情報が保存されている（中身は人間には解読不可能）
   

3. 同じ thread_id で名前を尋ね直す

   for chunk in graph.stream(
       {"messages": [{"role": "user", "content": "私の名前はなんですか?"}]},
       config,
       stream_mode="values"
   ):
       chunk["messages"][-1].pretty_print()
   

   

   → ちゃんと名前を覚えておいてくれました！

わかっていない部分

短期記憶への接続と Model Registry／Serving Endpoint との兼ね合い

• LangGraph では接続情報を コンパイル時 に渡す必要がある
  → 接続が切れるたびに再コンパイル = モデル再デプロイに相当
• 現状、Lakebase の OAuth トークン有効期限は 1 時間
  → 1 時間ごとに Model Registry へ再登録、Serving Endpoint を再サーブ？

このあたりのスマートな運用方法はまだ模索中です……。
有識者の皆さま、ぜひご意見をお寄せください！ ><

まとめ

• Lakebase により Databricks でもネイティブな Postgres OLTP が利用可能になり、AI エージェントのメモリ や リアルタイム特徴量ストア など幅広い用途に活用できる
• LangGraph + Lakebase を組み合わせることで、対話履歴を永続化するエージェントをシンプルに構築できた
• 運用面では トークン期限と再デプロイ の課題があるので、アイデア募集中です！

参考になった記事など

• Lakebase周り
  • リリースブログ
  • What Is a Lakebase?
  • Lakebaseとは何ですか?
  • Databricks Lakebase
  • Lakebase徹底考察、彼はなぜDatabricksに生まれたのか、、、？
  • 【Databricks】OLTPデータベースであるLakebaseを触ってみた
• エージェント周り
  • LangChainとLangGraphによるRAG・AIエージェント［実践］入門
  • Understanding Short-Term Memory in LangGraph: A Hands-On Guide
  • Using PostgreSQL with LangGraph for State Management and Vector Storage
  • Checkpoints