19c0e04 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

feat: make it possible for api key to be programatic

19c0e04Kujtim Hoxha1 year ago

cspell.json+1−39

compare with main

unifiedsplit

@@ -1,39 +1 @@

1−{

1+{"language":"en","flagWords":[],"version":"0.2","words":["afero","alecthomas","bubbletea","charmbracelet","charmtone","Charple","crush","diffview","Emph","filepicker","Focusable","fsext","GROQ","Guac","imageorient","Lanczos","lipgloss","lsps","lucasb","nfnt","oksvg","Preproc","rasterx","rivo","Sourcegraph","srwiley","Strikethrough","termenv","textinput","trashhalo","uniseg","Unticked","genai"]}

2− "flagWords": [],

3− "words": [

4− "afero",

5− "alecthomas",

6− "bubbletea",

7− "charmbracelet",

8− "charmtone",

9− "Charple",

10− "crush",

11− "diffview",

12− "Emph",

13− "filepicker",

14− "Focusable",

15− "fsext",

16− "GROQ",

17− "Guac",

18− "imageorient",

19− "Lanczos",

20− "lipgloss",

21− "lsps",

22− "lucasb",

23− "nfnt",

24− "oksvg",

25− "Preproc",

26− "rasterx",

27− "rivo",

28− "Sourcegraph",

29− "srwiley",

30− "Strikethrough",

31− "termenv",

32− "textinput",

33− "trashhalo",

34− "uniseg",

35− "Unticked"

36− ],

37− "version": "0.2",

38− "language": "en"

39−}

internal/config/config.go+3−3

@@ -592,11 +592,11 @@ func mergeProviderConfigs(base, global, local *Config) {

592 if cfg == nil {

593 continue

594 }

595− for providerName, globalProvider := range cfg.Providers {

595+ for providerName, p := range cfg.Providers {

596 if _, ok := base.Providers[providerName]; !ok {

597− base.Providers[providerName] = globalProvider

597+ base.Providers[providerName] = p

598 } else {

599− base.Providers[providerName] = mergeProviderConfig(providerName, base.Providers[providerName], globalProvider)

599+ base.Providers[providerName] = mergeProviderConfig(providerName, base.Providers[providerName], p)

600 }

601 }

602 }

internal/config/shell.gonew file+74

@@ -0,0 +1,74 @@

1+package config

3+import (

4+ "context"

5+ "fmt"

6+ "os"

7+ "strings"

8+ "time"

10+ "github.com/charmbracelet/crush/internal/llm/tools/shell"

11+ "github.com/charmbracelet/crush/internal/logging"

12+)

13+

14+// ExecuteCommand executes a shell command and returns the output

15+// This is a shared utility that can be used by both provider config and tools

16+func ExecuteCommand(ctx context.Context, command string, workingDir string) (string, error) {

17+ if workingDir == "" {

18+ workingDir = WorkingDirectory()

19+ }

20+

21+ persistentShell := shell.GetPersistentShell(workingDir)

22+

23+ stdout, stderr, err := persistentShell.Exec(ctx, command)

24+ if err != nil {

25+ logging.Debug("Command execution failed", "command", command, "error", err, "stderr", stderr)

26+ return "", fmt.Errorf("command execution failed: %w", err)

27+ }

28+

29+ return strings.TrimSpace(stdout), nil

30+}

31+

32+// ResolveAPIKey resolves an API key that can be either:

33+// - A direct string value

34+// - An environment variable (prefixed with $)

35+// - A shell command (wrapped in $(...))

36+func ResolveAPIKey(apiKey string) (string, error) {

37+ if !strings.HasPrefix(apiKey, "$") {

38+ return apiKey, nil

39+ }

40+

41+ if strings.HasPrefix(apiKey, "$(") && strings.HasSuffix(apiKey, ")") {

42+ command := strings.TrimSuffix(strings.TrimPrefix(apiKey, "$("), ")")

43+ logging.Debug("Resolving API key from command", "command", command)

44+ return resolveCommandAPIKey(command)

45+ }

46+

47+ envVar := strings.TrimPrefix(apiKey, "$")

48+ if value := os.Getenv(envVar); value != "" {

49+ logging.Debug("Resolved environment variable", "envVar", envVar, "value", value)

50+ return value, nil

51+ }

52+

53+ logging.Debug("Environment variable not found", "envVar", envVar)

54+

55+ return "", fmt.Errorf("environment variable %s not found", envVar)

56+}

57+

58+// resolveCommandAPIKey executes a command to get an API key, with caching support

59+func resolveCommandAPIKey(command string) (string, error) {

60+ ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)

61+ defer cancel()

62+

63+ logging.Debug("Executing command for API key", "command", command)

64+

65+ workingDir := WorkingDirectory()

66+

67+ result, err := ExecuteCommand(ctx, command, workingDir)

68+ if err != nil {

69+ return "", fmt.Errorf("failed to execute API key command: %w", err)

70+ }

71+ logging.Debug("Command executed successfully", "command", command, "result", result)

72+ return result, nil

73+}

