この記事は以下の英語記事を元にしています。
なお、元の記事は少し情報が古いため、最新の情報に合わせて書き換えています。

1年ほど前にExtismについての記事を書きましたが、開発が活発でどんどんアップデートされています。

ChatGPTを呼び出すプラグインをRustで作りながら、改めてExtismに入門する記事にしました。

準備

必須ではないですが、extsim CLIという便利なツールがあるのでインストールします。(この記事では使用します。)
https://extism.org/docs/install/
これはライブラリのインストールや、プラグインの呼び出しをコマンドラインからすることができます。

プロジェクトの作成

mkdir chatgpt-plugin
cd chatgpt-plugin
cargo init --lib
cargo add serde serde_json extism-pdk@1.0.0-rc1

作成されたCargo.tomlを以下のように書き加えます。

[lib]
crate_type = ["cdylib"]

プラグインの作成

src/lib.rsを以下のように書き換えていきます。
まずは、必要なライブラリをインポートします。

use extism_pdk::*;
use serde::Deserialize;
use serde_json::json;
use std::str::from_utf8;

次に、レスポンスの型を定義します。

#[derive(Deserialize)]
struct ChatMessage {
    content: String,
}

#[derive(Deserialize)]
struct ChatChoice {
    message: ChatMessage,
}

#[derive(Deserialize)]
struct ChatResult {
    choices: Vec<ChatChoice>,
}

#[plugin_fn]
pub fn call_chat_gpt(input: String) -> FnResult<String> {
    let api_key = config::get("open_ai_key")?
        .unwrap_or("Could not find config key 'open_ai_key'".to_string());

    let req = HttpRequest::new("https://api.openai.com/v1/chat/completions")
        .with_header("Authorization", format!("Bearer {}", api_key))
        .with_header("Content-Type", "application/json")
        .with_method("POST");

    let req_body = json!({
      "model": "gpt-3.5-turbo",
      "temperature": 0.7,
      "messages": [
        {
          "role": "user",
          "content": input,
        }
      ],
    });

    let res = http::request::<String>(&req, Some(req_body.to_string()))?;
    let body = res.body();
    let body = from_utf8(&body)?;
    let body: ChatResult = serde_json::from_str(body)?;

    Ok(body.choices[0].message.content.clone())
}

プラグインのビルド

ビルドすると、target/wasm32-unknown-unknown/release/chatgpt_plugin.wasmが生成されます。

cargo build --release --target wasm32-unknown-unknown

プラグインの実行

extism cliを使って、プラグインを実行します。
OpenAIのAPIキーを使用するので、環境変数に設定しておくか、直接入力して下さい。

extism call \
  target/wasm32-unknown-unknown/release/chatgpt_plugin.wasm \
  call_chat_gpt \
  --allow-host "*" \
  --set-config="{\"open_ai_key\": \"$OPENAI_API_KEY\"}" \
  --input="Please write me a haiku about Wasm"

output:

WebAssembly code,
Runs in all major browsers,
Fast and efficient.%

ホストアプリケーションからの呼び出し

せっかくなので、Pythonでホストアプリケーションを作って、プラグインを呼び出してみます。
Pythonには、OpenAIのライブラリがあるので以下のコードの使い道は全くないですが...

mkdir host
cd host
python3 -m venv venv
source venv/bin/activate
pip3 install extism==1.0.0rc0 --pre

コード

import json
import os
import extism

filename = "../target/wasm32-unknown-unknown/release/chatgpt_plugin.wasm"
input_text = "Pleease introduce yourself."

api_key = os.environ.get("OPENAI_API_KEY", "")


manifest = {
    "wasm": [{"path": filename}],
    "allowed_hosts": [
        "api.openai.com",
    ],
    "config": {
        "open_ai_key": api_key,
    },
}

with extism.Plugin(manifest, wasi=True) as plugin:
    call_chat_gpt = plugin.call(
        "call_chat_gpt",
        input_text.encode("utf-8"),
        parse=lambda output: bytes(output).decode("utf-8"),
    )

print(call_chat_gpt)

実行

python3 host.py

output:

Hello! I am an AI language model developed by OpenAI. I have been trained on a wide range of topics and can assist you with various tasks, answer questions, engage in conversations, and provide information on different subjects. How may I assist you today?

付録(元記事からの変更点)

config::getの返り値の変更

extism-pdk v0.3.4 以前、config::getの返り値はOption<String>だったが、Result<Option<String>, Error>になった。
Original

let api_key = config::get("open_ai_key").expect("Could not find config key 'open_ai_key'");

Change

let api_key = config::get("open_ai_key")?
    .unwrap_or("Could not find config key 'open_ai_key'".to_string());

--allow-hostの追加

extism runtime v0.4.0から、HTTPコールはデフォルトで許可されなくなった。
Original

extism call \
  target/wasm32-unknown-unknown/release/chatgpt_plugin.wasm \
  call_chat_gpt \
  --set-config="{\"open_ai_key\": \"$OPEN_AI_KEY\"}" \
  --input="Please write me a haiku about Wasm"

Change

extism call \
  target/wasm32-unknown-unknown/release/chatgpt_plugin.wasm \
  call_chat_gpt \
  --allow-host "*" \
  --set-config="{\"open_ai_key\": \"$OPENAI_API_KEY\"}" \
  --input="Please write me a haiku about Wasm"

Discussion