bedrocktool

package module
v0.7.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Sep 17, 2024 License: MIT Imports: 12 Imported by: 0

README

bedrocktool

Utility for Bedrock Converse in Golang

GoDoc Go Report Card License

Example

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"math/rand"
	"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/bedrockruntime"
	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types"
	"github.com/mashiike/bedrocktool"
)

func clockToolFunc(ctx context.Context, _ bedrocktool.EmptyWorkerInput) (types.ToolResultBlock, error) {
	slog.InfoContext(ctx, "call clock tool")
	return types.ToolResultBlock{
		Content: []types.ToolResultContentBlock{
			&types.ToolResultContentBlockMemberText{
				Value: time.Now().Format(time.RFC3339),
			},
		},
	}, nil
}

// json schema generated by github.com/invopop/jsonschema
// jsonschema tag specifies see as: https://github.com/invopop/jsonschema?tab=readme-ov-file
type weatherInput struct {
	City string `json:"city" jsonschema:"description=都市名 (例: 横浜,東京),default=東京, required=true"`
	When string `json:"when" jsonschema:"description=日時 RFC3339 (例: 2022-01-01T00:00:00Z), required=false"`
}

func weatherToolFunc(ctx context.Context, input weatherInput) (types.ToolResultBlock, error) {
	slog.InfoContext(ctx, "call weather tool", "city", input.City, "when", input.When)
	whethers := []string{"晴れ", "曇り", "雨", "雪"}
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	randomIndex := r.Intn(len(whethers))
	return types.ToolResultBlock{
		Content: []types.ToolResultContentBlock{
			&types.ToolResultContentBlockMemberText{
				Value: whethers[randomIndex],
			},
		},
	}, nil
}

func main() {
	ctx := context.Background()
	opts := make([]func(*config.LoadOptions) error, 0)
	awsCfg, err := config.LoadDefaultConfig(ctx, opts...)
	if err != nil {
		panic("configuration error, " + err.Error())
	}
	d := bedrocktool.NewFromConfig(awsCfg)
	type clockInput struct{}
	d.Register(
		"clock",
		"他のツールとの連携で、現在の時刻が必要な場合に使用します。このツールが有効な場合は、暗黙的な時刻の情報は事前学習知識ではなく、都度このツールによって現在時刻を取得してください。",
		bedrocktool.NewWorker(clockToolFunc),
	)
	d.Register(
		"weather",
		"指定された都市の指定された日時の天気を取得します。暗黙的に時間が必要な場合は、clockツールを使用してください。",
		bedrocktool.NewWorker(weatherToolFunc),
	)
	output, err := d.Converse(ctx, &bedrockruntime.ConverseInput{
		ModelId: aws.String("anthropic.claude-3-sonnet-20240229-v1:0"),
		Messages: []types.Message{
			{
				Role: types.ConversationRoleUser,
				Content: []types.ContentBlock{
					&types.ContentBlockMemberText{
						Value: "今日の鎌倉の天気を調べて。",
					},
				},
			},
		},
	})
	if err != nil {
		panic("converse error, " + err.Error())
	}
	toolUseByID := make(map[string]types.ToolUseBlock)
	for _, msg := range output {
		for _, content := range msg.Content {
			switch c := content.(type) {
			case *types.ContentBlockMemberText:
				fmt.Printf("[%s]:\t%s\n", msg.Role, c.Value)
			case *types.ContentBlockMemberImage:
				fmt.Printf("[%s]:\t<image %s>\n", msg.Role, c.Value.Format)
			case *types.ContentBlockMemberToolUse:
				toolUseByID[*c.Value.ToolUseId] = c.Value
			case *types.ContentBlockMemberToolResult:
				toolUse, ok := toolUseByID[*c.Value.ToolUseId]
				if !ok {
					fmt.Printf("tool use not found: %s\n", *c.Value.ToolUseId)
					continue
				}
				bs, err := json.Marshal(c.Value.Content)
				if err != nil {
					fmt.Printf("failed to marshal tool result: %s\n", err)
					continue
				}
				fmt.Printf("[%s]:\t%s\n", *toolUse.Name, string(bs))
			default:
			}
		}
	}
	// Example output:
	// 2024/06/05 16:15:38 INFO call clock tool
	// 2024/06/05 16:15:41 INFO call weather tool city=鎌倉 when=2024-06-05T16:15:38+09:00
	// [assistant]:    はい、鎌倉の今日の天気を調べましょう。
	// [clock]:        [{"Value":"2024-06-05T16:15:38+09:00"}]
	// [weather]:      [{"Value":"雪"}]
	// [assistant]:    今日(2024年6月5日)の鎌倉の天気は雪とのことですね。6月に雪というのは非常に珍しい気象状況だと思われます。鎌倉は温暖な気候が一般的なので、どこかの山間部など局所的な雪降りの可能性があります。雪が積もれば交通機関へも影響が出るかもしれません。防寒対策が必要かもしれませんので、外出の際は暖かい服装を心がけましょう。
}
bedrocktool.Dispacher