74+

internal/llm/agent/agent.go+2−2

@@ -253,10 +253,10 @@ func (a *agent) IsBusy() bool {

253 if cancelFunc, ok := value.(context.CancelFunc); ok {

254 if cancelFunc != nil {

255 busy = true

256− return false // Stop iterating

256+ return false

257 }

258 }

259− return true // Continue iterating

259+ return true

260 })

261 return busy

262}

internal/llm/provider/anthropic.go+22−10

@@ -21,12 +21,20 @@ import (

22type anthropicClient struct {

23 providerOptions providerClientOptions

24+ useBedrock bool

24 client anthropic.Client

25 client anthropic.Client

25}

26}

27type AnthropicClient ProviderClient

28type AnthropicClient ProviderClient

29func newAnthropicClient(opts providerClientOptions, useBedrock bool) AnthropicClient {

30func newAnthropicClient(opts providerClientOptions, useBedrock bool) AnthropicClient {

31+ return &anthropicClient{

32+ providerOptions: opts,

33+ client: createAnthropicClient(opts, useBedrock),

34+ }

35+}

36+

37+func createAnthropicClient(opts providerClientOptions, useBedrock bool) anthropic.Client {

30 anthropicClientOptions := []option.RequestOption{}

38 anthropicClientOptions := []option.RequestOption{}

31 if opts.apiKey != "" {

39 if opts.apiKey != "" {

32 anthropicClientOptions = append(anthropicClientOptions, option.WithAPIKey(opts.apiKey))

40 anthropicClientOptions = append(anthropicClientOptions, option.WithAPIKey(opts.apiKey))

@@ -34,12 +42,7 @@ func newAnthropicClient(opts providerClientOptions, useBedrock bool) AnthropicCl

34 if useBedrock {

42 if useBedrock {

35 anthropicClientOptions = append(anthropicClientOptions, bedrock.WithLoadDefaultConfig(context.Background()))

43 anthropicClientOptions = append(anthropicClientOptions, bedrock.WithLoadDefaultConfig(context.Background()))

36 }

44 }

37−

38− client := anthropic.NewClient(anthropicClientOptions...)

45+ return anthropic.NewClient(anthropicClientOptions...)

39− return &anthropicClient{

40− providerOptions: opts,

41− client: client,

42− }

43}

46}

45func (a *anthropicClient) convertMessages(messages []message.Message) (anthropicMessages []anthropic.MessageParam) {

48func (a *anthropicClient) convertMessages(messages []message.Message) (anthropicMessages []anthropic.MessageParam) {

@@ -385,12 +388,21 @@ func (a *anthropicClient) stream(ctx context.Context, messages []message.Message

385}

388}

386

389

387func (a *anthropicClient) shouldRetry(attempts int, err error) (bool, int64, error) {

390func (a *anthropicClient) shouldRetry(attempts int, err error) (bool, int64, error) {

388− var apierr *anthropic.Error

391+ var apiErr *anthropic.Error

389− if !errors.As(err, &apierr) {

392+ if !errors.As(err, &apiErr) {

390 return false, 0, err

393 return false, 0, err

391 }

394 }

392

395

396+ if apiErr.StatusCode == 401 {

397+ a.providerOptions.apiKey, err = config.ResolveAPIKey(a.providerOptions.config.APIKey)

398+ if err != nil {

399+ return false, 0, fmt.Errorf("failed to resolve API key: %w", err)

400+ }

401+ a.client = createAnthropicClient(a.providerOptions, a.useBedrock)

402+ return true, 0, nil

403+ }

404+

393− if apierr.StatusCode != 429 && apierr.StatusCode != 529 {

405+ if apiErr.StatusCode != 429 && apiErr.StatusCode != 529 {

394 return false, 0, err

406 return false, 0, err

395 }

407 }

396

408

@@ -399,7 +411,7 @@ func (a *anthropicClient) shouldRetry(attempts int, err error) (bool, int64, err

399 }

411 }

400

412

401 retryMs := 0

413 retryMs := 0

402− retryAfterValues := apierr.Response.Header.Values("Retry-After")

414+ retryAfterValues := apiErr.Response.Header.Values("Retry-After")

403

415

404 backoffMs := 2000 * (1 << (attempts - 1))

416 backoffMs := 2000 * (1 << (attempts - 1))

405 jitterMs := int(float64(backoffMs) * 0.2)

417 jitterMs := int(float64(backoffMs) * 0.2)

internal/llm/provider/gemini.go+22−1

@@ -25,7 +25,7 @@ type geminiClient struct {

25type GeminiClient ProviderClient

27func newGeminiClient(opts providerClientOptions) GeminiClient {

28− client, err := genai.NewClient(context.Background(), &genai.ClientConfig{APIKey: opts.apiKey, Backend: genai.BackendGeminiAPI})

28+ client, err := createGeminiClient(opts)

29 if err != nil {

30 logging.Error("Failed to create Gemini client", "error", err)

31 return nil

@@ -37,6 +37,14 @@ func newGeminiClient(opts providerClientOptions) GeminiClient {

37 }

38}

40+func createGeminiClient(opts providerClientOptions) (*genai.Client, error) {

41+ client, err := genai.NewClient(context.Background(), &genai.ClientConfig{APIKey: opts.apiKey, Backend: genai.BackendGeminiAPI})

42+ if err != nil {

43+ return nil, err

44+ }

45+ return client, nil

46+}

47+

40func (g *geminiClient) convertMessages(messages []message.Message) []*genai.Content {

48func (g *geminiClient) convertMessages(messages []message.Message) []*genai.Content {

41 var history []*genai.Content

49 var history []*genai.Content

42 for _, msg := range messages {

50 for _, msg := range messages {

@@ -414,6 +422,19 @@ func (g *geminiClient) shouldRetry(attempts int, err error) (bool, int64, error)

414 errMsg := err.Error()

422 errMsg := err.Error()

415 isRateLimit := contains(errMsg, "rate limit", "quota exceeded", "too many requests")

423 isRateLimit := contains(errMsg, "rate limit", "quota exceeded", "too many requests")

416

424

425+ // Check for token expiration (401 Unauthorized)

426+ if contains(errMsg, "unauthorized", "invalid api key", "api key expired") {

427+ g.providerOptions.apiKey, err = config.ResolveAPIKey(g.providerOptions.config.APIKey)

428+ if err != nil {

429+ return false, 0, fmt.Errorf("failed to resolve API key: %w", err)

430+ }

431+ g.client, err = createGeminiClient(g.providerOptions)

432+ if err != nil {

433+ return false, 0, fmt.Errorf("failed to create Gemini client after API key refresh: %w", err)

434+ }

435+ return true, 0, nil

436+ }

437+

417 // Check for common rate limit error messages

438 // Check for common rate limit error messages

418

439

419 if !isRateLimit {

440 if !isRateLimit {

internal/llm/provider/openai.go+22−9

@@ -26,6 +26,13 @@ type openaiClient struct {

26type OpenAIClient ProviderClient

28func newOpenAIClient(opts providerClientOptions) OpenAIClient {

29+ return &openaiClient{

30+ providerOptions: opts,

31+ client: createOpenAIClient(opts),

32+ }

33+}

34+

35+func createOpenAIClient(opts providerClientOptions) openai.Client {

29 openaiClientOptions := []option.RequestOption{}

36 openaiClientOptions := []option.RequestOption{}

30 if opts.apiKey != "" {

37 if opts.apiKey != "" {

31 openaiClientOptions = append(openaiClientOptions, option.WithAPIKey(opts.apiKey))

38 openaiClientOptions = append(openaiClientOptions, option.WithAPIKey(opts.apiKey))

@@ -40,11 +47,7 @@ func newOpenAIClient(opts providerClientOptions) OpenAIClient {

40 }

47 }

41 }

48 }

43− client := openai.NewClient(openaiClientOptions...)

50+ return openai.NewClient(openaiClientOptions...)

44− return &openaiClient{

45− providerOptions: opts,

46− client: client,

47− }

48}

51}

50func (o *openaiClient) convertMessages(messages []message.Message) (openaiMessages []openai.ChatCompletionMessageParamUnion) {

53func (o *openaiClient) convertMessages(messages []message.Message) (openaiMessages []openai.ChatCompletionMessageParamUnion) {

@@ -339,12 +342,22 @@ func (o *openaiClient) stream(ctx context.Context, messages []message.Message, t

339}

342}

340

343

341func (o *openaiClient) shouldRetry(attempts int, err error) (bool, int64, error) {

344func (o *openaiClient) shouldRetry(attempts int, err error) (bool, int64, error) {

342− var apierr *openai.Error

345+ var apiErr *openai.Error

343− if !errors.As(err, &apierr) {

346+ if !errors.As(err, &apiErr) {

344 return false, 0, err

347 return false, 0, err

345 }

348 }

346

349

350+ // Check for token expiration (401 Unauthorized)

351+ if apiErr.StatusCode == 401 {

352+ o.providerOptions.apiKey, err = config.ResolveAPIKey(o.providerOptions.config.APIKey)

353+ if err != nil {

354+ return false, 0, fmt.Errorf("failed to resolve API key: %w", err)

355+ }

356+ o.client = createOpenAIClient(o.providerOptions)

357+ return true, 0, nil

358+ }

359+

347− if apierr.StatusCode != 429 && apierr.StatusCode != 500 {

360+ if apiErr.StatusCode != 429 && apiErr.StatusCode != 500 {

348 return false, 0, err

361 return false, 0, err

349 }

362 }

350

363

@@ -353,7 +366,7 @@ func (o *openaiClient) shouldRetry(attempts int, err error) (bool, int64, error)

353 }

366 }

354

367

355 retryMs := 0

368 retryMs := 0

356− retryAfterValues := apierr.Response.Header.Values("Retry-After")

369+ retryAfterValues := apiErr.Response.Header.Values("Retry-After")

357

370

358 backoffMs := 2000 * (1 << (attempts - 1))

371 backoffMs := 2000 * (1 << (attempts - 1))

359 jitterMs := int(float64(backoffMs) * 0.2)

372 jitterMs := int(float64(backoffMs) * 0.2)

internal/llm/provider/provider.go+8−1

@@ -60,6 +60,7 @@ type Provider interface {

61type providerClientOptions struct {

62 baseURL string

63+ config config.ProviderConfig

63 apiKey string

64 apiKey string

64 modelType config.ModelType

65 modelType config.ModelType

65 model func(config.ModelType) config.Model

66 model func(config.ModelType) config.Model

@@ -134,9 +135,15 @@ func WithMaxTokens(maxTokens int64) ProviderClientOption {

134}

135}

135

136

136func NewProvider(cfg config.ProviderConfig, opts ...ProviderClientOption) (Provider, error) {

137func NewProvider(cfg config.ProviderConfig, opts ...ProviderClientOption) (Provider, error) {

138+ resolvedAPIKey, err := config.ResolveAPIKey(cfg.APIKey)

139+ if err != nil {

140+ return nil, fmt.Errorf("failed to resolve API key for provider %s: %w", cfg.ID, err)

141+ }

142+

137 clientOptions := providerClientOptions{

143 clientOptions := providerClientOptions{

138 baseURL: cfg.BaseURL,

144 baseURL: cfg.BaseURL,

139− apiKey: cfg.APIKey,

145+ config: cfg,

146+ apiKey: resolvedAPIKey,

140 extraHeaders: cfg.ExtraHeaders,

147 extraHeaders: cfg.ExtraHeaders,

141 model: func(tp config.ModelType) config.Model {

148 model: func(tp config.ModelType) config.Model {

142 return config.GetModel(tp)

149 return config.GetModel(tp)