Skip to content
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

feat: add credential management #65

Merged
merged 15 commits into from
Sep 19, 2024
27 changes: 27 additions & 0 deletions credentials.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
package gptscript

import "time"

type CredentialType string

const (
CredentialTypeTool CredentialType = "tool"
CredentialTypeModelProvider CredentialType = "modelProvider"
)

type Credential struct {
Context string `json:"context"`
ToolName string `json:"toolName"`
Type CredentialType `json:"type"`
Env map[string]string `json:"env"`
Ephemeral bool `json:"ephemeral,omitempty"`
ExpiresAt *time.Time `json:"expiresAt"`
RefreshToken string `json:"refreshToken"`
}

type CredentialRequest struct {
Content string `json:"content"`
AllContexts bool `json:"allContexts"`
Context []string `json:"context"`
Name string `json:"name"`
}
7 changes: 6 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,20 @@ module github.com/gptscript-ai/go-gptscript

go 1.23.0

require github.com/getkin/kin-openapi v0.124.0
require (
github.com/getkin/kin-openapi v0.124.0
github.com/stretchr/testify v1.8.4
Copy link
Collaborator

Choose a reason for hiding this comment

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

nit: I would prefer to not have this.

Copy link
Contributor

Choose a reason for hiding this comment

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

nit: why? doesn't everyone use testify? What else would you use?

)

require (
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/go-openapi/jsonpointer v0.20.2 // indirect
github.com/go-openapi/swag v0.22.8 // indirect
github.com/invopop/yaml v0.2.0 // indirect
github.com/josharian/intern v1.0.0 // indirect
github.com/mailru/easyjson v0.7.7 // indirect
github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect
github.com/perimeterx/marshmallow v1.1.5 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
61 changes: 61 additions & 0 deletions gptscript.go
Original file line number Diff line number Diff line change
Expand Up @@ -307,6 +307,67 @@ func (g *GPTScript) PromptResponse(ctx context.Context, resp PromptResponse) err
return err
}

type ListCredentialsOptions struct {
credCtxs []string
allContexts bool
Copy link
Collaborator

Choose a reason for hiding this comment

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

These need to be updated.

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

}

func (g *GPTScript) ListCredentials(ctx context.Context, opts ListCredentialsOptions) ([]Credential, error) {
req := CredentialRequest{}
if opts.allContexts {
req.AllContexts = true
} else if len(opts.credCtxs) > 0 {
req.Context = opts.credCtxs
} else {
req.Context = []string{"default"}
}

out, err := g.runBasicCommand(ctx, "credentials", req)
if err != nil {
return nil, err
}

var creds []Credential
if err = json.Unmarshal([]byte(out), &creds); err != nil {
return nil, err
}
return creds, nil
}

func (g *GPTScript) CreateCredential(ctx context.Context, cred Credential) error {
credJSON, err := json.Marshal(cred)
if err != nil {
return fmt.Errorf("failed to marshal credential: %w", err)
}

_, err = g.runBasicCommand(ctx, "credentials/create", CredentialRequest{Content: string(credJSON)})
return err
}

func (g *GPTScript) RevealCredential(ctx context.Context, credCtxs []string, name string) (Credential, error) {
out, err := g.runBasicCommand(ctx, "credentials/reveal", CredentialRequest{
Context: credCtxs,
Name: name,
})
if err != nil {
return Credential{}, err
}

var cred Credential
if err = json.Unmarshal([]byte(out), &cred); err != nil {
return Credential{}, err
}
return cred, nil
}

func (g *GPTScript) DeleteCredential(ctx context.Context, credCtx, name string) error {
_, err := g.runBasicCommand(ctx, "credentials/delete", CredentialRequest{
Context: []string{credCtx}, // Only one context can be specified for delete operations
Name: name,
})
return err
}

