Skip to content

Commit 91a600f

Browse files
thedadamsthedadams
authored
feat: add workspace API support (#72)
This change also reorganizes some of the code. Signed-off-by: Donnie Adams <donnie@acorn.io>
1 parent b701ad4 commit 91a600f

File tree

7 files changed

+564
-214
lines changed

7 files changed

+564
-214
lines changed

‎datasets.go‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,12 @@
11
package gptscript
22

3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
"os"
8+
)
9+
310
type DatasetElementMeta struct {
411
Name string `json:"name"`
512
Description string `json:"description"`
@@ -48,3 +55,146 @@ type getDatasetElementArgs struct {
4855
DatasetID string `json:"datasetID"`
4956
Element string `json:"element"`
5057
}
58+
59+
func (g *GPTScript) ListDatasets(ctx context.Context, workspace string) ([]DatasetMeta, error) {
60+
if workspace == "" {
61+
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
62+
}
63+
64+
out, err := g.runBasicCommand(ctx, "datasets", datasetRequest{
65+
Input: "{}",
66+
Workspace: workspace,
67+
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
68+
})
69+
if err != nil {
70+
return nil, err
71+
}
72+
73+
var datasets []DatasetMeta
74+
if err = json.Unmarshal([]byte(out), &datasets); err != nil {
75+
return nil, err
76+
}
77+
return datasets, nil
78+
}
79+
80+
func (g *GPTScript) CreateDataset(ctx context.Context, workspace, name, description string) (Dataset, error) {
81+
if workspace == "" {
82+
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
83+
}
84+
85+
args := createDatasetArgs{
86+
Name: name,
87+
Description: description,
88+
}
89+
argsJSON, err := json.Marshal(args)
90+
if err != nil {
91+
return Dataset{}, fmt.Errorf("failed to marshal dataset args: %w", err)
92+
}
93+
94+
out, err := g.runBasicCommand(ctx, "datasets/create", datasetRequest{
95+
Input: string(argsJSON),
96+
Workspace: workspace,
97+
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
98+
})
99+
if err != nil {
100+
return Dataset{}, err
101+
}
102+
103+
var dataset Dataset
104+
if err = json.Unmarshal([]byte(out), &dataset); err != nil {
105+
return Dataset{}, err
106+
}
107+
return dataset, nil
108+
}
109+
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")
113+
}
114+
115+
args := addDatasetElementArgs{
116+
DatasetID: datasetID,
117+
ElementName: elementName,
118+
ElementDescription: elementDescription,
119+
ElementContent: elementContent,
120+
}
121+
argsJSON, err := json.Marshal(args)
122+
if err != nil {
123+
return DatasetElementMeta{}, fmt.Errorf("failed to marshal element args: %w", err)
124+
}
125+
126+
out, err := g.runBasicCommand(ctx, "datasets/add-element", datasetRequest{
127+
Input: string(argsJSON),
128+
Workspace: workspace,
129+
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
130+
})
131+
if err != nil {
132+
return DatasetElementMeta{}, err
133+
}
134+
135+
var element DatasetElementMeta
136+
if err = json.Unmarshal([]byte(out), &element); err != nil {
137+
return DatasetElementMeta{}, err
138+
}
139+
return element, nil
140+
}
141+
142+
func (g *GPTScript) ListDatasetElements(ctx context.Context, workspace, datasetID string) ([]DatasetElementMeta, error) {
143+
if workspace == "" {
144+
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
145+
}
146+
147+
args := listDatasetElementArgs{
148+
DatasetID: datasetID,
149+
}
150+
argsJSON, err := json.Marshal(args)
151+
if err != nil {
152+
return nil, fmt.Errorf("failed to marshal element args: %w", err)
153+
}
154+
155+
out, err := g.runBasicCommand(ctx, "datasets/list-elements", datasetRequest{
156+
Input: string(argsJSON),
157+
Workspace: workspace,
158+
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
159+
})
160+
if err != nil {
161+
return nil, err
162+
}
163+
164+
var elements []DatasetElementMeta
165+
if err = json.Unmarshal([]byte(out), &elements); err != nil {
166+
return nil, err
167+
}
168+
return elements, nil
169+
}
170+
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+
176+
args := getDatasetElementArgs{
177+
DatasetID: datasetID,
178+
Element: elementName,
179+
}
180+
argsJSON, err := json.Marshal(args)
181+
if err != nil {
182+
return DatasetElement{}, fmt.Errorf("failed to marshal element args: %w", err)
183+
}
184+
185+
out, err := g.runBasicCommand(ctx, "datasets/get-element", datasetRequest{
186+
Input: string(argsJSON),
187+
Workspace: workspace,
188+
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
189+
})
190+
if err != nil {
191+
return DatasetElement{}, err
192+
}
193+
194+
var element DatasetElement
195+
if err = json.Unmarshal([]byte(out), &element); err != nil {
196+
return DatasetElement{}, err
197+
}
198+
199+
return element, nil
200+
}

