エージェントの永続化

エージェントの永続化（Agent Persistency）は、KoogフレームワークにおけるAIエージェントにチェックポイント機能を提供する機能です。 これにより、実行中の特定ポイントでエージェントの状態を保存・復元できるようになり、以下の機能が利用可能になります。

• 特定ポイントからのエージェント実行の再開
• 以前の状態へのロールバック
• セッションをまたいだエージェント状態の永続化

主要な概念

チェックポイント

チェックポイントは、エージェント実行中の特定時点における完全な状態を捕捉します。これには以下が含まれます。

• メッセージ履歴（ユーザー、システム、アシスタント、ツール間のすべてのやり取り）
• 現在実行中のノード
• 現在のノードに対する入力データ
• 作成日時

チェックポイントは一意のIDで識別され、特定のエージェントに関連付けられます。

前提条件

エージェントの永続化機能を利用するには、エージェントのストラテジー内のすべてのノードが一意の名前を持つ必要があります。 これは、機能がインストールされる際に強制されます。

require(ctx.strategy.metadata.uniqueNames) {
    "Checkpoint feature requires unique node names in the strategy metadata"
}

グラフ内のノードに一意の名前を設定してください。

インストール

エージェントの永続化機能を使用するには、エージェントの設定に追加します。

val agent = AIAgent(
    executor = executor,
    llmModel = OllamaModels.Meta.LLAMA_3_2,
) {
    install(Persistency) {
        // スナップショットにインメモリストレージを使用
        storage = InMemoryPersistencyStorageProvider("in-memory-storage")
        // 自動永続化を有効にする
        enableAutomaticPersistency = true
    }
}

設定オプション

エージェントの永続化機能には、主に2つの設定オプションがあります。

• ストレージプロバイダー: チェックポイントの保存と取得に使用されるプロバイダー。
• 継続的な永続化: 各ノードの実行後にチェックポイントを自動的に作成する機能。

ストレージプロバイダー

チェックポイントの保存と取得に使用されるストレージプロバイダーを設定します。

install(Persistency) {
    storage = InMemoryPersistencyStorageProvider("in-memory-storage")
}

フレームワークには、以下の組み込みプロバイダーが含まれています。

• InMemoryPersistencyStorageProvider: チェックポイントをメモリに保存します（アプリケーションの再起動時に失われます）。
• FilePersistencyStorageProvider: チェックポイントをファイルシステムに永続化します。
• NoPersistencyStorageProvider: チェックポイントを保存しないno-op実装です。これはデフォルトのプロバイダーです。

PersistencyStorageProviderインターフェースを実装することで、カスタムストレージプロバイダーを実装することもできます。 詳細については、カスタムストレージプロバイダーを参照してください。

継続的な永続化

継続的な永続化とは、各ノードの実行後にチェックポイントが自動的に作成されることを意味します。 継続的な永続化を有効にするには、以下のコードを使用します。

install(Persistency) {
    enableAutomaticPersistency = true
}

有効化すると、エージェントは各ノードの実行後に自動的にチェックポイントを作成し、きめ細かいリカバリを可能にします。

基本的な使用方法

チェックポイントの作成

エージェントの実行中の特定ポイントでチェックポイントを作成する方法については、以下のコードサンプルを参照してください。

suspend fun example(context: AIAgentContextBase) {
    // 現在の状態でチェックポイントを作成
    val checkpoint = context.persistency().createCheckpoint(
        agentContext = context,
        nodeId = "current-node-id",
        lastInput = inputData,
        lastInputType = inputType,
        checkpointId = context.runId,
    )

    // チェックポイントIDは後で使用するために保存できます
    val checkpointId = checkpoint?.checkpointId
}

チェックポイントからの復元

特定チェックポイントからエージェントの状態を復元するには、以下のコードサンプルに従ってください。

suspend fun example(context: AIAgentContextBase, checkpointId: String) {
    // 特定のチェックポイントにロールバック
    context.persistency().rollbackToCheckpoint(checkpointId, context)

    // または最新のチェックポイントにロールバック
    context.persistency().rollbackToLatestCheckpoint(context)
}

拡張関数の使用

エージェントの永続化機能は、チェックポイントを操作するための便利な拡張関数を提供します。

suspend fun example(context: AIAgentContextBase) {
    // チェックポイント機能にアクセス
    val checkpointFeature = context.persistency()

    // またはチェックポイント機能でアクションを実行
    context.withPersistency(context) { ctx ->
        // 'this'はチェックポイント機能です
        createCheckpoint(
            agentContext = ctx,
            nodeId = "current-node-id",
            lastInput = inputData,
            lastInputType = inputType,
            checkpointId = ctx.runId,
        )
    }
}

高度な使用方法

カスタムストレージプロバイダー

PersistencyStorageProviderインターフェースを実装することで、カスタムストレージプロバイダーを実装できます。

class MyCustomStorageProvider : PersistencyStorageProvider {
    override suspend fun getCheckpoints(agentId: String): List<AgentCheckpointData> {
        // 実装
    }
    
    override suspend fun saveCheckpoint(agentCheckpointData: AgentCheckpointData) {
        // 実装
    }
    
    override suspend fun getLatestCheckpoint(agentId: String): AgentCheckpointData? {
        // 実装
    }
}

機能設定でカスタムプロバイダーを使用するには、エージェントでエージェント永続化機能を設定する際に、それをストレージとして設定します。

install(Persistency) {
    storage = MyCustomStorageProvider()
}

実行ポイントの設定

高度な制御のために、エージェントの実行ポイントを直接設定できます。

fun example(context: AIAgentContextBase) {
    context.persistency().setExecutionPoint(
        agentContext = context,
        nodeId = "target-node-id",
        messageHistory = customMessageHistory,
        input = customInput
    )
}

これにより、チェックポイントからの復元だけでなく、エージェントの状態をよりきめ細かく制御できるようになります。