refactor: add context support to client initialization functions

- Add context parameter to `GetClient` function in `commit.go`
- Import `context` package in `openai.go`
- Add context parameter to `NewGemini` function in `openai.go`
- Add context parameter to `GetClient` function in `openai.go`
- Add context parameter to `New` function in `gemini.go`
- Use provided context in `genai.NewClient` initialization in `gemini.go`

Signed-off-by: Bo-Yi Wu <[email protected]>

Author: appleboy

cmd/commit.go

@@ -95,7 +95,7 @@ var commitCmd = &cobra.Command{
 
 		// check provider
 		provider := core.Platform(viper.GetString("openai.provider"))
-		client, err := GetClient(provider)
+		client, err := GetClient(cmd.Context(), provider)
 		if err != nil && !promptOnly {
 			return err
 		}

cmd/openai.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 
 	"github.com/appleboy/CodeGPT/core"
@@ -32,8 +33,9 @@ func NewOpenAI() (*openai.Client, error) {
 }
 
 // NewGemini returns a new Gemini client
-func NewGemini() (*gemini.Client, error) {
+func NewGemini(ctx context.Context) (*gemini.Client, error) {
 	return gemini.New(
+		ctx,
 		gemini.WithToken(viper.GetString("openai.api_key")),
 		gemini.WithModel(viper.GetString("openai.model")),
 		gemini.WithMaxTokens(viper.GetInt32("openai.max_tokens")),
@@ -43,10 +45,10 @@ func NewGemini() (*gemini.Client, error) {
 }
 
 // GetClient returns the generative client based on the platform
-func GetClient(p core.Platform) (core.Generative, error) {
+func GetClient(ctx context.Context, p core.Platform) (core.Generative, error) {
 	switch p {
 	case core.Gemini:
-		return NewGemini()
+		return NewGemini(ctx)
 	case core.OpenAI, core.Azure:
 		return NewOpenAI()
 	}

cmd/review.go

@@ -58,7 +58,7 @@ var reviewCmd = &cobra.Command{
 
 		// check provider
 		provider := core.Platform(viper.GetString("openai.provider"))
-		client, err := GetClient(provider)
+		client, err := GetClient(cmd.Context(), provider)
 		if err != nil {
 			return err
 		}

gemini/gemini.go

@@ -84,7 +84,7 @@ func (c *Client) GetSummaryPrefix(ctx context.Context, content string) (*core.Re
 	return r, nil
 }
 
-func New(opts ...Option) (c *Client, err error) {
+func New(ctx context.Context, opts ...Option) (c *Client, err error) {
 	// Create a new config object with the given options.
 	cfg := newConfig(opts...)
 
@@ -100,7 +100,7 @@ func New(opts ...Option) (c *Client, err error) {
 		temperature: cfg.temperature,
 	}
 
-	client, err := genai.NewClient(context.Background(), option.WithAPIKey(cfg.token))
+	client, err := genai.NewClient(ctx, option.WithAPIKey(cfg.token))
 	if err != nil {
 		return nil, err
 	}