Skip to content

Commit

Permalink
enhance: use byte slice for dataset element contents (#81)
Browse files Browse the repository at this point in the history
Signed-off-by: Grant Linville <[email protected]>
  • Loading branch information
g-linville authored Nov 1, 2024
1 parent 1fce3cc commit 3a3a70c
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 8 deletions.
7 changes: 4 additions & 3 deletions datasets.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package gptscript

import (
"context"
"encoding/base64"
"encoding/json"
"fmt"
"os"
Expand All @@ -14,7 +15,7 @@ type DatasetElementMeta struct {

type DatasetElement struct {
DatasetElementMeta `json:",inline"`
Contents string `json:"contents"`
Contents []byte `json:"contents"`
}

type DatasetMeta struct {
Expand Down Expand Up @@ -115,7 +116,7 @@ func (g *GPTScript) CreateDataset(ctx context.Context, workspaceID, name, descri
return dataset, nil
}

func (g *GPTScript) AddDatasetElement(ctx context.Context, workspaceID, datasetID, elementName, elementDescription, elementContent string) (DatasetElementMeta, error) {
func (g *GPTScript) AddDatasetElement(ctx context.Context, workspaceID, datasetID, elementName, elementDescription string, elementContent []byte) (DatasetElementMeta, error) {
if workspaceID == "" {
workspaceID = os.Getenv("GPTSCRIPT_WORKSPACE_ID")
}
Expand All @@ -124,7 +125,7 @@ func (g *GPTScript) AddDatasetElement(ctx context.Context, workspaceID, datasetI
DatasetID: datasetID,
ElementName: elementName,
ElementDescription: elementDescription,
ElementContent: elementContent,
ElementContent: base64.StdEncoding.EncodeToString(elementContent),
}
argsJSON, err := json.Marshal(args)
if err != nil {
Expand Down
10 changes: 5 additions & 5 deletions datasets_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ func TestDatasets(t *testing.T) {
require.Equal(t, 0, len(dataset.Elements))

// Add an element
elementMeta, err := g.AddDatasetElement(context.Background(), workspaceID, dataset.ID, "test-element", "This is a test element", "This is the content")
elementMeta, err := g.AddDatasetElement(context.Background(), workspaceID, dataset.ID, "test-element", "This is a test element", []byte("This is the content"))
require.NoError(t, err)
require.Equal(t, "test-element", elementMeta.Name)
require.Equal(t, "This is a test element", elementMeta.Description)
Expand All @@ -35,14 +35,14 @@ func TestDatasets(t *testing.T) {
Name: "test-element-2",
Description: "This is a test element 2",
},
Contents: "This is the content 2",
Contents: []byte("This is the content 2"),
},
{
DatasetElementMeta: DatasetElementMeta{
Name: "test-element-3",
Description: "This is a test element 3",
},
Contents: "This is the content 3",
Contents: []byte("This is the content 3"),
},
})
require.NoError(t, err)
Expand All @@ -52,14 +52,14 @@ func TestDatasets(t *testing.T) {
require.NoError(t, err)
require.Equal(t, "test-element", element.Name)
require.Equal(t, "This is a test element", element.Description)
require.Equal(t, "This is the content", element.Contents)
require.Equal(t, []byte("This is the content"), element.Contents)

// Get the third element
element, err = g.GetDatasetElement(context.Background(), workspaceID, dataset.ID, "test-element-3")
require.NoError(t, err)
require.Equal(t, "test-element-3", element.Name)
require.Equal(t, "This is a test element 3", element.Description)
require.Equal(t, "This is the content 3", element.Contents)
require.Equal(t, []byte("This is the content 3"), element.Contents)

// List elements in the dataset
elements, err := g.ListDatasetElements(context.Background(), workspaceID, dataset.ID)
Expand Down

0 comments on commit 3a3a70c

Please sign in to comment.