bedrocktool.Dispacher is a utility for Bedrock Converse in Golang. It provides a simple way to implement a tool for Bedrock Converse.

d := bedrocktool.NewFromConfig(awsCfg)
d.Register(toolName, toolDescription, bedrocktool.NewWorker(toolWorkerFunc))
output, err := d.Converse(ctx, &bedrockruntime.ConverseInput{
//... 
// this is the same input as bedrockruntime.ConverseInput
// but, ToolConfig is generated by bedrocktool.Dispacher, so you don't need to specify it.
})
Hooks

if you wont to track API Call: for example logging Usage Metrics, you can use bedrocktool.OnAfterModelCall

d := bedrocktool.NewFromConfig(awsCfg)
d.OnAfterModelCall = func(ctx context.Context, input *bedrockruntime.ConverseInput, output *bedrockruntime.ConverseOutput, err error) {
    // any thing
}
Middleware

if you want to add some common processing to all tools, you can use bedrocktool.Middleware

// implement bedrocktool.Middleware interface
type middleware struct {}
func (m middleware) HandleInputSchema(next bedrocktool.Worker) document.Interface {
    // any thing...
    return next()
}
func (m middleware) HandleExecute(ctx context.Context, in types.ToolUseBlock, next Worker) (types.ToolResultBlock, error) {
    // any thing...
    return next(ctx, in)
}

d := bedrocktool.NewFromConfig(awsCfg)
d.Use(middleware{})

helper types bedrocktool.InputSchemaMiddlewareFunc and bedrocktool.ExecuteMiddlewareFunc for easy implementation.

d := bedrocktool.NewFromConfig(awsCfg)
d.Use(bedrocktool.InputSchemaMiddlewareFunc(func(next bedrocktool.Worker) document.Interface {
    // any thing...
    return next()
}))
d.Use(bedrocktool.ExecuteMiddlewareFunc(func(ctx context.Context, in types.ToolUseBlock, next Worker) (types.ToolResultBlock, error) {
    // any thing...
    return next(ctx, in)
}))
Context-Bound ToolSet

Tools are grouped together in units called ToolSet. If you want to use a specific ToolSet only within a certain context, you can do so as follows:

ctx, ts := bedrocktool.WithToolSet(ctx)
ts.Register(toolName, toolDescription, bedrocktool.NewWorker(toolWorkerFunc))

d.Converse(ctx, &bedrockruntime.ConverseInput{
    //... 
})

This approach allows you to activate specific tools only within the scope of a given context.

Documentation

Index

Constants

This section is empty.

Variables

View Source 
var NoPanicOnRegisterError = false

flag of panic behavior, when the error occurred during the registration of the tool.

Functions

func GenerateInputSchemaDocument 

func GenerateInputSchemaDocument[T any]() (document.Interface, error)

func MarshalJSONSchema added in v0.3.1 

func MarshalJSONSchema(v any) ([]byte, error)

func NewContext 

func NewContext(parent context.Context, cc *ConverseContext) context.Context

func ToolName 

func ToolName(ctx context.Context) string

func ToolUseID 

func ToolUseID(ctx context.Context) string

func ValidateToolName added in v0.2.0 

func ValidateToolName(name string) error

Types

type BedrockConverseAPIClient 

type BedrockConverseAPIClient interface {
	Converse(ctx context.Context, params *bedrockruntime.ConverseInput, optFns ...func(*bedrockruntime.Options)) (*bedrockruntime.ConverseOutput, error)
}

