Skip to content

enhance: add support for args and aliases to credential tools #433

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 9 commits into from
Jun 6, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
59 changes: 50 additions & 9 deletions docs/docs/03-tools/04-credentials.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ directly from user input) and conveniently set them in the environment before ru
A credential provider tool looks just like any other GPTScript, with the following caveats:
- It cannot call the LLM and must run a command.
- It must print contents to stdout in the format `{"env":{"ENV_VAR_1":"value1","ENV_VAR_2":"value2"}}`.
- Any args defined on the tool will be ignored.

Here is a simple example of a credential provider tool that uses the builtin `sys.prompt` to ask the user for some input:

Expand Down Expand Up @@ -50,6 +49,30 @@ credentials: credential-tool-1.gpt, credential-tool-2.gpt
(tool stuff here)
```

## Credential Tool Arguments

A credential tool may define arguments. Here is an example:

```yaml
name: my-credential-tool
args: env: the environment variable to set
args: val: the value to set it to

#!/usr/bin/env bash

echo "{\"env\":{\"$ENV\":\"$VAL\"}}"
```

When you reference this credential tool in another file, you can use syntax like this to set both arguments:

```yaml
credential: my-credential-tool.gpt with MY_ENV_VAR as env and "my value" as val

(tool stuff here)
```

In this example, the tool's output would be `{"env":{"MY_ENV_VAR":"my value"}}`

## Storing Credentials

By default, credentials are automatically stored in a config file at `$XDG_CONFIG_HOME/gptscript/config.json`.
Expand All @@ -67,16 +90,30 @@ is called `gptscript-credential-wincred.exe`.)
There will likely be support added for other credential stores in the future.

:::note
Credentials received from credential provider tools that are not on GitHub (such as a local file) will not be stored
in the credentials store.
Credentials received from credential provider tools that are not on GitHub (such as a local file) and do not have an alias
will not be stored in the credentials store.
:::

## Credential Aliases

When you reference a credential tool in your script, you can give it an alias using the `as` keyword like this:

```yaml
credentials: my-credential-tool.gpt as myAlias

(tool stuff here)
```

This will store the resulting credential with the name `myAlias`.
This is useful when you want to reference the same credential tool in scripts that need to handle different credentials,
or when you want to store credentials that were provided by a tool that is not on GitHub.

## Credential Contexts

Each stored credential is uniquely identified by the name of its provider tool and the name of its context. A credential
context is basically a namespace for credentials. If you have multiple credentials from the same provider tool, you can
switch between them by defining them in different credential contexts. The default context is called `default`, and this
is used if none is specified.
Each stored credential is uniquely identified by the name of its provider tool (or alias, if one was specified) and the name of its context.
A credential context is basically a namespace for credentials. If you have multiple credentials from the same provider tool,
you can switch between them by defining them in different credential contexts. The default context is called `default`,
and this is used if none is specified.

You can set the credential context to use with the `--credential-context` flag when running GPTScript. For
example:
Expand All @@ -97,13 +134,17 @@ credentials in all contexts with `--all-contexts`.
You can delete a credential by running the following command:

```bash
gptscript credential delete --credential-context <credential context> <credential tool name>
gptscript credential delete --credential-context <credential context> <credential 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
## Credential Overrides (Advanced)

:::note
The syntax for this will change at some point in the future.
:::

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
Expand Down
2 changes: 1 addition & 1 deletion pkg/parser/parser.go
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ func isParam(line string, tool *types.Tool) (_ bool, err error) {
return false, err
}
case "credentials", "creds", "credential", "cred":
tool.Parameters.Credentials = append(tool.Parameters.Credentials, csv(strings.ToLower(value))...)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There was no need to be converting it to lowercase here.

tool.Parameters.Credentials = append(tool.Parameters.Credentials, csv(value)...)
default:
return false, nil
}
Expand Down
53 changes: 38 additions & 15 deletions pkg/runner/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -816,17 +816,27 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
continue
}

toolName, credentialAlias, args, err := types.ParseCredentialArgs(credToolName, callCtx.Input)
if err != nil {
return nil, fmt.Errorf("failed to parse credential tool %q: %w", credToolName, err)
}

var (
cred *credentials.Credential
exists bool
err error
)

// Only try to look up the cred if the tool is on GitHub.
if isGitHubTool(credToolName) {
cred, exists, err = store.Get(credToolName)
// 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 == "" {
cred, exists, err = store.Get(toolName)
if err != nil {
return nil, fmt.Errorf("failed to get credentials for tool %s: %w", credToolName, err)
return nil, fmt.Errorf("failed to get credentials for tool %s: %w", toolName, err)
}
} else if credentialAlias != "" {
cred, exists, err = store.Get(credentialAlias)
if err != nil {
return nil, fmt.Errorf("failed to get credentials for tool %s: %w", credentialAlias, err)
}
}

