From 0680146246ac4fb31981e68e58bf93bf5d9919c9 Mon Sep 17 00:00:00 2001 From: Donnie Adams Date: Thu, 17 Oct 2024 08:41:38 -0400 Subject: [PATCH] feat: simplify workspace API and add support for s3 Signed-off-by: Donnie Adams --- opts.go | 24 ++++--- workspace.go | 147 +++++++++----------------------------- workspace_test.go | 178 +++++++++++++++++++++++++++++++++------------- 3 files changed, 175 insertions(+), 174 deletions(-) diff --git a/opts.go b/opts.go index 6c9e1ab..4b71932 100644 --- a/opts.go +++ b/opts.go @@ -3,16 +3,17 @@ package gptscript // GlobalOptions allows specification of settings that are used for every call made. // These options can be overridden by the corresponding Options. type GlobalOptions struct { - URL string `json:"url"` - Token string `json:"token"` - OpenAIAPIKey string `json:"APIKey"` - OpenAIBaseURL string `json:"BaseURL"` - DefaultModel string `json:"DefaultModel"` - DefaultModelProvider string `json:"DefaultModelProvider"` - CacheDir string `json:"CacheDir"` - Env []string `json:"env"` - DatasetToolRepo string `json:"DatasetToolRepo"` - WorkspaceTool string `json:"WorkspaceTool"` + URL string `json:"url"` + Token string `json:"token"` + OpenAIAPIKey string `json:"APIKey"` + OpenAIBaseURL string `json:"BaseURL"` + DefaultModel string `json:"DefaultModel"` + DefaultModelProvider string `json:"DefaultModelProvider"` + CacheDir string `json:"CacheDir"` + Env []string `json:"env"` + DatasetToolRepo string `json:"DatasetToolRepo"` + WorkspaceTool string `json:"WorkspaceTool"` + WorkspaceDirectoryDataHome string `json:"WorkspaceDirectoryDataHome"` } func (g GlobalOptions) toEnv() []string { @@ -29,6 +30,9 @@ func (g GlobalOptions) toEnv() []string { if g.DefaultModelProvider != "" { args = append(args, "GPTSCRIPT_SDKSERVER_DEFAULT_MODEL_PROVIDER="+g.DefaultModelProvider) } + if g.WorkspaceDirectoryDataHome != "" { + args = append(args, "GPTSCRIPT_WORKSPACE_DIR="+g.WorkspaceDirectoryDataHome) + } return args } diff --git a/workspace.go b/workspace.go index 86b8c2e..2b73893 100644 --- a/workspace.go +++ b/workspace.go @@ -3,14 +3,16 @@ package gptscript import ( "context" "encoding/base64" - "encoding/json" "strings" ) -func (g *GPTScript) CreateWorkspace(ctx context.Context, providerType string) (string, error) { +func (g *GPTScript) CreateWorkspace(ctx context.Context, providerType string, fromWorkspaces ...string) (string, error) { out, err := g.runBasicCommand(ctx, "workspaces/create", map[string]any{ - "provider": providerType, - "workspaceTool": g.globalOpts.WorkspaceTool, + "providerType": providerType, + "fromWorkspaceIDs": fromWorkspaces, + "workspaceTool": g.globalOpts.WorkspaceTool, + "directoryDataHome": g.globalOpts.WorkspaceDirectoryDataHome, + "env": g.globalOpts.Env, }) if err != nil { return "", err @@ -19,156 +21,72 @@ func (g *GPTScript) CreateWorkspace(ctx context.Context, providerType string) (s return strings.TrimSpace(out), nil } -type DeleteWorkspaceOptions struct { - IgnoreNotFound bool -} - -func (g *GPTScript) DeleteWorkspace(ctx context.Context, workspaceID string, opts ...DeleteWorkspaceOptions) error { - var opt DeleteWorkspaceOptions - for _, o := range opts { - opt.IgnoreNotFound = opt.IgnoreNotFound || o.IgnoreNotFound - } +func (g *GPTScript) DeleteWorkspace(ctx context.Context, workspaceID string) error { _, err := g.runBasicCommand(ctx, "workspaces/delete", map[string]any{ - "id": workspaceID, - "ignoreNotFound": opt.IgnoreNotFound, - "workspaceTool": g.globalOpts.WorkspaceTool, - }) - - return err -} - -type CreateDirectoryInWorkspaceOptions struct { - IgnoreExists bool -} - -func (g *GPTScript) CreateDirectoryInWorkspace(ctx context.Context, workspaceID, dir string, opts ...CreateDirectoryInWorkspaceOptions) error { - var opt CreateDirectoryInWorkspaceOptions - for _, o := range opts { - opt.IgnoreExists = opt.IgnoreExists || o.IgnoreExists - } - - _, err := g.runBasicCommand(ctx, "workspaces/mkdir", map[string]any{ "id": workspaceID, - "directoryName": dir, - "ignoreExists": opt.IgnoreExists, "workspaceTool": g.globalOpts.WorkspaceTool, - }) - - return err -} - -type DeleteDirectoryInWorkspaceOptions struct { - IgnoreNotFound bool - MustBeEmpty bool -} - -func (g *GPTScript) DeleteDirectoryInWorkspace(ctx context.Context, workspaceID, dir string, opts ...DeleteDirectoryInWorkspaceOptions) error { - var opt DeleteDirectoryInWorkspaceOptions - for _, o := range opts { - o.IgnoreNotFound = opt.IgnoreNotFound || o.IgnoreNotFound - o.MustBeEmpty = opt.MustBeEmpty || o.MustBeEmpty - } - - _, err := g.runBasicCommand(ctx, "workspaces/rmdir", map[string]any{ - "id": workspaceID, - "directoryName": dir, - "ignoreNotFound": opt.IgnoreNotFound, - "mustBeEmpty": opt.MustBeEmpty, - "workspaceTool": g.globalOpts.WorkspaceTool, + "env": g.globalOpts.Env, }) return err } type ListFilesInWorkspaceOptions struct { - SubDir string - NonRecursive bool - ExcludeHidden bool -} - -type WorkspaceContent struct { - ID, Path, FileName string - Children []WorkspaceContent + Prefix string } -func (g *GPTScript) ListFilesInWorkspace(ctx context.Context, workspaceID string, opts ...ListFilesInWorkspaceOptions) (*WorkspaceContent, error) { +func (g *GPTScript) ListFilesInWorkspace(ctx context.Context, workspaceID string, opts ...ListFilesInWorkspaceOptions) ([]string, error) { var opt ListFilesInWorkspaceOptions for _, o := range opts { - if o.SubDir != "" { - opt.SubDir = o.SubDir + if o.Prefix != "" { + opt.Prefix = o.Prefix } - opt.NonRecursive = opt.NonRecursive || o.NonRecursive - opt.ExcludeHidden = opt.ExcludeHidden || o.ExcludeHidden } out, err := g.runBasicCommand(ctx, "workspaces/list", map[string]any{ "id": workspaceID, - "subDir": opt.SubDir, - "excludeHidden": opt.ExcludeHidden, - "nonRecursive": opt.NonRecursive, + "prefix": opt.Prefix, "workspaceTool": g.globalOpts.WorkspaceTool, - "json": true, + "env": g.globalOpts.Env, }) if err != nil { return nil, err } - var content []WorkspaceContent - err = json.Unmarshal([]byte(out), &content) - if err != nil { - return nil, err - } - - if len(content) == 0 { - return &WorkspaceContent{ID: workspaceID}, nil - } - - return &content[0], nil + // The first line of the output is the workspace ID, ignore it. + return strings.Split(strings.TrimSpace(out), "\n")[1:], nil } -type CreateFileInWorkspaceOptions struct { - MustNotExist bool - WithoutCreate bool - CreateDirs bool -} +func (g *GPTScript) RemoveAllWithPrefix(ctx context.Context, workspaceID, prefix string) error { + _, err := g.runBasicCommand(ctx, "workspaces/remove-all-with-prefix", map[string]any{ + "id": workspaceID, + "prefix": prefix, + "workspaceTool": g.globalOpts.WorkspaceTool, + "env": g.globalOpts.Env, + }) -func (g *GPTScript) WriteFileInWorkspace(ctx context.Context, workspaceID, filePath string, contents []byte, opts ...CreateFileInWorkspaceOptions) error { - var opt CreateFileInWorkspaceOptions - for _, o := range opts { - opt.MustNotExist = opt.MustNotExist || o.MustNotExist - opt.WithoutCreate = opt.WithoutCreate || o.WithoutCreate - opt.CreateDirs = opt.CreateDirs || o.CreateDirs - } + return err +} +func (g *GPTScript) WriteFileInWorkspace(ctx context.Context, workspaceID, filePath string, contents []byte) error { _, err := g.runBasicCommand(ctx, "workspaces/write-file", map[string]any{ "id": workspaceID, "contents": base64.StdEncoding.EncodeToString(contents), "filePath": filePath, - "mustNotExist": opt.MustNotExist, - "withoutCreate": opt.WithoutCreate, - "createDirs": opt.CreateDirs, "workspaceTool": g.globalOpts.WorkspaceTool, "base64EncodedInput": true, + "env": g.globalOpts.Env, }) return err } -type DeleteFileInWorkspaceOptions struct { - IgnoreNotFound bool -} - -func (g *GPTScript) DeleteFileInWorkspace(ctx context.Context, workspaceID, filePath string, opts ...DeleteFileInWorkspaceOptions) error { - var opt DeleteFileInWorkspaceOptions - for _, o := range opts { - opt.IgnoreNotFound = opt.IgnoreNotFound || o.IgnoreNotFound - } - +func (g *GPTScript) DeleteFileInWorkspace(ctx context.Context, workspaceID, filePath string) error { _, err := g.runBasicCommand(ctx, "workspaces/delete-file", map[string]any{ - "id": workspaceID, - "filePath": filePath, - "ignoreNotFound": opt.IgnoreNotFound, - "workspaceTool": g.globalOpts.WorkspaceTool, + "id": workspaceID, + "filePath": filePath, + "workspaceTool": g.globalOpts.WorkspaceTool, + "env": g.globalOpts.Env, }) return err @@ -180,6 +98,7 @@ func (g *GPTScript) ReadFileInWorkspace(ctx context.Context, workspaceID, filePa "filePath": filePath, "workspaceTool": g.globalOpts.WorkspaceTool, "base64EncodeOutput": true, + "env": g.globalOpts.Env, }) if err != nil { return nil, err diff --git a/workspace_test.go b/workspace_test.go index cba683f..6df5c3e 100644 --- a/workspace_test.go +++ b/workspace_test.go @@ -3,6 +3,7 @@ package gptscript import ( "bytes" "context" + "os" "testing" ) @@ -18,7 +19,7 @@ func TestCreateAndDeleteWorkspace(t *testing.T) { } } -func TestCreateDirectory(t *testing.T) { +func TestWriteReadAndDeleteFileFromWorkspace(t *testing.T) { id, err := g.CreateWorkspace(context.Background(), "directory") if err != nil { t.Fatalf("Error creating workspace: %v", err) @@ -31,18 +32,27 @@ func TestCreateDirectory(t *testing.T) { } }) - err = g.CreateDirectoryInWorkspace(context.Background(), id, "test") + err = g.WriteFileInWorkspace(context.Background(), id, "test.txt", []byte("test")) + if err != nil { + t.Fatalf("Error creating file: %v", err) + } + + content, err := g.ReadFileInWorkspace(context.Background(), id, "test.txt") if err != nil { - t.Fatalf("Error creating directory: %v", err) + t.Errorf("Error reading file: %v", err) } - err = g.DeleteDirectoryInWorkspace(context.Background(), id, "test") + if !bytes.Equal(content, []byte("test")) { + t.Errorf("Unexpected content: %s", content) + } + + err = g.DeleteFileInWorkspace(context.Background(), id, "test.txt") if err != nil { - t.Errorf("Error listing files: %v", err) + t.Errorf("Error deleting file: %v", err) } } -func TestWriteReadAndDeleteFileFromWorkspace(t *testing.T) { +func TestLsComplexWorkspace(t *testing.T) { id, err := g.CreateWorkspace(context.Background(), "directory") if err != nil { t.Fatalf("Error creating workspace: %v", err) @@ -55,6 +65,96 @@ func TestWriteReadAndDeleteFileFromWorkspace(t *testing.T) { } }) + err = g.WriteFileInWorkspace(context.Background(), id, "test/test1.txt", []byte("hello1")) + if err != nil { + t.Fatalf("Error creating file: %v", err) + } + + err = g.WriteFileInWorkspace(context.Background(), id, "test1/test2.txt", []byte("hello2")) + if err != nil { + t.Fatalf("Error creating file: %v", err) + } + + err = g.WriteFileInWorkspace(context.Background(), id, "test1/test3.txt", []byte("hello3")) + if err != nil { + t.Fatalf("Error creating file: %v", err) + } + + err = g.WriteFileInWorkspace(context.Background(), id, ".hidden.txt", []byte("hidden")) + if err != nil { + t.Fatalf("Error creating hidden file: %v", err) + } + + // List all files + content, err := g.ListFilesInWorkspace(context.Background(), id) + if err != nil { + t.Fatalf("Error listing files: %v", err) + } + + if len(content) != 4 { + t.Errorf("Unexpected number of files: %d", len(content)) + } + + // List files in subdirectory + content, err = g.ListFilesInWorkspace(context.Background(), id, ListFilesInWorkspaceOptions{Prefix: "test1"}) + if err != nil { + t.Fatalf("Error listing files: %v", err) + } + + if len(content) != 2 { + t.Errorf("Unexpected number of files: %d", len(content)) + } + + // Remove all files with test1 prefix + err = g.RemoveAllWithPrefix(context.Background(), id, "test1") + if err != nil { + t.Fatalf("Error removing files: %v", err) + } + + // List files in subdirectory + content, err = g.ListFilesInWorkspace(context.Background(), id) + if err != nil { + t.Fatalf("Error listing files: %v", err) + } + + if len(content) != 2 { + t.Errorf("Unexpected number of files: %d", len(content)) + } +} + +func TestCreateAndDeleteWorkspaceS3(t *testing.T) { + if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" || os.Getenv("WORKSPACE_PROVIDER_S3_BUCKET") == "" { + t.Skip("Skipping test because AWS credentials are not set") + } + + id, err := g.CreateWorkspace(context.Background(), "s3") + if err != nil { + t.Fatalf("Error creating workspace: %v", err) + } + + err = g.DeleteWorkspace(context.Background(), id) + if err != nil { + t.Errorf("Error deleting workspace: %v", err) + } +} + +func TestWriteReadAndDeleteFileFromWorkspaceS3(t *testing.T) { + if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" || os.Getenv("WORKSPACE_PROVIDER_S3_BUCKET") == "" { + t.Skip("Skipping test because AWS credentials are not set") + } + + id, err := g.CreateWorkspace(context.Background(), "s3") + if err != nil { + t.Fatalf("Error creating workspace: %v", err) + } + + t.Cleanup(func() { + err := g.DeleteWorkspace(context.Background(), id) + if err != nil { + t.Errorf("Error deleting workspace: %v", err) + } + }) + err = g.WriteFileInWorkspace(context.Background(), id, "test.txt", []byte("test")) if err != nil { t.Fatalf("Error creating file: %v", err) @@ -75,8 +175,12 @@ func TestWriteReadAndDeleteFileFromWorkspace(t *testing.T) { } } -func TestLsComplexWorkspace(t *testing.T) { - id, err := g.CreateWorkspace(context.Background(), "directory") +func TestLsComplexWorkspaceS3(t *testing.T) { + if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" || os.Getenv("WORKSPACE_PROVIDER_S3_BUCKET") == "" { + t.Skip("Skipping test because AWS credentials are not set") + } + + id, err := g.CreateWorkspace(context.Background(), "s3") if err != nil { t.Fatalf("Error creating workspace: %v", err) } @@ -88,29 +192,19 @@ func TestLsComplexWorkspace(t *testing.T) { } }) - err = g.CreateDirectoryInWorkspace(context.Background(), id, "test") - if err != nil { - t.Fatalf("Error creating directory: %v", err) - } - err = g.WriteFileInWorkspace(context.Background(), id, "test/test1.txt", []byte("hello1")) if err != nil { t.Fatalf("Error creating file: %v", err) } - err = g.WriteFileInWorkspace(context.Background(), id, "test1/test2.txt", []byte("hello2"), CreateFileInWorkspaceOptions{CreateDirs: true}) + err = g.WriteFileInWorkspace(context.Background(), id, "test1/test2.txt", []byte("hello2")) if err != nil { t.Fatalf("Error creating file: %v", err) } - err = g.WriteFileInWorkspace(context.Background(), id, "test1/test2.txt", []byte("hello-2"), CreateFileInWorkspaceOptions{MustNotExist: true}) - if err == nil { - t.Fatalf("Expected error creating file that must not exist") - } - - err = g.WriteFileInWorkspace(context.Background(), id, "test1/test3.txt", []byte("hello3"), CreateFileInWorkspaceOptions{WithoutCreate: true}) - if err == nil { - t.Fatalf("Expected error creating file that doesn't exist") + err = g.WriteFileInWorkspace(context.Background(), id, "test1/test3.txt", []byte("hello3")) + if err != nil { + t.Fatalf("Error creating file: %v", err) } err = g.WriteFileInWorkspace(context.Background(), id, ".hidden.txt", []byte("hidden")) @@ -124,49 +218,33 @@ func TestLsComplexWorkspace(t *testing.T) { t.Fatalf("Error listing files: %v", err) } - if content.ID != id { - t.Errorf("Unexpected ID: %s", content.ID) - } - - if content.Path != "" { - t.Errorf("Unexpected path: %s", content.Path) - } - - if content.FileName != "" { - t.Errorf("Unexpected filename: %s", content.FileName) - } - - if len(content.Children) != 3 { - t.Errorf("Unexpected number of files: %d", len(content.Children)) + if len(content) != 4 { + t.Errorf("Unexpected number of files: %d", len(content)) } // List files in subdirectory - content, err = g.ListFilesInWorkspace(context.Background(), id, ListFilesInWorkspaceOptions{SubDir: "test1"}) + content, err = g.ListFilesInWorkspace(context.Background(), id, ListFilesInWorkspaceOptions{Prefix: "test1"}) if err != nil { t.Fatalf("Error listing files: %v", err) } - if len(content.Children) != 1 { - t.Errorf("Unexpected number of files: %d", len(content.Children)) + if len(content) != 2 { + t.Errorf("Unexpected number of files: %d", len(content)) } - // Exclude hidden files - content, err = g.ListFilesInWorkspace(context.Background(), id, ListFilesInWorkspaceOptions{ExcludeHidden: true}) + // Remove all files with test1 prefix + err = g.RemoveAllWithPrefix(context.Background(), id, "test1") if err != nil { - t.Fatalf("Error listing files: %v", err) + t.Fatalf("Error removing files: %v", err) } - if len(content.Children) != 2 { - t.Errorf("Unexpected number of files when listing without hidden: %d", len(content.Children)) - } - - // List non-recursive - content, err = g.ListFilesInWorkspace(context.Background(), id, ListFilesInWorkspaceOptions{NonRecursive: true}) + // List files in subdirectory + content, err = g.ListFilesInWorkspace(context.Background(), id) if err != nil { t.Fatalf("Error listing files: %v", err) } - if len(content.Children) != 1 { - t.Errorf("Unexpected number of files when listing non-recursive: %d", len(content.Children)) + if len(content) != 2 { + t.Errorf("Unexpected number of files: %d", len(content)) } }