BedrockConverseAPIClient is a client for BedrockConverseAPI. see: https://pkg.go.dev/github.com/aws/aws-sdk-go-v2/service/bedrockruntime#Client.Converse

type ConverseContext 

type ConverseContext struct {
	// contains filtered or unexported fields
}

func FromContext 

func FromContext(ctx context.Context) (*ConverseContext, bool)

func (*ConverseContext) InputMessages 

func (cc *ConverseContext) InputMessages() []types.Message

func (*ConverseContext) ModelID 

func (cc *ConverseContext) ModelID() string

func (*ConverseContext) OutputMessages 

func (cc *ConverseContext) OutputMessages() []types.Message

func (*ConverseContext) SetModelID 

func (cc *ConverseContext) SetModelID(modelID string)

SetModelID sets the model ID for the conversation. for in tool use. model id upgrade/downgrade.

func (*ConverseContext) System 

func (cc *ConverseContext) System() []types.SystemContentBlock

type Dispacher 

type Dispacher struct {
	// contains filtered or unexported fields
}

Dispacher is a tool use dispacher. It is used to send messages to the specified Amazon Bedrock model.

func New 

func New(options bedrockruntime.Options, optFns ...func(*bedrockruntime.Options)) *Dispacher

New creates a new instance of the Bedrock Tool Use Dispacher.

func NewFromConfig 

func NewFromConfig(cfg aws.Config, optFns ...func(*bedrockruntime.Options)) *Dispacher

NewFromConfig creates a new instance of the Bedrock Tool Use Dispacher.

func NewWithClient 

func NewWithClient(client BedrockConverseAPIClient) *Dispacher

NewWithClient creates a new instance of the Bedrock Tool Use Dispacher.

func (*Dispacher) Converse 

func (d *Dispacher) Converse(ctx context.Context, params *bedrockruntime.ConverseInput, optFns ...func(*bedrockruntime.Options)) ([]types.Message, error)

Converse sends messages to the specified Amazon Bedrock model. input same as https://pkg.go.dev/github.com/aws/aws-sdk-go-v2/service/bedrockruntime#Client.Converse but output is different. because this function is multiple call Converse API. if you need track api call, use OnBeforeModelCall and OnAfterModelCall.

func (*Dispacher) GetError 

func (d *Dispacher) GetError() error

GetError returns the error that occurred during the registration of the tool.

func (*Dispacher) NewToolConfiguration 

func (d *Dispacher) NewToolConfiguration(ctx context.Context) *types.ToolConfiguration

NewToolConfiguration returns a new ToolConfiguration.

func (*Dispacher) OnAfterModelCall 

func (d *Dispacher) OnAfterModelCall(f func(context.Context, *bedrockruntime.ConverseInput, *bedrockruntime.ConverseOutput))

OnAfterModelCall sets the function to be called after the model is called. After Call Bedrock Converse API.

func (*Dispacher) OnAfterToolUse 

func (d *Dispacher) OnAfterToolUse(f func(context.Context, *types.ContentBlockMemberToolUse, types.ToolResultBlock))

OnAfterToolUse sets the function to be called after the tool is used. After Call Worker.Execute.

func (*Dispacher) OnBeforeModelCall 

func (d *Dispacher) OnBeforeModelCall(f func(context.Context, *bedrockruntime.ConverseInput))

OnBeforeModelCall sets the function to be called before the model is called. Before Call Bedrock Converse API.

func (*Dispacher) OnBeforeToolUse 

func (d *Dispacher) OnBeforeToolUse(f func(context.Context, *types.ContentBlockMemberToolUse))

OnBeforeToolUse sets the function to be called before the tool is used. Before Call Worker.Execute.

func (*Dispacher) Register 

func (d *Dispacher) Register(name string, description string, worker Worker, opts ...RegisterOption)

Register registers a tool with the specified name and description. if occurs error during the registration, it will panic. if you want to handle the error, use GetError and NoPanicOnRegisterError=false.

func (*Dispacher) RegisterTool added in v0.2.0 

func (d *Dispacher) RegisterTool(tool Tool, opts ...RegisterOption)

RegisterTool registers a tool. this is syntactic sugar for Register.

func (*Dispacher) ResolveWorker added in v0.2.0 

