LayerX AI Agent ブログリレー 8 日目の記事です。
バクラク事業 CTO @yyoshiki41 です。
この記事では、AI エージェントでのメモリを Vercel AI SDK が提供するジェネリックなミドルウェア機構である Language Model Middleware を使って実装する方法を紹介します。
Introduction
メモリレイヤは、AI エージェントの能力を拡張させる技術要素です。
Short-Term Memory, Long-Term Memory に分類され、以下のような整理が一般的です。いずれも環境情報やプロンプトとして扱われ、コンテキストとして機能します。
| Short-Term Memory | Long-Term Memory | |
|---|---|---|
| ストアされるデータ | 直近のメッセージや履歴 | セッションを越えて共有されるユーザー固有の情報や事実、過去の手順 |
| スコープ | 同一セッションとして、区切られたスコープ | セッションや時間限定もされない |
| 参照手段 | ・加工せずそのままコンテキストに含める ・コンテキストウィンドウを超える場合、古い情報を落とす場合もある |
・情報圧縮 (Compaction) した状態から、コンテキストに含める ・要約と外部ストレージの参照先をコンテキストに含め、LLM に利用判断させる |
| 例 | チャットでのスレッド | 既知のユーザー属性、セッション外の過去の履歴を参照して、解答する |
Short-Term Memory が AI エージェントの状態の一部として扱われるのに対し、Long-Term Memory はあくまで外部記憶のためコンテキストに含められます。
その加工方法も複数のパターンがあります。[1]
メモリレイヤのプレイヤー
AI エージェントの開発者向けにメモリレイヤを提供するサービスは、mem0, Memobase, Zep など複数の選択肢があります。
いずれもクラウドサービス版と OSS として公開され、ストレージをセルフホストして利用するオプションがあり、クイックに試せるものが多いです。
ユースケースに応じた選択が必要で、ユーザープロファイルへの特化や Knowledge Graph を構築し、AI エージェントに活用させるようなものもあります。
時間認識型の Knowledge Graph の詳細については、ブログリレー1日目の Temporal Knowledge Graphで作る！時間変化するナレッジを扱うAI Agentの世界 を覗いてみてください。
ai sdk でメモリレイヤを導入する
ai sdk は、LLM Model プロバイダーを抽象化し、OpenAI, Anthropic, Groq など複数のモデルも同じインターフェイスで扱えるようにしてくれます。
また、AI SDK v5 では Tool コールのライフサイクルイベントのフックが強化されたり、モデルの呼び出し部分をインタセプトする Language Model Middleware が提供されています。モデルを中心としつつ、そこへの依存は排除した設計を開発者に提供してくれています。
Language Model Middleware のドキュメント 内では、logging や cache、プロンプトのセーフティガードなどが紹介されています。
今回はこのミドルウェア機構を使って、メモリレイヤを汎用的に組み込むサンプルを紹介していきます。
Language Model Middleware の実装
実装インターフェイス
Language Model Middleware は、以下の3つの関数を実装することで提供できます。
-
transformParams: モデル呼び出しのパラメータを変換する関数 -
wrapGenerate: モデル呼び出しでアウトプットを生成する実行をラップする関数 -
wrapStream: モデルのストリーミング出力実行をラップする関数
ai@5.0.17 からは LanguageModelV2Middleware にエイリアスが貼られる形となり、LanguageModelMiddleware として型定義のインポートが可能になっています。
import type { LanguageModelMiddleware } from "ai";
ライフサイクル
モデルへのリクエスト部分で、コンテキスト補完するのがメモリの役割です。
ユーザー入力に応じたメモリ補完とメモリ更新のライフサイクル例が下記です。
今回、ユーザーインプット部分でメモリを検索し、システムメッセージに注入する部分を Language Model Middleware で実装します。
サンプルコード
1. メモリストアの実装
今回サンプルとして InMemory で動くシンプルなもので動かします。
ストレージレイヤに依存しないような形でインターフェイスを定義します。
export interface MemoryStore {
search(query: MemoryQuery): Promise<MemoryEntry[]>;
add(entries: MemoryEntry[]): Promise<void>;
}
export type MemoryEntry = {
id: string;
userId: string;
text: string;
vector?: number[];
createdAt: Date;
metadata?: Record<string, unknown>;
};
export type MemoryQuery = {
userId: string;
text: string;
topK?: number;
};
検索も InMemory で、cosine 類似度でリランクするような簡易的な実装です。
import { openai } from "@ai-sdk/openai";
import { embed, cosineSimilarity } from "ai";
import type { MemoryEntry, MemoryQuery, MemoryStore } from "./types";
const embeddingModel = openai.textEmbeddingModel("text-embedding-3-small");
export class InMemoryStore implements MemoryStore {
private items: MemoryEntry[] = [];
async add(entries: MemoryEntry[]): Promise<void> {
for (const entry of entries) {
const { embedding } = await embed({
model: embeddingModel,
value: entry.text,
});
const normalized: MemoryEntry = {
...entry,
vector: embedding,
};
this.items = this.items.filter((item) => item.id !== normalized.id);
this.items.push(normalized);
}
}
async search(query: MemoryQuery): Promise<MemoryEntry[]> {
const candidates = this.items.filter((entry) => {
if (entry.userId === query.userId) {
return true;
}
return false;
});
if (candidates.length === 0) {
return [];
}
const { embedding } = await embed({
model: embeddingModel,
value: query.text,
});
const topK = query.topK ?? 8;
return candidates
.map((entry) => ({
entry,
score: entry.vector
? cosineSimilarity(embedding, entry.vector)
: Number.NEGATIVE_INFINITY,
}))
.sort((a, b) => b.score - a.score)
.slice(0, topK)
.map((item) => item.entry);
}
}
2. MemoryMiddleware の実装
メインとなるミドルウェア部分です。今回、transformParams でユーザー入力をもとにメモリストアを検索し、システムメッセージとして注入しています。
プロンプトもナイーブに抽出されたメモリ内容を与えて、利用されるかも LLM に判断させます。
`You have access to persistent user memory.\n` +
`Use it when relevant and ignore when not applicable.\n\n` +
`### User Memory\n${formattedMemory}\n\n`,
const memoryMiddleware = ({
store,
resolveUser,
}: MemoryMiddlewareContext): LanguageModelMiddleware => ({
async transformParams({ params }) {
const user = resolveUser(params);
if (!user?.userId) {
return params;
}
const originalPrompt = params.prompt;
const userText = extractLatestUserText(originalPrompt);
if (userText === "") {
return params;
}
let memoryLines: string[] = [];
try {
const hits = await store.search({
userId: user.userId,
text: userText,
topK: 5,
});
memoryLines = hits.map((entry) => entry.text).filter(Boolean);
} catch {
// If the memory store is unavailable, continue without injecting memory.
return params;
}
if (memoryLines.length === 0) {
return params;
}
const prompt = injectMemory(originalPrompt, memoryLines);
return {
...params,
prompt,
};
},
});
function extractLatestUserText(prompt: LanguageModelV2Prompt): string {
for (let i = prompt.length - 1; i >= 0; i -= 1) {
const message = prompt[i];
if (!message || message.role !== "user") {
continue;
}
return message.content
.filter((part): part is LanguageModelV2TextPart => part.type === "text")
.map((part) => part.text)
.join("\n")
.trim();
}
return "";
}
function injectMemory(
prompt: LanguageModelV2Prompt,
memoryLines: string[]
): LanguageModelV2Prompt {
const formattedMemory = memoryLines.map((line) => `- ${line}`).join("\n");
const sysMessage: LanguageModelV2Message = {
role: "system",
content:
`You have access to persistent user memory.\n` +
`Use it when relevant and ignore when not applicable.\n\n` +
`### User Memory\n${formattedMemory}\n\n`,
};
return [sysMessage, ...prompt];
}
type MemoryMiddlewareContext = {
store: MemoryStore;
resolveUser: ResolveUserFn;
};
type ResolveUserResult = {
userId: string;
sessionId?: string;
};
type ResolveUserFn = (
params: LanguageModelV2CallOptions
) => ResolveUserResult | null | undefined;
ai sdk では wrapLanguageModel というモデル部分をラップし、メタデータやミドルウェアを注入できるユーティリティ関数が提供されており、これを利用します。
import type { LanguageModelV2, } from "@ai-sdk/provider";
import type { LanguageModelMiddleware } from "ai";
import { wrapLanguageModel } from "ai";
import type { MemoryStore } from "./memory/types";
export function wrapWithMemory({
model,
store,
resolveUser,
}: {
model: LanguageModelV2;
store: MemoryStore;
resolveUser: ResolveUserFn;
}) {
return wrapLanguageModel({
model,
middleware: memoryMiddleware({ store, resolveUser }),
});
}
3. モデル呼び出し部分で wrapLanguageModel を利用する
const userId = context.user.id;
const modelWithMemory = wrapWithMemory({
model,
store: memoryStore,
resolveUser: () => ({
userId,
}),
});
const result = streamText({
model: modelWithMemory,
...rest
});
デモ
今回のサンプル実装をそのまま動かしたものではありませんが、メモリについては今年8月に行った Bet AI Day でも取り上げています。
おわりに
Language Model Middleware によって、より低いレイヤでメモリにある情報をインジェクトする実装を紹介しました。実際にプロダクションに適用する場合には、メモリ抽出や更新方法、利用する情報の取捨選択など、ユースケースに応じて設計が必要です。
利用するメモリストアやサービスのAPI、モデルにも依存せず実装出来る ai sdk の抽象化は、メモリレイヤの導入ハードルを下げてくれます。
ai sdk は Web 技術との親和性高く、TypeScript で実装することができ、今後も AI エージェント開発の主力ツールとして利用していきます。
ここまでの速度で技術が発展し、全く新たなプロダクトを生み出せる機会というのは10年に一度しかないと思います。この激動の時代をまたとない機会と捉えて集中しつつも、全力で楽しんでLLMに向き合える方をお待ちしています！
募集
バクラク事業部では定期的な技術発信を行っています！
AI エージェント限らず、開発チームエネブルメント、プロダクト開発に興味のある方は、カジュアル面談もお待ちしております！
Discussion