diff --git a/src/gptscript.ts b/src/gptscript.ts index 4f1cf6e..7067a84 100644 --- a/src/gptscript.ts +++ b/src/gptscript.ts @@ -12,6 +12,7 @@ export interface GlobalOpts { BaseURL?: string DefaultModel?: string DefaultModelProvider?: string + DatasetToolRepo?: string Env?: string[] } @@ -390,6 +391,84 @@ export class GPTScript { await r.text() } + // Dataset methods + + async listDatasets(workspace: string): Promise> { + if (workspace == "") { + workspace = process.env.GPTSCRIPT_WORKSPACE_DIR ?? "" + } + + const r: Run = new RunSubcommand("datasets", "", {URL: this.opts.URL, Token: this.opts.Token}) + r.request({input: "{}", workspace: workspace, datasetToolRepo: this.opts.DatasetToolRepo ?? ""}) + const result = await r.text() + return JSON.parse(result) as Array + } + + async createDataset(workspace: string, name: string, description: string): Promise { + if (workspace == "") { + workspace = process.env.GPTSCRIPT_WORKSPACE_DIR ?? "" + } + + const r: Run = new RunSubcommand("datasets/create", "", {URL: this.opts.URL, Token: this.opts.Token}) + r.request({ + input: JSON.stringify({datasetName: name, datasetDescription: description}), + workspace: workspace, + datasetToolRepo: this.opts.DatasetToolRepo ?? "" + }) + const result = await r.text() + return JSON.parse(result) as Dataset + } + + async addDatasetElement(workspace: string, datasetID: string, elementName: string, elementDescription: string, elementContent: string): Promise { + if (workspace == "") { + workspace = process.env.GPTSCRIPT_WORKSPACE_DIR ?? "" + } + + const r: Run = new RunSubcommand("datasets/add-element", "", {URL: this.opts.URL, Token: this.opts.Token}) + r.request({ + input: JSON.stringify({ + datasetID, + elementName, + elementDescription, + elementContent + }), + workspace: workspace, + datasetToolRepo: this.opts.DatasetToolRepo ?? "" + }) + const result = await r.text() + return JSON.parse(result) as DatasetElementMeta + } + + async listDatasetElements(workspace: string, datasetID: string): Promise> { + if (workspace == "") { + workspace = process.env.GPTSCRIPT_WORKSPACE_DIR ?? "" + } + + const r: Run = new RunSubcommand("datasets/list-elements", "", {URL: this.opts.URL, Token: this.opts.Token}) + r.request({ + input: JSON.stringify({datasetID}), + workspace: workspace, + datasetToolRepo: this.opts.DatasetToolRepo ?? "" + }) + const result = await r.text() + return JSON.parse(result) as Array + } + + async getDatasetElement(workspace: string, datasetID: string, elementName: string): Promise { + if (workspace == "") { + workspace = process.env.GPTSCRIPT_WORKSPACE_DIR ?? "" + } + + const r: Run = new RunSubcommand("datasets/get-element", "", {URL: this.opts.URL, Token: this.opts.Token}) + r.request({ + input: JSON.stringify({datasetID, element: elementName}), + workspace: workspace, + datasetToolRepo: this.opts.DatasetToolRepo ?? "" + }) + const result = await r.text() + return JSON.parse(result) as DatasetElement + } + /** * Helper method to handle the common logic for loading. * @@ -1103,3 +1182,29 @@ function jsonToCredential(cred: string): Credential { refreshToken: c.refreshToken } } + +// Dataset types + +export interface DatasetElementMeta { + name: string + description: string +} + +export interface DatasetElement { + name: string + description: string + contents: string +} + +export interface DatasetMeta { + id: string + name: string + description: string +} + +export interface Dataset { + id: string + name: string + description: string + elements: Record +} diff --git a/tests/gptscript.test.ts b/tests/gptscript.test.ts index 5703985..94cc0bd 100644 --- a/tests/gptscript.test.ts +++ b/tests/gptscript.test.ts @@ -1,7 +1,7 @@ import * as gptscript from "../src/gptscript" import { ArgumentSchemaType, - CredentialType, + CredentialType, Dataset, getEnv, PropertyType, RunEventType, @@ -13,6 +13,7 @@ import path from "path" import {fileURLToPath} from "url" import * as fs from "node:fs" import {randomBytes} from "node:crypto" +import {tmpdir} from "node:os"; let gFirst: gptscript.GPTScript let g: gptscript.GPTScript @@ -885,4 +886,82 @@ describe("gptscript module", () => { throw new Error("failed to verify deletion: " + e) } }, 20000) + + test("dataset operations", async () => { + const datasetName = "test-" + randomBytes(10).toString("hex") + const workspace = fs.mkdtempSync(path.join(tmpdir(), "node-gptscript-")) + let datasetID: string + + // Create + try { + const dataset = await g.createDataset(workspace, 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 + } catch (e) { + throw new Error("failed to create dataset: " + e) + } + + // Add elements + try { + const e1 = await g.addDatasetElement( + workspace, + datasetID, + "element1", + "", + "this is element 1 contents" + ) + expect(e1.name).toEqual("element1") + expect(e1.description).toEqual("") + + const e2 = await g.addDatasetElement( + workspace, + datasetID, + "element2", + "a description", + "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) + } + + // Get elements + try { + const e1 = await g.getDatasetElement(workspace, datasetID, "element1") + expect(e1.name).toEqual("element1") + expect(e1.description).toBeUndefined() + expect(e1.contents).toEqual("this is element 1 contents") + + const e2 = await g.getDatasetElement(workspace, datasetID, "element2") + expect(e2.name).toEqual("element2") + expect(e2.description).toEqual("a description") + expect(e2.contents).toEqual("this is element 2 contents") + } catch (e) { + throw new Error("failed to get elements: " + e) + } + + // List the elements in the dataset + try { + const elements = await g.listDatasetElements(workspace, datasetID) + expect(elements.length).toEqual(2) + expect(elements.map(e => e.name)).toContain("element1") + expect(elements.map(e => e.name)).toContain("element2") + } catch (e) { + throw new Error("failed to list elements: " + e) + } + + // List datasets + try { + const datasets = await g.listDatasets(workspace) + expect(datasets.length).toBeGreaterThan(0) + expect(datasets.map(d => d.name)).toContain(datasetName) + } catch (e) { + throw new Error("failed to list datasets: " + e) + } + }, 20000) })