Agents for Bedrock で時間がかかったり重すぎたりする Actions を RETURN_CONTROL して Go クライアントでハンドリングする

2024-08-27 aws llm

Agents for Bedrock は Actions として Lambda 関数を登録して呼び出させることができる。ただ、その処理に時間がかる場合、Lambda のタイムアウトやリソース上限に当たったりすることが考えられる。また、重い処理を並列で実行したり、進行状況をユーザーに通知したりすることを考えると、Agent から呼び出すには都合が悪いことがある。それを解決するのが ReturnControls で Agent 側で Actions のハンドリングを行うのではなく、クライアントに呼び出すべき Action とその入力を返し、クライアントはその結果を Agent に渡すことができる。

CDK で Agents for Bedrock を作成し入力に基づいて Lambda 関数が呼び出されることを確認する - sambaiz-net

actionGroupExecutor に Lambda の ARN ではなく、customControl: RETURN_CONTROL を指定する。

return {
  actionGroupName: "TestActionGroup",

  actionGroupExecutor: {
    customControl: "RETURN_CONTROL",
  },
  functionSchema: {
    functions: [
      {
        name: "TestActionGroupFunction1",
        description: "Return a greeting message",
        parameters: {
          name: {
            type: "string",
            description: "Name of the user",
            required: true,
          },
        },
      },
    ],
  },
}

aws-sdk-go-v2 で Agent を呼び出す。Alias には DRAFT バージョンに対応する TSTALIASID を指定している。sessionID は任意のユニークな文字列で良いようだ。

package main

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime"
	"github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime/types"
)

func newBedrockAgentRuntimeClinet() *bedrockagentruntime.Client {
	cfg, err := config.LoadDefaultConfig(context.TODO(), config.WithRegion("us-east-1"))
	if err != nil {
		log.Fatalf("Failed to load configuration, %v", err)
	}

	return bedrockagentruntime.NewFromConfig(cfg)
}

func invokeAgent(client *bedrockagentruntime.Client, sessionID string, inputText *string, sessionState *types.SessionState) (*types.ResponseStreamMemberReturnControl, error) {
	out, err := client.InvokeAgent(context.TODO(), &bedrockagentruntime.InvokeAgentInput{
		AgentId:      aws.String(os.Getenv("AGENT_ID")),
		AgentAliasId: aws.String("TSTALIASID"),
		SessionId:    aws.String(sessionID),
		InputText:    inputText,
		SessionState: sessionState,
	})
	if err != nil {
		return nil, err
	}

	stream := out.GetStream()
	defer stream.Close()

	ch := stream.Events()
	for event := range ch {
		switch ev := event.(type) {
		case *types.ResponseStreamMemberChunk:
			fmt.Println("Chunk:", string(ev.Value.Bytes))

		case *types.ResponseStreamMemberReturnControl:
			return ev, nil
		}
	}

	return nil, nil
}

func main() {
	bedrockAgentRuntimeClient := newBedrockAgentRuntimeClinet()

	sessionID := fmt.Sprintf("session-%d", time.Now().Unix())

	fmt.Println("Invoking agent with input text")
	returnControl, err := invokeAgent(bedrockAgentRuntimeClient, sessionID, aws.String("Hi, I am Tom."), nil)
	if err != nil {
		log.Fatalf("Failed to invoke agent with input text, %v", err)
	}

	var (
		funcInput    *types.InvocationInputMemberMemberFunctionInvocationInput
		funcResponse string
	)
	if returnControl.Value.InvocationId == nil {
		fmt.Println("No invocation ID")
		return
	}
	for _, input := range returnControl.Value.InvocationInputs {
		in, ok := input.(*types.InvocationInputMemberMemberFunctionInvocationInput)
		if !ok {
			fmt.Println("No function invocation input")
			return
		}

		funcInput = in
		fmt.Println("Function:", *funcInput.Value.Function)
		fmt.Println("Parameters:")
		for _, param := range funcInput.Value.Parameters {
			if param.Name != nil && param.Value != nil {
				fmt.Println(*param.Name, *param.Value)
				funcResponse = fmt.Sprintf("おはこんハロチャオー！ %sさん!", *param.Value)
			}
		}
	}

	fmt.Println("Invoking agent with function responses")
	_, err = invokeAgent(bedrockAgentRuntimeClient, sessionID, nil, &types.SessionState{
		InvocationId: returnControl.Value.InvocationId,
		ReturnControlInvocationResults: []types.InvocationResultMember{
			&types.InvocationResultMemberMemberFunctionResult{
				Value: types.FunctionResult{
					ActionGroup: funcInput.Value.ActionGroup,
					Function:    funcInput.Value.Function,
					ResponseBody: map[string]types.ContentBody{
						"TEXT": {
							Body: aws.String(funcResponse),
						},
					},
				},
			},
		},
	})
	if err != nil {
		log.Fatalf("Failed to invoke agent with function responses, %v", err)
	}
}

結果、1回目の呼び出しで Function と Parameters が返り、2回目の呼び出しで SessionState として呼び出し結果を渡すと、Agent から Actions を呼び出したときと同様なレスポンスが得られた。

Invoking agent with input text
Function: TestActionGroupFunction1
Parameters:
name Tom
Invoking agent with function responses
Chunk: おはこんハロチャオー! Tomさん!

参考

サンプルコードで理解する Agents for Amazon Bedrock の Return of Control - Qiita