Skip to content

Commit bce1378

Browse files
g-linvilleg-linville
committed
chore: add Add Elements method and update for workspace provider
Signed-off-by: Grant Linville <[email protected]>
1 parent 3901872 commit bce1378

File tree

2 files changed

+80
-34
lines changed

2 files changed

+80
-34
lines changed

‎datasets.go

+44-22
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,10 @@ type Dataset struct {
3030
}
3131

3232
type datasetRequest struct {
33-
Input string `json:"input"`
34-
Workspace string `json:"workspace"`
35-
DatasetToolRepo string `json:"datasetToolRepo"`
33+
Input string `json:"input"`
34+
WorkspaceID string `json:"workspaceID"`
35+
DatasetToolRepo string `json:"datasetToolRepo"`
36+
Env []string `json:"env"`
3637
}
3738

3839
type createDatasetArgs struct {
@@ -47,6 +48,11 @@ type addDatasetElementArgs struct {
4748
ElementContent string `json:"elementContent"`
4849
}
4950

51+
type addDatasetElementsArgs struct {
52+
DatasetID string `json:"datasetID"`
53+
Elements []DatasetElement `json:"elements"`
54+
}
55+
5056
type listDatasetElementArgs struct {
5157
DatasetID string `json:"datasetID"`
5258
}
@@ -63,8 +69,9 @@ func (g *GPTScript) ListDatasets(ctx context.Context, workspace string) ([]Datas
6369

6470
out, err := g.runBasicCommand(ctx, "datasets", datasetRequest{
6571
Input: "{}",
66-
Workspace: workspace,
72+
WorkspaceID: workspace,
6773
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
74+
Env: g.globalOpts.Env,
6875
})
6976
if err != nil {
7077
return nil, err
@@ -77,9 +84,9 @@ func (g *GPTScript) ListDatasets(ctx context.Context, workspace string) ([]Datas
7784
return datasets, nil
7885
}
7986

80-
func (g *GPTScript) CreateDataset(ctx context.Context, workspace, name, description string) (Dataset, error) {
81-
if workspace == "" {
82-
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
87+
func (g *GPTScript) CreateDataset(ctx context.Context, workspaceID, name, description string) (Dataset, error) {
88+
if workspaceID == "" {
89+
workspaceID = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
8390
}
8491

8592
args := createDatasetArgs{
@@ -93,8 +100,9 @@ func (g *GPTScript) CreateDataset(ctx context.Context, workspace, name, descript
93100

94101
out, err := g.runBasicCommand(ctx, "datasets/create", datasetRequest{
95102
Input: string(argsJSON),
96-
Workspace: workspace,
103+
WorkspaceID: workspaceID,
97104
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
105+
Env: g.globalOpts.Env,
98106
})
99107
if err != nil {
100108
return Dataset{}, err
@@ -107,9 +115,9 @@ func (g *GPTScript) CreateDataset(ctx context.Context, workspace, name, descript
107115
return dataset, nil
108116
}
109117

110-
func (g *GPTScript) AddDatasetElement(ctx context.Context, workspace, datasetID, elementName, elementDescription, elementContent string) (DatasetElementMeta, error) {
111-
if workspace == "" {
112-
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
118+
func (g *GPTScript) AddDatasetElement(ctx context.Context, workspaceID, datasetID, elementName, elementDescription, elementContent string) (DatasetElementMeta, error) {
119+
if workspaceID == "" {
120+
workspaceID = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
113121
}
114122

115123
args := addDatasetElementArgs{
@@ -125,8 +133,9 @@ func (g *GPTScript) AddDatasetElement(ctx context.Context, workspace, datasetID,
125133

126134
out, err := g.runBasicCommand(ctx, "datasets/add-element", datasetRequest{
127135
Input: string(argsJSON),
128-
Workspace: workspace,
136+
WorkspaceID: workspaceID,
129137
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
138+
Env: g.globalOpts.Env,
130139
})
131140
if err != nil {
132141
return DatasetElementMeta{}, err
@@ -139,11 +148,26 @@ func (g *GPTScript) AddDatasetElement(ctx context.Context, workspace, datasetID,
139148
return element, nil
140149
}
141150

142-
func (g *GPTScript) ListDatasetElements(ctx context.Context, workspace, datasetID string) ([]DatasetElementMeta, error) {
143-
if workspace == "" {
144-
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
151+
func (g *GPTScript) AddDatasetElements(ctx context.Context, workspaceID, datasetID string, elements []DatasetElement) error {
152+
args := addDatasetElementsArgs{
153+
DatasetID: datasetID,
154+
Elements: elements,
145155
}
156+
argsJSON, err := json.Marshal(args)
157+
if err != nil {
158+
return fmt.Errorf("failed to marshal element args: %w", err)
159+
}
160+
161+
_, err = g.runBasicCommand(ctx, "datasets/add-elements", datasetRequest{
162+
Input: string(argsJSON),
163+
WorkspaceID: workspaceID,
164+
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
165+
Env: g.globalOpts.Env,
166+
})
167+
return err
168+
}
146169

170+
func (g *GPTScript) ListDatasetElements(ctx context.Context, workspaceID, datasetID string) ([]DatasetElementMeta, error) {
147171
args := listDatasetElementArgs{
148172
DatasetID: datasetID,
149173
}
@@ -154,8 +178,9 @@ func (g *GPTScript) ListDatasetElements(ctx context.Context, workspace, datasetI
154178

155179
out, err := g.runBasicCommand(ctx, "datasets/list-elements", datasetRequest{
156180
Input: string(argsJSON),
157-
Workspace: workspace,
181+
WorkspaceID: workspaceID,
158182
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
183+
Env: g.globalOpts.Env,
159184
})
160185
if err != nil {
161186
return nil, err
@@ -168,11 +193,7 @@ func (g *GPTScript) ListDatasetElements(ctx context.Context, workspace, datasetI
168193
return elements, nil
169194
}
170195

171-
func (g *GPTScript) GetDatasetElement(ctx context.Context, workspace, datasetID, elementName string) (DatasetElement, error) {
172-
if workspace == "" {
173-
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
174-
}
175-
196+
func (g *GPTScript) GetDatasetElement(ctx context.Context, workspaceID, datasetID, elementName string) (DatasetElement, error) {
176197
args := getDatasetElementArgs{
177198
DatasetID: datasetID,
178199
Element: elementName,
@@ -184,8 +205,9 @@ func (g *GPTScript) GetDatasetElement(ctx context.Context, workspace, datasetID,
184205

185206
out, err := g.runBasicCommand(ctx, "datasets/get-element", datasetRequest{
186207
Input: string(argsJSON),
187-
Workspace: workspace,
208+
WorkspaceID: workspaceID,
188209
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
210+
Env: g.globalOpts.Env,
189211
})
190212
if err != nil {
191213
return DatasetElement{}, err

‎datasets_test.go

+36-12
Original file line numberDiff line numberDiff line change
@@ -2,48 +2,72 @@ package gptscript
22

33
import (
44
"context"
5-
"os"
65
"testing"
76

87
"github.com/stretchr/testify/require"
98
)
109

1110
func TestDatasets(t *testing.T) {
12-
workspace, err := os.MkdirTemp("", "go-gptscript-test")
11+
workspaceID, err := g.CreateWorkspace(context.Background(), "directory")
1312
require.NoError(t, err)
13+
1414
defer func() {
15-
_ = os.RemoveAll(workspace)
15+
_ = g.DeleteWorkspace(context.Background(), DeleteWorkspaceOptions{WorkspaceID: workspaceID})
1616
}()
1717

1818
// Create a dataset
19-
dataset, err := g.CreateDataset(context.Background(), workspace, "test-dataset", "This is a test dataset")
19+
dataset, err := g.CreateDataset(context.Background(), workspaceID, "test-dataset", "This is a test dataset")
2020
require.NoError(t, err)
2121
require.Equal(t, "test-dataset", dataset.Name)
2222
require.Equal(t, "This is a test dataset", dataset.Description)
2323
require.Equal(t, 0, len(dataset.Elements))
2424

2525
// Add an element
26-
elementMeta, err := g.AddDatasetElement(context.Background(), workspace, dataset.ID, "test-element", "This is a test element", "This is the content")
26+
elementMeta, err := g.AddDatasetElement(context.Background(), workspaceID, dataset.ID, "test-element", "This is a test element", "This is the content")
2727
require.NoError(t, err)
2828
require.Equal(t, "test-element", elementMeta.Name)
2929
require.Equal(t, "This is a test element", elementMeta.Description)
3030

31-
// Get the element
32-
element, err := g.GetDatasetElement(context.Background(), workspace, dataset.ID, "test-element")
31+
// Add two more
32+
err = g.AddDatasetElements(context.Background(), workspaceID, dataset.ID, []DatasetElement{
33+
{
34+
DatasetElementMeta: DatasetElementMeta{
35+
Name: "test-element-2",
36+
Description: "This is a test element 2",
37+
},
38+
Contents: "This is the content 2",
39+
},
40+
{
41+
DatasetElementMeta: DatasetElementMeta{
42+
Name: "test-element-3",
43+
Description: "This is a test element 3",
44+
},
45+
Contents: "This is the content 3",
46+
},
47+
})
48+
require.NoError(t, err)
49+
50+
// Get the first element
51+
element, err := g.GetDatasetElement(context.Background(), workspaceID, dataset.ID, "test-element")
3352
require.NoError(t, err)
3453
require.Equal(t, "test-element", element.Name)
3554
require.Equal(t, "This is a test element", element.Description)
3655
require.Equal(t, "This is the content", element.Contents)
3756

57+
// Get the third element
58+
element, err = g.GetDatasetElement(context.Background(), workspaceID, dataset.ID, "test-element-3")
59+
require.NoError(t, err)
60+
require.Equal(t, "test-element-3", element.Name)
61+
require.Equal(t, "This is a test element 3", element.Description)
62+
require.Equal(t, "This is the content 3", element.Contents)
63+
3864
// List elements in the dataset
39-
elements, err := g.ListDatasetElements(context.Background(), workspace, dataset.ID)
65+
elements, err := g.ListDatasetElements(context.Background(), workspaceID, dataset.ID)
4066
require.NoError(t, err)
41-
require.Equal(t, 1, len(elements))
42-
require.Equal(t, "test-element", elements[0].Name)
43-
require.Equal(t, "This is a test element", elements[0].Description)
67+
require.Equal(t, 3, len(elements))
4468

4569
// List datasets
46-
datasets, err := g.ListDatasets(context.Background(), workspace)
70+
datasets, err := g.ListDatasets(context.Background(), workspaceID)
4771
require.NoError(t, err)
4872
require.Equal(t, 1, len(datasets))
4973
require.Equal(t, "test-dataset", datasets[0].Name)

0 commit comments

Comments
 (0)