git.iko.soy/crush
on main help sign in
add bedrock support
commit939ae03f42e61d0944da80381219e6bbdfc2d850 authorKujtim Hoxha <kujtimii.h@gmail.com> dateWed 9 Apr 2025, 17:45 (2 years ago) parentfde04bb
8 files changed+217−7
compare with main
unifiedsplit
go.mod+14 go.sum+28 internal/config/config.go+7 internal/llm/agent/agent.go+23 internal/llm/models/models.go+16 internal/llm/provider/anthropic.go+28−5 internal/llm/provider/bedrock.go+87 internal/tui/components/repl/editor.go+14−2
go.mod+14
@@ -47,6 +47,20 @@ require (
4747 github.com/alecthomas/chroma/v2 v2.15.0 // indirect
4848 github.com/andybalholm/cascadia v1.3.2 // indirect
4949 github.com/atotto/clipboard v0.1.4 // indirect
50+ github.com/aws/aws-sdk-go-v2 v1.30.3 // indirect
51+ github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3 // indirect
52+ github.com/aws/aws-sdk-go-v2/config v1.27.27 // indirect
53+ github.com/aws/aws-sdk-go-v2/credentials v1.17.27 // indirect
54+ github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11 // indirect
55+ github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15 // indirect
56+ github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15 // indirect
57+ github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect
58+ github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3 // indirect
59+ github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17 // indirect
60+ github.com/aws/aws-sdk-go-v2/service/sso v1.22.4 // indirect
61+ github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 // indirect
62+ github.com/aws/aws-sdk-go-v2/service/sts v1.30.3 // indirect
63+ github.com/aws/smithy-go v1.20.3 // indirect
5064 github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
5165 github.com/aymerick/douceur v0.2.0 // indirect
5266 github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect
go.sum+28
@@ -28,6 +28,34 @@ github.com/anthropics/anthropic-sdk-go v0.2.0-beta.2 h1:h7qxtumNjKPWFv1QM/HJy60M
2828github.com/anthropics/anthropic-sdk-go v0.2.0-beta.2/go.mod h1:AapDW22irxK2PSumZiQXYUFvsdQgkwIWlpESweWZI/c=
2929github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
3030github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
31+github.com/aws/aws-sdk-go-v2 v1.30.3 h1:jUeBtG0Ih+ZIFH0F4UkmL9w3cSpaMv9tYYDbzILP8dY=
32+github.com/aws/aws-sdk-go-v2 v1.30.3/go.mod h1:nIQjQVp5sfpQcTc9mPSr1B0PaWK5ByX9MOoDadSN4lc=
33+github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3 h1:tW1/Rkad38LA15X4UQtjXZXNKsCgkshC3EbmcUmghTg=
34+github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3/go.mod h1:UbnqO+zjqk3uIt9yCACHJ9IVNhyhOCnYk8yA19SAWrM=
35+github.com/aws/aws-sdk-go-v2/config v1.27.27 h1:HdqgGt1OAP0HkEDDShEl0oSYa9ZZBSOmKpdpsDMdO90=
36+github.com/aws/aws-sdk-go-v2/config v1.27.27/go.mod h1:MVYamCg76dFNINkZFu4n4RjDixhVr51HLj4ErWzrVwg=
37+github.com/aws/aws-sdk-go-v2/credentials v1.17.27 h1:2raNba6gr2IfA0eqqiP2XiQ0UVOpGPgDSi0I9iAP+UI=
38+github.com/aws/aws-sdk-go-v2/credentials v1.17.27/go.mod h1:gniiwbGahQByxan6YjQUMcW4Aov6bLC3m+evgcoN4r4=
39+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11 h1:KreluoV8FZDEtI6Co2xuNk/UqI9iwMrOx/87PBNIKqw=
40+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11/go.mod h1:SeSUYBLsMYFoRvHE0Tjvn7kbxaUhl75CJi1sbfhMxkU=
41+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15 h1:SoNJ4RlFEQEbtDcCEt+QG56MY4fm4W8rYirAmq+/DdU=
42+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15/go.mod h1:U9ke74k1n2bf+RIgoX1SXFed1HLs51OgUSs+Ph0KJP8=
43+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15 h1:C6WHdGnTDIYETAm5iErQUiVNsclNx9qbJVPIt03B6bI=
44+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15/go.mod h1:ZQLZqhcu+JhSrA9/NXRm8SkDvsycE+JkV3WGY41e+IM=
45+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 h1:hT8rVHwugYE2lEfdFE0QWVo81lF7jMrYJVDWI+f+VxU=
46+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY=
47+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3 h1:dT3MqvGhSoaIhRseqw2I0yH81l7wiR2vjs57O51EAm8=
48+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3/go.mod h1:GlAeCkHwugxdHaueRr4nhPuY+WW+gR8UjlcqzPr1SPI=
49+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17 h1:HGErhhrxZlQ044RiM+WdoZxp0p+EGM62y3L6pwA4olE=
50+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17/go.mod h1:RkZEx4l0EHYDJpWppMJ3nD9wZJAa8/0lq9aVC+r2UII=
51+github.com/aws/aws-sdk-go-v2/service/sso v1.22.4 h1:BXx0ZIxvrJdSgSvKTZ+yRBeSqqgPM89VPlulEcl37tM=
52+github.com/aws/aws-sdk-go-v2/service/sso v1.22.4/go.mod h1:ooyCOXjvJEsUw7x+ZDHeISPMhtwI3ZCB7ggFMcFfWLU=
53+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 h1:yiwVzJW2ZxZTurVbYWA7QOrAaCYQR72t0wrSBfoesUE=
54+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4/go.mod h1:0oxfLkpz3rQ/CHlx5hB7H69YUpFiI1tql6Q6Ne+1bCw=
55+github.com/aws/aws-sdk-go-v2/service/sts v1.30.3 h1:ZsDKRLXGWHk8WdtyYMoGNO7bTudrvuKpDKgMVRlepGE=
56+github.com/aws/aws-sdk-go-v2/service/sts v1.30.3/go.mod h1:zwySh8fpFyXp9yOr/KVzxOl8SRqgf/IDw5aUt9UKFcQ=
57+github.com/aws/smithy-go v1.20.3 h1:ryHwveWzPV5BIof6fyDvor6V3iUL7nTfiTKXHiW05nE=
58+github.com/aws/smithy-go v1.20.3/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E=
3159github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
3260github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
3361github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWpi6yML8=
internal/config/config.go+7
@@ -36,6 +36,11 @@ type Model struct {
3636 // TODO: Maybe support multiple models for different purposes
3737}
3838
39+type AnthropicConfig struct {
40+ DisableCache bool `json:"disableCache"`
41+ UseBedrock bool `json:"useBedrock"`
42+}
43+
3944type Provider struct {
4045 APIKey string `json:"apiKey"`
4146 Enabled bool `json:"enabled"`
@@ -130,6 +135,8 @@ func Load(debug bool) error {
130135 defaultModelSet = true
131136 }
132137 }
138+
139+ viper.SetDefault("providers.bedrock.enabled", true)
133140 // TODO: add more providers
134141 cfg = &Config{}
135142
internal/llm/agent/agent.go+23
@@ -380,6 +380,29 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid
380380 return nil, nil, err
381381 }
382382
383+ case models.ProviderBedrock:
384+ var err error
385+ agentProvider, err = provider.NewBedrockProvider(
386+ provider.WithBedrockSystemMessage(
387+ prompt.CoderAnthropicSystemPrompt(),
388+ ),
389+ provider.WithBedrockMaxTokens(maxTokens),
390+ provider.WithBedrockModel(model),
391+ )
392+ if err != nil {
393+ return nil, nil, err
394+ }
395+ titleGenerator, err = provider.NewBedrockProvider(
396+ provider.WithBedrockSystemMessage(
397+ prompt.TitlePrompt(),
398+ ),
399+ provider.WithBedrockMaxTokens(maxTokens),
400+ provider.WithBedrockModel(model),
401+ )
402+ if err != nil {
403+ return nil, nil, err
404+ }
405+
383406 }
384407
385408 return agentProvider, titleGenerator, nil
internal/llm/models/models.go+16
@@ -31,11 +31,15 @@ const (
3131
3232 // GROQ
3333 QWENQwq ModelID = "qwen-qwq"
34+
35+ // Bedrock
36+ BedrockClaude37Sonnet ModelID = "bedrock.claude-3.7-sonnet"
3437)
3538
3639const (
3740 ProviderOpenAI ModelProvider = "openai"
3841 ProviderAnthropic ModelProvider = "anthropic"
42+ ProviderBedrock ModelProvider = "bedrock"
3943 ProviderGemini ModelProvider = "gemini"
4044 ProviderGROQ ModelProvider = "groq"
4145)
@@ -119,4 +123,16 @@ var SupportedModels = map[ModelID]Model{
119123 CostPer1MOutCached: 0,
120124 CostPer1MOut: 0,
121125 },
126+
127+ // Bedrock
128+ BedrockClaude37Sonnet: {
129+ ID: BedrockClaude37Sonnet,
130+ Name: "Bedrock: Claude 3.7 Sonnet",
131+ Provider: ProviderBedrock,
132+ APIModel: "anthropic.claude-3-7-sonnet-20250219-v1:0",
133+ CostPer1MIn: 3.0,
134+ CostPer1MInCached: 3.75,
135+ CostPer1MOutCached: 0.30,
136+ CostPer1MOut: 15.0,
137+ },
122138}
internal/llm/provider/anthropic.go+28−5
@@ -9,6 +9,7 @@ import (
99 "time"
1010
1111 "github.com/anthropics/anthropic-sdk-go"
12+ "github.com/anthropics/anthropic-sdk-go/bedrock"
1213 "github.com/anthropics/anthropic-sdk-go/option"
1314 "github.com/kujtimiihoxha/termai/internal/llm/models"
1415 "github.com/kujtimiihoxha/termai/internal/llm/tools"
@@ -21,6 +22,8 @@ type anthropicProvider struct {
2122 maxTokens int64
2223 apiKey string
2324 systemMessage string
25+ useBedrock bool
26+ disableCache bool
2427}
2528
2629type AnthropicOption func(*anthropicProvider)
@@ -49,6 +52,18 @@ func WithAnthropicKey(apiKey string) AnthropicOption {
4952 }
5053}
5154
55+func WithAnthropicBedrock() AnthropicOption {
56+ return func(a *anthropicProvider) {
57+ a.useBedrock = true
58+ }
59+}
60+
61+func WithAnthropicDisableCache() AnthropicOption {
62+ return func(a *anthropicProvider) {
63+ a.disableCache = true
64+ }
65+}
66+
5267func NewAnthropicProvider(opts ...AnthropicOption) (Provider, error) {
5368 provider := &anthropicProvider{
5469 maxTokens: 1024,
@@ -62,7 +77,16 @@ func NewAnthropicProvider(opts ...AnthropicOption) (Provider, error) {
6277 return nil, errors.New("system message is required")
6378 }
6479
65− provider.client = anthropic.NewClient(option.WithAPIKey(provider.apiKey))
80+ anthropicOptions := []option.RequestOption{}
81+
82+ if provider.apiKey != "" {
83+ anthropicOptions = append(anthropicOptions, option.WithAPIKey(provider.apiKey))
84+ }
85+ if provider.useBedrock {
86+ anthropicOptions = append(anthropicOptions, bedrock.WithLoadDefaultConfig(context.Background()))
87+ }
88+
89+ provider.client = anthropic.NewClient(anthropicOptions...)
6690 return provider, nil
6791}
6892
@@ -338,7 +362,7 @@ func (a *anthropicProvider) convertToAnthropicTools(tools []tools.BaseTool) []an
338362 },
339363 }
340364
341− if i == len(tools)-1 {
365+ if i == len(tools)-1 && !a.disableCache {
342366 toolParam.CacheControl = anthropic.CacheControlEphemeralParam{
343367 Type: "ephemeral",
344368 }
@@ -358,7 +382,7 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag
358382 switch msg.Role {
359383 case message.User:
360384 content := anthropic.NewTextBlock(msg.Content().String())
361− if cachedBlocks < 2 {
385+ if cachedBlocks < 2 && !a.disableCache {
362386 content.OfRequestTextBlock.CacheControl = anthropic.CacheControlEphemeralParam{
363387 Type: "ephemeral",
364388 }
@@ -370,7 +394,7 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag
370394 blocks := []anthropic.ContentBlockParamUnion{}
371395 if msg.Content().String() != "" {
372396 content := anthropic.NewTextBlock(msg.Content().String())
373− if cachedBlocks < 2 {
397+ if cachedBlocks < 2 && !a.disableCache {
374398 content.OfRequestTextBlock.CacheControl = anthropic.CacheControlEphemeralParam{
375399 Type: "ephemeral",
376400 }
@@ -404,4 +428,3 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag
404428
405429 return anthropicMessages
406430}
407−
internal/llm/provider/bedrock.gonew file+87
@@ -0,0 +1,87 @@
1+package provider
2+
3+import (
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+
15+type bedrockProvider struct {
16+ childProvider Provider
17+ model models.Model
18+ maxTokens int64
19+ systemMessage string
20+}
21+
22+func (b *bedrockProvider) SendMessages(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
23+ return b.childProvider.SendMessages(ctx, messages, tools)
24+}
25+
26+func (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+
30+func 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+ return nil, errors.New("AWS_REGION or AWS_DEFAULT_REGION environment variable is required")
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+
69+type BedrockOption func(*bedrockProvider)
70+
71+func WithBedrockSystemMessage(message string) BedrockOption {
72+ return func(a *bedrockProvider) {
73+ a.systemMessage = message
74+ }
75+}
76+
77+func WithBedrockMaxTokens(maxTokens int64) BedrockOption {
78+ return func(a *bedrockProvider) {
79+ a.maxTokens = maxTokens
80+ }
81+}
82+
83+func WithBedrockModel(model models.Model) BedrockOption {
84+ return func(a *bedrockProvider) {
85+ a.model = model
86+ }
87+}
internal/tui/components/repl/editor.go+14−2
@@ -1,6 +1,7 @@
11package repl
22
33import (
4+ "log"
45 "strings"
56
67 "github.com/charmbracelet/bubbles/key"
@@ -138,11 +139,22 @@ func (m *editorCmp) SetSize(width int, height int) {
138139
139140func (m *editorCmp) Send() tea.Cmd {
140141 return func() tea.Msg {
141− messages, _ := m.app.Messages.List(m.sessionID)
142+ messages, err := m.app.Messages.List(m.sessionID)
143+ log.Printf("error: %v", err)
144+ log.Printf("messages: %v", messages)
145+
146+ if err != nil {
147+ return util.ReportError(err)
148+ }
142149 if hasUnfinishedMessages(messages) {
143150 return util.ReportWarn("Assistant is still working on the previous message")
144151 }
145− a, _ := agent.NewCoderAgent(m.app)
152+ a, err := agent.NewCoderAgent(m.app)
153+ log.Printf("error: %v", err)
154+ log.Printf("agent: %v", a)
155+ if err != nil {
156+ return util.ReportError(err)
157+ }
146158
147159 content := strings.Join(m.editor.GetBuffer().Lines(), "\n")
148160 go a.Generate(m.sessionID, content)