Agents for Amazon Bedrockで、作成済みのKnowledge baseを元にしたエージェントを作成してみた

この記事は、2023年のAWS re:inventでのAmazon Bedrockのアップデートについての内容です。Bedrockを使用してKnowledge baseを作成し、作成したKnowledge baseやモデルを使用してAgentを作成する方法について詳しく説明されています。Agentの作成手順やテスト方法、さらにAWS SDKを使用してAgentを実行する方法についても示されています。記事では、Bedrock Agentsの機能や利点についても触れており、今後のワークフローの改善に期待が寄せられています。

この記事は、「Amazon Bedrock Advent Calendar 2023」5日目の記事です。

2023年のAWS re:inventでは、Amazon Bedrockに関するアップデートが複数登場しました。昨日はBedrockを利用してKnowledge baseを作成する機能に挑戦してみました。今日は作成したKnowledge baseやモデルを利用してAgentを作る機能を試してみます。

Amazon Bedrock管理画面から、Agentsを作成する

Bedrock管理画面にアクセスすると、OrchestrationメニューにAgentsが増えています。新しくAgentsを作成するには、Create Agentsをクリックしましょう。

作成のためのウィザードが立ち上がります、1つ目のステップでは、ユーザーからの入力を受け付けるかなどの設定を行います。

また、IAMとかKMSまわりのお馴染みの設定系も、このステップで行います。

2番目のStepでは、Agentが利用する言語モデルを選べます。が、今の所AnthropicのClaude V2のみが利用できる様子で、AmazonのTitan Expressが今後サポートされる予定の様子です。

モデルを選択する際には、前提条件を伝えるテキストを1200文字（おそらく英語換算）まで入れることができます。今回は「あなたはAPI設計のエキスパートです。APIに関する質問にのみ回答してください。回答はMarkdown形式で行います。」と登録してみました。

Step3では、処理の中で呼び出すLambdaを登録できます。この辺りからOptionalな選択が増えてきました。Agentが処理を行う中で、必要に応じてAWS Lambdaを実行する仕組みがある様子で、LangChainの持つAgentのAWS版という印象を受けました。

次のステップでKnowledge baseを選べます。今回は前の記事で作成したものを指定します。

最後に確認画面が表示され、作成を行います。作成に成功すると、Agentの詳細画面に移動します。

作成したAgentをブラウザ上でテストする

右側のウィジェットでテストができます。ただしStream形式でレスポンスを返すわけではないため、返答が出てくるまで少し時間がかかります。あと、英語で返事が返ってきました。

Agentにエイリアスを設定してデプロイする

エイリアスが作れます。デフォルトはWorking draftですので、v1とかliveとかつけるとよさそうかなと思います。

AWS SDK (Node.js v3)からAgentを利用しようとしてみた

実際に組み込むことを想定して、AWS SDKからも呼び出してみました。まずはAWS SDKをインストールします。

npm i @aws-sdk/client-bedrock-agent-runtime

その後、BedrockAgentRuntimeClientとInvokeAgentCommandを利用することでエージェントを実行できる様子です。

import { BedrockAgentRuntimeClient, InvokeAgentCommand  } from '@aws-sdk/client-bedrock-agent-runtime'

const app = new Hono()
app.get('/invoke', async c => {
    const client = new BedrockAgentRuntimeClient()
    const sessionId = Math.random().toString(36).substring(2, 15)
    const command = new InvokeAgentCommand({
        agentId: 'xxxx',
        agentAliasId: 'xxxx',
        sessionId: sessionId,
        inputText: 'Stripeで請求書を作る方法を教えて'
    })
    const response = await client.send(command)
    return c.json(response)
})

実行結果がこんな感じです。テキストが返ってきていないですが、Streamのためでしょうか。一応動いてはいるように見えます。

{
  "$metadata": {
    "httpStatusCode": 200,
    "requestId": "1767da0b-5896-4c5c-9637-3a83ecf0d61a",
    "attempts": 1,
    "totalRetryDelay": 0
  },
  "contentType": "application/json",
  "sessionId": "1zq6hn578xo",
  "completion": {
    "options": {
      "messageStream": {
        "options": {
          "inputStream": {},
          "decoder": {
            "headerMarshaller": {},
            "messageBuffer": [],
            "isEndOfStream": false
          }
        }
      }
    }
  }
}

RetrieveCommandの方も試してみました。

import { BedrockAgentRuntimeClient, RetrieveCommand  } from '@aws-sdk/client-bedrock-agent-runtime'

const app = new Hono()
app.get('/invoke', async c => {
    const client = new BedrockAgentRuntimeClient()
    const sessionId = Math.random().toString(36).substring(2, 15)
    const command = new RetrieveCommand({ // RetrieveRequest
        knowledgeBaseId: "xxxx", // required
        retrievalQuery: { // KnowledgeBaseQuery
          text: "Stripeで請求書を作る方法を教えて", // required
        },
        retrievalConfiguration: { // KnowledgeBaseRetrievalConfiguration
          vectorSearchConfiguration: { // KnowledgeBaseVectorSearchConfiguration
            numberOfResults: Number(5), // required
          },
        },
    })
    const response = await client.send(command)
    return c.json(response)
})

こちらの場合は、回答に利用できるKnowledge base上のデータを配列で返してくれている様子です。

{
  "$metadata": {
    "httpStatusCode": 200,
    "requestId": "9259e8e0-711b-49cb-98f9-f11d73025cac",
    "attempts": 1,
    "totalRetryDelay": 0
  },
  "retrievalResults": [
    {
      "content": {
        "text": "Depending on                     the purchase, they can pay with four interest-free payments                     (Split Pay) or pay over a longer term (Installments), which                     might include interest. Check this                     [page](https://stripe.com/docs/payments/affirm) for more                     details like country availability.                   properties:                     display_preference:                       properties:                         preference:                           enum:                             - none                             - 'off'                             - 'on'                           type: string                       title: display_preference_param                       type: object                   title: payment_method_param                   type: object                 afterpay_clearpay:                   description: >-                     Afterpay gives your customers a way to pay for purchases in                     installments, check this                     [page](https://stripe.com/docs/payments/afterpay-clearpay)                     for more details like country availability. Afterpay is                     particularly popular among businesses selling fashion,                     beauty, and sports products.                   properties:                     display_preference:                       properties:                         preference:                           enum:                             - none                             - 'off'                             - 'on'                           type: string                       title: display_preference_param                       type: object                   title: payment_method_param                   type: object                 alipay:                   description: >-                     Alipay is a digital wallet in China that has more than a                     billion active users worldwide. Alipay users can pay on the                     web or on a mobile device using login credentials or their                     Alipay app."
      },
      "location": {
        "s3Location": {
          "uri": "s3://hidetaka-lab-tmp/stripe-openapi-spec3.yaml"
        },
        "type": "S3"
      },
      "score": 0.5854195952415466
    },

２つのAPIを組み合わせると、「回答」と「参考にしたデータソース」の2つをユーザーに提示できるということでしょうか。

Bedrock agents ≒ Managed LangChain

ワークショップに載っている図をみる限り、「LangChainみたいなやつを、AWS上で作れる」くらいのざっくり理解から入るのがよさそうに見えます。タスクの実行などもLambda経由で行えるので、AWS内のリソースの状況やStripe / Shopifyなどのデータと連携できると、いろいろ捗りそうな予感がします。