Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feat: add tools module #621

Merged
merged 23 commits into from
Mar 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions packages/wasm-tools/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
## Usage

```ts
import { TestTool } from "@llamaindex/wasm-tools";
const testTool = new TestTool();
testTool.call("1"); // get post has id = 1 (url: https://jsonplaceholder.typicode.com/todos?id=1)
```
55 changes: 55 additions & 0 deletions packages/wasm-tools/assembly/base.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
export class ToolParameterProperty {
type: string;
description: string | null = null;

constructor(type: string, description: string | null = null) {
this.type = type;
this.description = description;
}
}

// Because AssemblyScript does not support Record<string, ToolParameterProperty> yet,
// we have to use an array of key-value pairs instead.
// When loading the metadata in application, we will convert
// the array ToolParameterPropertyRecord[] to Record<string, ToolParameterProperty>.
export class ToolParameterPropertyRecord {
key: string;
value: ToolParameterProperty;

constructor(key: string, value: ToolParameterProperty) {
this.key = key;
this.value = value;
}
}

export class ToolParameters {
type: string;
properties: ToolParameterPropertyRecord[];
required: string[] | null = null;

constructor(
type: string,
properties: ToolParameterPropertyRecord[],
required: string[] | null = null,
) {
this.type = type;
this.properties = properties;
this.required = required;
}
}

export class ToolMetadata {
name: string;
description: string;
parameters: ToolParameters | null = null;

constructor(
name: string,
description: string,
parameters: ToolParameters | null = null,
) {
this.name = name;
this.description = description;
this.parameters = parameters;
}
}
1 change: 1 addition & 0 deletions packages/wasm-tools/assembly/http.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
export declare function get(url: string, headersString: string): void;
27 changes: 27 additions & 0 deletions packages/wasm-tools/assembly/test-tool/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
import {
ToolMetadata,
ToolParameterProperty,
ToolParameterPropertyRecord,
ToolParameters,
} from "../base";
import * as http from "../http";
export * from "../base";

export const defaultMetadata: ToolMetadata = new ToolMetadata(
"Test Tool",
"This is a test tool",
new ToolParameters(
"object",
[
new ToolParameterPropertyRecord(
"query",
new ToolParameterProperty("string", "The text query to search"),
),
],
["query"],
),
);

export function call(id: string): void {
http.get(`https://jsonplaceholder.typicode.com/todos?id=${id}`, "");
}
4 changes: 4 additions & 0 deletions packages/wasm-tools/assembly/tsconfig.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"extends": "assemblyscript/std/assembly.json",
"include": ["./**/*.ts"]
}
17 changes: 17 additions & 0 deletions packages/wasm-tools/bin/compile.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
import { execSync } from "child_process";
import { readdirSync } from "fs";

// get list of tools from folder names inside assembly folder
const tools = readdirSync("assembly").filter((dir) => !dir.includes("."));

// loop through each tool, compile it to wasm and verify it
tools.forEach((tool) => {
try {
execSync(
`asc assembly/${tool}/index.ts -b dist/${tool}.wasm -t dist/${tool}.wat --exportRuntime --sourceMap --optimize`,
);
} catch (error) {
console.error(`Error compiling module ${tool}:`, error.message);
process.exit(1);
}
});
59 changes: 59 additions & 0 deletions packages/wasm-tools/package.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
{
"name": "@llamaindex/wasm-tools",
"version": "0.0.1",
"license": "MIT",
"type": "module",
"dependencies": {
"@types/node": "^18.19.14",
"@assemblyscript/loader": "^0.19.9"
},
"devDependencies": {
"assemblyscript": "^0.19.9",
"@swc/cli": "^0.3.9",
"@swc/core": "^1.4.2",
"typescript": "^5.3.3"
},
"engines": {
"node": ">=18.0.0"
},
"types": "./dist/index.d.ts",
"main": "./dist/cjs/index.js",
"exports": {
".": {
"import": {
"types": "./dist/index.d.ts",
"default": "./dist/index.js"
},
"require": {
"types": "./dist/index.d.ts",
"default": "./dist/cjs/index.js"
}
},
"./*": {
"import": {
"types": "./dist/*.d.ts",
"default": "./dist/*.js"
},
"require": {
"types": "./dist/*.d.ts",
"default": "./dist/cjs/*.js"
}
}
},
"files": [
"dist",
"CHANGELOG.md"
],
"repository": {
"type": "git",
"url": "https://github.com/run-llama/LlamaIndexTS.git",
"directory": "packages/tools"
},
"scripts": {
"build": "rm -rf ./dist && pnpm run build:esm && pnpm run build:cjs && pnpm run build:type && pnpm run build:wasm",
"build:esm": "swc src -d dist --strip-leading-paths --config-file ../../.swcrc",
"build:cjs": "swc src -d dist/cjs --strip-leading-paths --config-file ../../.cjs.swcrc",
"build:type": "tsc -p tsconfig.json",
"build:wasm": "node bin/compile.js"
}
}
155 changes: 155 additions & 0 deletions packages/wasm-tools/src/factory.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,155 @@
// @ts-ignore
import loader from "@assemblyscript/loader";
import fs from "fs";
import type { BaseTool, ToolMetadata } from "./types.js";
import { arrayKVtoObject, transformObject } from "./utils/object.js";

