Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 657dcee

Browse files
thedadamsthedadams
authoredOct 17, 2024··
feat: simplify workspace API and add support for s3 (#73)
Signed-off-by: Donnie Adams <[email protected]>
1 parent 91a600f commit 657dcee

File tree

3 files changed

+175
-174
lines changed

3 files changed

+175
-174
lines changed
 

‎opts.go

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,17 @@ package gptscript
33
// GlobalOptions allows specification of settings that are used for every call made.
44
// These options can be overridden by the corresponding Options.
55
type GlobalOptions struct {
6-
URL string `json:"url"`
7-
Token string `json:"token"`
8-
OpenAIAPIKey string `json:"APIKey"`
9-
OpenAIBaseURL string `json:"BaseURL"`
10-
DefaultModel string `json:"DefaultModel"`
11-
DefaultModelProvider string `json:"DefaultModelProvider"`
12-
CacheDir string `json:"CacheDir"`
13-
Env []string `json:"env"`
14-
DatasetToolRepo string `json:"DatasetToolRepo"`
15-
WorkspaceTool string `json:"WorkspaceTool"`
6+
URL string `json:"url"`
7+
Token string `json:"token"`
8+
OpenAIAPIKey string `json:"APIKey"`
9+
OpenAIBaseURL string `json:"BaseURL"`
10+
DefaultModel string `json:"DefaultModel"`
11+
DefaultModelProvider string `json:"DefaultModelProvider"`
12+
CacheDir string `json:"CacheDir"`
13+
Env []string `json:"env"`
14+
DatasetToolRepo string `json:"DatasetToolRepo"`
15+
WorkspaceTool string `json:"WorkspaceTool"`
16+
WorkspaceDirectoryDataHome string `json:"WorkspaceDirectoryDataHome"`
1617
}
1718

1819
func (g GlobalOptions) toEnv() []string {
@@ -29,6 +30,9 @@ func (g GlobalOptions) toEnv() []string {
2930
if g.DefaultModelProvider != "" {
3031
args = append(args, "GPTSCRIPT_SDKSERVER_DEFAULT_MODEL_PROVIDER="+g.DefaultModelProvider)
3132
}
33+
if g.WorkspaceDirectoryDataHome != "" {
34+
args = append(args, "GPTSCRIPT_WORKSPACE_DIR="+g.WorkspaceDirectoryDataHome)
35+
}
3236

3337
return args
3438
}

‎workspace.go

Lines changed: 33 additions & 114 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,16 @@ package gptscript
33
import (
44
"context"
55
"encoding/base64"
6-
"encoding/json"
76
"strings"
87
)
98

10-
func (g *GPTScript) CreateWorkspace(ctx context.Context, providerType string) (string, error) {
9+
func (g *GPTScript) CreateWorkspace(ctx context.Context, providerType string, fromWorkspaces ...string) (string, error) {
1110
out, err := g.runBasicCommand(ctx, "workspaces/create", map[string]any{
12-
"provider": providerType,
13-
"workspaceTool": g.globalOpts.WorkspaceTool,
11+
"providerType": providerType,
12+
"fromWorkspaceIDs": fromWorkspaces,
13+
"workspaceTool": g.globalOpts.WorkspaceTool,
14+
"directoryDataHome": g.globalOpts.WorkspaceDirectoryDataHome,
15+
"env": g.globalOpts.Env,
1416
})
1517
if err != nil {
1618
return "", err
@@ -19,156 +21,72 @@ func (g *GPTScript) CreateWorkspace(ctx context.Context, providerType string) (s
1921
return strings.TrimSpace(out), nil
2022
}
2123

22-
type DeleteWorkspaceOptions struct {
23-
IgnoreNotFound bool
24-
}
25-
26-
func (g *GPTScript) DeleteWorkspace(ctx context.Context, workspaceID string, opts ...DeleteWorkspaceOptions) error {
27-
var opt DeleteWorkspaceOptions
28-
for _, o := range opts {
29-
opt.IgnoreNotFound = opt.IgnoreNotFound || o.IgnoreNotFound
30-
}
24+
func (g *GPTScript) DeleteWorkspace(ctx context.Context, workspaceID string) error {
3125
_, err := g.runBasicCommand(ctx, "workspaces/delete", map[string]any{
32-
"id": workspaceID,
33-
"ignoreNotFound": opt.IgnoreNotFound,
34-
"workspaceTool": g.globalOpts.WorkspaceTool,
35-
})
36-
37-
return err
38-
}
39-
40-
type CreateDirectoryInWorkspaceOptions struct {
41-
IgnoreExists bool
42-
}
43-
44-
func (g *GPTScript) CreateDirectoryInWorkspace(ctx context.Context, workspaceID, dir string, opts ...CreateDirectoryInWorkspaceOptions) error {
45-
var opt CreateDirectoryInWorkspaceOptions
46-
for _, o := range opts {
47-
opt.IgnoreExists = opt.IgnoreExists || o.IgnoreExists
48-
}
49-
50-
_, err := g.runBasicCommand(ctx, "workspaces/mkdir", map[string]any{
5126
"id": workspaceID,
52-
"directoryName": dir,
53-
"ignoreExists": opt.IgnoreExists,
5427
"workspaceTool": g.globalOpts.WorkspaceTool,
55-
})
56-
57-
return err
58-
}
59-
60-
type DeleteDirectoryInWorkspaceOptions struct {
61-
IgnoreNotFound bool
62-
MustBeEmpty bool
63-
}
64-
65-
func (g *GPTScript) DeleteDirectoryInWorkspace(ctx context.Context, workspaceID, dir string, opts ...DeleteDirectoryInWorkspaceOptions) error {
66-
var opt DeleteDirectoryInWorkspaceOptions
67-
for _, o := range opts {
68-
o.IgnoreNotFound = opt.IgnoreNotFound || o.IgnoreNotFound
69-
o.MustBeEmpty = opt.MustBeEmpty || o.MustBeEmpty
70-
}
71-
72-
_, err := g.runBasicCommand(ctx, "workspaces/rmdir", map[string]any{
73-
"id": workspaceID,
74-
"directoryName": dir,
75-
"ignoreNotFound": opt.IgnoreNotFound,
76-
"mustBeEmpty": opt.MustBeEmpty,
77-
"workspaceTool": g.globalOpts.WorkspaceTool,
28+
"env": g.globalOpts.Env,
7829
})
7930

8031
return err
8132
}
8233

8334
type ListFilesInWorkspaceOptions struct {
84-
SubDir string
85-
NonRecursive bool
86-
ExcludeHidden bool
87-
}
88-
89-
type WorkspaceContent struct {
90-
ID, Path, FileName string
91-
Children []WorkspaceContent
35+
Prefix string
9236
}
9337

94-
func (g *GPTScript) ListFilesInWorkspace(ctx context.Context, workspaceID string, opts ...ListFilesInWorkspaceOptions) (*WorkspaceContent, error) {
38+
func (g *GPTScript) ListFilesInWorkspace(ctx context.Context, workspaceID string, opts ...ListFilesInWorkspaceOptions) ([]string, error) {
9539
var opt ListFilesInWorkspaceOptions
9640
for _, o := range opts {
97-
if o.SubDir != "" {
98-
opt.SubDir = o.SubDir
41+
if o.Prefix != "" {
42+
opt.Prefix = o.Prefix
9943
}
100-
opt.NonRecursive = opt.NonRecursive || o.NonRecursive
101-
opt.ExcludeHidden = opt.ExcludeHidden || o.ExcludeHidden
10244
}
10345

10446
out, err := g.runBasicCommand(ctx, "workspaces/list", map[string]any{
10547
"id": workspaceID,
106-
"subDir": opt.SubDir,
107-
"excludeHidden": opt.ExcludeHidden,
108-
"nonRecursive": opt.NonRecursive,
48+
"prefix": opt.Prefix,
10949
"workspaceTool": g.globalOpts.WorkspaceTool,
110-
"json": true,
50+
"env": g.globalOpts.Env,
11151
})
11252
if err != nil {
11353
return nil, err
11454
}
11555

116-
var content []WorkspaceContent
117-
err = json.Unmarshal([]byte(out), &content)
118-
if err != nil {
119-
return nil, err
120-
}
121-
122-
if len(content) == 0 {
123-
return &WorkspaceContent{ID: workspaceID}, nil
124-
}
125-
126-
return &content[0], nil
56+
// The first line of the output is the workspace ID, ignore it.
57+
return strings.Split(strings.TrimSpace(out), "\n")[1:], nil
12758
}
12859

129-
type CreateFileInWorkspaceOptions struct {
130-
MustNotExist bool
131-
WithoutCreate bool
132-
CreateDirs bool
133-
}
60+
func (g *GPTScript) RemoveAllWithPrefix(ctx context.Context, workspaceID, prefix string) error {
61+
_, err := g.runBasicCommand(ctx, "workspaces/remove-all-with-prefix", map[string]any{
62+
"id": workspaceID,
63+
"prefix": prefix,
64+
"workspaceTool": g.globalOpts.WorkspaceTool,
65+
"env": g.globalOpts.Env,
66+
})
13467

135-
func (g *GPTScript) WriteFileInWorkspace(ctx context.Context, workspaceID, filePath string, contents []byte, opts ...CreateFileInWorkspaceOptions) error {
136-
var opt CreateFileInWorkspaceOptions
137-
for _, o := range opts {
138-
opt.MustNotExist = opt.MustNotExist || o.MustNotExist
139-
opt.WithoutCreate = opt.WithoutCreate || o.WithoutCreate
140-
opt.CreateDirs = opt.CreateDirs || o.CreateDirs
141-
}
68+
return err
69+
}
14270

71+
func (g *GPTScript) WriteFileInWorkspace(ctx context.Context, workspaceID, filePath string, contents []byte) error {
14372
_, err := g.runBasicCommand(ctx, "workspaces/write-file", map[string]any{
14473
"id": workspaceID,
14574
"contents": base64.StdEncoding.EncodeToString(contents),
14675
"filePath": filePath,
147-
"mustNotExist": opt.MustNotExist,
148-
"withoutCreate": opt.WithoutCreate,
149-
"createDirs": opt.CreateDirs,
15076
"workspaceTool": g.globalOpts.WorkspaceTool,
15177
"base64EncodedInput": true,
78+
"env": g.globalOpts.Env,
15279
})
15380

15481
return err
15582
}
15683

157-
type DeleteFileInWorkspaceOptions struct {
158-
IgnoreNotFound bool
159-
}
160-
161-
func (g *GPTScript) DeleteFileInWorkspace(ctx context.Context, workspaceID, filePath string, opts ...DeleteFileInWorkspaceOptions) error {
162-
var opt DeleteFileInWorkspaceOptions
163-
for _, o := range opts {
164-
opt.IgnoreNotFound = opt.IgnoreNotFound || o.IgnoreNotFound
165-
}
166-
84+
func (g *GPTScript) DeleteFileInWorkspace(ctx context.Context, workspaceID, filePath string) error {
16785
_, err := g.runBasicCommand(ctx, "workspaces/delete-file", map[string]any{
168-
"id": workspaceID,
169-
"filePath": filePath,
170-
"ignoreNotFound": opt.IgnoreNotFound,
171-
"workspaceTool": g.globalOpts.WorkspaceTool,
86+
"id": workspaceID,
87+
"filePath": filePath,
88+
"workspaceTool": g.globalOpts.WorkspaceTool,
89+
"env": g.globalOpts.Env,
17290
})
17391

17492
return err
@@ -180,6 +98,7 @@ func (g *GPTScript) ReadFileInWorkspace(ctx context.Context, workspaceID, filePa
18098
"filePath": filePath,
18199
"workspaceTool": g.globalOpts.WorkspaceTool,
182100
"base64EncodeOutput": true,
101+
"env": g.globalOpts.Env,
183102
})
184103
if err != nil {
185104
return nil, err

‎workspace_test.go

Lines changed: 128 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package gptscript
33
import (
44
"bytes"
55
"context"
6+
"os"
67
"testing"
78
)
89

@@ -18,7 +19,7 @@ func TestCreateAndDeleteWorkspace(t *testing.T) {
1819
}
1920
}
2021

21-
func TestCreateDirectory(t *testing.T) {
22+
func TestWriteReadAndDeleteFileFromWorkspace(t *testing.T) {
2223
id, err := g.CreateWorkspace(context.Background(), "directory")
2324
if err != nil {
2425
t.Fatalf("Error creating workspace: %v", err)
@@ -31,18 +32,27 @@ func TestCreateDirectory(t *testing.T) {
3132
}
3233
})
3334

34-
err = g.CreateDirectoryInWorkspace(context.Background(), id, "test")
35+
err = g.WriteFileInWorkspace(context.Background(), id, "test.txt", []byte("test"))
36+
if err != nil {
37+
t.Fatalf("Error creating file: %v", err)
38+
}
39+
40+
content, err := g.ReadFileInWorkspace(context.Background(), id, "test.txt")
3541
if err != nil {
36-
t.Fatalf("Error creating directory: %v", err)
42+
t.Errorf("Error reading file: %v", err)
3743
}
3844

39-
err = g.DeleteDirectoryInWorkspace(context.Background(), id, "test")
45+
if !bytes.Equal(content, []byte("test")) {
46+
t.Errorf("Unexpected content: %s", content)
47+
}
48+
49+
err = g.DeleteFileInWorkspace(context.Background(), id, "test.txt")
4050
if err != nil {
41-
t.Errorf("Error listing files: %v", err)
51+
t.Errorf("Error deleting file: %v", err)
4252
}
4353
}
4454

45-
func TestWriteReadAndDeleteFileFromWorkspace(t *testing.T) {
55+
func TestLsComplexWorkspace(t *testing.T) {
4656
id, err := g.CreateWorkspace(context.Background(), "directory")
4757
if err != nil {
4858
t.Fatalf("Error creating workspace: %v", err)
@@ -55,6 +65,96 @@ func TestWriteReadAndDeleteFileFromWorkspace(t *testing.T) {
5565
}
5666
})
5767

68+
err = g.WriteFileInWorkspace(context.Background(), id, "test/test1.txt", []byte("hello1"))
69+
if err != nil {
70+
t.Fatalf("Error creating file: %v", err)
71+
}
72+
73+
err = g.WriteFileInWorkspace(context.Background(), id, "test1/test2.txt", []byte("hello2"))
74+
if err != nil {
75+
t.Fatalf("Error creating file: %v", err)
76+
}
77+
78+
err = g.WriteFileInWorkspace(context.Background(), id, "test1/test3.txt", []byte("hello3"))
79+
if err != nil {
80+
t.Fatalf("Error creating file: %v", err)
81+
}
82+
83+
err = g.WriteFileInWorkspace(context.Background(), id, ".hidden.txt", []byte("hidden"))
84+
if err != nil {
85+
t.Fatalf("Error creating hidden file: %v", err)
86+
}
87+
88+
// List all files
89+
content, err := g.ListFilesInWorkspace(context.Background(), id)
90+
if err != nil {
91+
t.Fatalf("Error listing files: %v", err)
92+
}
93+
94+
if len(content) != 4 {
95+
t.Errorf("Unexpected number of files: %d", len(content))
96+
}
97+
98+
// List files in subdirectory
99+
content, err = g.ListFilesInWorkspace(context.Background(), id, ListFilesInWorkspaceOptions{Prefix: "test1"})
100+
if err != nil {
101+
t.Fatalf("Error listing files: %v", err)
102+
}
103+
104+
if len(content) != 2 {
105+
t.Errorf("Unexpected number of files: %d", len(content))
106+
}
107+
108+
// Remove all files with test1 prefix
109+
err = g.RemoveAllWithPrefix(context.Background(), id, "test1")
110+
if err != nil {
111+
t.Fatalf("Error removing files: %v", err)
112+
}
113+
114+
// List files in subdirectory
115+
content, err = g.ListFilesInWorkspace(context.Background(), id)
116+
if err != nil {
117+
t.Fatalf("Error listing files: %v", err)
118+
}
119+
120+
if len(content) != 2 {
121+
t.Errorf("Unexpected number of files: %d", len(content))
122+
}
123+
}
124+
125+
func TestCreateAndDeleteWorkspaceS3(t *testing.T) {
126+
if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" || os.Getenv("WORKSPACE_PROVIDER_S3_BUCKET") == "" {
127+
t.Skip("Skipping test because AWS credentials are not set")
128+
}
129+
130+
id, err := g.CreateWorkspace(context.Background(), "s3")
131+
if err != nil {
132+
t.Fatalf("Error creating workspace: %v", err)
133+
}
134+
135+
err = g.DeleteWorkspace(context.Background(), id)
136+
if err != nil {
137+
t.Errorf("Error deleting workspace: %v", err)
138+
}
139+
}
140+
141+
func TestWriteReadAndDeleteFileFromWorkspaceS3(t *testing.T) {
142+
if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" || os.Getenv("WORKSPACE_PROVIDER_S3_BUCKET") == "" {
143+
t.Skip("Skipping test because AWS credentials are not set")
144+
}
145+
146+
id, err := g.CreateWorkspace(context.Background(), "s3")
147+
if err != nil {
148+
t.Fatalf("Error creating workspace: %v", err)
149+
}
150+
151+
t.Cleanup(func() {
152+
err := g.DeleteWorkspace(context.Background(), id)
153+
if err != nil {
154+
t.Errorf("Error deleting workspace: %v", err)
155+
}
156+
})
157+
58158
err = g.WriteFileInWorkspace(context.Background(), id, "test.txt", []byte("test"))
59159
if err != nil {
60160
t.Fatalf("Error creating file: %v", err)
@@ -75,8 +175,12 @@ func TestWriteReadAndDeleteFileFromWorkspace(t *testing.T) {
75175
}
76176
}
77177

78-
func TestLsComplexWorkspace(t *testing.T) {
79-
id, err := g.CreateWorkspace(context.Background(), "directory")
178+
func TestLsComplexWorkspaceS3(t *testing.T) {
179+
if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" || os.Getenv("WORKSPACE_PROVIDER_S3_BUCKET") == "" {
180+
t.Skip("Skipping test because AWS credentials are not set")
181+
}
182+
183+
id, err := g.CreateWorkspace(context.Background(), "s3")
80184
if err != nil {
81185
t.Fatalf("Error creating workspace: %v", err)
82186
}
@@ -88,29 +192,19 @@ func TestLsComplexWorkspace(t *testing.T) {
88192
}
89193
})
90194

91-
err = g.CreateDirectoryInWorkspace(context.Background(), id, "test")
92-
if err != nil {
93-
t.Fatalf("Error creating directory: %v", err)
94-
}
95-
96195
err = g.WriteFileInWorkspace(context.Background(), id, "test/test1.txt", []byte("hello1"))
97196
if err != nil {
98197
t.Fatalf("Error creating file: %v", err)
99198
}
100199

101-
err = g.WriteFileInWorkspace(context.Background(), id, "test1/test2.txt", []byte("hello2"), CreateFileInWorkspaceOptions{CreateDirs: true})
200+
err = g.WriteFileInWorkspace(context.Background(), id, "test1/test2.txt", []byte("hello2"))
102201
if err != nil {
103202
t.Fatalf("Error creating file: %v", err)
104203
}
105204

106-
err = g.WriteFileInWorkspace(context.Background(), id, "test1/test2.txt", []byte("hello-2"), CreateFileInWorkspaceOptions{MustNotExist: true})
107-
if err == nil {
108-
t.Fatalf("Expected error creating file that must not exist")
109-
}
110-
111-
err = g.WriteFileInWorkspace(context.Background(), id, "test1/test3.txt", []byte("hello3"), CreateFileInWorkspaceOptions{WithoutCreate: true})
112-
if err == nil {
113-
t.Fatalf("Expected error creating file that doesn't exist")
205+
err = g.WriteFileInWorkspace(context.Background(), id, "test1/test3.txt", []byte("hello3"))
206+
if err != nil {
207+
t.Fatalf("Error creating file: %v", err)
114208
}
115209

116210
err = g.WriteFileInWorkspace(context.Background(), id, ".hidden.txt", []byte("hidden"))
@@ -124,49 +218,33 @@ func TestLsComplexWorkspace(t *testing.T) {
124218
t.Fatalf("Error listing files: %v", err)
125219
}
126220

127-
if content.ID != id {
128-
t.Errorf("Unexpected ID: %s", content.ID)
129-
}
130-
131-
if content.Path != "" {
132-
t.Errorf("Unexpected path: %s", content.Path)
133-
}
134-
135-
if content.FileName != "" {
136-
t.Errorf("Unexpected filename: %s", content.FileName)
137-
}
138-
139-
if len(content.Children) != 3 {
140-
t.Errorf("Unexpected number of files: %d", len(content.Children))
221+
if len(content) != 4 {
222+
t.Errorf("Unexpected number of files: %d", len(content))
141223
}
142224

143225
// List files in subdirectory
144-
content, err = g.ListFilesInWorkspace(context.Background(), id, ListFilesInWorkspaceOptions{SubDir: "test1"})
226+
content, err = g.ListFilesInWorkspace(context.Background(), id, ListFilesInWorkspaceOptions{Prefix: "test1"})
145227
if err != nil {
146228
t.Fatalf("Error listing files: %v", err)
147229
}
148230

149-
if len(content.Children) != 1 {
150-
t.Errorf("Unexpected number of files: %d", len(content.Children))
231+
if len(content) != 2 {
232+
t.Errorf("Unexpected number of files: %d", len(content))
151233
}
152234

153-
// Exclude hidden files
154-
content, err = g.ListFilesInWorkspace(context.Background(), id, ListFilesInWorkspaceOptions{ExcludeHidden: true})
235+
// Remove all files with test1 prefix
236+
err = g.RemoveAllWithPrefix(context.Background(), id, "test1")
155237
if err != nil {
156-
t.Fatalf("Error listing files: %v", err)
238+
t.Fatalf("Error removing files: %v", err)
157239
}
158240

159-
if len(content.Children) != 2 {
160-
t.Errorf("Unexpected number of files when listing without hidden: %d", len(content.Children))
161-
}
162-
163-
// List non-recursive
164-
content, err = g.ListFilesInWorkspace(context.Background(), id, ListFilesInWorkspaceOptions{NonRecursive: true})
241+
// List files in subdirectory
242+
content, err = g.ListFilesInWorkspace(context.Background(), id)
165243
if err != nil {
166244
t.Fatalf("Error listing files: %v", err)
167245
}
168246

169-
if len(content.Children) != 1 {
170-
t.Errorf("Unexpected number of files when listing non-recursive: %d", len(content.Children))
247+
if len(content) != 2 {
248+
t.Errorf("Unexpected number of files: %d", len(content))
171249
}
172250
}

0 commit comments

Comments
 (0)
Please sign in to comment.