func (d *Dispacher) ResolveWorker(ctx context.Context, name string) (Worker, bool)

ResolveWorker returns the worker registered with the specified name. diff of Worker: consider to temporary tool set.

func (*Dispacher) SetLogger 

func (d *Dispacher) SetLogger(logger *slog.Logger)

SetLogger sets the logger for the dispacher.

func (*Dispacher) SetToolChoice 

func (d *Dispacher) SetToolChoice(tc types.ToolChoice)

SetToolChoice sets the tool choice for the dispacher.

func (*Dispacher) SubToolSet 

func (d *Dispacher) SubToolSet() *ToolSet

func (*Dispacher) Use 

func (d *Dispacher) Use(middlewares ...Middleware)

func (*Dispacher) Worker 

func (d *Dispacher) Worker(name string) (Worker, bool)

Worker returns the worker registered with the specified name.

type EmptyWorkerInput 

type EmptyWorkerInput struct{}

type ExecuteMiddlewareFunc 

type ExecuteMiddlewareFunc func(context.Context, types.ToolUseBlock, Worker) (types.ToolResultBlock, error)

func (ExecuteMiddlewareFunc) HandleExecute 

func (f ExecuteMiddlewareFunc) HandleExecute(ctx context.Context, in types.ToolUseBlock, next Worker) (types.ToolResultBlock, error)

func (ExecuteMiddlewareFunc) HandleInputSchema 

func (f ExecuteMiddlewareFunc) HandleInputSchema(next Worker) document.Interface

type InputSchemaMiddlewareFunc 

type InputSchemaMiddlewareFunc func(Worker) document.Interface

func (InputSchemaMiddlewareFunc) HandleExecute 

func (f InputSchemaMiddlewareFunc) HandleExecute(ctx context.Context, in types.ToolUseBlock, next Worker) (types.ToolResultBlock, error)

func (InputSchemaMiddlewareFunc) HandleInputSchema 

func (f InputSchemaMiddlewareFunc) HandleInputSchema(next Worker) document.Interface

type Middleware 

type Middleware interface {
	HandleInputSchema(next Worker) document.Interface
	HandleExecute(ctx context.Context, in types.ToolUseBlock, next Worker) (types.ToolResultBlock, error)
}

type RegisterOption 

type RegisterOption func(*workerEntry)

RegisterOption is an option for registering a tool.

func WithToolEnabler 

func WithToolEnabler(f func(context.Context) bool) RegisterOption

WithToolEnabler sets the function to determine whether the tool is enabled. this function is called before the first time Bedorck Converse API If return false, not enabled the tool in this conversation.

type Registory 

type Registory interface {
	Register(name string, description string, worker Worker, opts ...RegisterOption)
}

type Tool added in v0.2.0 

type Tool interface {
	Name() string
	Description() string
	Worker() Worker
}

type ToolSet 

type ToolSet struct {
	// contains filtered or unexported fields
}

func WithToolSet added in v0.2.0 

func WithToolSet(ctx context.Context) (context.Context, *ToolSet)

func (*ToolSet) Exists 

func (ts *ToolSet) Exists(name string) bool

func (*ToolSet) GetError 

func (ts *ToolSet) GetError() error

func (*ToolSet) Register 

func (ts *ToolSet) Register(name string, description string, worker Worker, opts ...RegisterOption)

func (*ToolSet) RegisterTool added in v0.2.0 

func (ts *ToolSet) RegisterTool(tool Tool, opts ...RegisterOption)

func (*ToolSet) SubToolSet 

func (ts *ToolSet) SubToolSet() *ToolSet

func (*ToolSet) Tools 

func (ts *ToolSet) Tools(ctx context.Context) []types.Tool

func (*ToolSet) Use 

func (ts *ToolSet) Use(middlewares ...Middleware)

func (*ToolSet) Worker 

func (ts *ToolSet) Worker(name string) (Worker, bool)

type Worker 

type Worker interface {
	InputSchema() document.Interface
	Execute(context.Context, types.ToolUseBlock) (types.ToolResultBlock, error)
}

func NewWorker 

func NewWorker[T any](f func(context.Context, T) (types.ToolResultBlock, error)) Worker

Source Files

View all Source files

Directories

Path Synopsis
tool
remote
urlreader

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.