From d4b222c34e757b62429dc43ae9cc44aa3f6c9700 Mon Sep 17 00:00:00 2001 From: Grant Linville Date: Wed, 6 Nov 2024 17:07:34 -0500 Subject: [PATCH] chore: update for dataset rewrite (#102) Signed-off-by: Grant Linville --- src/gptscript.ts | 112 +++++++++++----------------------------- tests/gptscript.test.ts | 104 ++++++++++++++----------------------- 2 files changed, 67 insertions(+), 149 deletions(-) diff --git a/src/gptscript.ts b/src/gptscript.ts index 406f702..d241cb0 100644 --- a/src/gptscript.ts +++ b/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,98 +386,51 @@ export class GPTScript { }) } - // Dataset methods - - async listDatasets(workspaceID: string): Promise> { - if (workspaceID == "") { - workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? "" - } - + // returns an array of dataset IDs + async listDatasets(): Promise> { 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 } - async createDataset(workspaceID: string, name: string, description: string): Promise { - 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 { - 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) { - if (workspaceID === "") { - workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? "" - } - + async addDatasetElements(elements: Array, opts: {name?: string, description?: string, datasetID?: string}): Promise { 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> { - if (workspaceID == "") { - workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? "" - } - + async listDatasetElements(datasetID: string): Promise> { 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 } - async getDatasetElement(workspaceID: string, datasetID: string, elementName: string): Promise { - if (workspaceID == "") { - workspaceID = process.env.GPTSCRIPT_WORKSPACE_ID ?? "" - } - + async getDatasetElement(datasetID: string, elementName: string): Promise { 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 }) @@ -485,7 +438,8 @@ export class GPTScript { 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 + contents?: string + binaryContents?: ArrayBuffer } diff --git a/tests/gptscript.test.ts b/tests/gptscript.test.ts index fa3dec9..077107c 100644 --- a/tests/gptscript.test.ts +++ b/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) }