Expand All @@ -838,12 +848,21 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
return nil, fmt.Errorf("failed to find ID for tool %s", credToolName)
}

subCtx, err := callCtx.SubCall(callCtx.Ctx, "", credToolRefs[0].ToolID, "", engine.CredentialToolCategory) // leaving callID as "" will cause it to be set by the engine
input := ""
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
input := ""
var input string

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed

if args != nil {
inputBytes, err := json.Marshal(args)
if err != nil {
return nil, fmt.Errorf("failed to marshal args for tool %s: %w", credToolName, err)
}
input = string(inputBytes)
}

subCtx, err := callCtx.SubCall(callCtx.Ctx, input, credToolRefs[0].ToolID, "", engine.CredentialToolCategory) // leaving callID as "" will cause it to be set by the engine
if err != nil {
return nil, fmt.Errorf("failed to create subcall context for tool %s: %w", credToolName, err)
}

res, err := r.call(subCtx, monitor, env, "")
res, err := r.call(subCtx, monitor, env, input)
if err != nil {
return nil, fmt.Errorf("failed to run credential tool %s: %w", credToolName, err)
}
Expand All @@ -860,9 +879,13 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
}

cred = &credentials.Credential{
ToolName: credToolName,
Type: credentials.CredentialTypeTool,
Env: envMap.Env,
Type: credentials.CredentialTypeTool,
Env: envMap.Env,
}
if credentialAlias != "" {
cred.ToolName = credentialAlias
} else {
cred.ToolName = toolName
}

isEmpty := true
Expand All @@ -873,15 +896,15 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
}
}

// Only store the credential if the tool is on GitHub, and the credential is non-empty.
if isGitHubTool(credToolName) && callCtx.Program.ToolSet[credToolRefs[0].ToolID].Source.Repo != nil {
// Only store the credential if the tool is on GitHub or has an alias, and the credential is non-empty.
if (isGitHubTool(toolName) && callCtx.Program.ToolSet[credToolRefs[0].ToolID].Source.Repo != nil) || credentialAlias != "" {
if isEmpty {
log.Warnf("Not saving empty credential for tool %s", credToolName)
log.Warnf("Not saving empty credential for tool %s", toolName)
} else if err := store.Add(*cred); err != nil {
return nil, fmt.Errorf("failed to add credential for tool %s: %w", credToolName, err)
return nil, fmt.Errorf("failed to add credential for tool %s: %w", toolName, err)
}
} else {
log.Warnf("Not saving credential for local tool %s - credentials will only be saved for tools from GitHub.", credToolName)
log.Warnf("Not saving credential for tool %s - credentials will only be saved for tools from GitHub, or tools that use aliases.", toolName)
}
}

Expand Down
153 changes: 153 additions & 0 deletions pkg/types/credential_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
package types

import "testing"

func TestParseCredentialArgs(t *testing.T) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💯