export default class ToolFactory {
/**
* Transform the metadata from the assemblyscript raw format to the application format
* Convert asm string to ts string
* Convert asm array to ts array
* Convert the properties from an Array to a Record<string, { type: string; description?: string }>
* Convert the argsKwargs from an Array to a Record<string, any>
*/
private static wasmInstanceToConfigs(wasmInstance: any): BaseTool {
const { __pin, __unpin, __getString, __getArray, __newString } =
wasmInstance.exports;

const getObjectByAddress = (address: any, classWrapper: any) => {
const object = classWrapper.wrap(__pin(address));
__unpin(address);
return object;
};

const {
ToolMetadata,
ToolParameters,
ToolParameterPropertyRecord,
ToolParameterProperty,
} = wasmInstance.exports;

const { defaultMetadata, call } = wasmInstance.exports;
const metadata = transformObject(
getObjectByAddress(defaultMetadata, ToolMetadata),
{
name: __getString,
description: __getString,
parameters: (parameters) => {
if (!parameters) return null;
const parametersObj = getObjectByAddress(parameters, ToolParameters);
return transformObject(parametersObj, {
type: __getString,
required: (required) => {
const requiredArray = __getArray(required);
return requiredArray.map(__getString);
},
properties: (properties) => {
const propertiesArray = __getArray(properties);
const arr = propertiesArray.map((property: any) => {
return transformObject(
getObjectByAddress(property, ToolParameterPropertyRecord),
{
key: __getString,
value: (value) => {
return transformObject(
getObjectByAddress(value, ToolParameterProperty),
{
type: __getString,
description: __getString,
},
);
},
},
);
});
return arrayKVtoObject(arr);
},
});
},
},
) as ToolMetadata;

// Wrap assemblyscript function to a ts function
const callFunction = (...args: string[]): string => {
const argsString = args.map((arg) => __pin(__newString(arg)));
return __getString(call(...argsString));
};

return {
metadata,
call: callFunction,
};
}

private static initWasmInstanceFromFile = (filePath: string) => {
const wasmFile = fs.readFileSync(
`node_modules/@llamaindex/tools/dist/${filePath}.wasm`,
);

const wasmInstance = loader.instantiateSync(wasmFile, {
http: {
// import fetch from JavaScript and use it in WebAssembly
get(url: string, headersString: string) {
const stringHeaders = wasmInstance.exports
.__getString(headersString)
.split(",,,,");

stringHeaders.pop();
const headers: Record<string, string> = {};
for (let i = 0; i < stringHeaders.length; i++) {
headers[stringHeaders[i]] = stringHeaders[i + 1];
i++;
}
fetch(wasmInstance.exports.__getString(url), {
headers: {
...headers,
},
mode: "no-cors",
method: "GET",
})
.then((fetched) => {
fetched.json().then((data) => {
console.log("Response from API call: ", data);
// Add callback to handle data if needed
return wasmInstance.exports.__newString(JSON.stringify(data));
});
})
.catch((err) => {
console.error(wasmInstance.exports.__newString(err.message));
});
},
},
});

return wasmInstance;
};

private static getToolConfigs = (filePath: string): BaseTool => {
const wasmInstance = this.initWasmInstanceFromFile(filePath);
const toolConfigs = this.wasmInstanceToConfigs(wasmInstance);
return toolConfigs;
};

private static configsToToolClass = (toolConfigs: BaseTool) => {
return class implements BaseTool {
call = toolConfigs.call;
metadata: ToolMetadata;
constructor(metadata: ToolMetadata) {
this.metadata = metadata || toolConfigs.metadata;
}
};
};

public static get toolList(): string[] {
return fs
.readdirSync("node_modules/@llamaindex/tools/dist")
.filter((file) => file.endsWith(".wasm"))
.map((file) => file.replace(".wasm", ""));
}

public static toClass = (tool: string) => {
const toolConfigs = this.getToolConfigs(tool);
return this.configsToToolClass(toolConfigs);
};
}
1 change: 1 addition & 0 deletions packages/wasm-tools/src/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
export * from "./tools.js";
3 changes: 3 additions & 0 deletions packages/wasm-tools/src/tools.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
import ToolFactory from "./factory.js";

export const TestTool = ToolFactory.toClass("test-tool");
16 changes: 16 additions & 0 deletions packages/wasm-tools/src/types.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
export type ToolParameters = {
type: string | "object";
properties: Record<string, { type: string; description?: string }>;
required?: string[];
};

export interface ToolMetadata {
description: string;
name: string;
parameters?: ToolParameters;
}

export interface BaseTool {
call?: (...args: any[]) => any;
metadata: ToolMetadata;
}
23 changes: 23 additions & 0 deletions packages/wasm-tools/src/utils/object.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
export const transformObject = (
obj: any,
transfomer: Record<string, (value: any) => any>,
) => {
const newObj: Record<string, any> = {};
for (const key in transfomer) {
newObj[key] = transfomer[key](obj[key]);
}
return newObj;
};

export const arrayKVtoObject = (
array: {
key: string;
value: any;
}[],
) => {
const obj: Record<string, any> = {};
for (const item of array) {
obj[item.key] = item.value;
}
return obj;
};
Loading
Loading