939ae03 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

add bedrock support

939ae03Kujtim Hoxha2 years ago

go.mod+14

compare with main

unifiedsplit

@@ -47,6 +47,20 @@ require (

47 github.com/alecthomas/chroma/v2 v2.15.0 // indirect

48 github.com/andybalholm/cascadia v1.3.2 // indirect

49 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

50 github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect

64 github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect

51 github.com/aymerick/douceur v0.2.0 // indirect

65 github.com/aymerick/douceur v0.2.0 // indirect

52 github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect

66 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

28github.com/anthropics/anthropic-sdk-go v0.2.0-beta.2/go.mod h1:AapDW22irxK2PSumZiQXYUFvsdQgkwIWlpESweWZI/c=

29github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=

30github.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=

31github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=

59github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=

32github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=

60github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=

33github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWpi6yML8=

61github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWpi6yML8=

internal/config/config.go+7

@@ -36,6 +36,11 @@ type Model struct {

36 // TODO: Maybe support multiple models for different purposes

37}

39+type AnthropicConfig struct {

40+ DisableCache bool `json:"disableCache"`

41+ UseBedrock bool `json:"useBedrock"`

42+}

43+

39type Provider struct {

44type Provider struct {

40 APIKey string `json:"apiKey"`

45 APIKey string `json:"apiKey"`

41 Enabled bool `json:"enabled"`

46 Enabled bool `json:"enabled"`

@@ -130,6 +135,8 @@ func Load(debug bool) error {

130 defaultModelSet = true

135 defaultModelSet = true

131 }

136 }

132 }

137 }

138+

139+ viper.SetDefault("providers.bedrock.enabled", true)

133 // TODO: add more providers

140 // TODO: add more providers

134 cfg = &Config{}

141 cfg = &Config{}

135

142

internal/llm/agent/agent.go+23

@@ -380,6 +380,29 @@ func getAgentProviders(ctx context.Context, model models.Model) (provider.Provid

380 return nil, nil, err

381 }

382

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+

383 }

406 }

384

407

385 return agentProvider, titleGenerator, nil

408 return agentProvider, titleGenerator, nil

internal/llm/models/models.go+16

@@ -31,11 +31,15 @@ const (

32 // GROQ

33 QWENQwq ModelID = "qwen-qwq"

34+

35+ // Bedrock

36+ BedrockClaude37Sonnet ModelID = "bedrock.claude-3.7-sonnet"

34)

37)

36const (

39const (

37 ProviderOpenAI ModelProvider = "openai"

40 ProviderOpenAI ModelProvider = "openai"

38 ProviderAnthropic ModelProvider = "anthropic"

41 ProviderAnthropic ModelProvider = "anthropic"

42+ ProviderBedrock ModelProvider = "bedrock"

39 ProviderGemini ModelProvider = "gemini"

43 ProviderGemini ModelProvider = "gemini"

40 ProviderGROQ ModelProvider = "groq"

44 ProviderGROQ ModelProvider = "groq"

41)

45)

@@ -119,4 +123,16 @@ var SupportedModels = map[ModelID]Model{

119 CostPer1MOutCached: 0,

123 CostPer1MOutCached: 0,

120 CostPer1MOut: 0,

124 CostPer1MOut: 0,

121 },

125 },

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+ },

122}

138}

internal/llm/provider/anthropic.go+28−5

@@ -9,6 +9,7 @@ import (

9 "time"

11 "github.com/anthropics/anthropic-sdk-go"

12+ "github.com/anthropics/anthropic-sdk-go/bedrock"

12 "github.com/anthropics/anthropic-sdk-go/option"

13 "github.com/anthropics/anthropic-sdk-go/option"

13 "github.com/kujtimiihoxha/termai/internal/llm/models"

14 "github.com/kujtimiihoxha/termai/internal/llm/models"

14 "github.com/kujtimiihoxha/termai/internal/llm/tools"

15 "github.com/kujtimiihoxha/termai/internal/llm/tools"

@@ -21,6 +22,8 @@ type anthropicProvider struct {

21 maxTokens int64

22 maxTokens int64

22 apiKey string

23 apiKey string

23 systemMessage string

24 systemMessage string

25+ useBedrock bool

26+ disableCache bool

24}

27}

26type AnthropicOption func(*anthropicProvider)

