From 638f64506b324065b55998c616633ad1d5888326 Mon Sep 17 00:00:00 2001 From: Donnie Adams Date: Wed, 26 Jun 2024 09:23:38 -0400 Subject: [PATCH] feat: build credential helpers just in time Instead of building the credential helpers on startup, this change builds them when they are needed. This speeds up start time for the SDK. Signed-off-by: Donnie Adams --- pkg/cli/credential.go | 16 +++++++++++--- pkg/cli/credential_delete.go | 19 ++++++++++++---- pkg/cli/credential_show.go | 19 ++++++++++++---- pkg/cli/eval.go | 2 +- pkg/cli/gptscript.go | 2 +- pkg/credentials/noop.go | 10 +++++---- pkg/credentials/store.go | 43 ++++++++++++++++++++++-------------- pkg/gptscript/gptscript.go | 13 +++++------ pkg/openai/client.go | 4 ++-- pkg/prompt/credential.go | 4 ++-- pkg/remote/remote.go | 4 ++-- pkg/runner/runner.go | 6 ++--- pkg/sdkserver/run.go | 2 +- pkg/sdkserver/server.go | 2 +- 14 files changed, 94 insertions(+), 52 deletions(-) diff --git a/pkg/cli/credential.go b/pkg/cli/credential.go index 14f832eb..b0c4a30a 100644 --- 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 index db1f97b8..9c986c54 100644 --- 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 index 92911dde..ccfe3675 100644 --- 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 index 7beaa8a0..2cd4b1b5 100644 --- 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 index 90af42d7..a8561abf 100644 --- 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 index 1dfa61e8..5f3cc5ad 100644 --- 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 index 287c1aa6..3940184b 100644 --- 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 index 63cce7ba..f8264a44 100644 --- 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,11 +91,6 @@ 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()) } @@ -103,11 +98,13 @@ func New(o ...Options) (*GPTScript, error) { 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 index 42ff83bc..6e3c9c60 100644 --- a/pkg/openai/client.go +++ b/pkg/openai/client.go @@ -93,7 +93,7 @@ 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 @@ -101,7 +101,7 @@ func NewClient(credStore credentials.CredentialStore, opts ...Options) (*Client, // 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 index 9202ed49..a47a9168 100644 --- 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 index 3837879a..7d7cff6e 100644 --- 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, diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index 466eddb7..cecc7afd 100644 --- a/pkg/runner/runner.go +++ b/pkg/runner/runner.go @@ -873,12 +873,12 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env // Only try to look up the cred if the tool is on GitHub or has an alias. // If it is a GitHub tool and has an alias, the alias overrides the tool name, so we use it as the credential name. if isGitHubTool(toolName) && credentialAlias == "" { - c, exists, err = r.credStore.Get(toolName) + c, exists, err = r.credStore.Get(callCtx.Ctx, toolName) if err != nil { return nil, fmt.Errorf("failed to get credentials for tool %s: %w", toolName, err) } } else if credentialAlias != "" { - c, exists, err = r.credStore.Get(credentialAlias) + c, exists, err = r.credStore.Get(callCtx.Ctx, credentialAlias) if err != nil { return nil, fmt.Errorf("failed to get credentials for tool %s: %w", credentialAlias, err) } @@ -942,7 +942,7 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env if (isGitHubTool(toolName) && callCtx.Program.ToolSet[credToolRefs[0].ToolID].Source.Repo != nil) || credentialAlias != "" { if isEmpty { log.Warnf("Not saving empty credential for tool %s", toolName) - } else if err := r.credStore.Add(*c); err != nil { + } else if err := r.credStore.Add(callCtx.Ctx, *c); err != nil { return nil, fmt.Errorf("failed to add credential for tool %s: %w", toolName, err) } } else { diff --git a/pkg/sdkserver/run.go b/pkg/sdkserver/run.go index 321e76c1..dc155557 100644 --- a/pkg/sdkserver/run.go +++ b/pkg/sdkserver/run.go @@ -17,7 +17,7 @@ import ( type loaderFunc func(context.Context, string, string, ...loader.Options) (types.Program, error) func (s *server) execAndStream(ctx context.Context, programLoader loaderFunc, logger mvl.Logger, w http.ResponseWriter, opts gptscript.Options, chatState, input, subTool string, toolDef fmt.Stringer) { - g, err := gptscript.New(s.gptscriptOpts, opts) + g, err := gptscript.New(ctx, s.gptscriptOpts, opts) if err != nil { writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to initialize gptscript: %w", err)) return diff --git a/pkg/sdkserver/server.go b/pkg/sdkserver/server.go index c51e21e4..4556f69e 100644 --- a/pkg/sdkserver/server.go +++ b/pkg/sdkserver/server.go @@ -53,7 +53,7 @@ func Start(ctx context.Context, opts Options) error { // prompt server because it is only used for fmt, parse, etc. opts.Env = append(opts.Env, fmt.Sprintf("%s=%s", types.PromptTokenEnvVar, token)) - g, err := gptscript.New(opts.Options) + g, err := gptscript.New(ctx, opts.Options) if err != nil { return err }