feat: build wasm with extism

.github/workflows/preview.yml

@@ -18,6 +18,11 @@ jobs:
           node-version-file: ".nvmrc"
           cache: "pnpm"
 
+      - name: Install Extism
+        run: |
+          curl -O https://raw.githubusercontent.com/extism/js-pdk/main/install.sh
+          sh install.sh
+
       - name: Install dependencies
         run: pnpm install
 

.github/workflows/publish_github.yml

@@ -20,6 +20,11 @@ jobs:
           node-version-file: ".nvmrc"
           cache: "pnpm"
 
+      - name: Install Extism
+        run: |
+          curl -O https://raw.githubusercontent.com/extism/js-pdk/main/install.sh
+          sh install.sh
+
       - name: Install dependencies
         run: pnpm install
 

.github/workflows/release.yml

@@ -23,6 +23,11 @@ jobs:
           node-version-file: ".nvmrc"
           cache: "pnpm"
 
+      - name: Install Extism
+        run: |
+          curl -O https://raw.githubusercontent.com/extism/js-pdk/main/install.sh
+          sh install.sh
+
       - name: Install dependencies
         run: pnpm install
 

.github/workflows/test.yml

@@ -51,6 +51,10 @@ jobs:
         with:
           node-version: ${{ matrix.node-version }}
           cache: "pnpm"
+      - name: Install Extism
+        run: |
+          curl -O https://raw.githubusercontent.com/extism/js-pdk/main/install.sh
+          sh install.sh
       - name: Install dependencies
         run: pnpm install
       - name: Run tests
@@ -66,6 +70,10 @@ jobs:
         with:
           node-version-file: ".nvmrc"
           cache: "pnpm"
+      - name: Install Extism
+        run: |
+          curl -O https://raw.githubusercontent.com/extism/js-pdk/main/install.sh
+          sh install.sh
       - name: Install dependencies
         run: pnpm install
       - name: Build
@@ -103,6 +111,10 @@ jobs:
         with:
           node-version-file: ".nvmrc"
           cache: "pnpm"
+      - name: Install Extism
+        run: |
+          curl -O https://raw.githubusercontent.com/extism/js-pdk/main/install.sh
+          sh install.sh
       - name: Install dependencies
         run: pnpm install
       - name: Build llamaindex
@@ -122,6 +134,10 @@ jobs:
         with:
           node-version-file: ".nvmrc"
           cache: "pnpm"
+      - name: Install Extism
+        run: |
+          curl -O https://raw.githubusercontent.com/extism/js-pdk/main/install.sh
+          sh install.sh
       - name: Install dependencies
         run: pnpm install
       - name: Build

packages/extism-tools/README.md

