chore: add addDatasetElements function (#99)

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

Hide file tree

Showing 2 changed files with 48 additions and 1 deletion.
diff --git a/src/gptscript.ts b/src/gptscript.ts
@@ -434,6 +434,19 @@ export class GPTScript {
         return JSON.parse(result) as DatasetElementMeta
     }
 
+    async addDatasetElements(workspaceID: string, datasetID: string, elements: Array<DatasetElement>) {
+        if (workspaceID === "") {
+            workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""
+        }
+
+        return await this.runBasicCommand("datasets/add-elements", {
+            input: JSON.stringify({datasetID, elements}),
+            workspaceID: workspaceID,
+            datasetToolRepo: this.opts.DatasetToolRepo ?? "",
+            env: this.opts.Env,
+        })
+    }
+
     async listDatasetElements(workspaceID: string, datasetID: string): Promise<Array<DatasetElementMeta>> {
         if (workspaceID == "") {
             workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? ""

diff --git a/tests/gptscript.test.ts b/tests/gptscript.test.ts
@@ -929,6 +929,28 @@ describe("gptscript module", () => {
             throw new Error("failed to add elements: " + e)
         }
 
+        // Add two elements at once.
+        try {
+            await g.addDatasetElements(
+              workspaceID,
+              datasetID,
+              [
+                {
+                  name: "element3",
+                  description: "a description",
+                  contents: "this is element 3 contents"
+                },
+                {
+                  name: "element4",
+                  description: "a description",
+                  contents: "this is element 4 contents"
+                }
+              ]
+            )
+        } catch (e) {
+            throw new Error("failed to add elements: " + e)
+        }
+
         // Get elements
         try {
             const e1 = await g.getDatasetElement(workspaceID, datasetID, "element1")
@@ -940,16 +962,28 @@ describe("gptscript module", () => {
             expect(e2.name).toEqual("element2")
             expect(e2.description).toEqual("a description")
             expect(e2.contents).toEqual("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")
+
+            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")
         } 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(2)
+            expect(elements.length).toEqual(4)
             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)
         }