1package provider
2
3import (
4 "context"
5 "errors"
6 "fmt"
7 "os"
8 "strings"
9
10 "github.com/kujtimiihoxha/termai/internal/llm/models"
11 "github.com/kujtimiihoxha/termai/internal/llm/tools"
12 "github.com/kujtimiihoxha/termai/internal/message"
13)
14
15type bedrockProvider struct {
16 childProvider Provider
17 model models.Model
18 maxTokens int64
19 systemMessage string
20}
21
22func (b *bedrockProvider) SendMessages(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
23 return b.childProvider.SendMessages(ctx, messages, tools)
24}
25
26func (b *bedrockProvider) StreamResponse(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (<-chan ProviderEvent, error) {
27 return b.childProvider.StreamResponse(ctx, messages, tools)
28}
29
30func NewBedrockProvider(opts ...BedrockOption) (Provider, error) {
31 provider := &bedrockProvider{}
32 for _, opt := range opts {
33 opt(provider)
34 }
35
36 // based on the AWS region prefix the model name with, us, eu, ap, sa, etc.
37 region := os.Getenv("AWS_REGION")
38 if region == "" {
39 region = os.Getenv("AWS_DEFAULT_REGION")
40 }
41
42 if region == "" {
43 region = "us-east-1" // default region
44 }
45 if len(region) < 2 {
46 return nil, errors.New("AWS_REGION or AWS_DEFAULT_REGION environment variable is invalid")
47 }
48 regionPrefix := region[:2]
49 provider.model.APIModel = fmt.Sprintf("%s.%s", regionPrefix, provider.model.APIModel)
50
51 if strings.Contains(string(provider.model.APIModel), "anthropic") {
52 anthropic, err := NewAnthropicProvider(
53 WithAnthropicModel(provider.model),
54 WithAnthropicMaxTokens(provider.maxTokens),
55 WithAnthropicSystemMessage(provider.systemMessage),
56 WithAnthropicBedrock(),
57 WithAnthropicDisableCache(),
58 )
59 provider.childProvider = anthropic
60 if err != nil {
61 return nil, err
62 }
63 } else {
64 return nil, errors.New("unsupported model for bedrock provider")
65 }
66 return provider, nil
67}
68
69type BedrockOption func(*bedrockProvider)
70
71func WithBedrockSystemMessage(message string) BedrockOption {
72 return func(a *bedrockProvider) {
73 a.systemMessage = message
74 }
75}
76
77func WithBedrockMaxTokens(maxTokens int64) BedrockOption {
78 return func(a *bedrockProvider) {
79 a.maxTokens = maxTokens
80 }
81}
82
83func WithBedrockModel(model models.Model) BedrockOption {
84 return func(a *bedrockProvider) {
85 a.model = model
86 }
87}