Merge pull request #561 from thedadams/deferred-cred-helper-build

feat: build credential helpers just in time
gptscript-ai · Jun 26, 2024 · 6055f37 · 6055f37
2 parents ead5eaf + 638f645
commit 6055f37
Show file tree

Hide file tree

Showing 14 changed files with 94 additions and 52 deletions.
diff --git a/pkg/cli/credential.go b/pkg/cli/credential.go
@@ -12,6 +12,8 @@ import (
 	"github.com/gptscript-ai/gptscript/pkg/cache"
 	"github.com/gptscript-ai/gptscript/pkg/config"
 	"github.com/gptscript-ai/gptscript/pkg/credentials"
+	"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
+	"github.com/gptscript-ai/gptscript/pkg/runner"
 	"github.com/spf13/cobra"
 )
 
@@ -35,7 +37,7 @@ func (c *Credential) Customize(cmd *cobra.Command) {
 	cmd.AddCommand(cmd2.Command(&Show{root: c.root}))
 }
 
-func (c *Credential) Run(_ *cobra.Command, _ []string) error {
+func (c *Credential) Run(cmd *cobra.Command, _ []string) error {
 	cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
 	if err != nil {
 		return fmt.Errorf("failed to read CLI config: %w", err)
@@ -51,14 +53,22 @@ func (c *Credential) Run(_ *cobra.Command, _ []string) error {
 		return err
 	}
 	opts.Cache = cache.Complete(opts.Cache)
+	opts.Runner = runner.Complete(opts.Runner)
+	if opts.Runner.RuntimeManager == nil {
+		opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
+	}
+
+	if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg, opts.Env); err != nil {
+		return err
+	}
 
 	// Initialize the credential store and get all the credentials.
-	store, err := credentials.NewStore(cfg, ctx, opts.Cache.CacheDir)
+	store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, ctx, opts.Cache.CacheDir)
 	if err != nil {
 		return fmt.Errorf("failed to get credentials store: %w", err)
 	}
 
-	creds, err := store.List()
+	creds, err := store.List(cmd.Context())
 	if err != nil {
 		return fmt.Errorf("failed to list credentials: %w", err)
 	}

diff --git a/pkg/cli/credential_delete.go b/pkg/cli/credential_delete.go
@@ -6,6 +6,8 @@ import (
 	"github.com/gptscript-ai/gptscript/pkg/cache"
 	"github.com/gptscript-ai/gptscript/pkg/config"
 	"github.com/gptscript-ai/gptscript/pkg/credentials"
+	"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
+	"github.com/gptscript-ai/gptscript/pkg/runner"
 	"github.com/spf13/cobra"
 )
 
@@ -21,24 +23,33 @@ func (c *Delete) Customize(cmd *cobra.Command) {
 	cmd.Args = cobra.ExactArgs(1)
 }
 
-func (c *Delete) Run(_ *cobra.Command, args []string) error {
+func (c *Delete) Run(cmd *cobra.Command, args []string) error {
 	opts, err := c.root.NewGPTScriptOpts()
 	if err != nil {
 		return err
 	}
-	opts.Cache = cache.Complete(opts.Cache)
 
 	cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
 	if err != nil {
 		return fmt.Errorf("failed to read CLI config: %w", err)
 	}
 
-	store, err := credentials.NewStore(cfg, c.root.CredentialContext, opts.Cache.CacheDir)
+	opts.Cache = cache.Complete(opts.Cache)
+	opts.Runner = runner.Complete(opts.Runner)
+	if opts.Runner.RuntimeManager == nil {
+		opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
+	}
+
+	if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg, opts.Env); err != nil {
+		return err
+	}
+
+	store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, c.root.CredentialContext, opts.Cache.CacheDir)
 	if err != nil {
 		return fmt.Errorf("failed to get credentials store: %w", err)
 	}
 
-	if err = store.Remove(args[0]); err != nil {
+	if err = store.Remove(cmd.Context(), args[0]); err != nil {
 		return fmt.Errorf("failed to remove credential: %w", err)
 	}
 	return nil

diff --git a/pkg/cli/credential_show.go b/pkg/cli/credential_show.go
@@ -8,6 +8,8 @@ import (
 	"github.com/gptscript-ai/gptscript/pkg/cache"
 	"github.com/gptscript-ai/gptscript/pkg/config"
 	"github.com/gptscript-ai/gptscript/pkg/credentials"
+	"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
+	"github.com/gptscript-ai/gptscript/pkg/runner"
 	"github.com/spf13/cobra"
 )
 
@@ -23,24 +25,33 @@ func (c *Show) Customize(cmd *cobra.Command) {
 	cmd.Args = cobra.ExactArgs(1)
 }
 
-func (c *Show) Run(_ *cobra.Command, args []string) error {
+func (c *Show) Run(cmd *cobra.Command, args []string) error {
 	opts, err := c.root.NewGPTScriptOpts()
 	if err != nil {
 		return err
 	}
-	opts.Cache = cache.Complete(opts.Cache)
 
 	cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
 	if err != nil {
 		return fmt.Errorf("failed to read CLI config: %w", err)
 	}
 
-	store, err := credentials.NewStore(cfg, c.root.CredentialContext, opts.Cache.CacheDir)
+	opts.Cache = cache.Complete(opts.Cache)
+	opts.Runner = runner.Complete(opts.Runner)
+	if opts.Runner.RuntimeManager == nil {
+		opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
+	}
+
+	if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg, opts.Env); err != nil {
+		return err
+	}
+
+	store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, c.root.CredentialContext, opts.Cache.CacheDir)
 	if err != nil {
 		return fmt.Errorf("failed to get credentials store: %w", err)
 	}
 
-	cred, exists, err := store.Get(args[0])
+	cred, exists, err := store.Get(cmd.Context(), args[0])
 	if err != nil {
 		return fmt.Errorf("failed to get credential: %w", err)
 	}

diff --git a/pkg/cli/eval.go b/pkg/cli/eval.go
@@ -56,7 +56,7 @@ func (e *Eval) Run(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	runner, err := gptscript.New(opts)
+	runner, err := gptscript.New(cmd.Context(), opts)
 	if err != nil {
 		return err
 	}

diff --git a/pkg/cli/gptscript.go b/pkg/cli/gptscript.go
@@ -380,7 +380,7 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) (retErr error) {
 
 	ctx := cmd.Context()
 
-	gptScript, err := gptscript.New(gptOpt)
+	gptScript, err := gptscript.New(ctx, gptOpt)
 	if err != nil {
 		return err
 	}

diff --git a/pkg/credentials/noop.go b/pkg/credentials/noop.go
@@ -1,19 +1,21 @@
 package credentials
 
+import "context"
+
 type NoopStore struct{}
 
-func (s NoopStore) Get(_ string) (*Credential, bool, error) {
+func (s NoopStore) Get(context.Context, string) (*Credential, bool, error) {
 	return nil, false, nil
 }
 
-func (s NoopStore) Add(_ Credential) error {
+func (s NoopStore) Add(context.Context, Credential) error {
 	return nil
 }
 
-func (s NoopStore) Remove(_ string) error {
+func (s NoopStore) Remove(context.Context, string) error {
 	return nil
 }
 
-func (s NoopStore) List() ([]Credential, error) {
+func (s NoopStore) List(context.Context) ([]Credential, error) {
 	return nil, nil
 }
diff --git a/pkg/credentials/store.go b/pkg/credentials/store.go
@@ -1,6 +1,7 @@
 package credentials
 
 import (
+	"context"
 	"fmt"
 	"path/filepath"
 	"regexp"
@@ -10,32 +11,38 @@ import (
 	"github.com/gptscript-ai/gptscript/pkg/config"
 )
 
+type CredentialBuilder interface {
+	EnsureCredentialHelpers(ctx context.Context) error
+}
+
 type CredentialStore interface {
-	Get(toolName string) (*Credential, bool, error)
-	Add(cred Credential) error
-	Remove(toolName string) error
-	List() ([]Credential, error)
+	Get(ctx context.Context, toolName string) (*Credential, bool, error)
+	Add(ctx context.Context, cred Credential) error
+	Remove(ctx context.Context, toolName string) error
+	List(ctx context.Context) ([]Credential, error)
 }
 
 type Store struct {
 	credCtx        string
+	credBuilder    CredentialBuilder
 	credHelperDirs CredentialHelperDirs
 	cfg            *config.CLIConfig
 }
 
-func NewStore(cfg *config.CLIConfig, credCtx, cacheDir string) (CredentialStore, error) {
+func NewStore(cfg *config.CLIConfig, credentialBuilder CredentialBuilder, credCtx, cacheDir string) (CredentialStore, error) {
 	if err := validateCredentialCtx(credCtx); err != nil {
 		return nil, err
 	}
 	return Store{
 		credCtx:        credCtx,
+		credBuilder:    credentialBuilder,
 		credHelperDirs: GetCredentialHelperDirs(cacheDir),
 		cfg:            cfg,
 	}, nil
 }
 
-func (s Store) Get(toolName string) (*Credential, bool, error) {
-	store, err := s.getStore()
+func (s Store) Get(ctx context.Context, toolName string) (*Credential, bool, error) {
+	store, err := s.getStore(ctx)
 	if err != nil {
 		return nil, false, err
 	}
@@ -57,9 +64,9 @@ func (s Store) Get(toolName string) (*Credential, bool, error) {
 	return &cred, true, nil
 }
 
-func (s Store) Add(cred Credential) error {
+func (s Store) Add(ctx context.Context, cred Credential) error {
 	cred.Context = s.credCtx
-	store, err := s.getStore()
+	store, err := s.getStore(ctx)
 	if err != nil {
 		return err
 	}
@@ -70,16 +77,16 @@ func (s Store) Add(cred Credential) error {
 	return store.Store(auth)
 }
 
-func (s Store) Remove(toolName string) error {
-	store, err := s.getStore()
+func (s Store) Remove(ctx context.Context, toolName string) error {
+	store, err := s.getStore(ctx)
 	if err != nil {
 		return err
 	}
 	return store.Erase(toolNameWithCtx(toolName, s.credCtx))
 }
 
-func (s Store) List() ([]Credential, error) {
-	store, err := s.getStore()
+func (s Store) List(ctx context.Context) ([]Credential, error) {
+	store, err := s.getStore(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -106,17 +113,21 @@ func (s Store) List() ([]Credential, error) {
 	return creds, nil
 }
 
-func (s *Store) getStore() (credentials.Store, error) {
-	return s.getStoreByHelper(config.GPTScriptHelperPrefix + s.cfg.CredentialsStore)
+func (s *Store) getStore(ctx context.Context) (credentials.Store, error) {
+	return s.getStoreByHelper(ctx, config.GPTScriptHelperPrefix+s.cfg.CredentialsStore)
 }
 
-func (s *Store) getStoreByHelper(helper string) (credentials.Store, error) {
+func (s *Store) getStoreByHelper(ctx context.Context, helper string) (credentials.Store, error) {
 	if helper == "" || helper == config.GPTScriptHelperPrefix+"file" {
 		return credentials.NewFileStore(s.cfg), nil
 	}
 
 	// If the helper is referencing one of the credential helper programs, then reference the full path.
 	if strings.HasPrefix(helper, "gptscript-credential-") {
+		if err := s.credBuilder.EnsureCredentialHelpers(ctx); err != nil {
+			return nil, err
+		}
+
 		helper = filepath.Join(s.credHelperDirs.BinDir, helper)
 	}
 

diff --git a/pkg/gptscript/gptscript.go b/pkg/gptscript/gptscript.go
@@ -77,7 +77,7 @@ func complete(opts ...Options) Options {
 	return result
 }
 
-func New(o ...Options) (*GPTScript, error) {
+func New(ctx context.Context, o ...Options) (*GPTScript, error) {
 	opts := complete(o...)
 	registry := llm.NewRegistry()
 
@@ -91,23 +91,20 @@ func New(o ...Options) (*GPTScript, error) {
 		return nil, err
 	}
 
-	credStore, err := credentials.NewStore(cliCfg, opts.CredentialContext, cacheClient.CacheDir())
-	if err != nil {
-		return nil, err
-	}
-
 	if opts.Runner.RuntimeManager == nil {
 		opts.Runner.RuntimeManager = runtimes.Default(cacheClient.CacheDir())
 	}
 
 	if err := opts.Runner.RuntimeManager.SetUpCredentialHelpers(context.Background(), cliCfg, opts.Env); err != nil {
 		return nil, err
 	}
-	if err := opts.Runner.RuntimeManager.EnsureCredentialHelpers(context.Background()); err != nil {
+
+	credStore, err := credentials.NewStore(cliCfg, opts.Runner.RuntimeManager, opts.CredentialContext, cacheClient.CacheDir())
+	if err != nil {
 		return nil, err
 	}
 
-	oaiClient, err := openai.NewClient(credStore, opts.OpenAI, openai.Options{
+	oaiClient, err := openai.NewClient(ctx, credStore, opts.OpenAI, openai.Options{
 		Cache:   cacheClient,
 		SetSeed: true,
 	})

diff --git a/pkg/openai/client.go b/pkg/openai/client.go
@@ -93,15 +93,15 @@ func complete(opts ...Options) (Options, error) {
 	return result, err
 }
 
-func NewClient(credStore credentials.CredentialStore, opts ...Options) (*Client, error) {
+func NewClient(ctx context.Context, credStore credentials.CredentialStore, opts ...Options) (*Client, error) {
 	opt, err := complete(opts...)
 	if err != nil {
 		return nil, err
 	}
 
 	// If the API key is not set, try to get it from the cred store
 	if opt.APIKey == "" && opt.BaseURL == "" {
-		cred, exists, err := credStore.Get(BuiltinCredName)
+		cred, exists, err := credStore.Get(ctx, BuiltinCredName)
 		if err != nil {
 			return nil, err
 		}

diff --git a/pkg/prompt/credential.go b/pkg/prompt/credential.go
@@ -9,7 +9,7 @@ import (
 )
 
 func GetModelProviderCredential(ctx context.Context, credStore credentials.CredentialStore, credName, env, message string, envs []string) (string, error) {
-	cred, exists, err := credStore.Get(credName)
+	cred, exists, err := credStore.Get(ctx, credName)
 	if err != nil {
 		return "", err
 	}
@@ -25,7 +25,7 @@ func GetModelProviderCredential(ctx context.Context, credStore credentials.Crede
 		}
 
 		k = gjson.Get(result, "key").String()
-		if err := credStore.Add(credentials.Credential{
+		if err := credStore.Add(ctx, credentials.Credential{
 			ToolName: credName,
 			Type:     credentials.CredentialTypeModelProvider,
 			Env: map[string]string{

diff --git a/pkg/remote/remote.go b/pkg/remote/remote.go
@@ -116,7 +116,7 @@ func (c *Client) clientFromURL(ctx context.Context, apiURL string) (*openai.Clie
 		}
 	}
 
-	return openai.NewClient(c.credStore, openai.Options{
+	return openai.NewClient(ctx, c.credStore, openai.Options{
 		BaseURL: apiURL,
 		Cache:   c.cache,
 		APIKey:  key,
@@ -163,7 +163,7 @@ func (c *Client) load(ctx context.Context, toolName string) (*openai.Client, err
 		url += "/v1"
 	}
 
-	client, err = openai.NewClient(c.credStore, openai.Options{
+	client, err = openai.NewClient(ctx, c.credStore, openai.Options{
 		BaseURL:  url,
 		Cache:    c.cache,
 		CacheKey: prg.EntryToolID,