Feat: add tools module (#621)

packages/wasm-tools/README.md

+## 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)
+```

packages/wasm-tools/assembly/base.ts

+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;
+  }
+}

packages/wasm-tools/assembly/http.ts

+export declare function get(url: string, headersString: string): void;

packages/wasm-tools/assembly/test-tool/index.ts

+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}`, "");
+}

packages/wasm-tools/assembly/tsconfig.json

+{
+  "extends": "assemblyscript/std/assembly.json",
+  "include": ["./**/*.ts"]
+}

packages/wasm-tools/bin/compile.js

+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);
+  }
+});

packages/wasm-tools/package.json

+{
+  "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"
+  }
+}

packages/wasm-tools/src/factory.ts

+// @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);
+  };
+}

packages/wasm-tools/src/index.ts

+export * from "./tools.js";

packages/wasm-tools/src/tools.ts

+import ToolFactory from "./factory.js";
+
+export const TestTool = ToolFactory.toClass("test-tool");

packages/wasm-tools/src/types.ts

+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;
+}

packages/wasm-tools/src/utils/object.ts

+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;
+};

packages/wasm-tools/tsconfig.json

+{
+  "extends": "../../tsconfig.json",
+  "compilerOptions": {
+    "rootDir": "./src",
+    "outDir": "./dist",
+    "tsBuildInfoFile": "./dist/.tsbuildinfo",
+    "emitDeclarationOnly": true,
+    "module": "node16",
+    "moduleResolution": "node16",
+    "skipLibCheck": true,
+    "strict": true
+  },
+  "include": ["./src"],
+  "exclude": ["node_modules"],
+  "references": [
+    {
+      "path": "../env/tsconfig.json"
+    }
+  ]
+}

pnpm-lock.yaml

@@ -388,6 +388,28 @@ importers:
   packages/tsconfig: {}
+  packages/wasm-tools:
+    dependencies:
+      '@assemblyscript/loader':
+        specifier: ^0.19.9
+        version: 0.19.23
+      '@types/node':
+        specifier: ^18.19.14
+        version: 18.19.14
+    devDependencies:
+      '@swc/cli':
+        specifier: ^0.3.9
+        version: 0.3.9(@swc/core@1.4.2)
+      '@swc/core':
+        specifier: ^1.4.2
+        version: 1.4.2
+      assemblyscript:
+        specifier: ^0.19.9
+        version: 0.19.23
+      typescript:
+        specifier: ^5.3.3
+        version: 5.3.3
+
 packages:
   /@aashutoshrathi/word-wrap@1.2.6:
@@ -555,6 +577,10 @@ packages:
       - encoding
     dev: false
+  /@assemblyscript/loader@0.19.23:
+    resolution: {integrity: sha512-ulkCYfFbYj01ie1MDOyxv2F6SpRN1TOj7fQxbP07D6HmeR+gr2JLSmINKjga2emB+b1L2KGrFKBTc+e00p54nw==}
+    dev: false
+
   /@aws-crypto/sha256-js@5.2.0:
     resolution: {integrity: sha512-FFQQyu7edu4ufvIZ+OadFpHHOt+eSTBaYaki44c+akjg7qZg9oOQeLlk77F6tSYqjDAFClrHJk9tMf0HdVyOvA==}
     engines: {node: '>=16.0.0'}
@@ -5122,6 +5148,15 @@ packages:
       - utf-8-validate
     dev: false
+  /assemblyscript@0.19.23:
+    resolution: {integrity: sha512-fwOQNZVTMga5KRsfY80g7cpOl4PsFQczMwHzdtgoqLXaYhkhavufKb0sB0l3T1DUxpAufA0KNhlbpuuhZUwxMA==}
+    hasBin: true
+    dependencies:
+      binaryen: 102.0.0-nightly.20211028
+      long: 5.2.3
+      source-map-support: 0.5.21
+    dev: true
+
   /assertion-error@1.1.0:
     resolution: {integrity: sha512-jgsaNduz+ndvGyFt3uSuWqvy4lCnIJiovtouQN5JZHOKCS2QuhEdbcQHFhVksz2N2U9hXJo8odG7ETyWlEeuDw==}
     dev: true
@@ -5315,6 +5350,11 @@ packages:
     resolution: {integrity: sha512-jDctJ/IVQbZoJykoeHbhXpOlNBqGNcwXJKJog42E5HDPUwQTSdjCHdihjj0DlnheQ7blbT6dHOafNAiS8ooQKA==}
     engines: {node: '>=8'}
+  /binaryen@102.0.0-nightly.20211028:
+    resolution: {integrity: sha512-GCJBVB5exbxzzvyt8MGDv/MeUjs6gkXDvf4xOIItRBptYl0Tz5sm1o/uG95YK0L0VeG5ajDu3hRtkBP2kzqC5w==}
+    hasBin: true
+    dev: true
+
   /bl@4.1.0:
     resolution: {integrity: sha512-1W07cM9gS6DcLperZfFSj+bWLtaPGSOHWhPiGzXmvVJbRLdG82sH/Kn8EtW1VqWVA54AKf2h5k5BbnIbwF3h6w==}
     dependencies:
@@ -8344,7 +8384,7 @@ packages:
     dependencies:
       array-union: 2.1.0
       dir-glob: 3.0.1
-      fast-glob: 3.3.2
+      fast-glob: 3.3.1
       ignore: 5.3.1
       merge2: 1.4.1
       slash: 3.0.0