Skip to content

Commit c8c2d47

Browse files
g-linvilleg-linville
committed
chore: update for dataset rewrite
Signed-off-by: Grant Linville <[email protected]>
1 parent df4a8c9 commit c8c2d47

File tree

3 files changed

+88
-169
lines changed

3 files changed

+88
-169
lines changed

‎datasets.go

+45-139
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package gptscript
22

33
import (
44
"context"
5-
"encoding/base64"
65
"encoding/json"
76
"fmt"
87
"os"
@@ -15,181 +14,92 @@ type DatasetElementMeta struct {
1514

1615
type DatasetElement struct {
1716
DatasetElementMeta `json:",inline"`
18-
Contents []byte `json:"contents"`
19-
}
20-
21-
type DatasetMeta struct {
22-
ID string `json:"id"`
23-
Name string `json:"name"`
24-
Description string `json:"description"`
17+
Contents string `json:"contents"`
18+
BinaryContents []byte `json:"binaryContents"`
2519
}
2620

2721
type Dataset struct {
28-
DatasetMeta `json:",inline"`
29-
BaseDir string `json:"baseDir,omitempty"`
30-
Elements map[string]DatasetElementMeta `json:"elements"`
22+
ID string `json:"id"`
23+
BaseDir string `json:"baseDir,omitempty"`
24+
Elements map[string]DatasetElementMeta `json:"elements"`
3125
}
3226

3327
type datasetRequest struct {
34-
Input string `json:"input"`
35-
WorkspaceID string `json:"workspaceID"`
36-
DatasetToolRepo string `json:"datasetToolRepo"`
37-
Env []string `json:"env"`
38-
}
39-
40-
type createDatasetArgs struct {
41-
Name string `json:"datasetName"`
42-
Description string `json:"datasetDescription"`
43-
}
44-
45-
type addDatasetElementArgs struct {
46-
DatasetID string `json:"datasetID"`
47-
ElementName string `json:"elementName"`
48-
ElementDescription string `json:"elementDescription"`
49-
ElementContent string `json:"elementContent"`
28+
Input string `json:"input"`
29+
DatasetTool string `json:"datasetTool"`
30+
Env []string `json:"env"`
5031
}
5132

5233
type addDatasetElementsArgs struct {
53-
DatasetID string `json:"datasetID"`
54-
Elements []DatasetElement `json:"elements"`
34+
WorkspaceID string `json:"workspaceID"`
35+
DatasetID string `json:"datasetID"`
36+
Elements []DatasetElement `json:"elements"`
5537
}
5638

5739
type listDatasetElementArgs struct {
58-
DatasetID string `json:"datasetID"`
40+
WorkspaceID string `json:"workspaceID"`
41+
DatasetID string `json:"datasetID"`
5942
}
6043

6144
type getDatasetElementArgs struct {
62-
DatasetID string `json:"datasetID"`
63-
Element string `json:"element"`
45+
WorkspaceID string `json:"workspaceID"`
46+
DatasetID string `json:"datasetID"`
47+
Element string `json:"name"`
6448
}
6549

66-
func (g *GPTScript) ListDatasets(ctx context.Context, workspaceID string) ([]DatasetMeta, error) {
67-
if workspaceID == "" {
68-
workspaceID = os.Getenv("GPTSCRIPT_WORKSPACE_ID")
69-
}
70-
50+
func (g *GPTScript) ListDatasets(ctx context.Context) ([]string, error) {
7151
out, err := g.runBasicCommand(ctx, "datasets", datasetRequest{
72-
Input: "{}",
73-
WorkspaceID: workspaceID,
74-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
75-
Env: g.globalOpts.Env,
52+
Input: fmt.Sprintf(`{"workspaceID": %q}`, os.Getenv("GPTSCRIPT_WORKSPACE_ID")),
53+
DatasetTool: g.globalOpts.DatasetTool,
54+
Env: g.globalOpts.Env,
7655
})
7756
if err != nil {
7857
return nil, err
7958
}
8059

81-
var datasets []DatasetMeta
60+
var datasets []string
8261
if err = json.Unmarshal([]byte(out), &datasets); err != nil {
8362
return nil, err
8463
}
8564
return datasets, nil
8665
}
8766

88-
func (g *GPTScript) CreateDataset(ctx context.Context, workspaceID, name, description string) (Dataset, error) {
89-
if workspaceID == "" {
90-
workspaceID = os.Getenv("GPTSCRIPT_WORKSPACE_ID")
91-
}
92-
93-
args := createDatasetArgs{
94-
Name: name,
95-
Description: description,
96-
}
97-
argsJSON, err := json.Marshal(args)
98-
if err != nil {
99-
return Dataset{}, fmt.Errorf("failed to marshal dataset args: %w", err)
100-
}
101-
102-
out, err := g.runBasicCommand(ctx, "datasets/create", datasetRequest{
103-
Input: string(argsJSON),
104-
WorkspaceID: workspaceID,
105-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
106-
Env: g.globalOpts.Env,
107-
})
108-
if err != nil {
109-
return Dataset{}, err
110-
}
111-
112-
var dataset Dataset
113-
if err = json.Unmarshal([]byte(out), &dataset); err != nil {
114-
return Dataset{}, err
115-
}
116-
return dataset, nil
117-
}
118-
119-
func (g *GPTScript) AddDatasetElement(ctx context.Context, workspaceID, datasetID, elementName, elementDescription string, elementContent []byte) (DatasetElementMeta, error) {
120-
if workspaceID == "" {
121-
workspaceID = os.Getenv("GPTSCRIPT_WORKSPACE_ID")
122-
}
123-
124-
args := addDatasetElementArgs{
125-
DatasetID: datasetID,
126-
ElementName: elementName,
127-
ElementDescription: elementDescription,
128-
ElementContent: base64.StdEncoding.EncodeToString(elementContent),
129-
}
130-
argsJSON, err := json.Marshal(args)
131-
if err != nil {
132-
return DatasetElementMeta{}, fmt.Errorf("failed to marshal element args: %w", err)
133-
}
134-
135-
out, err := g.runBasicCommand(ctx, "datasets/add-element", datasetRequest{
136-
Input: string(argsJSON),
137-
WorkspaceID: workspaceID,
138-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
139-
Env: g.globalOpts.Env,
140-
})
141-
if err != nil {
142-
return DatasetElementMeta{}, err
143-
}
144-
145-
var element DatasetElementMeta
146-
if err = json.Unmarshal([]byte(out), &element); err != nil {
147-
return DatasetElementMeta{}, err
148-
}
149-
return element, nil
67+
func (g *GPTScript) CreateDatasetWithElements(ctx context.Context, elements []DatasetElement) (string, error) {
68+
return g.AddDatasetElements(ctx, "", elements)
15069
}
15170

152-
func (g *GPTScript) AddDatasetElements(ctx context.Context, workspaceID, datasetID string, elements []DatasetElement) error {
153-
if workspaceID == "" {
154-
workspaceID = os.Getenv("GPTSCRIPT_WORKSPACE_ID")
155-
}
156-
71+
func (g *GPTScript) AddDatasetElements(ctx context.Context, datasetID string, elements []DatasetElement) (string, error) {
15772
args := addDatasetElementsArgs{
158-
DatasetID: datasetID,
159-
Elements: elements,
73+
WorkspaceID: os.Getenv("GPTSCRIPT_WORKSPACE_ID"),
74+
DatasetID: datasetID,
75+
Elements: elements,
16076
}
16177
argsJSON, err := json.Marshal(args)
16278
if err != nil {
163-
return fmt.Errorf("failed to marshal element args: %w", err)
79+
return "", fmt.Errorf("failed to marshal element args: %w", err)
16480
}
16581

166-
_, err = g.runBasicCommand(ctx, "datasets/add-elements", datasetRequest{
167-
Input: string(argsJSON),
168-
WorkspaceID: workspaceID,
169-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
170-
Env: g.globalOpts.Env,
82+
return g.runBasicCommand(ctx, "datasets/add-elements", datasetRequest{
83+
Input: string(argsJSON),
84+
DatasetTool: g.globalOpts.DatasetTool,
85+
Env: g.globalOpts.Env,
17186
})
172-
return err
17387
}
17488

175-
func (g *GPTScript) ListDatasetElements(ctx context.Context, workspaceID, datasetID string) ([]DatasetElementMeta, error) {
176-
if workspaceID == "" {
177-
workspaceID = os.Getenv("GPTSCRIPT_WORKSPACE_ID")
178-
}
179-
89+
func (g *GPTScript) ListDatasetElements(ctx context.Context, datasetID string) ([]DatasetElementMeta, error) {
18090
args := listDatasetElementArgs{
181-
DatasetID: datasetID,
91+
WorkspaceID: os.Getenv("GPTSCRIPT_WORKSPACE_ID"),
92+
DatasetID: datasetID,
18293
}
18394
argsJSON, err := json.Marshal(args)
18495
if err != nil {
18596
return nil, fmt.Errorf("failed to marshal element args: %w", err)
18697
}
18798

18899
out, err := g.runBasicCommand(ctx, "datasets/list-elements", datasetRequest{
189-
Input: string(argsJSON),
190-
WorkspaceID: workspaceID,
191-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
192-
Env: g.globalOpts.Env,
100+
Input: string(argsJSON),
101+
DatasetTool: g.globalOpts.DatasetTool,
102+
Env: g.globalOpts.Env,
193103
})
194104
if err != nil {
195105
return nil, err
@@ -202,25 +112,21 @@ func (g *GPTScript) ListDatasetElements(ctx context.Context, workspaceID, datase
202112
return elements, nil
203113
}
204114

205-
func (g *GPTScript) GetDatasetElement(ctx context.Context, workspaceID, datasetID, elementName string) (DatasetElement, error) {
206-
if workspaceID == "" {
207-
workspaceID = os.Getenv("GPTSCRIPT_WORKSPACE_ID")
208-
}
209-
115+
func (g *GPTScript) GetDatasetElement(ctx context.Context, datasetID, elementName string) (DatasetElement, error) {
210116
args := getDatasetElementArgs{
211-
DatasetID: datasetID,
212-
Element: elementName,
117+
WorkspaceID: os.Getenv("GPTSCRIPT_WORKSPACE_ID"),
118+
DatasetID: datasetID,
119+
Element: elementName,
213120
}
214121
argsJSON, err := json.Marshal(args)
215122
if err != nil {
216123
return DatasetElement{}, fmt.Errorf("failed to marshal element args: %w", err)
217124
}
218125

219126
out, err := g.runBasicCommand(ctx, "datasets/get-element", datasetRequest{
220-
Input: string(argsJSON),
221-
WorkspaceID: workspaceID,
222-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
223-
Env: g.globalOpts.Env,
127+
Input: string(argsJSON),
128+
DatasetTool: g.globalOpts.DatasetTool,
129+
Env: g.globalOpts.Env,
224130
})
225131
if err != nil {
226132
return DatasetElement{}, err

‎datasets_test.go

+41-28
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package gptscript
22

33
import (
44
"context"
5+
"os"
56
"testing"
67

78
"github.com/stretchr/testify/require"
@@ -11,66 +12,78 @@ func TestDatasets(t *testing.T) {
1112
workspaceID, err := g.CreateWorkspace(context.Background(), "directory")
1213
require.NoError(t, err)
1314

15+
require.NoError(t, os.Setenv("GPTSCRIPT_WORKSPACE_ID", workspaceID))
16+
1417
defer func() {
1518
_ = g.DeleteWorkspace(context.Background(), workspaceID)
1619
}()
1720

18-
// Create a dataset
19-
dataset, err := g.CreateDataset(context.Background(), workspaceID, "test-dataset", "This is a test dataset")
20-
require.NoError(t, err)
21-
require.Equal(t, "test-dataset", dataset.Name)
22-
require.Equal(t, "This is a test dataset", dataset.Description)
23-
require.Equal(t, 0, len(dataset.Elements))
24-
25-
// Add an element
26-
elementMeta, err := g.AddDatasetElement(context.Background(), workspaceID, dataset.ID, "test-element", "This is a test element", []byte("This is the content"))
21+
datasetID, err := g.CreateDatasetWithElements(context.Background(), []DatasetElement{
22+
{
23+
DatasetElementMeta: DatasetElementMeta{
24+
Name: "test-element-1",
25+
Description: "This is a test element 1",
26+
},
27+
Contents: "This is the content 1",
28+
},
29+
})
2730
require.NoError(t, err)
28-
require.Equal(t, "test-element", elementMeta.Name)
29-
require.Equal(t, "This is a test element", elementMeta.Description)
3031

31-
// Add two more
32-
err = g.AddDatasetElements(context.Background(), workspaceID, dataset.ID, []DatasetElement{
32+
// Add two more elements
33+
_, err = g.AddDatasetElements(context.Background(), datasetID, []DatasetElement{
3334
{
3435
DatasetElementMeta: DatasetElementMeta{
3536
Name: "test-element-2",
3637
Description: "This is a test element 2",
3738
},
38-
Contents: []byte("This is the content 2"),
39+
Contents: "This is the content 2",
3940
},
4041
{
4142
DatasetElementMeta: DatasetElementMeta{
4243
Name: "test-element-3",
4344
Description: "This is a test element 3",
4445
},
45-
Contents: []byte("This is the content 3"),
46+
Contents: "This is the content 3",
47+
},
48+
{
49+
DatasetElementMeta: DatasetElementMeta{
50+
Name: "binary-element",
51+
Description: "this element has binary contents",
52+
},
53+
BinaryContents: []byte("binary contents"),
4654
},
4755
})
4856
require.NoError(t, err)
4957

5058
// Get the first element
51-
element, err := g.GetDatasetElement(context.Background(), workspaceID, dataset.ID, "test-element")
59+
element, err := g.GetDatasetElement(context.Background(), datasetID, "test-element-1")
5260
require.NoError(t, err)
53-
require.Equal(t, "test-element", element.Name)
54-
require.Equal(t, "This is a test element", element.Description)
55-
require.Equal(t, []byte("This is the content"), element.Contents)
61+
require.Equal(t, "test-element-1", element.Name)
62+
require.Equal(t, "This is a test element 1", element.Description)
63+
require.Equal(t, "This is the content 1", element.Contents)
5664

5765
// Get the third element
58-
element, err = g.GetDatasetElement(context.Background(), workspaceID, dataset.ID, "test-element-3")
66+
element, err = g.GetDatasetElement(context.Background(), datasetID, "test-element-3")
5967
require.NoError(t, err)
6068
require.Equal(t, "test-element-3", element.Name)
6169
require.Equal(t, "This is a test element 3", element.Description)
62-
require.Equal(t, []byte("This is the content 3"), element.Contents)
70+
require.Equal(t, "This is the content 3", element.Contents)
71+
72+
// Get the binary element
73+
element, err = g.GetDatasetElement(context.Background(), datasetID, "binary-element")
74+
require.NoError(t, err)
75+
require.Equal(t, "binary-element", element.Name)
76+
require.Equal(t, "this element has binary contents", element.Description)
77+
require.Equal(t, []byte("binary contents"), element.BinaryContents)
6378

6479
// List elements in the dataset
65-
elements, err := g.ListDatasetElements(context.Background(), workspaceID, dataset.ID)
80+
elements, err := g.ListDatasetElements(context.Background(), datasetID)
6681
require.NoError(t, err)
67-
require.Equal(t, 3, len(elements))
82+
require.Equal(t, 4, len(elements))
6883

6984
// List datasets
70-
datasets, err := g.ListDatasets(context.Background(), workspaceID)
85+
datasetIDs, err := g.ListDatasets(context.Background())
7186
require.NoError(t, err)
72-
require.Equal(t, 1, len(datasets))
73-
require.Equal(t, "test-dataset", datasets[0].Name)
74-
require.Equal(t, "This is a test dataset", datasets[0].Description)
75-
require.Equal(t, dataset.ID, datasets[0].ID)
87+
require.Equal(t, 1, len(datasetIDs))
88+
require.Equal(t, datasetID, datasetIDs[0])
7689
}

‎opts.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ type GlobalOptions struct {
1111
DefaultModelProvider string `json:"DefaultModelProvider"`
1212
CacheDir string `json:"CacheDir"`
1313
Env []string `json:"env"`
14-
DatasetToolRepo string `json:"DatasetToolRepo"`
14+
DatasetTool string `json:"DatasetTool"`
1515
WorkspaceTool string `json:"WorkspaceTool"`
1616
}
1717

@@ -46,7 +46,7 @@ func completeGlobalOptions(opts ...GlobalOptions) GlobalOptions {
4646
result.OpenAIBaseURL = firstSet(opt.OpenAIBaseURL, result.OpenAIBaseURL)
4747
result.DefaultModel = firstSet(opt.DefaultModel, result.DefaultModel)
4848
result.DefaultModelProvider = firstSet(opt.DefaultModelProvider, result.DefaultModelProvider)
49-
result.DatasetToolRepo = firstSet(opt.DatasetToolRepo, result.DatasetToolRepo)
49+
result.DatasetTool = firstSet(opt.DatasetTool, result.DatasetTool)
5050
result.WorkspaceTool = firstSet(opt.WorkspaceTool, result.WorkspaceTool)
5151
result.Env = append(result.Env, opt.Env...)
5252
}

0 commit comments

Comments
 (0)