feat: build wasm with extism

.github/workflows/lint_on_push_or_pull.yml

@@ -19,6 +19,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: Run lint

.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,31 @@
+## Extism Tools
+
+### Prerequisites for Development
+
+- [Extism CLI](https://github.com/extism/cli?tab=readme-ov-file#installation)
+- [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
+```
+
+### Test WASM files
+
+```bash
+extism call wiki.wasm wikiCall --wasi --allow-host "*.wikipedia.org" --input='{"query": "Ho Chi Minh City"}'
+```
+
+### Run WASM files in Node.js using Extism SDK (https://github.com/extism/js-sdk)
+
+```bash
+node examples/wiki.js
+```
+
+### Run WASM files in Python using Extism SDK (https://github.com/extism/python-sdk)
+
+```bash
+python examples/wiki.py
+```

packages/extism-tools/bin/compile.js

@@ -0,0 +1,28 @@
+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 folder names inside assembly folder
+const tools = readdirSync(WASM_SRC_FOLDER).filter((dir) => !dir.includes("."));
+
+// 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}/${tool}.js -i ${WASM_SRC_FOLDER}/${tool}/${tool}.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/WikipediaTool.js

@@ -0,0 +1,17 @@
+import { OpenAIAgent } from "llamaindex";
+import { WikipediaTool } from "../../dist/WikipediaTool.js";
+
+const wikiTool = new WikipediaTool();
+void wikiTool.call({ query: "Ho Chi Minh City" }).then(console.log);
+
+async function main() {
+  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/tool/WikipediaTool.ts

@@ -0,0 +1,17 @@
+import { OpenAIAgent } from "llamaindex";
+import { WikipediaTool } from "../../dist/WikipediaTool.js";
+
+const wikiTool = new WikipediaTool();
+void wikiTool.call({ query: "Ho Chi Minh City" }).then(console.log);
+
+async function main() {
+  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,27 @@
+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:
+    data = plugin.call(
+        "wikiCall",
+        json.dumps({"query": "Ho Chi Minh City"}),
+        parse=lambda output: json.loads(bytes(output).decode("utf-8")),
+    )
+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/ExtismTool.ts

@@ -0,0 +1,22 @@
+import createPlugin, { type ExtismPluginOptions } from "@extism/extism";
+import type { BaseTool } from "@llamaindex/core/llms";
+
+export abstract class ExtismTool<T> implements BaseTool<T> {
+  readonly WASM_FILE_PATH = "./dist/wasm";
+  wasmFilename!: string;
+  metadata!: any;
+
+  get pluginConfig(): ExtismPluginOptions {
+    return {
+      useWasi: true,
+      runInWorker: true,
+    };
+  }
+
+  async getPlugin() {
+    return await createPlugin(
+      `${this.WASM_FILE_PATH}/${this.wasmFilename}`,
+      this.pluginConfig,
+    );
+  }
+}

packages/extism-tools/src/WikipediaTool.ts

@@ -0,0 +1,67 @@
+import { type ExtismPluginOptions } from "@extism/extism";
+import type { ToolMetadata } from "@llamaindex/core/llms";
+import type { JSONSchemaType } from "ajv";
+import { ExtismTool } from "./ExtismTool.js";
+
+type WikipediaParameter = {
+  query: string;
+};
+
+export type WikipediaToolParams = {
+  metadata?: ToolMetadata<JSONSchemaType<WikipediaParameter>>;
+};
+
+const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<WikipediaParameter>> = {
+  name: "wikipedia_tool",
+  description: "A tool that uses a query engine to search Wikipedia.",
+  parameters: {
+    type: "object",
+    properties: {
+      query: {
+        type: "string",
+        description: "The query to search for",
+      },
+    },
+    required: ["query"],
+  },
+};
+
+export class WikipediaTool extends ExtismTool<WikipediaParameter> {
+  readonly ALLOWED_HOSTS = ["*.wikipedia.org"];
+  readonly MAX_HTTP_RESPONSE_BYTES = 100 * 1024 * 1024;
+
+  wasmFilename: string = "wiki.wasm";
+  metadata: ToolMetadata<JSONSchemaType<WikipediaParameter>>;
+
+  constructor(params?: WikipediaToolParams) {
+    super();
+    this.metadata = params?.metadata || DEFAULT_META_DATA;
+  }
+
+  get pluginConfig(): ExtismPluginOptions {
+    return {
+      useWasi: true,
+      runInWorker: true,
+      allowedHosts: this.ALLOWED_HOSTS,
+      memory: { maxHttpResponseBytes: this.MAX_HTTP_RESPONSE_BYTES },
+    };
+  }
+
+  async runWikiSummary(query: string) {
+    const plugin = await this.getPlugin();
+    try {
+      const params = { query };
+      const data = await plugin.call("summary", JSON.stringify(params));
+      if (!data) return "No search results.";
+      return data.json().extract;
+    } catch (e) {
+      console.error(e);
+    } finally {
+      await plugin.close();
+    }
+  }
+
+  async call({ query }: WikipediaParameter): Promise<string> {
+    return await this.runWikiSummary(query);
+  }
+}