‎datasets_test.go‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
package gptscript
2+
3+
import (
4+
"context"
5+
"os"
6+
"testing"
7+
8+
"github.com/stretchr/testify/require"
9+
)
10+
11+
func TestDatasets(t *testing.T) {
12+
workspace, err := os.MkdirTemp("", "go-gptscript-test")
13+
require.NoError(t, err)
14+
defer func() {
15+
_ = os.RemoveAll(workspace)
16+
}()
17+
18+
// Create a dataset
19+
dataset, err := g.CreateDataset(context.Background(), workspace, "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(), workspace, dataset.ID, "test-element", "This is a test element", "This is the content")
27+
require.NoError(t, err)
28+
require.Equal(t, "test-element", elementMeta.Name)
29+
require.Equal(t, "This is a test element", elementMeta.Description)
30+
31+
// Get the element
32+
element, err := g.GetDatasetElement(context.Background(), workspace, dataset.ID, "test-element")
33+
require.NoError(t, err)
34+
require.Equal(t, "test-element", element.Name)
35+
require.Equal(t, "This is a test element", element.Description)
36+
require.Equal(t, "This is the content", element.Contents)
37+
38+
// List elements in the dataset
39+
elements, err := g.ListDatasetElements(context.Background(), workspace, dataset.ID)
40+
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)
44+
45+
// List datasets
46+
datasets, err := g.ListDatasets(context.Background(), workspace)
47+
require.NoError(t, err)
48+
require.Equal(t, 1, len(datasets))
49+
require.Equal(t, "test-dataset", datasets[0].Name)
50+
require.Equal(t, "This is a test dataset", datasets[0].Description)
51+
require.Equal(t, dataset.ID, datasets[0].ID)
52+
}

‎gptscript.go‎

Lines changed: 0 additions & 171 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import (
77
"context"
88
"encoding/base64"
99
"encoding/json"
10-
"errors"
1110
"fmt"
1211
"io"
1312
"log/slog"
@@ -389,176 +388,6 @@ func (g *GPTScript) DeleteCredential(ctx context.Context, credCtx, name string)
389388
return err
390389
}
391390

