mazyu36の日記

某SIer所属のクラウドエンジニアのブログ

Bedrock multi-agent collaboration を Generative AI Constructs だけで作りたかった

AWS CDK Amazon Bedrock

本記事は Amazon Bedrock Advent Calendar 2024 の 15日目の記事です。

re:Invent 2024期間にAmazon Bedrock の multi-agent collaboration 機能が発表されました（執筆時点ではプレビュー）。

aws.amazon.com

複数のAgentをオーケストレーションできるSupervisor Agentを作成し、Multi Agent を実現できるものです。

multi-agent collaboration については以下の Workshop が公開されています。

github.com

これの理解を深めるために、Generative AI Constructs を使って、CDKで構築してみようというのが今回のテーマです。

...が、タイトルにある通り執筆時点では全てを構築することはできませんでした。Supervisor Agentに必要な設定項目が CloudFormationでまだ設定不可能なためです。

そのため、本記事では個々の Agent は Generative AI Constructs で作成しつつ、Supervisor Agent のみ手動作成しています。

目次

アーキテクチャ

以下画像はWorkshopリポジトリからの引用です。

電力消費に関するMulti Agentです。Supervisor Agentの背後に以下の3つの Agent が居ます（リポジトリのREADMEを元にClaudeにまとめてもらったもの）

1. Forecasting Agent

  • 役割: エネルギー消費データの収集、予測、分析を行う
  • 主な機能:
    • 現在のエネルギー消費データの提供
    • 消費予測の提供
    • ユーザー消費統計の提供
    • コード解釈能力で予測データを分析

2. Solar Panel Agent

  • 役割: 太陽光発電システムの設置と運用を支援
  • 主な機能:
    • 設置ガイドラインと要件の提供
    • メンテナンスの指示とスケジュールの提供
    • サポートチケットの作成
    • 既存サポートチケットの追跡

3. Peak Load Manager Agent

  • 役割: 電力のピーク時負荷を最適化し、効率的な配電をサポート
  • 主な機能:
    • 不要なプロセスの特定
    • ピーク時と非ピーク時の使用状況の分析
    • グリッド割り当ての最適化

CDK実装

実装は以下にあります。

github.com

なお Generative AI Constructs を使った Agent の実装ですが、別記事で紹介しています。概ねこれ通りのため、以下記事から異なる部分を中心に紹介します。

mazyu36.hatenablog.com

1. Forecasting Agent

1つ目の Agent です。ざっくり特徴としては以下です。

  • Knowledge Baseがある。ベクターストアは OpenSearch Severless。
  • Agent から呼び出される Lambda は DynamoDB からデータを取得。
  • DynamoDB のデータは SageMaker が計算を行ったもの。ただし本 Workshop では計算結果が JSON で提供されているので、それを 何かしらの手段で DynamoDB に投入する。

※画像はWorkshopノートブックからの引用

CDKの実装は lib/constructs/1-energy-forecast.ts にあります。

前回の記事との実装の違いは主に以下3点のため、こちらを中心に紹介します。

  • Agent の Alias を設定
  • bedrock.S3DataSource で破壊的な変更あり
  • function schema を使用した Action Group の定義

Agent の Alias を設定

後のSupervisor Agent作成時に、各Agent の Alias が必要なため定義しておきます。

// Agentを定義
const agent = new bedrock.Agent(this, 'EnergyForecastAgent', {
  foundationModel: bedrock.BedrockFoundationModel.AMAZON_NOVA_PRO_V1,
  instruction,
  description,
  idleSessionTTL: Duration.seconds(1800),
  name: forecastAgentName,
});

// Aliasを定義。Alias 名は未指定の場合 latest になる。
new bedrock.AgentAlias(this, 'EnergyForecastAgentAlias', {
  agentId: agent.agentId,
});

bedrock.S3DataSource で破壊的な変更あり

chunkingStrategy を bedrock.ChunkingStrategy でまとめて定義する形になっているので、過去バージョンから流用する場合は修正が必要です。

new bedrock.S3DataSource(this, 'DataSourceForEnergyForecastAgent', {
  bucket: dataSourceBucket,
  knowledgeBase,
  // Chunking Strategyを定義
  chunkingStrategy: bedrock.ChunkingStrategy.fixedSize({
    maxTokens: 512,
    overlapPercentage: 20,
  }),
})

function schema を使用した Action Group の定義。

過去記事では OpenAPI を使用して、Agentが呼び出せるAPIを定義していましたが、今回は function schema を使用します。function schema 自体の説明は以下の記事が参考になります。

qiita.com

以下が実装例です。今後3ヶ月間のエネルギー消費量の予測を get_forecasted_consumption で取得できることを定義しています。

const actionGroup = new bedrock.AgentActionGroup(this, 'EnergyForecastActionGroup', {
  actionGroupExecutor: {
    lambda: actionGroupFunction,
  },
  actionGroupState: 'ENABLED',
  actionGroupName: 'forecast_consumption_actions',
  description: 'Function to get usage forecast for a user',
  // function schema の定義
  functionSchema: {
    functions: [
      {
        name: 'get_forecasted_consumption',
        description: 'Gets the next 3 months energy usage forecast',
        parameters: {
          customer_id: {
            description: 'Unique customer identifier',
            required: true,
            type: 'string'
          }
        }
      },
// 中略
    ]
  },
});

2. Solar Panel Agent

2つ目の Agent です。アーキテクチャ的には 1. Forecasting Agent とほぼ同じです。

※画像はWorkshopノートブックからの引用

実装は lib/constructs/2-solar-panel.ts にあります。

実装も 1. Forecasting Agentとほぼ同じため、実装説明は省略します。

3. Peak Load Manager Agent

3つ目のAgentです。こちらはさらにシンプルとなり、Knowledge Base が存在しません。

※画像はWorkshopノートブックからの引用

実装は lib/constructs/3-peak-load-manager.ts にあります。こちらも実装説明は省略します。

Supervisor Agent をマネコンで作っていく

ではSupervisor Agentを手で作っていきます。Agent作成時に Enable multi-agent collaboration にチェックを入れます。

description や instruction は Workshop 指定のものを入れていきます（私の実装リポジトリのREADMEにも転記してあります）。

また model について、Workshop では Claude を使っていますが、クレジットを使いたかったので Amazon Nova Pro にしています（問題なく動作します）。

あとは Collaborator として 3つの Agent を追加していきます。

  1. Forecasting Agent
  2. Solar Panel Agent
  3. Peak Load Manager Agent

Agent 名、Alias、Instructionなどを設定していきます。CDK 実装で Agent Alias を定義していたのは、ここで設定が必須なためです。

Instruction は Workshop指定のものを設定します。

以下画像は 1. Forecasting Agent ですが、2, 3についても同様に設定します。

設定後に保存することで Supervisor Agent の作成は完了です。

動かしてみる

では実際に動かしてみます。動かす前に以下2つが必要なのでやっておきましょう

準備ができたら Workshopのプロンプト で動かしてみます。

異なる Agent 向けの Prompt を投げたところ、ちゃんと Supervisor が捌いて各 Agent が動作し期待通りの結果が返ってきています。問題ないですね。

おわりに

思った以上に簡単に作れたので夢が広がりそうですね。

あとは multi-agent collaboration が CloudFormation にも対応し IaC で定義・管理できるようになると嬉しいです。対応した際は改めてブログ書きます。