Merge branch 'main' into feat/build-wasm-with-extism

.changeset/shy-bottles-shop.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/extism-tools": patch
+---
+
+Build WASM tools 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/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

@@ -42,6 +42,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 E2E Tests
@@ -62,6 +66,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
@@ -77,6 +85,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
@@ -114,6 +126,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
@@ -133,6 +149,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,24 @@
+## Extism Tools
+
+### Prerequisites for Development
+
+- [Extism PDK](https://github.com/extism/js-pdk?tab=readme-ov-file#linux-macos)
+
+### Build WASM files
+
+```bash
+pnpm run build
+```
+
+### Run WASM files in Node.js using Extism SDK (https://github.com/extism/js-sdk)
+
+```bash
+cd examples
+pnpm run test:wiki
+```
+
+### 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/package.json

@@ -0,0 +1,17 @@
+{
+  "name": "@llamaindex/extism-tools-examples",
+  "private": true,
+  "version": "1.0.0",
+  "type": "module",
+  "scripts": {
+    "test:wiki": "tsx ./src/wiki.ts",
+    "test:todo": "tsx ./src/todo.ts"
+  },
+  "dependencies": {
+    "@llamaindex/extism-tools": "workspace:*",
+    "llamaindex": "workspace:*"
+  },
+  "devDependencies": {
+    "tsx": "^4.19.0"
+  }
+}

packages/extism-tools/examples/src/todo.ts

@@ -0,0 +1,13 @@
+import { ExtismToolFactory } from "@llamaindex/extism-tools";
+import { OpenAI, OpenAIAgent, Settings } from "llamaindex";
+
+async function main() {
+  const TodoTool = await ExtismToolFactory.createToolClass("todo");
+  const todoTool = new TodoTool();
+  Settings.llm = new OpenAI();
+  const agent = new OpenAIAgent({ tools: [todoTool] });
+  const result = await agent.chat({ message: "Get first todo" });
+  console.log(result.message);
+}
+
+void main();

packages/extism-tools/examples/src/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/examples/src/wiki.ts

@@ -0,0 +1,13 @@
+import { ExtismToolFactory } from "@llamaindex/extism-tools";
+import { OpenAI, OpenAIAgent, Settings } from "llamaindex";
+
+async function main() {
+  const WikiTool = await ExtismToolFactory.createToolClass("wiki");
+  const wikiTool = new WikiTool();
+  Settings.llm = new OpenAI();
+  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/tsconfig.json

@@ -0,0 +1,9 @@
+{
+  "extends": "../../../tsconfig.json",
+  "compilerOptions": {
+    "outDir": "./lib",
+    "module": "node16",
+    "moduleResolution": "node16"
+  },
+  "include": ["./src"]
+}

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-rc8",
+    "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,82 @@
+import createPlugin, { type Plugin } from "@extism/extism";
+import type { JSONSchemaType } from "ajv";
+import type { BaseToolWithCall, ToolMetadata } from "llamaindex";
+import path from "path";
+import { fileURLToPath } from "url";
+
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = path.dirname(__filename);
+export const WASM_DIRECTORY = path.join(__dirname, "..", "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 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 const DEFAULT_TOOL_PARAMS: Omit<CreateToolClassParams, "wasmFilename"> =
+  {
+    allowedHosts: ["*"],
+    maxHttpResponseBytes: DEFAULT_MAX_HTTP_RESPONSE_BYTES,
+    transformResponse: (response: any) => response,
+  };
+
+export class ExtismToolFactory {
+  static async createToolClass(
+    toolName: string,
+    params: Omit<CreateToolClassParams, "wasmFilename"> = DEFAULT_TOOL_PARAMS,
+  ): Promise<new (params?: ToolClassParams) => BaseToolWithCall<ToolParams>> {
+    const config = { ...params, wasmFilename: `${toolName}.wasm` };
+    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 BaseToolWithCall<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/src/index.ts

@@ -0,0 +1 @@
+export { ExtismToolFactory } from "./ExtismToolFactory.js";

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

packages/extism-tools/wasm/todo.js

@@ -0,0 +1,34 @@
+function getMetadata() {
+  const metadata = {
+    name: "todo_tool",
+    description: "A tool helps search todo.",
+    parameters: {
+      type: "object",
+      properties: {
+        index: {
+          type: "number",
+          description: "The index of the todo to search for.",
+        },
+      },
+      required: ["index"],
+    },
+  };
+  Host.outputString(JSON.stringify(metadata));
+}
+
+function call() {
+  const params = JSON.parse(Host.inputString());
+  const index = params?.index;
+  if (!index) throw new Error("No index provided");
+  const request = {
+    method: "GET",
+    url: `https://jsonplaceholder.typicode.com/todos/${encodeURIComponent(index)}`,
+  };
+  const response = Http.request(request);
+  if (response.status != 200)
+    throw new Error(`Got non 200 response ${response.status}`);
+
+  Host.outputString(response.body);
+}
+
+module.exports = { getMetadata, call };

packages/extism-tools/wasm/wiki.js

@@ -0,0 +1,33 @@
+function getMetadata() {
+  const metadata = {
+    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"],
+    },
+  };
+  Host.outputString(JSON.stringify(metadata));
+}
+
+function call() {
+  const params = JSON.parse(Host.inputString());
+  const query = params?.query;
+  if (!query) throw new Error("No query provided");
+  const request = {
+    method: "GET",
+    url: `https://en.wikipedia.org/api/rest_v1/page/summary/${encodeURIComponent(query)}`,
+  };
+  const response = Http.request(request);
+  if (response.status != 200)
+    throw new Error(`Got non 200 response ${response.status}`);
+  Host.outputString(response.body);
+}
+
+module.exports = { getMetadata, call };

pnpm-lock.yaml

@@ -481,6 +481,50 @@ importers:
         specifier: ^1.1.2
         version: 1.1.2
 
+  packages/extism-tools:
+    dependencies:
+      '@extism/extism':
+        specifier: ^2.0.0-rc8
+        version: 2.0.0-rc8
+      '@llamaindex/core':
+        specifier: workspace:*
+        version: link:../core
+      '@llamaindex/env':
+        specifier: workspace:*
+        version: link:../env
+      ajv:
+        specifier: ^8.17.1
+        version: 8.17.1
+      llamaindex:
+        specifier: workspace:*
+        version: link:../llamaindex
+    devDependencies:
+      '@swc/cli':
+        specifier: ^0.4.0
+        version: 0.4.0(@swc/core@1.7.22(@swc/helpers@0.5.13))(chokidar@3.6.0)
+      '@swc/core':
+        specifier: ^1.7.22
+        version: 1.7.22(@swc/helpers@0.5.13)
+      '@types/node':
+        specifier: ^22.5.1
+        version: 22.5.4
+      typescript:
+        specifier: ^5.5.4
+        version: 5.6.2
+
+  packages/extism-tools/examples:
+    dependencies:
+      '@llamaindex/extism-tools':
+        specifier: workspace:*
+        version: link:..
+      llamaindex:
+        specifier: workspace:*
+        version: link:../../llamaindex
+    devDependencies:
+      tsx:
+        specifier: ^4.19.0
+        version: 4.19.0
+
   packages/llamaindex:
     dependencies:
       '@anthropic-ai/sdk':
@@ -2819,6 +2863,9 @@ packages:
     resolution: {integrity: sha512-autAXT203ixhqei9xt+qkYOvY8l6LAFIdT2UXc/RPNeUVfqRF1BV94GTJyVPFKT8nFM6MyVJhjLj9E8JWvf5zQ==}
     engines: {node: ^18.18.0 || ^20.9.0 || >=21.1.0}
 
+  '@extism/extism@2.0.0-rc8':
+    resolution: {integrity: sha512-a90ti82PDEiuANoDlFVBMjlW/sLnJl9zGSzvYY3eydX2cQyFA8AsdmC4JUWLK+NEP9eeta83+6N5B8UwYG33bw==}
+
   '@faker-js/faker@8.4.1':
     resolution: {integrity: sha512-XQ3cU+Q8Uqmrbf2e0cIC/QN43sTBSC8KF12u29Mb47tWrt2hAgBXSgpZMj4Ao8Uk0iJcU99QsOCaIL8934obCg==}
     engines: {node: ^14.17.0 || ^16.13.0 || >=18.0.0, npm: '>=6.14.13'}
@@ -14449,6 +14496,8 @@ snapshots:
       levn: 0.4.1
     optional: true
 
+  '@extism/extism@2.0.0-rc8': {}
+
   '@faker-js/faker@8.4.1': {}
 
   '@fastify/busboy@2.1.1': {}

pnpm-workspace.yaml

@@ -9,3 +9,4 @@ packages:
   - "packages/autotool/examples/*"
   - "examples/"
   - "examples/*"
+  - "packages/extism-tools/examples"

tsconfig.json

@@ -82,6 +82,12 @@
     },
     {
       "path": "./packages/experimental/tsconfig.json"
+    },
+    {
+      "path": "./packages/extism-tools/tsconfig.json"
+    },
+    {
+      "path": "./packages/extism-tools/examples/tsconfig.json"
     }
   ]
 }