enhance: add support for args and aliases to credential tools

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

Author: g-linville

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:
 
@@ -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`.
@@ -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:
@@ -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

pkg/parser/parser.go

@@ -136,7 +136,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, csv(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", 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", credToolName, err)
+				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
+			input := ""
+			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)
 			}
@@ -859,9 +878,11 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
 				return nil, fmt.Errorf("failed to unmarshal credential tool %s response: %w", credToolName, err)
 			}
 
-			cred = &credentials.Credential{
-				ToolName: credToolName,
-				Env:      envMap.Env,
+			cred = &credentials.Credential{Env: envMap.Env}
+			if credentialAlias != "" {
+				cred.ToolName = credentialAlias
+			} else {
+				cred.ToolName = toolName
 			}
 
 			isEmpty := true
@@ -872,15 +893,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"
+
+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 (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)
+				}
+			}
+		})
+	}
+}