An open API service indexing awesome lists of open source software.

https://github.com/mashiike/bedrocktool

Utility for Bedrock Converse in Golang
https://github.com/mashiike/bedrocktool

Last synced: 3 months ago
JSON representation

Utility for Bedrock Converse in Golang

Awesome Lists containing this project

README

        

# bedrocktool
Utility for Bedrock Converse in Golang

[![GoDoc](https://godoc.org/github.com/mashiike/bedrocktool?status.svg)](https://godoc.org/github.com/mashiike/bedrocktool)
[![Go Report Card](https://goreportcard.com/badge/github.com/mashiike/bedrocktool)](https://goreportcard.com/report/github.com/mashiike/bedrocktool)
[![License](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)

## Example

```go
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\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.

```go
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`
```go
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`
```go
// 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.

```go
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:

```go
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.