repo_description: Giving Kubernetes Superpowers to everyone
repo_topics: ai devops kubernetes llama openai sre tooling
code: package ai
  
  import (
  	"context"
  	"errors"
  	"fmt"
  	"github.com/aws/aws-sdk-go-v2/aws"
  	awsconfig "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"
  	"os"
  	"strings"
  )
  
  const amazonBedrockConverseClientName = "amazonbedrockconverse"
  
  type bedrockConverseAPI interface {
  	Converse(ctx context.Context, input *bedrockruntime.ConverseInput, optFns ...func(*bedrockruntime.Options)) (*bedrockruntime.ConverseOutput, error)
  }
  
  type AmazonBedrockConverseClient struct {
  	nopCloser
  
  	client        bedrockConverseAPI
  	model         string
  	temperature   float32
  	topP          float32
  	maxTokens     int
  	stopSequences []string
  }
  
  func getRegion(region string) string {
  	if os.Getenv("AWS_DEFAULT_REGION") != "" {
  		region = os.Getenv("AWS_DEFAULT_REGION")
  	}
  	// Return the supplied provider region if not overridden by environment variable
  	return region
  }
  
  func (a *AmazonBedrockConverseClient) getModelFromString(model string) (string, error) {
  	if model == "" {
  		return "", errors.New("model name cannot be empty")
  	}
  	model = strings.TrimSpace(model)
  
  	return model, nil
  }
  
  func processError(err error, modelId string) error {
  	errMsg := err.Error()
  	if strings.Contains(errMsg, "no such host") {
  		return fmt.Errorf(`the bedrock service is not available in the selected region.
                      please double-check the service availability for your region at
                      https://aws.amazon.com/about-aws/global-infrastructure/regional-product-services/`)
  	} else if strings.Contains(errMsg, "Could not resolve the foundation model") {
  		return fmt.Errorf(`could not resolve the foundation model from model identifier: \"%s\".
                      please verify that the requested model exists and is accessible
                      within the specified region`, modelId)
  	} else {
  		return fmt.Errorf("could not invoke model: \"%s\". here is why: %s", modelId, err)
  	}
  }
  
  func (a *AmazonBedrockConverseClient) Configure(config IAIConfig) error {
  	modelInput := config.GetModel()
  
  	var region = getRegion(config.GetProviderRegion())
  
  	// Only create AWS clients if they haven't been injected (for testing)
  	if a.client == nil {
  		cfg, err := awsconfig.LoadDefaultConfig(context.Background(),
  			awsconfig.WithRegion(region),
  		)
  		if err != nil {
  			if strings.Contains(err.Error(), "InvalidAccessKeyId") || strings.Contains(err.Error(), "SignatureDoesNotMatch") || strings.Contains(err.Error(), "NoCredentialProviders") {
  				return fmt.Errorf("aws credentials are invalid or missing. Please check your environment variables or aws config. details: %v", err)
  			}
  			return fmt.Errorf("failed to load aws config for region %s: %w", region, err)
  		}
  
  		a.client = bedrockruntime.NewFromConfig(cfg)
  	}
  
  	foundModel, err := a.getModelFromString(modelInput)
  	if err != nil {
  		return fmt.Errorf("failed to find model configuration for %s: %w", modelInput, err)
  	}
  	a.model = foundModel
  
  	// Set common configuration parameters
  	a.temperature = config.GetTemperature()
  	a.topP = config.GetTopP()
  	a.maxTokens = config.GetMaxTokens()
  	a.stopSequences = config.GetStopSequences()
  
  	return nil
  }
  
  func extractTextFromConverseOutput(output types.ConverseOutput, modelId string) (string, error) {
  	if output == nil {
  		return "", fmt.Errorf("empty response from model: %s", modelId)
  	}
  
  	msg, ok := output.(*types.ConverseOutputMemberMessage)
  	if !ok {
  		return "", fmt.Errorf("unexpected response type from model: %s", modelId)
  	}
  
  	if len(msg.Value.Content) == 0 {
  		return "", fmt.Errorf("no content returned from model: %s", modelId)
  	}
  
  	var builder strings.Builder
  
  	for _, block := range msg.Value.Content {
  		if textBlock, ok := block.(*types.ContentBlockMemberText); ok && textBlock != nil {
  			builder.WriteString(textBlock.Value)
  		}
  	}
  
  	if builder.Len() == 0 {
  		return "", fmt.Errorf("no text content returned from model: %s", modelId)
  	}
  
  	return builder.String(), nil
  }
  
  func (a *AmazonBedrockConverseClient) GetCompletion(ctx context.Context, prompt string) (string, error) {
  	var content = types.ContentBlockMemberText{
  		Value: prompt,
  	}
  	var message = types.Message{
  		Content: []types.ContentBlock{&content},
  		Role:    "user",
  	}
  	var converseInput = bedrockruntime.ConverseInput{
  		ModelId:  aws.String(a.model),
  		Messages: []types.Message{message},
  		InferenceConfig: &types.InferenceConfiguration{
  			Temperature:   aws.Float32(a.temperature),
  			TopP:          aws.Float32(a.topP),
  			MaxTokens:     aws.Int32(int32(a.maxTokens)),
  			StopSequences: a.stopSequences,
  		},
  	}
  	response, err := a.client.Converse(ctx, &converseInput)
  	if err != nil {
  		return "", processError(err, a.model)
  	}
  
  	text, err := extractTextFromConverseOutput(response.Output, a.model)
  	if err != nil {
  		return "", err
  	}
  
  	return text, nil
  }
  
  func (a *AmazonBedrockConverseClient) GetName() string {
  	return amazonBedrockConverseClientName
  }
  