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 }