tests := []struct {
name string
toolName string
input string
expectedName string
expectedAlias string
expectedArgs map[string]string
wantErr bool
}{
{
name: "empty",
toolName: "",
expectedName: "",
expectedAlias: "",
},
{
name: "tool name only",
toolName: "myCredentialTool",
expectedName: "myCredentialTool",
expectedAlias: "",
},
{
name: "tool name and alias",
toolName: "myCredentialTool as myAlias",
expectedName: "myCredentialTool",
expectedAlias: "myAlias",
},
{
name: "tool name with one arg",
toolName: "myCredentialTool with value1 as arg1",
expectedName: "myCredentialTool",
expectedAlias: "",
expectedArgs: map[string]string{
"arg1": "value1",
},
},
{
name: "tool name with two args",
toolName: "myCredentialTool with value1 as arg1 and value2 as arg2",
expectedName: "myCredentialTool",
expectedAlias: "",
expectedArgs: map[string]string{
"arg1": "value1",
"arg2": "value2",
},
},
{
name: "tool name with alias and one arg",
toolName: "myCredentialTool as myAlias with value1 as arg1",
expectedName: "myCredentialTool",
expectedAlias: "myAlias",
expectedArgs: map[string]string{
"arg1": "value1",
},
},
{
name: "tool name with alias and two args",
toolName: "myCredentialTool as myAlias with value1 as arg1 and value2 as arg2",
expectedName: "myCredentialTool",
expectedAlias: "myAlias",
expectedArgs: map[string]string{
"arg1": "value1",
"arg2": "value2",
},
},
{
name: "tool name with quoted args",
toolName: `myCredentialTool with "value one" as arg1 and "value two" as arg2`,
expectedName: "myCredentialTool",
expectedAlias: "",
expectedArgs: map[string]string{
"arg1": "value one",
"arg2": "value two",
},
},
{
name: "tool name with arg references",
toolName: `myCredentialTool with ${var1} as arg1 and ${var2} as arg2`,
input: `{"var1": "value1", "var2": "value2"}`,
expectedName: "myCredentialTool",
expectedAlias: "",
expectedArgs: map[string]string{
"arg1": "value1",
"arg2": "value2",
},
},
{
name: "tool name with alias but no 'as' (invalid)",
toolName: "myCredentialTool myAlias",
wantErr: true,
},
{
name: "tool name with 'as' but no alias (invalid)",
toolName: "myCredentialTool as",
wantErr: true,
},
{
name: "tool with 'with' but no args (invalid)",
toolName: "myCredentialTool with",
wantErr: true,
},
{
name: "tool with args but no 'with' (invalid)",
toolName: "myCredentialTool value1 as arg1",
wantErr: true,
},
{
name: "tool with trailing 'and' (invalid)",
toolName: "myCredentialTool with value1 as arg1 and",
wantErr: true,
},
{
name: "tool with quoted arg but the quote is unterminated (invalid)",
toolName: `myCredentialTool with "value one" as arg1 and "value two as arg2`,
wantErr: true,
},
{
name: "invalid input",
toolName: "myCredentialTool",
input: `{"asdf":"asdf"`,
wantErr: true,
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
originalName, alias, args, err := ParseCredentialArgs(tt.toolName, tt.input)
if (err != nil) != tt.wantErr {
t.Errorf("ParseCredentialArgs() error = %v, wantErr %v", err, tt.wantErr)
return
}
if originalName != tt.expectedName {
t.Errorf("ParseCredentialArgs() originalName = %v, expectedName %v", originalName, tt.expectedName)
}
if alias != tt.expectedAlias {
t.Errorf("ParseCredentialArgs() alias = %v, expectedAlias %v", alias, tt.expectedAlias)
}
if len(args) != len(tt.expectedArgs) {
t.Errorf("ParseCredentialArgs() args = %v, expectedArgs %v", args, tt.expectedArgs)
}
for k, v := range tt.expectedArgs {
if args[k] != v {
t.Errorf("ParseCredentialArgs() args[%s] = %v, expectedArgs[%s] %v", k, args[k], k, v)
}
}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider using the github.com/stretchr/testify/assert and github.com/stretchr/testify/require packages:

Suggested change
originalName, alias, args, err := ParseCredentialArgs(tt.toolName, tt.input)
if (err != nil) != tt.wantErr {
t.Errorf("ParseCredentialArgs() error = %v, wantErr %v", err, tt.wantErr)
return
}
if originalName != tt.expectedName {
t.Errorf("ParseCredentialArgs() originalName = %v, expectedName %v", originalName, tt.expectedName)
}
if alias != tt.expectedAlias {
t.Errorf("ParseCredentialArgs() alias = %v, expectedAlias %v", alias, tt.expectedAlias)
}
if len(args) != len(tt.expectedArgs) {
t.Errorf("ParseCredentialArgs() args = %v, expectedArgs %v", args, tt.expectedArgs)
}
for k, v := range tt.expectedArgs {
if args[k] != v {
t.Errorf("ParseCredentialArgs() args[%s] = %v, expectedArgs[%s] %v", k, args[k], k, v)
}
}
originalName, alias, args, err := ParseCredentialArgs(tt.toolName, tt.input)
// Require halts the test immediately when an assertion fails
if tt.wantErr {
require.Error(t, err, "expected an error but got none")
return
}
require.NoError(t, err, "did not expect an error but got one")
require.Equal(t, tt.expectedName, originalName, "unexpected original name")
require.Equal(t, tt.expectedAlias, alias, "unexpected alias")
require.Equal(t, len(tt.expectedArgs), len(args), "unexpected number of args")
for k, v := range tt.expectedArgs {
// Assert marks the test as failed when an assertion fails, but doesn't halt, allowing the test to progress further and identify more potential issues
assert.Equal(t, v, args[k], "unexpected value for args[%s]", k)
}

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh thanks lol. I let Copilot generate basically all of this stuff and did not bother to think how it could be improved. Lazy me.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed

})
}
}
Loading
Loading