enhance: add credential overrides (#263)

Signed-off-by: Grant Linville <[email protected]>
Co-authored-by: Nick Hale <[email protected]>

Author: g-linville

docs/docs/03-tools/04-credentials.md

@@ -99,3 +99,54 @@ You can delete a credential by running the following command:
 ```bash
 gptscript credential delete --credential-context <credential context> <credential tool name>
 ```
+
+The `--show-env-vars` argument will also display the names of the environment variables that are set by the credential.
+This is useful when working with credential overrides.
+
+## Credential Overrides
+
+You can bypass credential tools and stored credentials by setting the `--credential-override` argument (or the
+`GPTSCRIPT_CREDENTIAL_OVERRIDE` environment variable) when running GPTScript. To set up a credential override, you
+need to be aware of which environment variables the credential tool sets. You can find this out by running the
+`gptscript credential --show-env-vars` command.
+
+### Format
+
+The `--credential-override` argument must be formatted in one of the following three ways:
+
+#### 1. Key-Value Pairs
+
+`toolA:ENV_VAR_1=value1,ENV_VAR_2=value2;toolB:ENV_VAR_1=value3,ENV_VAR_2=value4`
+
+In this example, both `toolA` and `toolB` provide the variables `ENV_VAR_1` and `ENV_VAR_2`.
+This will set the environment variables `ENV_VAR_1` and `ENV_VAR_2` to the specific values provided for each tool.
+
+#### 2. Environment Variables
+
+`toolA:ENV_VAR_1,ENV_VAR_2;toolB:ENV_VAR_3,ENV_VAR_4`
+
+In this example, `toolA` provides the variables `ENV_VAR_1` and `ENV_VAR_2`, and `toolB` provides the variables `ENV_VAR_3` and `ENV_VAR_4`.
+This will read the values of `ENV_VAR_1` through `ENV_VAR_4` from the current environment and set them for each tool.
+This is a direct mapping of environment variable names. **This is not recommended when overriding credentials for
+multiple tools that use the same environment variable names.**
+
+#### 3. Environment Variable Mapping
+
+`toolA:ENV_VAR_1->TOOL_A_ENV_VAR_1,ENV_VAR_2->TOOL_A_ENV_VAR_2;toolB:ENV_VAR_1->TOOL_B_ENV_VAR_1,ENV_VAR_2->TOOL_B_ENV_VAR_2`
+
+In this example, `toolA` and `toolB` both provide the variables `ENV_VAR_1` and `ENV_VAR_2`.
+This will set the environment variables `ENV_VAR_1` and `ENV_VAR_2` to the values of `TOOL_A_ENV_VAR_1` and
+`TOOL_A_ENV_VAR_2` from the current environment for `toolA`. The same applies for `toolB`, but with the values of
+`TOOL_B_ENV_VAR_1` and `TOOL_B_ENV_VAR_2`. This is a mapping of one environment variable name to another.
+
+### Real-World Example
+
+Here is an example of how you can use a credential override to skip running the credential tool for the Brave Search tool:
+
+```bash
+gptscript --credential-override "github.com/gptscript-ai/search/brave-credential:GPTSCRIPT_BRAVE_SEARCH_TOKEN->MY_BRAVE_SEARCH_TOKEN" github.com/gptscript-ai/search/brave '{"q": "cute cats"}'
+```
+
+If you run this command, rather than being prompted by the credential tool for your token, GPTScript will read the contents
+of the environment variable `MY_BRAVE_SEARCH_TOKEN` and set that as the variable `GPTSCRIPT_BRAVE_SEARCH_TOKEN` when it runs
+the script.

pkg/cli/credential.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"sort"
+	"strings"
 	"text/tabwriter"
 
 	cmd2 "github.com/acorn-io/cmd"
@@ -15,6 +16,7 @@ import (
 type Credential struct {
 	root        *GPTScript
 	AllContexts bool `usage:"List credentials for all contexts" local:"true"`
+	ShowEnvVars bool `usage:"Show names of environment variables in each credential" local:"true"`
 }
 
 func (c *Credential) Customize(cmd *cobra.Command) {
@@ -57,18 +59,47 @@ func (c *Credential) Run(_ *cobra.Command, _ []string) error {
 
 		w := tabwriter.NewWriter(os.Stdout, 10, 1, 3, ' ', 0)
 		defer w.Flush()
-		_, _ = w.Write([]byte("CONTEXT\tTOOL\n"))
-		for _, cred := range creds {
-			_, _ = fmt.Fprintf(w, "%s\t%s\n", cred.Context, cred.ToolName)
+
+		if c.ShowEnvVars {
+			_, _ = w.Write([]byte("CONTEXT\tTOOL\tENVIRONMENT VARIABLES\n"))
+
+			for _, cred := range creds {
+				envVars := make([]string, 0, len(cred.Env))
+				for envVar := range cred.Env {
+					envVars = append(envVars, envVar)
+				}
+				sort.Strings(envVars)
+				_, _ = fmt.Fprintf(w, "%s\t%s\t%s\n", cred.Context, cred.ToolName, strings.Join(envVars, ", "))
+			}
+		} else {
+			_, _ = w.Write([]byte("CONTEXT\tTOOL\n"))
+			for _, cred := range creds {
+				_, _ = fmt.Fprintf(w, "%s\t%s\n", cred.Context, cred.ToolName)
+			}
 		}
 	} else {
 		// Sort credentials by tool name
 		sort.Slice(creds, func(i, j int) bool {
 			return creds[i].ToolName < creds[j].ToolName
 		})
 
-		for _, cred := range creds {
-			fmt.Println(cred.ToolName)
+		if c.ShowEnvVars {
+			w := tabwriter.NewWriter(os.Stdout, 10, 1, 3, ' ', 0)
+			defer w.Flush()
+			_, _ = w.Write([]byte("TOOL\tENVIRONMENT VARIABLES\n"))
+
+			for _, cred := range creds {
+				envVars := make([]string, 0, len(cred.Env))
+				for envVar := range cred.Env {
+					envVars = append(envVars, envVar)
+				}
+				sort.Strings(envVars)
+				_, _ = fmt.Fprintf(w, "%s\t%s\n", cred.ToolName, strings.Join(envVars, ", "))
+			}
+		} else {
+			for _, cred := range creds {
+				fmt.Println(cred.ToolName)
+			}
 		}
 	}
 

pkg/cli/gptscript.go

@@ -39,22 +39,23 @@ type GPTScript struct {
 	CacheOptions
 	OpenAIOptions
 	DisplayOptions
-	Color             *bool  `usage:"Use color in output (default true)" default:"true"`
-	Confirm           bool   `usage:"Prompt before running potentially dangerous commands"`
-	Debug             bool   `usage:"Enable debug logging"`
-	Quiet             *bool  `usage:"No output logging (set --quiet=false to force on even when there is no TTY)" short:"q"`
-	Output            string `usage:"Save output to a file, or - for stdout" short:"o"`
-	Input             string `usage:"Read input from a file (\"-\" for stdin)" short:"f"`
-	SubTool           string `usage:"Use tool of this name, not the first tool in file" local:"true"`
-	Assemble          bool   `usage:"Assemble tool to a single artifact, saved to --output" hidden:"true" local:"true"`
-	ListModels        bool   `usage:"List the models available and exit" local:"true"`
-	ListTools         bool   `usage:"List built-in tools and exit" local:"true"`
-	Server            bool   `usage:"Start server" local:"true"`
-	ListenAddress     string `usage:"Server listen address" default:"127.0.0.1:9090" local:"true"`
-	Chdir             string `usage:"Change current working directory" short:"C"`
-	Daemon            bool   `usage:"Run tool as a daemon" local:"true" hidden:"true"`
-	Ports             string `usage:"The port range to use for ephemeral daemon ports (ex: 11000-12000)" hidden:"true"`
-	CredentialContext string `usage:"Context name in which to store credentials" default:"default"`
+	Color              *bool  `usage:"Use color in output (default true)" default:"true"`
+	Confirm            bool   `usage:"Prompt before running potentially dangerous commands"`
+	Debug              bool   `usage:"Enable debug logging"`
+	Quiet              *bool  `usage:"No output logging (set --quiet=false to force on even when there is no TTY)" short:"q"`
+	Output             string `usage:"Save output to a file, or - for stdout" short:"o"`
+	Input              string `usage:"Read input from a file (\"-\" for stdin)" short:"f"`
+	SubTool            string `usage:"Use tool of this name, not the first tool in file" local:"true"`
+	Assemble           bool   `usage:"Assemble tool to a single artifact, saved to --output" hidden:"true" local:"true"`
+	ListModels         bool   `usage:"List the models available and exit" local:"true"`
+	ListTools          bool   `usage:"List built-in tools and exit" local:"true"`
+	Server             bool   `usage:"Start server" local:"true"`
+	ListenAddress      string `usage:"Server listen address" default:"127.0.0.1:9090" local:"true"`
+	Chdir              string `usage:"Change current working directory" short:"C"`
+	Daemon             bool   `usage:"Run tool as a daemon" local:"true" hidden:"true"`
+	Ports              string `usage:"The port range to use for ephemeral daemon ports (ex: 11000-12000)" hidden:"true"`
+	CredentialContext  string `usage:"Context name in which to store credentials" default:"default"`
+	CredentialOverride string `usage:"Credentials to override (ex: --credential-override github.com/example/cred-tool:API_TOKEN=1234)"`
 }
 
 func New() *cobra.Command {
@@ -134,6 +135,8 @@ func (r *GPTScript) NewGPTScriptOpts() (gptscript.Options, error) {
 		opts.Runner.EndPort = endNum
 	}
 
+	opts.Runner.CredentialOverride = r.CredentialOverride
+
 	return opts, nil
 }
 

pkg/runner/credentials.go

@@ -0,0 +1,55 @@
+package runner
+
+import (
+	"fmt"
+	"os"
+	"strings"
+)
+
+// parseCredentialOverrides parses a string of credential overrides that the user provided as a command line arg.
+// The format of credential overrides can be one of three things:
+// tool1:ENV1,ENV2;tool2:ENV1,ENV2 (direct mapping of environment variables)
+// tool1:ENV1=VALUE1,ENV2=VALUE2;tool2:ENV1=VALUE1,ENV2=VALUE2 (key-value pairs)
+// tool1:ENV1->OTHER_ENV1,ENV2->OTHER_ENV2;tool2:ENV1->OTHER_ENV1,ENV2->OTHER_ENV2 (mapping to other environment variables)
+//
+// This function turns it into a map[string]map[string]string like this:
+//
+//	{
+//	  "tool1": {
+//	    "ENV1": "VALUE1",
+//	    "ENV2": "VALUE2",
+//	  },
+//	  "tool2": {
+//	    "ENV1": "VALUE1",
+//	    "ENV2": "VALUE2",
+//	  },
+//	}
+func parseCredentialOverrides(override string) (map[string]map[string]string, error) {
+	credentialOverrides := make(map[string]map[string]string)
+
+	for _, o := range strings.Split(override, ";") {
+		toolName, envs, found := strings.Cut(o, ":")
+		if !found {
+			return nil, fmt.Errorf("invalid credential override: %s", o)
+		}
+		envMap := make(map[string]string)
+		for _, env := range strings.Split(envs, ",") {
+			key, value, found := strings.Cut(env, "=")
+			if !found {
+				var envVar string
+				key, envVar, found = strings.Cut(env, "->")
+				if found {
+					// User did a mapping of key -> other env var, so look up the value.
+					value = os.Getenv(envVar)
+				} else {
+					// User just passed an env var name as the key, so look up the value.
+					value = os.Getenv(key)
+				}
+			}
+			envMap[key] = value
+		}
+		credentialOverrides[toolName] = envMap
+	}
+
+	return credentialOverrides, nil
+}

pkg/runner/runner.go

@@ -30,10 +30,11 @@ type Monitor interface {
 type MonitorKey struct{}
 
 type Options struct {
-	MonitorFactory MonitorFactory        `usage:"-"`
-	RuntimeManager engine.RuntimeManager `usage:"-"`
-	StartPort      int64                 `usage:"-"`
-	EndPort        int64                 `usage:"-"`
+	MonitorFactory     MonitorFactory        `usage:"-"`
+	RuntimeManager     engine.RuntimeManager `usage:"-"`
+	StartPort          int64                 `usage:"-"`
+	EndPort            int64                 `usage:"-"`
+	CredentialOverride string                `usage:"-"`
 }
 
 func complete(opts ...Options) (result Options) {
@@ -42,6 +43,7 @@ func complete(opts ...Options) (result Options) {
 		result.RuntimeManager = types.FirstSet(opt.RuntimeManager, result.RuntimeManager)
 		result.StartPort = types.FirstSet(opt.StartPort, result.StartPort)
 		result.EndPort = types.FirstSet(opt.EndPort, result.EndPort)
+		result.CredentialOverride = types.FirstSet(opt.CredentialOverride, result.CredentialOverride)
 	}
 	if result.MonitorFactory == nil {
 		result.MonitorFactory = noopFactory{}
@@ -62,6 +64,7 @@ type Runner struct {
 	ports          engine.Ports
 	credCtx        string
 	credMutex      sync.Mutex
+	credOverrides  string
 }
 
 func New(client engine.Model, credCtx string, opts ...Options) (*Runner, error) {
@@ -73,6 +76,7 @@ func New(client engine.Model, credCtx string, opts ...Options) (*Runner, error)
 		runtimeManager: opt.RuntimeManager,
 		credCtx:        credCtx,
 		credMutex:      sync.Mutex{},
+		credOverrides:  opt.CredentialOverride,
 	}
 
 	if opt.StartPort != 0 {
@@ -336,6 +340,7 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
 	r.credMutex.Lock()
 	defer r.credMutex.Unlock()
 
+	// Set up the credential store.
 	c, err := config.ReadCLIConfig("")
 	if err != nil {
 		return nil, fmt.Errorf("failed to read CLI config: %w", err)
@@ -346,7 +351,24 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
 		return nil, fmt.Errorf("failed to create credentials store: %w", err)
 	}
 
+	// Parse the credential overrides from the command line argument, if there are any.
+	var credOverrides map[string]map[string]string
+	if r.credOverrides != "" {
+		credOverrides, err = parseCredentialOverrides(r.credOverrides)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse credential overrides: %w", err)
+		}
+	}
+
 	for _, credToolName := range callCtx.Tool.Credentials {
+		// Check whether the credential was overridden before we attempt to find it in the store or run the tool.
+		if override, exists := credOverrides[credToolName]; exists {
+			for k, v := range override {
+				env = append(env, fmt.Sprintf("%s=%s", k, v))
+			}
+			continue
+		}
+
 		var (
 			cred   *credentials.Credential
 			exists bool