Skip to content

Commit 41f7e13

Browse files
thedadamsthedadams
committed
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 <[email protected]>
1 parent ead5eaf commit 41f7e13

File tree

13 files changed

+88
-48
lines changed

13 files changed

+88
-48
lines changed

‎pkg/cli/credential.go

+13-3
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import (
1212
"github.com/gptscript-ai/gptscript/pkg/cache"
1313
"github.com/gptscript-ai/gptscript/pkg/config"
1414
"github.com/gptscript-ai/gptscript/pkg/credentials"
15+
"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
16+
"github.com/gptscript-ai/gptscript/pkg/runner"
1517
"github.com/spf13/cobra"
1618
)
1719

@@ -35,7 +37,7 @@ func (c *Credential) Customize(cmd *cobra.Command) {
3537
cmd.AddCommand(cmd2.Command(&Show{root: c.root}))
3638
}
3739

38-
func (c *Credential) Run(_ *cobra.Command, _ []string) error {
40+
func (c *Credential) Run(cmd *cobra.Command, _ []string) error {
3941
cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
4042
if err != nil {
4143
return fmt.Errorf("failed to read CLI config: %w", err)
@@ -51,14 +53,22 @@ func (c *Credential) Run(_ *cobra.Command, _ []string) error {
5153
return err
5254
}
5355
opts.Cache = cache.Complete(opts.Cache)
56+
opts.Runner = runner.Complete(opts.Runner)
57+
if opts.Runner.RuntimeManager == nil {
58+
opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
59+
}
60+
61+
if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg, opts.Env); err != nil {
62+
return err
63+
}
5464

5565
// Initialize the credential store and get all the credentials.
56-
store, err := credentials.NewStore(cfg, ctx, opts.Cache.CacheDir)
66+
store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, ctx, opts.Cache.CacheDir)
5767
if err != nil {
5868
return fmt.Errorf("failed to get credentials store: %w", err)
5969
}
6070

61-
creds, err := store.List()
71+
creds, err := store.List(cmd.Context())
6272
if err != nil {
6373
return fmt.Errorf("failed to list credentials: %w", err)
6474
}

‎pkg/cli/credential_delete.go

+15-4
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import (
66
"github.com/gptscript-ai/gptscript/pkg/cache"
77
"github.com/gptscript-ai/gptscript/pkg/config"
88
"github.com/gptscript-ai/gptscript/pkg/credentials"
9+
"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
10+
"github.com/gptscript-ai/gptscript/pkg/runner"
911
"github.com/spf13/cobra"
1012
)
1113

@@ -21,24 +23,33 @@ func (c *Delete) Customize(cmd *cobra.Command) {
2123
cmd.Args = cobra.ExactArgs(1)
2224
}
2325

24-
func (c *Delete) Run(_ *cobra.Command, args []string) error {
26+
func (c *Delete) Run(cmd *cobra.Command, args []string) error {
2527
opts, err := c.root.NewGPTScriptOpts()
2628
if err != nil {
2729
return err
2830
}
29-
opts.Cache = cache.Complete(opts.Cache)
3031

3132
cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
3233
if err != nil {
3334
return fmt.Errorf("failed to read CLI config: %w", err)
3435
}
3536

36-
store, err := credentials.NewStore(cfg, c.root.CredentialContext, opts.Cache.CacheDir)
37+
opts.Cache = cache.Complete(opts.Cache)
38+
opts.Runner = runner.Complete(opts.Runner)
39+
if opts.Runner.RuntimeManager == nil {
40+
opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
41+
}
42+
43+
if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg, opts.Env); err != nil {
44+
return err
45+
}
46+
47+
store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, c.root.CredentialContext, opts.Cache.CacheDir)
3748
if err != nil {
3849
return fmt.Errorf("failed to get credentials store: %w", err)
3950
}
4051

41-
if err = store.Remove(args[0]); err != nil {
52+
if err = store.Remove(cmd.Context(), args[0]); err != nil {
4253
return fmt.Errorf("failed to remove credential: %w", err)
4354
}
4455
return nil

‎pkg/cli/credential_show.go

+15-4
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import (
88
"github.com/gptscript-ai/gptscript/pkg/cache"
99
"github.com/gptscript-ai/gptscript/pkg/config"
1010
"github.com/gptscript-ai/gptscript/pkg/credentials"
11+
"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
12+
"github.com/gptscript-ai/gptscript/pkg/runner"
1113
"github.com/spf13/cobra"
1214
)
1315

@@ -23,24 +25,33 @@ func (c *Show) Customize(cmd *cobra.Command) {
2325
cmd.Args = cobra.ExactArgs(1)
2426
}
2527

26-
func (c *Show) Run(_ *cobra.Command, args []string) error {
28+
func (c *Show) Run(cmd *cobra.Command, args []string) error {
2729
opts, err := c.root.NewGPTScriptOpts()
2830
if err != nil {
2931
return err
3032
}
31-
opts.Cache = cache.Complete(opts.Cache)
3233

3334
cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
3435
if err != nil {
3536
return fmt.Errorf("failed to read CLI config: %w", err)
3637
}
3738

38-
store, err := credentials.NewStore(cfg, c.root.CredentialContext, opts.Cache.CacheDir)
39+
opts.Cache = cache.Complete(opts.Cache)
40+
opts.Runner = runner.Complete(opts.Runner)
41+
if opts.Runner.RuntimeManager == nil {
42+
opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
43+
}
44+
45+
if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg, opts.Env); err != nil {
46+
return err
47+
}
48+
49+
store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, c.root.CredentialContext, opts.Cache.CacheDir)
3950
if err != nil {
4051
return fmt.Errorf("failed to get credentials store: %w", err)
4152
}
4253

43-
cred, exists, err := store.Get(args[0])
54+
cred, exists, err := store.Get(cmd.Context(), args[0])
4455
if err != nil {
4556
return fmt.Errorf("failed to get credential: %w", err)
4657
}

‎pkg/cli/eval.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ func (e *Eval) Run(cmd *cobra.Command, args []string) error {
5656
return err
5757
}
5858

59-
runner, err := gptscript.New(opts)
59+
runner, err := gptscript.New(cmd.Context(), opts)
6060
if err != nil {
6161
return err
6262
}

‎pkg/cli/gptscript.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -380,7 +380,7 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) (retErr error) {
380380

381381
ctx := cmd.Context()
382382

383-
gptScript, err := gptscript.New(gptOpt)
383+
gptScript, err := gptscript.New(ctx, gptOpt)
384384
if err != nil {
385385
return err
386386
}

‎pkg/credentials/store.go

+27-16
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package credentials
22

33
import (
4+
"context"
45
"fmt"
56
"path/filepath"
67
"regexp"
@@ -10,32 +11,38 @@ import (
1011
"github.com/gptscript-ai/gptscript/pkg/config"
1112
)
1213

14+
type CredentialBuilder interface {
15+
EnsureCredentialHelpers(ctx context.Context) error
16+
}
17+
1318
type CredentialStore interface {
14-
Get(toolName string) (*Credential, bool, error)
15-
Add(cred Credential) error
16-
Remove(toolName string) error
17-
List() ([]Credential, error)
19+
Get(ctx context.Context, toolName string) (*Credential, bool, error)
20+
Add(ctx context.Context, cred Credential) error
21+
Remove(ctx context.Context, toolName string) error
22+
List(ctx context.Context) ([]Credential, error)
1823
}
1924

2025
type Store struct {
2126
credCtx string
27+
credBuilder CredentialBuilder
2228
credHelperDirs CredentialHelperDirs
2329
cfg *config.CLIConfig
2430
}
2531

26-
func NewStore(cfg *config.CLIConfig, credCtx, cacheDir string) (CredentialStore, error) {
32+
func NewStore(cfg *config.CLIConfig, credentialBuilder CredentialBuilder, credCtx, cacheDir string) (CredentialStore, error) {
2733
if err := validateCredentialCtx(credCtx); err != nil {
2834
return nil, err
2935
}
3036
return Store{
3137
credCtx: credCtx,
38+
credBuilder: credentialBuilder,
3239
credHelperDirs: GetCredentialHelperDirs(cacheDir),
3340
cfg: cfg,
3441
}, nil
3542
}
3643

37-
func (s Store) Get(toolName string) (*Credential, bool, error) {
38-
store, err := s.getStore()
44+
func (s Store) Get(ctx context.Context, toolName string) (*Credential, bool, error) {
45+
store, err := s.getStore(ctx)
3946
if err != nil {
4047
return nil, false, err
4148
}
@@ -57,9 +64,9 @@ func (s Store) Get(toolName string) (*Credential, bool, error) {
5764
return &cred, true, nil
5865
}
5966

60-
func (s Store) Add(cred Credential) error {
67+
func (s Store) Add(ctx context.Context, cred Credential) error {
6168
cred.Context = s.credCtx
62-
store, err := s.getStore()
69+
store, err := s.getStore(ctx)
6370
if err != nil {
6471
return err
6572
}
@@ -70,16 +77,16 @@ func (s Store) Add(cred Credential) error {
7077
return store.Store(auth)
7178
}
7279

73-
func (s Store) Remove(toolName string) error {
74-
store, err := s.getStore()
80+
func (s Store) Remove(ctx context.Context, toolName string) error {
81+
store, err := s.getStore(ctx)
7582
if err != nil {
7683
return err
7784
}
7885
return store.Erase(toolNameWithCtx(toolName, s.credCtx))
7986
}
8087

81-
func (s Store) List() ([]Credential, error) {
82-
store, err := s.getStore()
88+
func (s Store) List(ctx context.Context) ([]Credential, error) {
89+
store, err := s.getStore(ctx)
8390
if err != nil {
8491
return nil, err
8592
}
@@ -106,17 +113,21 @@ func (s Store) List() ([]Credential, error) {
106113
return creds, nil
107114
}
108115

109-
func (s *Store) getStore() (credentials.Store, error) {
110-
return s.getStoreByHelper(config.GPTScriptHelperPrefix + s.cfg.CredentialsStore)
116+
func (s *Store) getStore(ctx context.Context) (credentials.Store, error) {
117+
return s.getStoreByHelper(ctx, config.GPTScriptHelperPrefix+s.cfg.CredentialsStore)
111118
}
112119

113-
func (s *Store) getStoreByHelper(helper string) (credentials.Store, error) {
120+
func (s *Store) getStoreByHelper(ctx context.Context, helper string) (credentials.Store, error) {
114121
if helper == "" || helper == config.GPTScriptHelperPrefix+"file" {
115122
return credentials.NewFileStore(s.cfg), nil
116123
}
117124

118125
// If the helper is referencing one of the credential helper programs, then reference the full path.
119126
if strings.HasPrefix(helper, "gptscript-credential-") {
127+
if err := s.credBuilder.EnsureCredentialHelpers(ctx); err != nil {
128+
return nil, err
129+
}
130+
120131
helper = filepath.Join(s.credHelperDirs.BinDir, helper)
121132
}
122133

‎pkg/gptscript/gptscript.go

+5-8
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ func complete(opts ...Options) Options {
7777
return result
7878
}
7979

80-
func New(o ...Options) (*GPTScript, error) {
80+
func New(ctx context.Context, o ...Options) (*GPTScript, error) {
8181
opts := complete(o...)
8282
registry := llm.NewRegistry()
8383

@@ -91,23 +91,20 @@ func New(o ...Options) (*GPTScript, error) {
9191
return nil, err
9292
}
9393

94-
credStore, err := credentials.NewStore(cliCfg, opts.CredentialContext, cacheClient.CacheDir())
95-
if err != nil {
96-
return nil, err
97-
}
98-
9994
if opts.Runner.RuntimeManager == nil {
10095
opts.Runner.RuntimeManager = runtimes.Default(cacheClient.CacheDir())
10196
}
10297

10398
if err := opts.Runner.RuntimeManager.SetUpCredentialHelpers(context.Background(), cliCfg, opts.Env); err != nil {
10499
return nil, err
105100
}
106-
if err := opts.Runner.RuntimeManager.EnsureCredentialHelpers(context.Background()); err != nil {
101+
102+
credStore, err := credentials.NewStore(cliCfg, opts.Runner.RuntimeManager, opts.CredentialContext, cacheClient.CacheDir())
103+
if err != nil {
107104
return nil, err
108105
}
109106

110-
oaiClient, err := openai.NewClient(credStore, opts.OpenAI, openai.Options{
107+
oaiClient, err := openai.NewClient(ctx, credStore, opts.OpenAI, openai.Options{
111108
Cache: cacheClient,
112109
SetSeed: true,
113110
})

‎pkg/openai/client.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -93,15 +93,15 @@ func complete(opts ...Options) (Options, error) {
9393
return result, err
9494
}
9595

96-
func NewClient(credStore credentials.CredentialStore, opts ...Options) (*Client, error) {
96+
func NewClient(ctx context.Context, credStore credentials.CredentialStore, opts ...Options) (*Client, error) {
9797
opt, err := complete(opts...)
9898
if err != nil {
9999
return nil, err
100100
}
101101

102102
// If the API key is not set, try to get it from the cred store
103103
if opt.APIKey == "" && opt.BaseURL == "" {
104-
cred, exists, err := credStore.Get(BuiltinCredName)
104+
cred, exists, err := credStore.Get(ctx, BuiltinCredName)
105105
if err != nil {
106106
return nil, err
107107
}

‎pkg/prompt/credential.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import (
99
)
1010

1111
func GetModelProviderCredential(ctx context.Context, credStore credentials.CredentialStore, credName, env, message string, envs []string) (string, error) {
12-
cred, exists, err := credStore.Get(credName)
12+
cred, exists, err := credStore.Get(ctx, credName)
1313
if err != nil {
1414
return "", err
1515
}
@@ -25,7 +25,7 @@ func GetModelProviderCredential(ctx context.Context, credStore credentials.Crede
2525
}
2626

2727
k = gjson.Get(result, "key").String()
28-
if err := credStore.Add(credentials.Credential{
28+
if err := credStore.Add(ctx, credentials.Credential{
2929
ToolName: credName,
3030
Type: credentials.CredentialTypeModelProvider,
3131
Env: map[string]string{

‎pkg/remote/remote.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ func (c *Client) clientFromURL(ctx context.Context, apiURL string) (*openai.Clie
116116
}
117117
}
118118

119-
return openai.NewClient(c.credStore, openai.Options{
119+
return openai.NewClient(ctx, c.credStore, openai.Options{
120120
BaseURL: apiURL,
121121
Cache: c.cache,
122122
APIKey: key,
@@ -163,7 +163,7 @@ func (c *Client) load(ctx context.Context, toolName string) (*openai.Client, err
163163
url += "/v1"
164164
}
165165

166-
client, err = openai.NewClient(c.credStore, openai.Options{
166+
client, err = openai.NewClient(ctx, c.credStore, openai.Options{
167167
BaseURL: url,
168168
Cache: c.cache,
169169
CacheKey: prg.EntryToolID,

‎pkg/runner/runner.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -873,12 +873,12 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
873873
// Only try to look up the cred if the tool is on GitHub or has an alias.
874874
// If it is a GitHub tool and has an alias, the alias overrides the tool name, so we use it as the credential name.
875875
if isGitHubTool(toolName) && credentialAlias == "" {
876-
c, exists, err = r.credStore.Get(toolName)
876+
c, exists, err = r.credStore.Get(callCtx.Ctx, toolName)
877877
if err != nil {
878878
return nil, fmt.Errorf("failed to get credentials for tool %s: %w", toolName, err)
879879
}
880880
} else if credentialAlias != "" {
881-
c, exists, err = r.credStore.Get(credentialAlias)
881+
c, exists, err = r.credStore.Get(callCtx.Ctx, credentialAlias)
882882
if err != nil {
883883
return nil, fmt.Errorf("failed to get credentials for tool %s: %w", credentialAlias, err)
884884
}
@@ -942,7 +942,7 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
942942
if (isGitHubTool(toolName) && callCtx.Program.ToolSet[credToolRefs[0].ToolID].Source.Repo != nil) || credentialAlias != "" {
943943
if isEmpty {
944944
log.Warnf("Not saving empty credential for tool %s", toolName)
945-
} else if err := r.credStore.Add(*c); err != nil {
945+
} else if err := r.credStore.Add(callCtx.Ctx, *c); err != nil {
946946
return nil, fmt.Errorf("failed to add credential for tool %s: %w", toolName, err)
947947
}
948948
} else {

‎pkg/sdkserver/run.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import (
1717
type loaderFunc func(context.Context, string, string, ...loader.Options) (types.Program, error)
1818

1919
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) {
20-
g, err := gptscript.New(s.gptscriptOpts, opts)
20+
g, err := gptscript.New(ctx, s.gptscriptOpts, opts)
2121
if err != nil {
2222
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to initialize gptscript: %w", err))
2323
return

‎pkg/sdkserver/server.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ func Start(ctx context.Context, opts Options) error {
5353
// prompt server because it is only used for fmt, parse, etc.
5454
opts.Env = append(opts.Env, fmt.Sprintf("%s=%s", types.PromptTokenEnvVar, token))
5555

56-
g, err := gptscript.New(opts.Options)
56+
g, err := gptscript.New(ctx, opts.Options)
5757
if err != nil {
5858
return err
5959
}

0 commit comments

Comments
 (0)