enhance: add support for args and aliases to credential tools (gptscript-ai#433)

Signed-off-by: Grant Linville <[email protected]>

Author: g-linville

.github/workflows/test.yaml

@@ -27,12 +27,11 @@ jobs:
           cache: false
           go-version: "1.21"
       - name: Build UI
-        run: |
-          if [ ${{ matrix.os }} != 'windows-latest' ]; then
-            make build-ui
-          fi
+        if: matrix.os == 'ubuntu-22.04'
+        run: make build-ui
         shell: bash
       - name: Validate
+        if: matrix.os == 'ubuntu-22.04'
         run: make validate
       - name: Build
         run: make build

.golangci.yml

@@ -3,7 +3,7 @@ run:
 
 output:
   formats:
-  - github-actions
+  - format: colored-line-number
 
 linters:
   disable-all: true

Makefile

@@ -21,7 +21,7 @@ tidy:
 test:
 	go test -v ./...
 
-GOLANGCI_LINT_VERSION ?= v1.56.1
+GOLANGCI_LINT_VERSION ?= v1.59.0
 lint:
 	if ! command -v golangci-lint &> /dev/null; then \
   		echo "Could not find golangci-lint, installing version $(GOLANGCI_LINT_VERSION)."; \

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

@@ -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:
 
@@ -42,14 +41,39 @@ LLM about it or even tell the LLM about the tool.
 If GPTScript has called the credential provider tool in the same context (more on that later), then it will use the stored
 credential instead of fetching it again.
 
-You can also specify multiple credential tools for the same script:
+You can also specify multiple credential tools for the same script, but they must be on separate lines:
 
 ```yaml
-credentials: credential-tool-1.gpt, credential-tool-2.gpt
+credentials: credential-tool-1.gpt
+credentials: 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`.
@@ -67,16 +91,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:
@@ -97,13 +135,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

pkg/parser/parser.go

@@ -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))...)
+		tool.Parameters.Credentials = append(tool.Parameters.Credentials, value)
 	default:
 		return false, nil
 	}

pkg/runner/runner.go

@@ -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)
 			}
 		}
 
@@ -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
+			var input string
+			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)
 			}
@@ -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
@@ -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)
 			}
 		}
 

pkg/types/credential_test.go

@@ -0,0 +1,153 @@
+package types
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseCredentialArgs(t *testing.T) {
+	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 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.Equal(t, v, args[k], "unexpected value for args[%s]", k)
+			}
+		})
+	}
+}