Handle Actions that take too long or are too heavy for Agents for Bedrock by using RETURN_CONTROL and processing them in a Go client

2024-08-27 aws llm

Agents for Bedrock can register and invoke Lambda functions as Actions. However, for processes that take a long time, there’s a risk of hitting Lambda timeouts or resource limits. Moreover, when considering running heavy processes in parallel or notifying users of progress, it can be inconvenient to call these from the Agent. ReturnControls solves this by not handling Actions on the Agent side, but instead returning the Action that should be called and its inputs to the client. The client can then pass the results back to the Agent.

Create Agents for Bedrock using CDK and check that Lambda functions are called based on input - sambaiz-net

In the actionGroupExecutor, instead of specifying the Lambda ARN, you specify 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,
          },
        },
      },
    ],
  },
}

Call the Agent using aws-sdk-go-v2. For the Alias, we specify TSTALIASID corresponding to the DRAFT version. It seems that sessionID can be any unique string.

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)
	}
}

As a result, the first call returns the function and parameters, and when the call results are passed as SessionState in the second call, we receive a response similar to when the action is called directly from the agent.

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

Reference

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