Skip to main content

One post tagged with "rust"

View All Tags

· 3 min read
Benjamin Eckel

We’ve seen a few people try to create Extism plug-ins that make use of ChatGPT. The simplest way you can accomplish this is to use Extism’s built-in HTTP host function. A more advanced option would be to create your own host function.

Let’s go the first route and create the simplest possible plug-in that utilizes ChatGPT. We’ll do this in Rust, but the same is possible from our other plug-in languages.

Project Setup​

First let’s initialize our rust project:

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

Now tell cargo we wish to build a shared library by putting this in Cargo.toml

[lib]
crate_type = ["cdylib"]

Writing lib.rs​

First we import the code we need into lib.rs:

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

Then let’s create some structs for the response body. We’re just putting what we need for this example. See the completion docs for the whole response.

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

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

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

Now let’s write our code. We’re going to skip error handling and edge cases as an exercise for the user:

#[plugin_fn]
pub fn call_chat_gpt(input: String) -> FnResult<String> {
  // get API key from the plug-in config key open_ai_key
  let api_key = config::get("open_ai_key").expect("Could not find config key 'open_ai_key'");
  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");

  // We could make our own structs for the body
  // this is a quick way to make some unstructured JSON
  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())
}
caution

Think carefully about how you handle secrets with plug-ins. You should always follow best practices like keeping secrets encrypted at rest and with plug-ins, you must take extra care not to leak secrets across users. Currently there are two ways you can handle secrets:

  1. You can use plug-in configs and the runtime function to update them. If you do this you’ll probably want to assume all configs are secrets and take extra care not to leak configs across customers.
  2. You provide some kind of host function for the plug-in programmer to fetch and decrypt their secrets from a secret provider / manager.

In this example we are going to pass the key from the host via the config object and leaving safely managing it as an exercise to the reader.

Compiling and Running​

We can compile and run this plug-in like we always do:

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

Note we do not need a wasi target since we are using Extism’s HTTP host function.

Now we can use the extism CLI to test.

# copy/paste Open AI secret key in environment variable
read -s -p "Enter OpenAI API Key:" OPEN_AI_KEY

# optionally just use export:
# export OPEN_AI_KEY=<past-key-here>

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"

Wasm code compiled,
Runs in browser, fast and light,
Web apps flourish bright.