enhance: use ArrayBuffer for dataset element contents (#100)

Signed-off-by: Grant Linville <[email protected]>
gptscript-ai · Nov 1, 2024 · b30db3c · b30db3c
1 parent ea42fed
commit b30db3c
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 14 deletions.
diff --git a/src/gptscript.ts b/src/gptscript.ts
@@ -415,7 +415,7 @@ export class GPTScript {
         return JSON.parse(result) as Dataset
     }
 
-    async addDatasetElement(workspaceID: string, datasetID: string, elementName: string, elementDescription: string, elementContent: string): Promise<DatasetElementMeta> {
+    async addDatasetElement(workspaceID: string, datasetID: string, elementName: string, elementDescription: string, elementContent: ArrayBuffer): Promise<DatasetElementMeta> {
         if (workspaceID == "") {
             workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""
         }
@@ -425,7 +425,7 @@ export class GPTScript {
                 datasetID,
                 elementName: elementName,
                 elementDescription: elementDescription,
-                elementContent: elementContent
+                elementContent: Buffer.from(elementContent).toString("base64")
             }),
             workspaceID: workspaceID,
             datasetToolRepo: this.opts.DatasetToolRepo ?? "",
@@ -439,8 +439,16 @@ export class GPTScript {
             workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""
         }
 
+        const serializableElements = elements.map(e => {
+            return {
+                name: e.name,
+                description: e.description,
+                contents: Buffer.from(e.contents).toString("base64")
+            }
+        })
+
         return await this.runBasicCommand("datasets/add-elements", {
-            input: JSON.stringify({datasetID, elements}),
+            input: JSON.stringify({datasetID, elements: serializableElements}),
             workspaceID: workspaceID,
             datasetToolRepo: this.opts.DatasetToolRepo ?? "",
             env: this.opts.Env,
@@ -452,7 +460,6 @@ export class GPTScript {
             workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""
         }
 
-
         const result = await this.runBasicCommand("datasets/list-elements", {
             input: JSON.stringify({datasetID}),
             workspaceID: workspaceID,
@@ -473,7 +480,13 @@ export class GPTScript {
             datasetToolRepo: this.opts.DatasetToolRepo ?? "",
             env: this.opts.Env
         })
-        return JSON.parse(result) as DatasetElement
+
+        const element = JSON.parse(result)
+        return {
+            name: element.name,
+            description: element.description,
+            contents: Buffer.from(element.contents, "base64")
+        }
     }
 
     async createWorkspace(providerType: string, ...fromWorkspaces: string[]): Promise<string> {
@@ -1309,7 +1322,7 @@ export interface DatasetElementMeta {
 export interface DatasetElement {
     name: string
     description: string
-    contents: string
+    contents: ArrayBuffer
 }
 
 export interface DatasetMeta {

diff --git a/tests/gptscript.test.ts b/tests/gptscript.test.ts
@@ -911,7 +911,7 @@ describe("gptscript module", () => {
                 datasetID,
                 "element1",
                 "",
-                "this is element 1 contents"
+                Buffer.from("this is element 1 contents")
             )
             expect(e1.name).toEqual("element1")
             expect(e1.description).toEqual("")
@@ -921,7 +921,7 @@ describe("gptscript module", () => {
                 datasetID,
                 "element2",
                 "a description",
-                "this is element 2 contents"
+                Buffer.from("this is element 2 contents")
             )
             expect(e2.name).toEqual("element2")
             expect(e2.description).toEqual("a description")
@@ -938,12 +938,12 @@ describe("gptscript module", () => {
                 {
                   name: "element3",
                   description: "a description",
-                  contents: "this is element 3 contents"
+                  contents: Buffer.from("this is element 3 contents")
                 },
                 {
                   name: "element4",
                   description: "a description",
-                  contents: "this is element 4 contents"
+                  contents: Buffer.from("this is element 4 contents")
                 }
               ]
             )
@@ -956,22 +956,22 @@ describe("gptscript module", () => {
             const e1 = await g.getDatasetElement(workspaceID, datasetID, "element1")
             expect(e1.name).toEqual("element1")
             expect(e1.description).toBeUndefined()
-            expect(e1.contents).toEqual("this is element 1 contents")
+            expect(e1.contents).toEqual(Buffer.from("this is element 1 contents"))
 
             const e2 = await g.getDatasetElement(workspaceID, datasetID, "element2")
             expect(e2.name).toEqual("element2")
             expect(e2.description).toEqual("a description")
-            expect(e2.contents).toEqual("this is element 2 contents")
+            expect(e2.contents).toEqual(Buffer.from("this is element 2 contents"))
 
             const e3 = await g.getDatasetElement(workspaceID, datasetID, "element3")
             expect(e3.name).toEqual("element3")
             expect(e3.description).toEqual("a description")
-            expect(e3.contents).toEqual("this is element 3 contents")
+            expect(e3.contents).toEqual(Buffer.from("this is element 3 contents"))
 
             const e4 = await g.getDatasetElement(workspaceID, datasetID, "element4")
             expect(e4.name).toEqual("element4")
             expect(e4.description).toEqual("a description")
-            expect(e4.contents).toEqual("this is element 4 contents")
+            expect(e4.contents).toEqual(Buffer.from("this is element 4 contents"))
         } catch (e) {
             throw new Error("failed to get elements: " + e)
         }