chore: update for dataset rewrite (#102)

Signed-off-by: Grant Linville <[email protected]>

Author: g-linville

src/gptscript.ts

@@ -12,7 +12,7 @@ export interface GlobalOpts {
     BaseURL?: string
     DefaultModel?: string
     DefaultModelProvider?: string
-    DatasetToolRepo?: string
+    DatasetTool?: string
     WorkspaceTool?: string
     Env?: string[]
 }
@@ -386,106 +386,60 @@ export class GPTScript {
         })
     }
 
-    // Dataset methods
-
-    async listDatasets(workspaceID: string): Promise<Array<DatasetMeta>> {
-        if (workspaceID == "") {
-            workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""
-        }
-
+    // returns an array of dataset IDs
+    async listDatasets(): Promise<Array<DatasetMeta>> {
         const result = await this.runBasicCommand("datasets", {
-            workspaceID: workspaceID,
-            datasetToolRepo: this.opts.DatasetToolRepo ?? "",
+            input: "{}",
+            datasetTool: this.opts.DatasetTool ?? "",
             env: this.opts.Env
         })
         return JSON.parse(result) as Array<DatasetMeta>
     }
 
-    async createDataset(workspaceID: string, name: string, description: string): Promise<Dataset> {
-        if (workspaceID == "") {
-            workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""
-        }
-
-        const result = await this.runBasicCommand("datasets/create", {
-            input: JSON.stringify({datasetName: name, datasetDescription: description}),
-            workspaceID: workspaceID,
-            datasetToolRepo: this.opts.DatasetToolRepo ?? "",
-            env: this.opts.Env
-        })
-        return JSON.parse(result) as Dataset
-    }
-
-    async addDatasetElement(workspaceID: string, datasetID: string, elementName: string, elementDescription: string, elementContent: ArrayBuffer): Promise<DatasetElementMeta> {
-        if (workspaceID == "") {
-            workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""
-        }
-
-        const result = await this.runBasicCommand("datasets/add-element", {
-            input: JSON.stringify({
-                datasetID,
-                elementName: elementName,
-                elementDescription: elementDescription,
-                elementContent: Buffer.from(elementContent).toString("base64")
-            }),
-            workspaceID: workspaceID,
-            datasetToolRepo: this.opts.DatasetToolRepo ?? "",
-            env: this.opts.Env
-        })
-        return JSON.parse(result) as DatasetElementMeta
-    }
-
-    async addDatasetElements(workspaceID: string, datasetID: string, elements: Array<DatasetElement>) {
-        if (workspaceID === "") {
-            workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""
-        }
-
+    async addDatasetElements(elements: Array<DatasetElement>, opts: {name?: string, description?: string, datasetID?: string}): Promise<string> {
         const serializableElements = elements.map(e => {
             return {
                 name: e.name,
                 description: e.description,
-                contents: Buffer.from(e.contents).toString("base64")
+                contents: e.contents,
+                binaryContents: Buffer.from(e.binaryContents ?? Buffer.from("")).toString("base64")
             }
         })
 
         return await this.runBasicCommand("datasets/add-elements", {
-            input: JSON.stringify({datasetID, elements: serializableElements}),
-            workspaceID: workspaceID,
-            datasetToolRepo: this.opts.DatasetToolRepo ?? "",
-            env: this.opts.Env,
+            input: JSON.stringify({
+                name: opts.name ?? "",
+                description: opts.description ?? "",
+                datasetID: opts.datasetID ?? "",
+                elements: serializableElements
+            }),
+            datasetTool: this.opts.DatasetTool ?? "",
+            env: this.opts.Env
         })
     }
 
-    async listDatasetElements(workspaceID: string, datasetID: string): Promise<Array<DatasetElementMeta>> {
-        if (workspaceID == "") {
-            workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""
-        }
-
+    async listDatasetElements(datasetID: string): Promise<Array<DatasetElementMeta>> {
         const result = await this.runBasicCommand("datasets/list-elements", {
             input: JSON.stringify({datasetID}),
-            workspaceID: workspaceID,
-            datasetToolRepo: this.opts.DatasetToolRepo ?? "",
+            datasetTool: this.opts.DatasetTool ?? "",
             env: this.opts.Env
         })
         return JSON.parse(result) as Array<DatasetElementMeta>
     }
 
-    async getDatasetElement(workspaceID: string, datasetID: string, elementName: string): Promise<DatasetElement> {
-        if (workspaceID == "") {
-            workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""
-        }
-
+    async getDatasetElement(datasetID: string, elementName: string): Promise<DatasetElement> {
         const result = await this.runBasicCommand("datasets/get-element", {
-            input: JSON.stringify({datasetID, element: elementName}),
-            workspaceID: workspaceID,
-            datasetToolRepo: this.opts.DatasetToolRepo ?? "",
+            input: JSON.stringify({datasetID, name: elementName}),
+            datasetTool: this.opts.DatasetTool ?? "",
             env: this.opts.Env
         })
 
         const element = JSON.parse(result)
         return {
             name: element.name,
             description: element.description,
-            contents: Buffer.from(element.contents, "base64")
+            contents: element.contents,
+            binaryContents: Buffer.from(element.binaryContents ?? "", "base64")
         }
     }
 
@@ -1312,28 +1266,20 @@ function jsonToCredential(cred: string): Credential {
     }
 }
 
-// Dataset types
-
-export interface DatasetElementMeta {
-    name: string
-    description: string
-}
-
-export interface DatasetElement {
+export interface DatasetMeta {
+    id: string
     name: string
     description: string
-    contents: ArrayBuffer
 }
 
-export interface DatasetMeta {
-    id: string
+export interface DatasetElementMeta {
     name: string
     description: string
 }
 
-export interface Dataset {
-    id: string
+export interface DatasetElement {
     name: string
     description: string
-    elements: Record<string, DatasetElementMeta>
+    contents?: string
+    binaryContents?: ArrayBuffer
 }

tests/gptscript.test.ts

@@ -887,112 +887,84 @@ describe("gptscript module", () => {
     }, 20000)
 
     test("dataset operations", async () => {
-        const datasetName = "test-" + randomBytes(10).toString("hex")
-        const workspaceID = await g.createWorkspace("directory")
+        process.env.GPTSCRIPT_WORKSPACE_ID = await g.createWorkspace("directory")
+
+        const client = new gptscript.GPTScript({
+            APIKey: process.env.OPENAI_API_KEY,
+            Env: Object.entries(process.env).map(([k, v]) => `${k}=${v}`)
+        })
+
         let datasetID: string
 
-        // Create
+        // Create and add two elements
         try {
-            const dataset = await g.createDataset(workspaceID, datasetName, "a test dataset")
-            expect(dataset).toBeDefined()
-            expect(dataset.name).toEqual(datasetName)
-            expect(dataset.description).toEqual("a test dataset")
-            expect(dataset.id.length).toBeGreaterThan(0)
-            expect(dataset.elements).toEqual({})
-            datasetID = dataset.id
+            datasetID = await client.addDatasetElements([
+                {
+                    name: "element1",
+                    description: "",
+                    contents: "this is element 1 contents"
+                },
+                {
+                    name: "element2",
+                    description: "a description",
+                    binaryContents: Buffer.from("this is element 2 contents")
+                }
+            ], {name: "test-dataset", description: "a test dataset"})
         } catch (e) {
             throw new Error("failed to create dataset: " + e)
         }
 
-        // Add elements
-        try {
-            const e1 = await g.addDatasetElement(
-                workspaceID,
-                datasetID,
-                "element1",
-                "",
-                Buffer.from("this is element 1 contents")
-            )
-            expect(e1.name).toEqual("element1")
-            expect(e1.description).toEqual("")
-
-            const e2 = await g.addDatasetElement(
-                workspaceID,
-                datasetID,
-                "element2",
-                "a description",
-                Buffer.from("this is element 2 contents")
-            )
-            expect(e2.name).toEqual("element2")
-            expect(e2.description).toEqual("a description")
-        } catch (e) {
-            throw new Error("failed to add elements: " + e)
-        }
-
-        // Add two elements at once.
+        // Add another element
         try {
-            await g.addDatasetElements(
-              workspaceID,
-              datasetID,
-              [
+            await client.addDatasetElements([
                 {
-                  name: "element3",
-                  description: "a description",
-                  contents: Buffer.from("this is element 3 contents")
-                },
-                {
-                  name: "element4",
-                  description: "a description",
-                  contents: Buffer.from("this is element 4 contents")
+                    name: "element3",
+                    description: "a description",
+                    contents: "this is element 3 contents"
                 }
-              ]
-            )
+            ], {datasetID: datasetID})
         } catch (e) {
             throw new Error("failed to add elements: " + e)
         }
 
         // Get elements
         try {
-            const e1 = await g.getDatasetElement(workspaceID, datasetID, "element1")
+            const e1 = await client.getDatasetElement(datasetID, "element1")
             expect(e1.name).toEqual("element1")
             expect(e1.description).toBeUndefined()
-            expect(e1.contents).toEqual(Buffer.from("this is element 1 contents"))
+            expect(e1.contents).toEqual("this is element 1 contents")
 
-            const e2 = await g.getDatasetElement(workspaceID, datasetID, "element2")
+            const e2 = await client.getDatasetElement(datasetID, "element2")
             expect(e2.name).toEqual("element2")
             expect(e2.description).toEqual("a description")
-            expect(e2.contents).toEqual(Buffer.from("this is element 2 contents"))
+            expect(e2.binaryContents).toEqual(Buffer.from("this is element 2 contents"))
 
-            const e3 = await g.getDatasetElement(workspaceID, datasetID, "element3")
+            const e3 = await client.getDatasetElement(datasetID, "element3")
             expect(e3.name).toEqual("element3")
             expect(e3.description).toEqual("a description")
-            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(Buffer.from("this is element 4 contents"))
+            expect(e3.contents).toEqual("this is element 3 contents")
         } catch (e) {
             throw new Error("failed to get elements: " + e)
         }
 
         // List the elements in the dataset
         try {
-            const elements = await g.listDatasetElements(workspaceID, datasetID)
-            expect(elements.length).toEqual(4)
+            const elements = await client.listDatasetElements(datasetID)
+            expect(elements.length).toEqual(3)
             expect(elements.map(e => e.name)).toContain("element1")
             expect(elements.map(e => e.name)).toContain("element2")
             expect(elements.map(e => e.name)).toContain("element3")
-            expect(elements.map(e => e.name)).toContain("element4")
         } catch (e) {
             throw new Error("failed to list elements: " + e)
         }
 
         // List datasets
         try {
-            const datasets = await g.listDatasets(workspaceID)
+            const datasets = await client.listDatasets()
             expect(datasets.length).toBeGreaterThan(0)
-            expect(datasets.map(d => d.name)).toContain(datasetName)
+            expect(datasets[0].id).toEqual(datasetID)
+            expect(datasets[0].name).toEqual("test-dataset")
+            expect(datasets[0].description).toEqual("a test dataset")
         } catch (e) {
             throw new Error("failed to list datasets: " + e)
         }