diff --git a/.changeset/shy-bottles-shop.md b/.changeset/shy-bottles-shop.md new file mode 100644 index 0000000000..3f5822f44b --- /dev/null +++ b/.changeset/shy-bottles-shop.md @@ -0,0 +1,5 @@ +--- +"@llamaindex/extism-tools": patch +--- + +Build WASM tools with extism diff --git a/.github/workflows/lint_on_push_or_pull.yml b/.github/workflows/lint_on_push_or_pull.yml index a6ff03c13f..0ea7931f44 100644 --- a/.github/workflows/lint_on_push_or_pull.yml +++ b/.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 diff --git a/.github/workflows/preview.yml b/.github/workflows/preview.yml index c25886a413..021ae659ab 100644 --- a/.github/workflows/preview.yml +++ b/.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 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 0cd5f69b24..131027a925 100644 --- a/.github/workflows/release.yml +++ b/.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 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2e48f5e4b8..f3a11aa183 100644 --- a/.github/workflows/test.yml +++ b/.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 diff --git a/packages/extism-tools/README.md b/packages/extism-tools/README.md new file mode 100644 index 0000000000..eb39ad5794 --- /dev/null +++ b/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 +``` diff --git a/packages/extism-tools/bin/compile.js b/packages/extism-tools/bin/compile.js new file mode 100644 index 0000000000..82a6051576 --- /dev/null +++ b/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); + } +}); diff --git a/packages/extism-tools/examples/package.json b/packages/extism-tools/examples/package.json new file mode 100644 index 0000000000..2159d62e96 --- /dev/null +++ b/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" + } +} diff --git a/packages/extism-tools/examples/src/todo.ts b/packages/extism-tools/examples/src/todo.ts new file mode 100644 index 0000000000..50484aa958 --- /dev/null +++ b/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(); diff --git a/packages/extism-tools/examples/src/wiki.py b/packages/extism-tools/examples/src/wiki.py new file mode 100644 index 0000000000..fa8c7f6997 --- /dev/null +++ b/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) diff --git a/packages/extism-tools/examples/src/wiki.ts b/packages/extism-tools/examples/src/wiki.ts new file mode 100644 index 0000000000..b9673554c0 --- /dev/null +++ b/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(); diff --git a/packages/extism-tools/examples/tsconfig.json b/packages/extism-tools/examples/tsconfig.json new file mode 100644 index 0000000000..d474970090 --- /dev/null +++ b/packages/extism-tools/examples/tsconfig.json @@ -0,0 +1,9 @@ +{ + "extends": "../../../tsconfig.json", + "compilerOptions": { + "outDir": "./lib", + "module": "node16", + "moduleResolution": "node16" + }, + "include": ["./src"] +} diff --git a/packages/extism-tools/package.json b/packages/extism-tools/package.json new file mode 100644 index 0000000000..766d1306c7 --- /dev/null +++ b/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" + } +} diff --git a/packages/extism-tools/src/ExtismToolFactory.ts b/packages/extism-tools/src/ExtismToolFactory.ts new file mode 100644 index 0000000000..9af376ef0d --- /dev/null +++ b/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; + +export type ToolClassParams = { + metadata?: ToolMetadata>; +}; + +export type CreateToolClassParams = { + wasmFilename: string; + allowedHosts: string[]; + maxHttpResponseBytes: number; + transformResponse: (response: any) => any; +}; + +export const createPluginInstance = async ( + params: Omit, +): Promise => { + 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 = + { + allowedHosts: ["*"], + maxHttpResponseBytes: DEFAULT_MAX_HTTP_RESPONSE_BYTES, + transformResponse: (response: any) => response, + }; + +export class ExtismToolFactory { + static async createToolClass( + toolName: string, + params: Omit = DEFAULT_TOOL_PARAMS, + ): Promise BaseToolWithCall> { + 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 { + metadata: ToolMetadata>; + + constructor(params?: ToolClassParams) { + this.metadata = params?.metadata || defaultMetadata; + } + + async call(input: ToolParams): Promise { + 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(); + } + } +} diff --git a/packages/extism-tools/src/index.ts b/packages/extism-tools/src/index.ts new file mode 100644 index 0000000000..92155018f3 --- /dev/null +++ b/packages/extism-tools/src/index.ts @@ -0,0 +1 @@ +export { ExtismToolFactory } from "./ExtismToolFactory.js"; diff --git a/packages/extism-tools/tsconfig.json b/packages/extism-tools/tsconfig.json new file mode 100644 index 0000000000..dc0512c073 --- /dev/null +++ b/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" + } + ] +} diff --git a/packages/extism-tools/wasm/index.d.ts b/packages/extism-tools/wasm/index.d.ts new file mode 100644 index 0000000000..390d5752e4 --- /dev/null +++ b/packages/extism-tools/wasm/index.d.ts @@ -0,0 +1,4 @@ +declare module "main" { + export function getMetadata(): I32; + export function call(): I32; +} diff --git a/packages/extism-tools/wasm/todo.js b/packages/extism-tools/wasm/todo.js new file mode 100644 index 0000000000..6946b556e8 --- /dev/null +++ b/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 }; diff --git a/packages/extism-tools/wasm/wiki.js b/packages/extism-tools/wasm/wiki.js new file mode 100644 index 0000000000..93767e708c --- /dev/null +++ b/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 }; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index cbceaeb403..c5c0e196ef 100644 --- a/pnpm-lock.yaml +++ b/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': {} diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 78e2bf3501..04b5e37186 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -9,3 +9,4 @@ packages: - "packages/autotool/examples/*" - "examples/" - "examples/*" + - "packages/extism-tools/examples" diff --git a/tsconfig.json b/tsconfig.json index 11288b4933..39d7517f58 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -82,6 +82,12 @@ }, { "path": "./packages/experimental/tsconfig.json" + }, + { + "path": "./packages/extism-tools/tsconfig.json" + }, + { + "path": "./packages/extism-tools/examples/tsconfig.json" } ] }