29type AnthropicOption func(*anthropicProvider)

@@ -49,6 +52,18 @@ func WithAnthropicKey(apiKey string) AnthropicOption {

49 }

52 }

50}

53}

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+

52func NewAnthropicProvider(opts ...AnthropicOption) (Provider, error) {

67func NewAnthropicProvider(opts ...AnthropicOption) (Provider, error) {

53 provider := &anthropicProvider{

68 provider := &anthropicProvider{

54 maxTokens: 1024,

69 maxTokens: 1024,

@@ -62,7 +77,16 @@ func NewAnthropicProvider(opts ...AnthropicOption) (Provider, error) {

62 return nil, errors.New("system message is required")

77 return nil, errors.New("system message is required")

63 }

78 }

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+

65− provider.client = anthropic.NewClient(option.WithAPIKey(provider.apiKey))

89+ provider.client = anthropic.NewClient(anthropicOptions...)

66 return provider, nil

90 return provider, nil

67}

91}

@@ -338,7 +362,7 @@ func (a *anthropicProvider) convertToAnthropicTools(tools []tools.BaseTool) []an

338 },

362 },

339 }

363 }

340

364

341− if i == len(tools)-1 {

365+ if i == len(tools)-1 && !a.disableCache {

342 toolParam.CacheControl = anthropic.CacheControlEphemeralParam{

366 toolParam.CacheControl = anthropic.CacheControlEphemeralParam{

343 Type: "ephemeral",

367 Type: "ephemeral",

344 }

368 }

@@ -358,7 +382,7 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag

358 switch msg.Role {

382 switch msg.Role {

359 case message.User:

383 case message.User:

360 content := anthropic.NewTextBlock(msg.Content().String())

384 content := anthropic.NewTextBlock(msg.Content().String())

361− if cachedBlocks < 2 {

385+ if cachedBlocks < 2 && !a.disableCache {

362 content.OfRequestTextBlock.CacheControl = anthropic.CacheControlEphemeralParam{

386 content.OfRequestTextBlock.CacheControl = anthropic.CacheControlEphemeralParam{

363 Type: "ephemeral",

387 Type: "ephemeral",

364 }

388 }

@@ -370,7 +394,7 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag

370 blocks := []anthropic.ContentBlockParamUnion{}

394 blocks := []anthropic.ContentBlockParamUnion{}

371 if msg.Content().String() != "" {

395 if msg.Content().String() != "" {

372 content := anthropic.NewTextBlock(msg.Content().String())

396 content := anthropic.NewTextBlock(msg.Content().String())

373− if cachedBlocks < 2 {

397+ if cachedBlocks < 2 && !a.disableCache {

374 content.OfRequestTextBlock.CacheControl = anthropic.CacheControlEphemeralParam{

398 content.OfRequestTextBlock.CacheControl = anthropic.CacheControlEphemeralParam{

375 Type: "ephemeral",

399 Type: "ephemeral",

376 }

400 }

@@ -404,4 +428,3 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag

404

428

405 return anthropicMessages

429 return anthropicMessages

406}

430}

407−

internal/llm/provider/bedrock.gonew file+87

@@ -0,0 +1,87 @@

1+package provider

3+import (

4+ "context"

5+ "errors"

6+ "fmt"

7+ "os"

8+ "strings"

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 @@

1package repl

3import (

4+ "log"

4 "strings"

5 "strings"

6 "github.com/charmbracelet/bubbles/key"

7 "github.com/charmbracelet/bubbles/key"

@@ -138,11 +139,22 @@ func (m *editorCmp) SetSize(width int, height int) {

138

139

139func (m *editorCmp) Send() tea.Cmd {

140func (m *editorCmp) Send() tea.Cmd {

140 return func() tea.Msg {

141 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+ }

142 if hasUnfinishedMessages(messages) {

149 if hasUnfinishedMessages(messages) {

143 return util.ReportWarn("Assistant is still working on the previous message")

150 return util.ReportWarn("Assistant is still working on the previous message")

144 }

151 }

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+ }

146

158

147 content := strings.Join(m.editor.GetBuffer().Lines(), "\n")

159 content := strings.Join(m.editor.GetBuffer().Lines(), "\n")

148 go a.Generate(m.sessionID, content)

160 go a.Generate(m.sessionID, content)