Skip to content

enhance: add credential overrides #263

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions docs/docs/03-tools/04-credentials.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
41 changes: 36 additions & 5 deletions pkg/cli/credential.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"fmt"
"os"
"sort"
"strings"
"text/tabwriter"

cmd2 "github.com/acorn-io/cmd"
Expand All @@ -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) {
Expand Down Expand Up @@ -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)
}
}
}

Expand Down
35 changes: 19 additions & 16 deletions pkg/cli/gptscript.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -134,6 +135,8 @@ func (r *GPTScript) NewGPTScriptOpts() (gptscript.Options, error) {
opts.Runner.EndPort = endNum
}

opts.Runner.CredentialOverride = r.CredentialOverride

return opts, nil
}

Expand Down
55 changes: 55 additions & 0 deletions pkg/runner/credentials.go
Original file line number Diff line number Diff line change
@@ -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
}
30 changes: 26 additions & 4 deletions pkg/runner/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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{}
Expand All @@ -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) {
Expand All @@ -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 {
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down