392-
// Dataset methods
393-
394-
func (g *GPTScript) ListDatasets(ctx context.Context, workspace string) ([]DatasetMeta, error) {
395-
if workspace == "" {
396-
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
397-
}
398-
399-
out, err := g.runBasicCommand(ctx, "datasets", datasetRequest{
400-
Input: "{}",
401-
Workspace: workspace,
402-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
403-
})
404-
405-
if err != nil {
406-
return nil, err
407-
}
408-
409-
if strings.HasPrefix(out, "ERROR:") {
410-
return nil, errors.New(out)
411-
}
412-
413-
var datasets []DatasetMeta
414-
if err = json.Unmarshal([]byte(out), &datasets); err != nil {
415-
return nil, err
416-
}
417-
return datasets, nil
418-
}
419-
420-
func (g *GPTScript) CreateDataset(ctx context.Context, workspace, name, description string) (Dataset, error) {
421-
if workspace == "" {
422-
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
423-
}
424-
425-
args := createDatasetArgs{
426-
Name: name,
427-
Description: description,
428-
}
429-
argsJSON, err := json.Marshal(args)
430-
if err != nil {
431-
return Dataset{}, fmt.Errorf("failed to marshal dataset args: %w", err)
432-
}
433-
434-
out, err := g.runBasicCommand(ctx, "datasets/create", datasetRequest{
435-
Input: string(argsJSON),
436-
Workspace: workspace,
437-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
438-
})
439-
440-
if err != nil {
441-
return Dataset{}, err
442-
}
443-
444-
if strings.HasPrefix(out, "ERROR:") {
445-
return Dataset{}, errors.New(out)
446-
}
447-
448-
var dataset Dataset
449-
if err = json.Unmarshal([]byte(out), &dataset); err != nil {
450-
return Dataset{}, err
451-
}
452-
return dataset, nil
453-
}
454-
455-
func (g *GPTScript) AddDatasetElement(ctx context.Context, workspace, datasetID, elementName, elementDescription, elementContent string) (DatasetElementMeta, error) {
456-
if workspace == "" {
457-
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
458-
}
459-
460-
args := addDatasetElementArgs{
461-
DatasetID: datasetID,
462-
ElementName: elementName,
463-
ElementDescription: elementDescription,
464-
ElementContent: elementContent,
465-
}
466-
argsJSON, err := json.Marshal(args)
467-
if err != nil {
468-
return DatasetElementMeta{}, fmt.Errorf("failed to marshal element args: %w", err)
469-
}
470-
471-
out, err := g.runBasicCommand(ctx, "datasets/add-element", datasetRequest{
472-
Input: string(argsJSON),
473-
Workspace: workspace,
474-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
475-
})
476-
477-
if err != nil {
478-
return DatasetElementMeta{}, err
479-
}
480-
481-
if strings.HasPrefix(out, "ERROR:") {
482-
return DatasetElementMeta{}, errors.New(out)
483-
}
484-
485-
var element DatasetElementMeta
486-
if err = json.Unmarshal([]byte(out), &element); err != nil {
487-
return DatasetElementMeta{}, err
488-
}
489-
return element, nil
490-
}
491-
492-
func (g *GPTScript) ListDatasetElements(ctx context.Context, workspace, datasetID string) ([]DatasetElementMeta, error) {
493-
if workspace == "" {
494-
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
495-
}
496-
497-
args := listDatasetElementArgs{
498-
DatasetID: datasetID,
499-
}
500-
argsJSON, err := json.Marshal(args)
501-
if err != nil {
502-
return nil, fmt.Errorf("failed to marshal element args: %w", err)
503-
}
504-
505-
out, err := g.runBasicCommand(ctx, "datasets/list-elements", datasetRequest{
506-
Input: string(argsJSON),
507-
Workspace: workspace,
508-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
509-
})
510-
511-
if err != nil {
512-
return nil, err
513-
}
514-
515-
if strings.HasPrefix(out, "ERROR:") {
516-
return nil, errors.New(out)
517-
}
518-
519-
var elements []DatasetElementMeta
520-
if err = json.Unmarshal([]byte(out), &elements); err != nil {
521-
return nil, err
522-
}
523-
return elements, nil
524-
}
525-
526-
func (g *GPTScript) GetDatasetElement(ctx context.Context, workspace, datasetID, elementName string) (DatasetElement, error) {
527-
if workspace == "" {
528-
workspace = os.Getenv("GPTSCRIPT_WORKSPACE_DIR")
529-
}
530-
531-
args := getDatasetElementArgs{
532-
DatasetID: datasetID,
533-
Element: elementName,
534-
}
535-
argsJSON, err := json.Marshal(args)
536-
if err != nil {
537-
return DatasetElement{}, fmt.Errorf("failed to marshal element args: %w", err)
538-
}
539-
540-
out, err := g.runBasicCommand(ctx, "datasets/get-element", datasetRequest{
541-
Input: string(argsJSON),
542-
Workspace: workspace,
543-
DatasetToolRepo: g.globalOpts.DatasetToolRepo,
544-
})
545-
546-
if err != nil {
547-
return DatasetElement{}, err
548-
}
549-
550-
if strings.HasPrefix(out, "ERROR:") {
551-
return DatasetElement{}, errors.New(out)
552-
}
553-
554-
var element DatasetElement
555-
if err = json.Unmarshal([]byte(out), &element); err != nil {
556-
return DatasetElement{}, err
557-
}
558-
559-
return element, nil
560-
}
561-
562391
func (g *GPTScript) runBasicCommand(ctx context.Context, requestPath string, body any) (string, error) {
563392
run := &Run{
564393
url: g.globalOpts.URL,

0 commit comments

Comments
 (0)