Skip to content

feat: add credential management #65

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 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 {
CredentialContexts []string
AllContexts bool
}

func (g *GPTScript) ListCredentials(ctx context.Context, opts ListCredentialsOptions) ([]Credential, error) {
req := CredentialRequest{}
if opts.AllContexts {
req.AllContexts = true
} else if len(opts.CredentialContexts) > 0 {
req.Context = opts.CredentialContexts
} 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) {
Copy link
Contributor

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{
CredentialContexts: []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"`
CredentialContexts []string `json:"credentialContext"` // json tag is left singular to match SDKServer
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:
Copy link
Contributor

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
46 changes: 46 additions & 0 deletions run_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,13 @@ package gptscript

import (
"context"
"crypto/rand"
"encoding/hex"
"os"
"runtime"
"testing"

"github.com/stretchr/testify/require"
)

func TestRestartingErrorRun(t *testing.T) {
Expand Down Expand Up @@ -42,3 +47,44 @@ func TestRestartingErrorRun(t *testing.T) {
t.Errorf("executing run with input of 0 should not fail: %v", err)
}
}

func TestStackedContexts(t *testing.T) {
const name = "testcred"

wd, err := os.Getwd()
require.NoError(t, err)

bytes := make([]byte, 32)
_, err = rand.Read(bytes)
require.NoError(t, err)

context1 := hex.EncodeToString(bytes)[:16]
context2 := hex.EncodeToString(bytes)[16:]

run, err := g.Run(context.Background(), wd+"/test/credential.gpt", Options{
CredentialContexts: []string{context1, context2},
})
require.NoError(t, err)

_, err = run.Text()
require.NoError(t, err)

// The credential should exist in context1 now.
cred, err := g.RevealCredential(context.Background(), []string{context1, context2}, name)
require.NoError(t, err)
require.Equal(t, cred.Context, context1)

// Now change the context order and run the script again.
run, err = g.Run(context.Background(), wd+"/test/credential.gpt", Options{
CredentialContexts: []string{context2, context1},
})
require.NoError(t, err)

_, err = run.Text()
require.NoError(t, err)

// Now make sure the credential exists in context1 still.
cred, err = g.RevealCredential(context.Background(), []string{context2, context1}, name)
require.NoError(t, err)
require.Equal(t, cred.Context, context1)
}
13 changes: 13 additions & 0 deletions test/credential.gpt
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
name: echocred
credential: mycredentialtool as testcred

#!/usr/bin/env bash

echo $VALUE

---
name: mycredentialtool

#!sys.echo

{"env":{"VALUE":"hello"}}
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