func (g *GPTScript) runBasicCommand(ctx context.Context, requestPath string, body any) (string, error) {
run := &Run{
url: g.url,
Expand Down
45 changes: 45 additions & 0 deletions gptscript_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,18 @@ package gptscript

import (
"context"
"errors"
"fmt"
"math/rand"
"os"
"path/filepath"
"runtime"
"strconv"
"strings"
"testing"

"github.com/getkin/kin-openapi/openapi3"
"github.com/stretchr/testify/require"
)

var g *GPTScript
Expand Down Expand Up @@ -1448,3 +1452,44 @@ func TestLoadTools(t *testing.T) {
t.Errorf("Unexpected name: %s", prg.Name)
}
}

func TestCredentials(t *testing.T) {

Choose a reason for hiding this comment

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

👍

// We will test in the following order of create, list, reveal, delete.
name := "test-" + strconv.Itoa(rand.Int())
if len(name) > 20 {
name = name[:20]
}

// Create
err := g.CreateCredential(context.Background(), Credential{
Context: "testing",
ToolName: name,
Type: CredentialTypeTool,
Env: map[string]string{"ENV": "testing"},
RefreshToken: "my-refresh-token",
})
require.NoError(t, err)

// List
creds, err := g.ListCredentials(context.Background(), ListCredentialsOptions{
credCtxs: []string{"testing"},
})
require.NoError(t, err)
require.GreaterOrEqual(t, len(creds), 1)

// Reveal
cred, err := g.RevealCredential(context.Background(), []string{"testing"}, name)
require.NoError(t, err)
require.Contains(t, cred.Env, "ENV")
require.Equal(t, cred.Env["ENV"], "testing")
require.Equal(t, cred.RefreshToken, "my-refresh-token")

// Delete
err = g.DeleteCredential(context.Background(), "testing", name)
require.NoError(t, err)

// Delete again and make sure we get a NotFoundError
err = g.DeleteCredential(context.Background(), "testing", name)
require.Error(t, err)
require.True(t, errors.As(err, &ErrNotFound{}))
}
1 change: 1 addition & 0 deletions opts.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ type Options struct {
IncludeEvents bool `json:"includeEvents"`
Prompt bool `json:"prompt"`
CredentialOverrides []string `json:"credentialOverrides"`
CredentialContext []string `json:"credentialContext"`
Copy link
Collaborator

Choose a reason for hiding this comment

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

Let's make the field plural, but keep the json tag the same so that it is compatible with the SDK server. Maybe add a comment that this difference is intentional.

Also, it is possible to add a test for this?

Copy link
Member Author

Choose a reason for hiding this comment

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

👍

Copy link
Member Author

Choose a reason for hiding this comment

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

Done

Location string `json:"location"`
ForceSequential bool `json:"forceSequential"`
}
24 changes: 24 additions & 0 deletions run.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,14 @@ import (

var errAbortRun = errors.New("run aborted")

type ErrNotFound struct {
Message string
}

func (e ErrNotFound) Error() string {
return e.Message
}

type Run struct {
url, requestPath, toolPath string
tools []ToolDef
Expand All @@ -36,6 +44,7 @@ type Run struct {
output, errput string
events chan Frame
lock sync.Mutex
responseCode int
}

// Text returns the text output of the gptscript. It blocks until the output is ready.
Expand All @@ -60,6 +69,11 @@ func (r *Run) State() RunState {
// Err returns the error that caused the gptscript to fail, if any.
func (r *Run) Err() error {
if r.err != nil {
if r.responseCode == http.StatusNotFound {
return ErrNotFound{
Message: fmt.Sprintf("run encountered an error: %s", r.errput),
}
}
return fmt.Errorf("run encountered an error: %w with error output: %s", r.err, r.errput)
}
return nil
Expand Down Expand Up @@ -235,6 +249,7 @@ func (r *Run) request(ctx context.Context, payload any) (err error) {
return r.err
}

r.responseCode = resp.StatusCode
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusBadRequest {
r.state = Error
r.err = fmt.Errorf("run encountered an error")
Expand Down Expand Up @@ -325,6 +340,15 @@ func (r *Run) request(ctx context.Context, payload any) (err error) {

done, _ = out["done"].(bool)
r.rawOutput = out
case []any:

Choose a reason for hiding this comment

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

Is this checking if credential returns an array of objects? if so then it seems fine to me

b, err := json.Marshal(out)
if err != nil {
r.state = Error
r.err = fmt.Errorf("failed to process stdout: %w", err)
return
}

r.output = string(b)
Comment on lines +343 to +351
Copy link
Member Author

Choose a reason for hiding this comment

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

Is this correct? Maybe I should just return a string in the SDK instead of the full slice?

default:
r.state = Error
r.err = fmt.Errorf("failed to process stdout, invalid type: %T", out)
Expand Down
2 changes: 1 addition & 1 deletion test/global-tools.gpt
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ Runbook 3

---
Name: tool_1
Global Tools: github.com/gptscript-ai/knowledge, github.com/drpebcak/duckdb, github.com/gptscript-ai/browser, github.com/gptscript-ai/browser-search/google, github.com/gptscript-ai/browser-search/google-question-answerer
Global Tools: github.com/drpebcak/duckdb, github.com/gptscript-ai/browser, github.com/gptscript-ai/browser-search/google, github.com/gptscript-ai/browser-search/google-question-answerer

Say "Hello!"

Expand Down