Skip to content

Commit

Permalink
chore: update for dataset rewrite (#102)
Browse files Browse the repository at this point in the history 
Signed-off-by: Grant Linville <[email protected]>
  • Loading branch information
@g-linville
g-linville authored Nov 6, 2024
1 parent 56135fa commit d4b222c
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 149 deletions.
112 changes: 29 additions & 83 deletions src/gptscript.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ export interface GlobalOpts {
BaseURL?: string
DefaultModel?: string
DefaultModelProvider?: string
DatasetToolRepo?: string
DatasetTool?: string
WorkspaceTool?: string
Env?: string[]
}
Expand Down Expand Up @@ -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")
}
}

Expand Down Expand Up @@ -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
}
104 changes: 38 additions & 66 deletions tests/gptscript.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down

0 comments on commit d4b222c

Please sign in to comment.