@@ -0,0 +1,30 @@
+## Extism Tools
+
+### Prerequisites for Development
+
+- [Extism PDK](https://github.com/extism/js-pdk?tab=readme-ov-file#linux-macos)
+
+### Build WASM files
+
+```bash
+cd packages/wasm-tools/extism/wiki
+extism-js wiki.js -i wiki.d.ts -o wiki.wasm
+```
+
+### Run WASM files in Node.js using Extism SDK (https://github.com/extism/js-sdk)
+
+```bash
+node examples/wasm/wiki.js
+```
+
+### Test using WASM tool with OpenAI Agent
+
+```bash
+node examples/tool/WikipediaTool.js
+```
+
+### Run WASM files in Python using Extism SDK (https://github.com/extism/python-sdk)
+
+```bash
+python examples/wasm/wiki.py
+```

packages/extism-tools/bin/compile.js

@@ -0,0 +1,31 @@
+// needed as extism-js doesn't support compiling multiple files
+import { execSync } from "child_process";
+import { mkdirSync, readdirSync } from "fs";
+
+const WASM_SRC_FOLDER = "wasm";
+const WASM_OUTPUT_FOLDER = "dist/wasm";
+
+// get list of tools from files (except index.d.ts)
+const tools = readdirSync(WASM_SRC_FOLDER)
+  .filter((file) => !file.includes("index.d.ts"))
+  .map((file) => file.split(".")[0]);
+
+// create dist/wasm folder if it doesn't exist using fs
+try {
+  mkdirSync(WASM_OUTPUT_FOLDER, { recursive: true });
+} catch (error) {
+  console.error("Error creating dist/wasm folder:", error.message);
+  process.exit(1);
+}
+
+// loop through each tool, compile it to wasm using extism-js
+tools.forEach((tool) => {
+  try {
+    execSync(
+      `extism-js ${WASM_SRC_FOLDER}/${tool}.js -i ${WASM_SRC_FOLDER}/index.d.ts -o ${WASM_OUTPUT_FOLDER}/${tool}.wasm`,
+    );
+  } catch (error) {
+    console.error(`Error compiling module ${tool}:`, error.message);
+    process.exit(1);
+  }
+});

packages/extism-tools/examples/tool/wiki.js

@@ -0,0 +1,12 @@
+import { OpenAIAgent } from "llamaindex";
+import { ExtismToolFactory } from "../../dist/ExtismToolFactory.js";
+
+async function main() {
+  const WikiTool = await ExtismToolFactory.createToolClass("wiki");
+  const wikiTool = new WikiTool();
+  const agent = new OpenAIAgent({ tools: [wikiTool] });
+  const result = await agent.chat({ message: "Ho Chi Minh City" });
+  console.log(result.message);
+}
+
+void main();

packages/extism-tools/examples/wasm/wiki.js

@@ -0,0 +1,23 @@
+import createPlugin from "@extism/extism";
+
+async function main() {
+  const plugin = await createPlugin("./dist/wasm/wiki.wasm", {
+    useWasi: true,
+    functions: {},
+    runInWorker: true,
+    allowedHosts: ["*.wikipedia.org"],
+    memory: { maxHttpResponseBytes: 100 * 1024 * 1024 },
+  });
+
+  try {
+    const params = { query: "Ho Chi Minh City" };
+    const data = await plugin.call("summary", JSON.stringify(params));
+    console.log(data.json());
+  } catch (e) {
+    console.error(e);
+  } finally {
+    await plugin.close();
+  }
+}
+
+void main();

packages/extism-tools/examples/wasm/wiki.py

@@ -0,0 +1,34 @@
+import extism
+import json
+from os.path import join, dirname
+
+
+def read_local_wasm(file_name):
+    path = join(dirname(__file__), file_name)  # Change this to your wasm file path
+    with open(path, "rb") as wasm_file:
+        return wasm_file.read()
+
+
+def _manifest(file_name):
+    wasm = read_local_wasm(file_name)
+    return {
+        "wasm": [{"data": wasm}],
+        "allowed_hosts": ["*.wikipedia.org"],
+    }
+
+
+manifest = _manifest("wiki.wasm")
+with extism.Plugin(manifest, wasi=True) as plugin:
+    metadata = plugin.call(
+        "getMetadata",
+        "",
+        parse=lambda output: json.loads(bytes(output).decode("utf-8")),
+    )
+    data = plugin.call(
+        "call",
+        json.dumps({"query": "Ho Chi Minh City"}),
+        parse=lambda output: json.loads(bytes(output).decode("utf-8")),
+    )
+
+print(metadata)
+print(data)

packages/extism-tools/package.json

@@ -0,0 +1,62 @@
+{
+  "name": "@llamaindex/extism-tools",
+  "version": "0.0.1",
+  "license": "MIT",
+  "type": "module",
+  "dependencies": {
+    "@extism/extism": "^2.0.0-rc7",
+    "ajv": "^8.17.1",
+    "@llamaindex/core": "workspace:*",
+    "@llamaindex/env": "workspace:*",
+    "llamaindex": "workspace:*"
+  },
+  "devDependencies": {
+    "@swc/cli": "^0.4.0",
+    "@swc/core": "^1.7.22",
+    "typescript": "^5.5.4",
+    "@types/node": "^22.5.1"
+  },
+  "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/extism-tools"
+  },
+  "scripts": {
+    "build": "rm -rf ./dist && pnpm run build:wasm && pnpm run build:esm && pnpm run build:cjs && pnpm run build:type",
+    "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/extism-tools/src/ExtismToolFactory.ts

@@ -0,0 +1,85 @@
+import createPlugin, { type Plugin } from "@extism/extism";
+import type { JSONSchemaType } from "ajv";
+import type { BaseTool, ToolMetadata } from "llamaindex";
+
+export const WASM_DIRECTORY = "./dist/wasm";
+export const DEFAULT_MAX_HTTP_RESPONSE_BYTES = 100 * 1024 * 1024; // 100 MB
+
+export type ToolParams = Record<string, any>;
+
+export type ToolClassParams = {
+  metadata?: ToolMetadata<JSONSchemaType<ToolParams>>;
+};
+
+export type CreateToolClassParams = {
+  wasmFilename: string;
+  allowedHosts: string[];
+  maxHttpResponseBytes: number;
+  transformResponse: (response: any) => any;
+};
+
+export enum ExtismTool {
+  WIKI = "wiki",
+}
+
+export const ToolMap: Record<`${ExtismTool}`, CreateToolClassParams> = {
+  [ExtismTool.WIKI]: {
+    wasmFilename: "wiki.wasm",
+    allowedHosts: ["*.wikipedia.org"],
+    maxHttpResponseBytes: DEFAULT_MAX_HTTP_RESPONSE_BYTES,
+    transformResponse: (response: any) => response.extract,
+  },
+};
+
+export const createPluginInstance = async (
+  params: Omit<CreateToolClassParams, "transformResponse">,
+): Promise<Plugin> => {
+  const { wasmFilename, allowedHosts, maxHttpResponseBytes } = params;
+  const plugin = await createPlugin(`${WASM_DIRECTORY}/${wasmFilename}`, {
+    useWasi: true,
+    runInWorker: true,
+    allowedHosts,
+    memory: { maxHttpResponseBytes },
+  });
+  return plugin;
+};
+
+export class ExtismToolFactory {
+  static async createToolClass(
+    toolName: `${ExtismTool}`,
+  ): Promise<new (params: ToolClassParams) => BaseTool<ToolParams>> {
+    const config = ToolMap[toolName];
+    if (!config) throw new Error(`Tool ${toolName} not supported yet`);
+
+    const plugin = await createPluginInstance(config);
+    try {
+      const wasmMetadata = await plugin.call("getMetadata");
+      if (!wasmMetadata) {
+        throw new Error("The WASM plugin did not return metadata.");
+      }
+      const defaultMetadata = wasmMetadata.json();
+
+      return class implements BaseTool<ToolParams> {
+        metadata: ToolMetadata<JSONSchemaType<ToolParams>>;
+
+        constructor(params: ToolClassParams) {
+          this.metadata = params?.metadata || defaultMetadata;
+        }
+
+        async call(input: ToolParams): Promise<string> {
+          const pluginInstance = await createPluginInstance(config);
+          const data = await pluginInstance.call("call", JSON.stringify(input));
+          if (!data) return "No result";
+          const result = config.transformResponse(data.json());
+          await pluginInstance.close();
+          return result;
+        }
+      };
+    } catch (e) {
+      console.error(e);
+      throw new Error("Failed to create Tool instance.");
+    } finally {
+      await plugin.close();
+    }
+  }
+}

packages/extism-tools/tsconfig.json

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

packages/extism-tools/wasm/index.d.ts

@@ -0,0 +1,4 @@
+declare module "main" {
+  export function getMetadata(): I32;
+  export function call(): I32;
+}