eaa9a4a · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

feat: allow to paste bedrock api keys (#2407)

eaa9a4aAndrey Nering7 months ago

internal/agent/coordinator.go+9−5

compare with main

unifiedsplit

@@ -725,7 +725,7 @@ func (c *coordinator) buildAzureProvider(baseURL, apiKey string, headers map[str

725 return azure.New(opts...)

726}

727

728−func (c *coordinator) buildBedrockProvider(headers map[string]string) (fantasy.Provider, error) {

728+func (c *coordinator) buildBedrockProvider(apiKey string, headers map[string]string) (fantasy.Provider, error) {

729 var opts []bedrock.Option

730 if c.cfg.Config().Options.Debug {

731 httpClient := log.NewHTTPClient()

@@ -734,9 +734,13 @@ func (c *coordinator) buildBedrockProvider(headers map[string]string) (fantasy.P

734 if len(headers) > 0 {

735 opts = append(opts, bedrock.WithHeaders(headers))

736 }

737− bearerToken := os.Getenv("AWS_BEARER_TOKEN_BEDROCK")

737+ switch {

738− if bearerToken != "" {

738+ case apiKey != "":

739− opts = append(opts, bedrock.WithAPIKey(bearerToken))

739+ opts = append(opts, bedrock.WithAPIKey(apiKey))

740+ case os.Getenv("AWS_BEARER_TOKEN_BEDROCK") != "":

741+ opts = append(opts, bedrock.WithAPIKey(os.Getenv("AWS_BEARER_TOKEN_BEDROCK")))

742+ default:

743+ // Skip, let the SDK do authentication.

740 }

744 }

741 return bedrock.New(opts...)

745 return bedrock.New(opts...)

742}

746}

@@ -824,7 +828,7 @@ func (c *coordinator) buildProvider(providerCfg config.ProviderConfig, model con

824 case azure.Name:

828 case azure.Name:

825 return c.buildAzureProvider(baseURL, apiKey, headers, providerCfg.ExtraParams)

829 return c.buildAzureProvider(baseURL, apiKey, headers, providerCfg.ExtraParams)

826 case bedrock.Name:

830 case bedrock.Name:

827− return c.buildBedrockProvider(headers)

831+ return c.buildBedrockProvider(apiKey, headers)

828 case google.Name:

832 case google.Name:

829 return c.buildGoogleProvider(baseURL, apiKey, headers)

833 return c.buildGoogleProvider(baseURL, apiKey, headers)

830 case "google-vertex":

834 case "google-vertex":

internal/config/config.go+9

@@ -3,6 +3,7 @@ package config

3import (

4 "cmp"

5 "context"

6+ "errors"

6 "fmt"

7 "fmt"

7 "log/slog"

8 "log/slog"

8 "maps"

9 "maps"

@@ -584,6 +585,14 @@ func (c *ProviderConfig) TestConnection(resolver VariableResolver) error {

584 baseURL, _ := resolver.ResolveValue(c.BaseURL)

585 baseURL, _ := resolver.ResolveValue(c.BaseURL)

585 baseURL = cmp.Or(baseURL, "https://generativelanguage.googleapis.com")

586 baseURL = cmp.Or(baseURL, "https://generativelanguage.googleapis.com")

586 testURL = baseURL + "/v1beta/models?key=" + url.QueryEscape(apiKey)

587 testURL = baseURL + "/v1beta/models?key=" + url.QueryEscape(apiKey)

588+ case catwalk.TypeBedrock:

589+ // NOTE: Bedrock has a `/foundation-models` endpoint that we could in

590+ // theory use, but apparently the authorization is region-specific,

591+ // so it's not so trivial.

592+ if strings.HasPrefix(apiKey, "ABSK") { // Bedrock API keys

593+ return nil

594+ }

595+ return errors.New("not a valid bedrock api key")

587 }

596 }

588

597

589 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

598 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)