handle context

Signed-off-by: Grant Linville <[email protected]>

Author: g-linville

pkg/llm/registry.go

@@ -64,7 +64,7 @@ func (r *Registry) Call(ctx context.Context, messageRequest types.CompletionRequ
 
 	if len(errs) > 0 && oaiClient != nil {
 		// Prompt the user to enter their OpenAI API key and try again.
-		if err := oaiClient.RetrieveAPIKey(); err != nil {
+		if err := oaiClient.RetrieveAPIKey(ctx); err != nil {
 			return nil, err
 		}
 		ok, err := oaiClient.Supports(ctx, messageRequest.Model)

pkg/openai/client.go

@@ -318,7 +318,7 @@ func toMessages(request types.CompletionRequest, compat bool) (result []openai.C
 
 func (c *Client) Call(ctx context.Context, messageRequest types.CompletionRequest, status chan<- types.CompletionStatus) (*types.CompletionMessage, error) {
 	if err := c.ValidAuth(); err != nil {
-		if err := c.RetrieveAPIKey(); err != nil {
+		if err := c.RetrieveAPIKey(ctx); err != nil {
 			return nil, err
 		}
 	}
@@ -569,8 +569,8 @@ func (c *Client) call(ctx context.Context, request openai.ChatCompletionRequest,
 	}
 }
 
-func (c *Client) RetrieveAPIKey() error {
-	k, err := prompt.GetModelProviderCredential(BuiltinCredName, "OPENAI_API_KEY", "Please provide your OpenAI API key:", c.credCtx, c.envs, c.cliCfg)
+func (c *Client) RetrieveAPIKey(ctx context.Context) error {
+	k, err := prompt.GetModelProviderCredential(ctx, BuiltinCredName, "OPENAI_API_KEY", "Please provide your OpenAI API key:", c.credCtx, c.envs, c.cliCfg)
 	if err != nil {
 		return err
 	}

pkg/prompt/credential.go

@@ -9,7 +9,7 @@ import (
 	"github.com/tidwall/gjson"
 )
 
-func GetModelProviderCredential(credName, env, message, credCtx string, envs []string, cliCfg *config.CLIConfig) (string, error) {
+func GetModelProviderCredential(ctx context.Context, credName, env, message, credCtx string, envs []string, cliCfg *config.CLIConfig) (string, error) {
 	store, err := credentials.NewStore(cliCfg, credCtx)
 	if err != nil {
 		return "", err
@@ -24,8 +24,7 @@ func GetModelProviderCredential(credName, env, message, credCtx string, envs []s
 	if exists {
 		k = cred.Env[env]
 	} else {
-		// SysPrompt doesn't use its first two arguments, so we can safely pass whatever to them
-		result, err := SysPrompt(context.Background(), envs, fmt.Sprintf(`{"message":"%s","fields":"key","sensitive":"true"}`, message))
+		result, err := SysPrompt(ctx, envs, fmt.Sprintf(`{"message":"%s","fields":"key","sensitive":"true"}`, message))
 		if err != nil {
 			return "", err
 		}

pkg/remote/remote.go

@@ -101,7 +101,7 @@ func isHTTPURL(toolName string) bool {
 		strings.HasPrefix(toolName, "https://")
 }
 
-func (c *Client) clientFromURL(apiURL string) (*openai.Client, error) {
+func (c *Client) clientFromURL(ctx context.Context, apiURL string) (*openai.Client, error) {
 	parsed, err := url.Parse(apiURL)
 	if err != nil {
 		return nil, err
@@ -111,7 +111,7 @@ func (c *Client) clientFromURL(apiURL string) (*openai.Client, error) {
 
 	if key == "" {
 		var err error
-		key, err = c.retrieveAPIKey(env, apiURL)
+		key, err = c.retrieveAPIKey(ctx, env, apiURL)
 		if err != nil {
 			return nil, err
 		}
@@ -138,7 +138,7 @@ func (c *Client) load(ctx context.Context, toolName string) (*openai.Client, err
 	}
 
 	if isHTTPURL(toolName) {
-		remoteClient, err := c.clientFromURL(toolName)
+		remoteClient, err := c.clientFromURL(ctx, toolName)
 		if err != nil {
 			return nil, err
 		}
@@ -177,6 +177,6 @@ func (c *Client) load(ctx context.Context, toolName string) (*openai.Client, err
 	return client, nil
 }
 
-func (c *Client) retrieveAPIKey(env, url string) (string, error) {
-	return prompt.GetModelProviderCredential(url, env, fmt.Sprintf("Please provide your API key for %s", url), c.credCtx, c.envs, c.cliCfg)
+func (c *Client) retrieveAPIKey(ctx context.Context, env, url string) (string, error) {
+	return prompt.GetModelProviderCredential(ctx, url, env, fmt.Sprintf("Please provide your API key for %s", url), c.credCtx, c.envs, c.cliCfg)
 }