mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-06-12 18:08:54 +00:00
412705783d
* fix(pico): preserve image media across pico attachments and client * * fix ci * fix(pico): preserve text when client media parsing fails - Skip non-inline Pico attachment URLs instead of treating them as invalid inline media - Preserve pico_client text messages when malformed media payloads are received - Add regression coverage for media.create, download attachments, and invalid media payloads * fix lint
815 lines
23 KiB
Go
815 lines
23 KiB
Go
//go:build bedrock
|
|
|
|
// PicoClaw - Ultra-lightweight personal AI agent
|
|
// License: MIT
|
|
//
|
|
// Copyright (c) 2026 PicoClaw contributors
|
|
|
|
// Package bedrock implements the LLM provider interface for AWS Bedrock.
|
|
// It uses the Bedrock Runtime Converse API for unified access to multiple
|
|
// model families (Claude, Llama, Mistral, etc.) with tool/function calling support.
|
|
package bedrock
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"math"
|
|
"strings"
|
|
"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/document"
|
|
"github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types"
|
|
|
|
"github.com/sipeed/picoclaw/pkg/providers/common"
|
|
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
|
|
)
|
|
|
|
type (
|
|
ToolCall = protocoltypes.ToolCall
|
|
FunctionCall = protocoltypes.FunctionCall
|
|
LLMResponse = protocoltypes.LLMResponse
|
|
UsageInfo = protocoltypes.UsageInfo
|
|
Message = protocoltypes.Message
|
|
ToolDefinition = protocoltypes.ToolDefinition
|
|
ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition
|
|
)
|
|
|
|
// Provider implements the LLM provider interface for AWS Bedrock.
|
|
type Provider struct {
|
|
client *bedrockruntime.Client
|
|
region string
|
|
requestTimeout time.Duration
|
|
}
|
|
|
|
// Option configures the Bedrock Provider.
|
|
type Option func(*providerConfig)
|
|
|
|
type providerConfig struct {
|
|
region string
|
|
profile string
|
|
baseEndpoint string
|
|
requestTimeout time.Duration
|
|
}
|
|
|
|
// WithRegion sets the AWS region for Bedrock requests.
|
|
func WithRegion(region string) Option {
|
|
return func(c *providerConfig) {
|
|
c.region = region
|
|
}
|
|
}
|
|
|
|
// WithProfile sets the AWS profile to use for credentials.
|
|
func WithProfile(profile string) Option {
|
|
return func(c *providerConfig) {
|
|
c.profile = profile
|
|
}
|
|
}
|
|
|
|
// WithBaseEndpoint sets a custom Bedrock endpoint URL.
|
|
// Example: https://bedrock-runtime.us-east-1.amazonaws.com
|
|
func WithBaseEndpoint(endpoint string) Option {
|
|
return func(c *providerConfig) {
|
|
c.baseEndpoint = endpoint
|
|
}
|
|
}
|
|
|
|
// WithRequestTimeout sets the timeout for Bedrock API requests.
|
|
func WithRequestTimeout(timeout time.Duration) Option {
|
|
return func(c *providerConfig) {
|
|
c.requestTimeout = timeout
|
|
}
|
|
}
|
|
|
|
// NewProvider creates a new AWS Bedrock provider.
|
|
// It uses the default AWS credential chain (env vars, shared config, IAM roles, etc.).
|
|
func NewProvider(ctx context.Context, opts ...Option) (*Provider, error) {
|
|
pc := &providerConfig{}
|
|
for _, opt := range opts {
|
|
opt(pc)
|
|
}
|
|
|
|
// Build AWS config options
|
|
var configOpts []func(*config.LoadOptions) error
|
|
|
|
if pc.region != "" {
|
|
configOpts = append(configOpts, config.WithRegion(pc.region))
|
|
}
|
|
|
|
if pc.profile != "" {
|
|
configOpts = append(configOpts, config.WithSharedConfigProfile(pc.profile))
|
|
}
|
|
|
|
// Load AWS config with automatic credential discovery
|
|
cfg, err := config.LoadDefaultConfig(ctx, configOpts...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("loading AWS config: %w", err)
|
|
}
|
|
|
|
// Validate region is set - required for Bedrock request signing
|
|
if cfg.Region == "" {
|
|
return nil, fmt.Errorf(
|
|
"AWS region not configured: set AWS_REGION, AWS_DEFAULT_REGION, or use WithRegion option",
|
|
)
|
|
}
|
|
|
|
// Build client options
|
|
var clientOpts []func(*bedrockruntime.Options)
|
|
if pc.baseEndpoint != "" {
|
|
clientOpts = append(clientOpts, func(o *bedrockruntime.Options) {
|
|
o.BaseEndpoint = aws.String(pc.baseEndpoint)
|
|
})
|
|
}
|
|
|
|
client := bedrockruntime.NewFromConfig(cfg, clientOpts...)
|
|
|
|
return &Provider{
|
|
client: client,
|
|
region: cfg.Region,
|
|
requestTimeout: pc.requestTimeout,
|
|
}, nil
|
|
}
|
|
|
|
// converseParams holds the shared request parameters for Converse and ConverseStream.
|
|
type converseParams struct {
|
|
messages []types.Message
|
|
system []types.SystemContentBlock
|
|
inferenceConfig *types.InferenceConfiguration
|
|
toolConfig *types.ToolConfiguration
|
|
}
|
|
|
|
func buildConverseParams(messages []Message, tools []ToolDefinition, options map[string]any) converseParams {
|
|
bedrockMessages, systemPrompts := convertMessages(messages)
|
|
|
|
var inferenceConfig *types.InferenceConfiguration
|
|
|
|
if maxTokens, ok := common.AsInt(options["max_tokens"]); ok && maxTokens > 0 {
|
|
if inferenceConfig == nil {
|
|
inferenceConfig = &types.InferenceConfiguration{}
|
|
}
|
|
if maxTokens > math.MaxInt32 {
|
|
maxTokens = math.MaxInt32
|
|
}
|
|
inferenceConfig.MaxTokens = aws.Int32(int32(maxTokens))
|
|
}
|
|
|
|
if temp, ok := common.AsFloat(options["temperature"]); ok {
|
|
if inferenceConfig == nil {
|
|
inferenceConfig = &types.InferenceConfiguration{}
|
|
}
|
|
inferenceConfig.Temperature = aws.Float32(float32(temp))
|
|
}
|
|
|
|
var toolConfig *types.ToolConfiguration
|
|
if len(tools) > 0 {
|
|
tc := convertTools(tools)
|
|
if len(tc.Tools) > 0 {
|
|
toolConfig = tc
|
|
}
|
|
}
|
|
|
|
return converseParams{
|
|
messages: bedrockMessages,
|
|
system: systemPrompts,
|
|
inferenceConfig: inferenceConfig,
|
|
toolConfig: toolConfig,
|
|
}
|
|
}
|
|
|
|
// Chat sends messages to AWS Bedrock using the Converse API.
|
|
func (p *Provider) Chat(
|
|
ctx context.Context,
|
|
messages []Message,
|
|
tools []ToolDefinition,
|
|
model string,
|
|
options map[string]any,
|
|
) (*LLMResponse, error) {
|
|
effectiveTimeout := p.requestTimeout
|
|
if effectiveTimeout <= 0 {
|
|
effectiveTimeout = common.DefaultRequestTimeout
|
|
}
|
|
if _, hasDeadline := ctx.Deadline(); !hasDeadline {
|
|
var cancel context.CancelFunc
|
|
ctx, cancel = context.WithTimeout(ctx, effectiveTimeout)
|
|
defer cancel()
|
|
}
|
|
|
|
params := buildConverseParams(messages, tools, options)
|
|
input := &bedrockruntime.ConverseInput{
|
|
ModelId: aws.String(model),
|
|
Messages: params.messages,
|
|
InferenceConfig: params.inferenceConfig,
|
|
ToolConfig: params.toolConfig,
|
|
}
|
|
if len(params.system) > 0 {
|
|
input.System = params.system
|
|
}
|
|
|
|
output, err := p.client.Converse(ctx, input)
|
|
if err != nil {
|
|
if isSSOTokenError(err) {
|
|
return nil, fmt.Errorf(
|
|
"bedrock converse: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w",
|
|
err,
|
|
)
|
|
}
|
|
return nil, fmt.Errorf("bedrock converse: %w", err)
|
|
}
|
|
|
|
return parseResponse(output)
|
|
}
|
|
|
|
// ChatStream sends messages to AWS Bedrock using the ConverseStream API.
|
|
// It streams the accumulated text so far via the onChunk callback and returns the complete response.
|
|
func (p *Provider) ChatStream(
|
|
ctx context.Context,
|
|
messages []Message,
|
|
tools []ToolDefinition,
|
|
model string,
|
|
options map[string]any,
|
|
onChunk func(accumulated string),
|
|
) (*LLMResponse, error) {
|
|
if p.requestTimeout > 0 {
|
|
if _, hasDeadline := ctx.Deadline(); !hasDeadline {
|
|
var cancel context.CancelFunc
|
|
ctx, cancel = context.WithTimeout(ctx, p.requestTimeout)
|
|
defer cancel()
|
|
}
|
|
}
|
|
|
|
params := buildConverseParams(messages, tools, options)
|
|
input := &bedrockruntime.ConverseStreamInput{
|
|
ModelId: aws.String(model),
|
|
Messages: params.messages,
|
|
InferenceConfig: params.inferenceConfig,
|
|
ToolConfig: params.toolConfig,
|
|
}
|
|
if len(params.system) > 0 {
|
|
input.System = params.system
|
|
}
|
|
|
|
output, err := p.client.ConverseStream(ctx, input)
|
|
if err != nil {
|
|
if isSSOTokenError(err) {
|
|
return nil, fmt.Errorf(
|
|
"bedrock conversestream: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w",
|
|
err,
|
|
)
|
|
}
|
|
return nil, fmt.Errorf("bedrock conversestream: %w", err)
|
|
}
|
|
|
|
return parseStreamResponse(ctx, output.GetStream(), onChunk)
|
|
}
|
|
|
|
// converseStreamReader abstracts the Bedrock event stream so parseStreamResponse
|
|
// can be unit-tested with a mock event source.
|
|
type converseStreamReader interface {
|
|
Events() <-chan types.ConverseStreamOutput
|
|
Err() error
|
|
Close() error
|
|
}
|
|
|
|
// parseStreamResponse processes the ConverseStream event stream and accumulates the response.
|
|
func parseStreamResponse(
|
|
ctx context.Context,
|
|
stream converseStreamReader,
|
|
onChunk func(accumulated string),
|
|
) (resp *LLMResponse, err error) {
|
|
if stream == nil {
|
|
return nil, fmt.Errorf("bedrock conversestream: nil event stream")
|
|
}
|
|
defer func() {
|
|
if closeErr := stream.Close(); closeErr != nil {
|
|
if err == nil {
|
|
err = fmt.Errorf("bedrock conversestream: close event stream: %w", closeErr)
|
|
} else {
|
|
log.Printf("bedrock conversestream: close event stream: %v", closeErr)
|
|
}
|
|
}
|
|
}()
|
|
|
|
var textContent strings.Builder
|
|
finishReason := "stop"
|
|
var usage *UsageInfo
|
|
toolCalls := make([]ToolCall, 0)
|
|
|
|
// Track active tool use blocks by index
|
|
type toolAccum struct {
|
|
id string
|
|
name string
|
|
argsJSON strings.Builder
|
|
}
|
|
activeTools := map[int]*toolAccum{}
|
|
|
|
events := stream.Events()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case event, ok := <-events:
|
|
if !ok {
|
|
// Stream closed
|
|
goto done
|
|
}
|
|
|
|
switch e := event.(type) {
|
|
case *types.ConverseStreamOutputMemberContentBlockStart:
|
|
// New content block starting
|
|
if toolUse, ok := e.Value.Start.(*types.ContentBlockStartMemberToolUse); ok {
|
|
activeTools[int(aws.ToInt32(e.Value.ContentBlockIndex))] = &toolAccum{
|
|
id: aws.ToString(toolUse.Value.ToolUseId),
|
|
name: aws.ToString(toolUse.Value.Name),
|
|
}
|
|
}
|
|
|
|
case *types.ConverseStreamOutputMemberContentBlockDelta:
|
|
// Content delta
|
|
switch delta := e.Value.Delta.(type) {
|
|
case *types.ContentBlockDeltaMemberText:
|
|
textContent.WriteString(delta.Value)
|
|
if onChunk != nil {
|
|
onChunk(textContent.String())
|
|
}
|
|
case *types.ContentBlockDeltaMemberToolUse:
|
|
idx := int(aws.ToInt32(e.Value.ContentBlockIndex))
|
|
if tool, exists := activeTools[idx]; exists {
|
|
tool.argsJSON.WriteString(aws.ToString(delta.Value.Input))
|
|
}
|
|
}
|
|
|
|
case *types.ConverseStreamOutputMemberContentBlockStop:
|
|
// Content block finished - finalize tool if it was a tool use
|
|
idx := int(aws.ToInt32(e.Value.ContentBlockIndex))
|
|
if tool, exists := activeTools[idx]; exists {
|
|
args := make(map[string]any)
|
|
argsStr := tool.argsJSON.String()
|
|
if argsStr != "" {
|
|
if err := json.Unmarshal([]byte(argsStr), &args); err != nil {
|
|
log.Printf("bedrock: stream: failed to parse tool arguments for %q: %v", tool.name, err)
|
|
args = map[string]any{"raw": argsStr}
|
|
}
|
|
}
|
|
funcArgs := argsStr
|
|
if argsJSON, marshalErr := json.Marshal(args); marshalErr == nil {
|
|
funcArgs = string(argsJSON)
|
|
}
|
|
toolCalls = append(toolCalls, ToolCall{
|
|
ID: tool.id,
|
|
Name: tool.name,
|
|
Arguments: args,
|
|
Function: &FunctionCall{
|
|
Name: tool.name,
|
|
Arguments: funcArgs,
|
|
},
|
|
})
|
|
delete(activeTools, idx)
|
|
}
|
|
|
|
case *types.ConverseStreamOutputMemberMessageStop:
|
|
// Message complete
|
|
switch e.Value.StopReason {
|
|
case types.StopReasonToolUse:
|
|
finishReason = "tool_calls"
|
|
case types.StopReasonMaxTokens:
|
|
finishReason = "length"
|
|
case types.StopReasonEndTurn:
|
|
finishReason = "stop"
|
|
case types.StopReasonStopSequence:
|
|
finishReason = "stop"
|
|
case types.StopReasonContentFiltered:
|
|
finishReason = "content_filter"
|
|
default:
|
|
finishReason = "stop"
|
|
}
|
|
|
|
case *types.ConverseStreamOutputMemberMetadata:
|
|
// Usage metadata
|
|
if e.Value.Usage != nil {
|
|
usage = &UsageInfo{
|
|
PromptTokens: int(aws.ToInt32(e.Value.Usage.InputTokens)),
|
|
CompletionTokens: int(aws.ToInt32(e.Value.Usage.OutputTokens)),
|
|
TotalTokens: int(
|
|
aws.ToInt32(e.Value.Usage.InputTokens),
|
|
) + int(
|
|
aws.ToInt32(e.Value.Usage.OutputTokens),
|
|
),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
done:
|
|
if err := stream.Err(); err != nil {
|
|
return nil, fmt.Errorf("bedrock conversestream: %w", err)
|
|
}
|
|
|
|
return &LLMResponse{
|
|
Content: textContent.String(),
|
|
ToolCalls: toolCalls,
|
|
FinishReason: finishReason,
|
|
Usage: usage,
|
|
}, nil
|
|
}
|
|
|
|
// GetDefaultModel returns an empty string as Bedrock models are user-configured.
|
|
func (p *Provider) GetDefaultModel() string {
|
|
return ""
|
|
}
|
|
|
|
// Region returns the AWS region configured for this Provider.
|
|
func (p *Provider) Region() string {
|
|
return p.region
|
|
}
|
|
|
|
// convertMessages converts internal messages to Bedrock Converse format.
|
|
// Returns the conversation messages and any system prompts separately.
|
|
// Note: Bedrock requires all tool results for a given assistant turn to be in a single
|
|
// user message with multiple ToolResultBlock content blocks. This function merges
|
|
// consecutive tool result messages accordingly.
|
|
func convertMessages(messages []Message) ([]types.Message, []types.SystemContentBlock) {
|
|
var bedrockMessages []types.Message
|
|
var systemPrompts []types.SystemContentBlock
|
|
|
|
// Helper to check if a message is a tool result
|
|
isToolResult := func(msg Message) bool {
|
|
return (msg.Role == "tool" || (msg.Role == "user" && msg.ToolCallID != "")) && msg.ToolCallID != ""
|
|
}
|
|
|
|
// Helper to create a tool result content block
|
|
makeToolResultBlock := func(msg Message) types.ContentBlock {
|
|
return &types.ContentBlockMemberToolResult{
|
|
Value: types.ToolResultBlock{
|
|
ToolUseId: aws.String(msg.ToolCallID),
|
|
Content: []types.ToolResultContentBlock{
|
|
&types.ToolResultContentBlockMemberText{
|
|
Value: msg.Content,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
i := 0
|
|
for i < len(messages) {
|
|
msg := messages[i]
|
|
|
|
switch {
|
|
case msg.Role == "system":
|
|
// System messages go to the System field
|
|
systemPrompts = append(systemPrompts, &types.SystemContentBlockMemberText{
|
|
Value: msg.Content,
|
|
})
|
|
i++
|
|
|
|
case isToolResult(msg):
|
|
// Collect all consecutive tool results into a single user message
|
|
// Bedrock requires all tool results for a turn in one message
|
|
var toolResultBlocks []types.ContentBlock
|
|
for i < len(messages) && isToolResult(messages[i]) {
|
|
toolResultBlocks = append(toolResultBlocks, makeToolResultBlock(messages[i]))
|
|
i++
|
|
}
|
|
bedrockMessages = append(bedrockMessages, types.Message{
|
|
Role: types.ConversationRoleUser,
|
|
Content: toolResultBlocks,
|
|
})
|
|
|
|
case msg.Role == "user":
|
|
// Regular user message (no ToolCallID)
|
|
content := buildUserContent(msg)
|
|
bedrockMessages = append(bedrockMessages, types.Message{
|
|
Role: types.ConversationRoleUser,
|
|
Content: content,
|
|
})
|
|
i++
|
|
|
|
case msg.Role == "assistant":
|
|
content := buildAssistantContent(msg)
|
|
bedrockMessages = append(bedrockMessages, types.Message{
|
|
Role: types.ConversationRoleAssistant,
|
|
Content: content,
|
|
})
|
|
i++
|
|
|
|
case msg.Role == "tool" && msg.ToolCallID == "":
|
|
// Tool message without ToolCallID - treat as regular user message
|
|
content := buildUserContent(msg)
|
|
bedrockMessages = append(bedrockMessages, types.Message{
|
|
Role: types.ConversationRoleUser,
|
|
Content: content,
|
|
})
|
|
i++
|
|
|
|
default:
|
|
// Unknown role - skip
|
|
i++
|
|
}
|
|
}
|
|
|
|
return bedrockMessages, systemPrompts
|
|
}
|
|
|
|
// buildUserContent builds Bedrock content blocks for a user message.
|
|
func buildUserContent(msg Message) []types.ContentBlock {
|
|
var content []types.ContentBlock
|
|
|
|
// Add text content
|
|
if msg.Content != "" {
|
|
content = append(content, &types.ContentBlockMemberText{
|
|
Value: msg.Content,
|
|
})
|
|
}
|
|
|
|
// Add images from Media field
|
|
for _, mediaURL := range msg.Media {
|
|
if strings.HasPrefix(mediaURL, "data:image/") {
|
|
// Parse data URL: data:image/jpeg;base64,<data>
|
|
parts := strings.SplitN(mediaURL, ",", 2)
|
|
if len(parts) != 2 {
|
|
continue
|
|
}
|
|
|
|
// Extract media type from "data:image/jpeg;base64"
|
|
mediaType := ""
|
|
header := parts[0]
|
|
if idx := strings.Index(header, "/"); idx != -1 {
|
|
end := strings.Index(header[idx:], ";")
|
|
if end == -1 {
|
|
end = len(header) - idx
|
|
}
|
|
mediaType = header[idx+1 : idx+end]
|
|
}
|
|
|
|
// Verify this is base64 encoded
|
|
if !strings.Contains(header, ";base64") {
|
|
continue // Skip non-base64 encoded data
|
|
}
|
|
|
|
// Map media type to Bedrock format
|
|
var format types.ImageFormat
|
|
switch mediaType {
|
|
case "jpeg", "jpg":
|
|
format = types.ImageFormatJpeg
|
|
case "png":
|
|
format = types.ImageFormatPng
|
|
case "gif":
|
|
format = types.ImageFormatGif
|
|
case "webp":
|
|
format = types.ImageFormatWebp
|
|
default:
|
|
continue // Skip unsupported formats
|
|
}
|
|
|
|
// Check size before decoding to prevent excessive memory allocation
|
|
// Bedrock has a ~20MB request limit; cap decoded images at 10MB
|
|
const maxImageSize = 10 * 1024 * 1024
|
|
decodedLen := base64.StdEncoding.DecodedLen(len(parts[1]))
|
|
if decodedLen > maxImageSize {
|
|
log.Printf("bedrock: skipping image exceeding size limit (%d bytes > %d)", decodedLen, maxImageSize)
|
|
continue
|
|
}
|
|
|
|
// Decode base64 data
|
|
imageData, err := base64.StdEncoding.DecodeString(parts[1])
|
|
if err != nil {
|
|
log.Printf("bedrock: failed to decode base64 image data: %v", err)
|
|
continue
|
|
}
|
|
|
|
content = append(content, &types.ContentBlockMemberImage{
|
|
Value: types.ImageBlock{
|
|
Format: format,
|
|
Source: &types.ImageSourceMemberBytes{
|
|
Value: imageData,
|
|
},
|
|
},
|
|
})
|
|
}
|
|
}
|
|
|
|
// Bedrock requires at least one content block; add empty text if needed
|
|
if len(content) == 0 {
|
|
content = append(content, &types.ContentBlockMemberText{Value: ""})
|
|
}
|
|
|
|
return content
|
|
}
|
|
|
|
// buildAssistantContent builds Bedrock content blocks for an assistant message.
|
|
func buildAssistantContent(msg Message) []types.ContentBlock {
|
|
var content []types.ContentBlock
|
|
|
|
// Add text content if present
|
|
if msg.Content != "" {
|
|
content = append(content, &types.ContentBlockMemberText{
|
|
Value: msg.Content,
|
|
})
|
|
}
|
|
|
|
// Add tool use blocks
|
|
for _, tc := range msg.ToolCalls {
|
|
// Validate tool call ID - Bedrock requires non-empty ToolUseId
|
|
if strings.TrimSpace(tc.ID) == "" {
|
|
log.Printf("bedrock: skipping tool call with empty ID (name: %q)", tc.Name)
|
|
continue
|
|
}
|
|
|
|
// Resolve tool name: prefer tc.Name, fallback to tc.Function.Name
|
|
// (tc.Name/tc.Arguments are json:"-" and may be empty when from JSON)
|
|
toolName := tc.Name
|
|
if toolName == "" && tc.Function != nil {
|
|
toolName = tc.Function.Name
|
|
}
|
|
if strings.TrimSpace(toolName) == "" {
|
|
continue
|
|
}
|
|
|
|
// Resolve arguments: prefer tc.Arguments, fallback to parsing tc.Function.Arguments
|
|
args := tc.Arguments
|
|
if args == nil && tc.Function != nil && tc.Function.Arguments != "" {
|
|
if err := json.Unmarshal([]byte(tc.Function.Arguments), &args); err != nil {
|
|
log.Printf("bedrock: failed to parse Function.Arguments for tool %q: %v", toolName, err)
|
|
args = map[string]any{}
|
|
}
|
|
}
|
|
if args == nil {
|
|
args = map[string]any{}
|
|
}
|
|
|
|
// Convert arguments to a Bedrock document using NewLazyDocument
|
|
inputDoc := document.NewLazyDocument(args)
|
|
|
|
content = append(content, &types.ContentBlockMemberToolUse{
|
|
Value: types.ToolUseBlock{
|
|
ToolUseId: aws.String(tc.ID),
|
|
Name: aws.String(toolName),
|
|
Input: inputDoc,
|
|
},
|
|
})
|
|
}
|
|
|
|
// Bedrock requires at least one content block; add empty text if needed
|
|
if len(content) == 0 {
|
|
content = append(content, &types.ContentBlockMemberText{Value: ""})
|
|
}
|
|
|
|
return content
|
|
}
|
|
|
|
// convertTools converts tool definitions to Bedrock format.
|
|
func convertTools(tools []ToolDefinition) *types.ToolConfiguration {
|
|
bedrockTools := make([]types.Tool, 0, len(tools))
|
|
|
|
for _, tool := range tools {
|
|
// Skip tools with empty names
|
|
if strings.TrimSpace(tool.Function.Name) == "" {
|
|
continue
|
|
}
|
|
|
|
// Ensure parameters is not nil - default to minimal object schema
|
|
params := tool.Function.Parameters
|
|
if params == nil {
|
|
params = map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{},
|
|
}
|
|
}
|
|
|
|
// Convert parameters schema to a Bedrock document
|
|
inputSchema := document.NewLazyDocument(params)
|
|
|
|
bedrockTools = append(bedrockTools, &types.ToolMemberToolSpec{
|
|
Value: types.ToolSpecification{
|
|
Name: aws.String(tool.Function.Name),
|
|
Description: aws.String(tool.Function.Description),
|
|
InputSchema: &types.ToolInputSchemaMemberJson{
|
|
Value: inputSchema,
|
|
},
|
|
},
|
|
})
|
|
}
|
|
|
|
return &types.ToolConfiguration{
|
|
Tools: bedrockTools,
|
|
}
|
|
}
|
|
|
|
// parseResponse converts Bedrock Converse output to LLMResponse.
|
|
func parseResponse(output *bedrockruntime.ConverseOutput) (*LLMResponse, error) {
|
|
var content strings.Builder
|
|
toolCalls := make([]ToolCall, 0)
|
|
|
|
// Process output content blocks
|
|
if output.Output != nil {
|
|
if msgOutput, ok := output.Output.(*types.ConverseOutputMemberMessage); ok {
|
|
for _, block := range msgOutput.Value.Content {
|
|
switch b := block.(type) {
|
|
case *types.ContentBlockMemberText:
|
|
content.WriteString(b.Value)
|
|
|
|
case *types.ContentBlockMemberToolUse:
|
|
// Unmarshal the document interface to a map
|
|
args := make(map[string]any)
|
|
if b.Value.Input != nil {
|
|
if err := b.Value.Input.UnmarshalSmithyDocument(&args); err != nil {
|
|
log.Printf("bedrock: failed to unmarshal tool input for tool %q (id %q): %v",
|
|
aws.ToString(b.Value.Name),
|
|
aws.ToString(b.Value.ToolUseId),
|
|
err,
|
|
)
|
|
args = make(map[string]any)
|
|
}
|
|
}
|
|
|
|
// Serialize arguments to JSON string for FunctionCall
|
|
argsJSON, err := json.Marshal(args)
|
|
if err != nil {
|
|
log.Printf("bedrock: failed to marshal tool arguments for tool %q (id %q): %v",
|
|
aws.ToString(b.Value.Name),
|
|
aws.ToString(b.Value.ToolUseId),
|
|
err,
|
|
)
|
|
argsJSON = []byte("{}")
|
|
}
|
|
|
|
toolCalls = append(toolCalls, ToolCall{
|
|
ID: aws.ToString(b.Value.ToolUseId),
|
|
Name: aws.ToString(b.Value.Name),
|
|
Arguments: args,
|
|
Function: &FunctionCall{
|
|
Name: aws.ToString(b.Value.Name),
|
|
Arguments: string(argsJSON),
|
|
},
|
|
})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Map stop reason
|
|
finishReason := "stop"
|
|
switch output.StopReason {
|
|
case types.StopReasonToolUse:
|
|
finishReason = "tool_calls"
|
|
case types.StopReasonMaxTokens:
|
|
finishReason = "length"
|
|
case types.StopReasonEndTurn:
|
|
finishReason = "stop"
|
|
case types.StopReasonStopSequence:
|
|
finishReason = "stop"
|
|
case types.StopReasonContentFiltered:
|
|
finishReason = "content_filter"
|
|
}
|
|
|
|
// Build usage info
|
|
var usage *UsageInfo
|
|
if output.Usage != nil {
|
|
usage = &UsageInfo{
|
|
PromptTokens: int(aws.ToInt32(output.Usage.InputTokens)),
|
|
CompletionTokens: int(aws.ToInt32(output.Usage.OutputTokens)),
|
|
TotalTokens: int(aws.ToInt32(output.Usage.InputTokens)) + int(aws.ToInt32(output.Usage.OutputTokens)),
|
|
}
|
|
}
|
|
|
|
return &LLMResponse{
|
|
Content: content.String(),
|
|
ToolCalls: toolCalls,
|
|
FinishReason: finishReason,
|
|
Usage: usage,
|
|
}, nil
|
|
}
|
|
|
|
// isSSOTokenError checks if the error is related to expired or invalid AWS SSO tokens.
|
|
// This helps provide actionable guidance when SSO credentials need to be refreshed.
|
|
// Only matches SSO-specific error patterns to avoid misclassifying other AWS credential errors.
|
|
func isSSOTokenError(err error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
lower := strings.ToLower(err.Error())
|
|
|
|
// Check for specific SSO token expiration/refresh-related error patterns (case-insensitive)
|
|
// Avoid matching generic patterns that could match non-SSO AWS errors (e.g., STS ExpiredToken)
|
|
if strings.Contains(lower, "refresh cached sso token") {
|
|
return true
|
|
}
|
|
if strings.Contains(lower, "read cached sso token") {
|
|
return true
|
|
}
|
|
if strings.Contains(lower, "sso oidc") {
|
|
return true
|
|
}
|
|
if strings.Contains(lower, "invalidgrantexception") {
|
|
return true
|
|
}
|
|
|
|
return false
|
|
}
|