From d083bcd7e2752f535b08a65f64b38d08c2ac7a0c Mon Sep 17 00:00:00 2001
From: Grant Linville <grant@acorn.io>
Date: Thu, 31 Oct 2024 16:28:44 -0400
Subject: [PATCH] enhance: use byte slice for dataset element contents

Signed-off-by: Grant Linville <grant@acorn.io>
---
 datasets.go      |  7 ++++---
 datasets_test.go | 10 +++++-----
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/datasets.go b/datasets.go
index 7b9f377..0b39edc 100644
--- a/datasets.go
+++ b/datasets.go
@@ -2,6 +2,7 @@ package gptscript
 
 import (
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"os"
@@ -14,7 +15,7 @@ type DatasetElementMeta struct {
 
 type DatasetElement struct {
 	DatasetElementMeta `json:",inline"`
-	Contents           string `json:"contents"`
+	Contents           []byte `json:"contents"`
 }
 
 type DatasetMeta struct {
@@ -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")
 	}
@@ -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 {
diff --git a/datasets_test.go b/datasets_test.go
index 44d72ab..bd08d3f 100644
--- a/datasets_test.go
+++ b/datasets_test.go
@@ -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)
@@ -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)
@@ -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)