-
Notifications
You must be signed in to change notification settings - Fork 15
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: add support for credential auth for non-tool-references
Support authentication for agent and workflow tools that use non-tool-references in their definition. That is, a GitHub hosted tool or a local tool. During the implementation and testing for non-tool-references, I also spotted a gap: nested tools. That is, if an agent uses a tool that uses a tool that has a credential, then the authentication for that tool would not be processed. After this change, it will. Signed-off-by: Donnie Adams <[email protected]>
- Loading branch information
Showing
5 changed files
with
279 additions
and
179 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,173 @@ | ||
| package creds | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "net/url" | ||
| "path" | ||
| "slices" | ||
| "strings" | ||
|
|
||
| "github.com/gptscript-ai/go-gptscript" | ||
| gtypes "github.com/gptscript-ai/gptscript/pkg/types" | ||
| "github.com/obot-platform/obot/pkg/system" | ||
| ) | ||
|
|
||
| func DetermineCredsAndCredNames(prg *gptscript.Program, tool gptscript.Tool, name string) ([]string, []string, error) { | ||
| // The available tool references from this tool are the tool itself and any tool this tool exports. | ||
| toolRefs := make([]toolRef, 0, len(tool.Export)+len(tool.Tools)+1) | ||
| toolRefs = append(toolRefs, toolRef{ | ||
| ToolReference: gptscript.ToolReference{ | ||
| Reference: name, | ||
| ToolID: prg.EntryToolID, | ||
| }, | ||
| name: name, | ||
| }) | ||
| for _, t := range tool.Tools { | ||
| for _, ref := range tool.ToolMapping[t] { | ||
| toolRefs = append(toolRefs, toolRef{ | ||
| ToolReference: ref, | ||
| name: t, | ||
| }) | ||
| } | ||
| } | ||
|
|
||
| credentials := make([]string, 0, len(tool.Credentials)+len(tool.Export)+len(tool.Tools)) | ||
| credentialNames := make([]string, 0, len(tool.Credentials)+len(tool.Export)+len(tool.Tools)) | ||
| seen := make(map[string]struct{}) | ||
| for len(toolRefs) > 0 { | ||
| ref := toolRefs[0] | ||
| toolRefs = toolRefs[1:] | ||
|
|
||
| if _, ok := seen[ref.ToolID]; ok { | ||
| continue | ||
| } | ||
| seen[ref.ToolID] = struct{}{} | ||
|
|
||
| t := prg.ToolSet[ref.ToolID] | ||
|
|
||
| // Add the tools that this tool exports if we haven't already seen them. | ||
| for _, e := range t.Export { | ||
| refs := t.ToolMapping[e] | ||
| for _, r := range refs { | ||
| if _, ok := seen[r.ToolID]; !ok { | ||
| toolRefs = append(toolRefs, toolRef{ | ||
| ToolReference: r, | ||
| name: ref.name, | ||
| }) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| for _, cred := range t.Credentials { | ||
| parsedCred := cred | ||
| credToolName, credSubTool := gtypes.SplitToolRef(cred) | ||
| if strings.HasPrefix(credToolName, ".") { | ||
| toolName, _ := gtypes.SplitToolRef(ref.Reference) | ||
| if !path.IsAbs(toolName) { | ||
| if !strings.HasPrefix(toolName, ".") { | ||
| toolName, _ = gtypes.SplitToolRef(ref.name) | ||
| } else { | ||
| toolName = path.Join(ref.name, toolName) | ||
| } | ||
| } | ||
|
|
||
| refURL, err := url.Parse(toolName) | ||
| if err != nil { | ||
| continue | ||
| } | ||
|
|
||
| if strings.HasSuffix(refURL.Path, ".gpt") { | ||
| refURL.Path = path.Dir(refURL.Path) | ||
| } | ||
|
|
||
| refURL.Path = path.Join(refURL.Path, credToolName) | ||
| parsedCred = refURL.String() | ||
| if refURL.Host == "" { | ||
| // This is only a path, so url unescape it. | ||
| // No need to check the error here, we would have errored when parsing. | ||
| parsedCred, _ = url.PathUnescape(parsedCred) | ||
| } | ||
|
|
||
| if credSubTool != "" { | ||
| parsedCred = fmt.Sprintf("%s from %s", credSubTool, parsedCred) | ||
| } | ||
| } | ||
|
|
||
| if parsedCred != "" && !slices.Contains(credentials, parsedCred) { | ||
| credentials = append(credentials, parsedCred) | ||
| } | ||
|
|
||
| credNames, err := determineCredentialNames(prg, prg.ToolSet[ref.ToolID], cred) | ||
| if err != nil { | ||
| return credentials, credentialNames, err | ||
| } | ||
|
|
||
| for _, n := range credNames { | ||
| if !slices.Contains(credentialNames, n) { | ||
| credentialNames = append(credentialNames, n) | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| return credentials, credentialNames, nil | ||
| } | ||
|
|
||
| func determineCredentialNames(prg *gptscript.Program, tool gptscript.Tool, toolName string) ([]string, error) { | ||
| if toolName == system.ModelProviderCredential { | ||
| return []string{system.ModelProviderCredential}, nil | ||
| } | ||
|
|
||
| var subTool string | ||
| parsedToolName, alias, args, err := gtypes.ParseCredentialArgs(toolName, "") | ||
| if err != nil { | ||
| parsedToolName, subTool = gtypes.SplitToolRef(toolName) | ||
| parsedToolName, alias, args, err = gtypes.ParseCredentialArgs(parsedToolName, "") | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| } | ||
|
|
||
| if alias != "" { | ||
| return []string{alias}, nil | ||
| } | ||
|
|
||
| if args == nil { | ||
| // This is a tool and not the credential format. Parse the tool from the program to determine the alias | ||
| toolNames := make([]string, 0, len(tool.Credentials)) | ||
| if subTool == "" { | ||
| toolName = parsedToolName | ||
| } | ||
| for _, cred := range tool.Credentials { | ||
| if cred == toolName { | ||
| if len(tool.ToolMapping[cred]) == 0 { | ||
| return nil, fmt.Errorf("cannot find credential name for tool %q", toolName) | ||
| } | ||
|
|
||
| for _, ref := range tool.ToolMapping[cred] { | ||
| for _, c := range prg.ToolSet[ref.ToolID].ExportCredentials { | ||
| names, err := determineCredentialNames(prg, prg.ToolSet[ref.ToolID], c) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| toolNames = append(toolNames, names...) | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| if len(toolNames) > 0 { | ||
| return toolNames, nil | ||
| } | ||
|
|
||
| return nil, fmt.Errorf("tool %q not found in program", toolName) | ||
| } | ||
|
|
||
| return []string{toolName}, nil | ||
| } | ||
|
|
||
| type toolRef struct { | ||
| gptscript.ToolReference | ||
| name string | ||
| } |
Oops, something went wrong.