2 posts tagged with "ChatGPT"

In our last post we showed how to call ChatGPT from inside an Extism plug-in. Although it’s a useful capability for plug-in developers, there isn’t anything specially suited for Wasm here. So we’ve been thinking about what role Wasm might be able to play in this LLM enabled future. In this post we will explore one idea around code generation.

Code Generation​

One key touted feature of LLMs has been their ability to generate code. Many programmers have interacted with this through tools like Copilot or by just asking ChatGPT programming questions. Programmers generally copy-paste these code snippets or review it in some way before they integrate it into their apps. Directly running code generated by an LLM has some of the same risks that Extism deals with. The code should be considered untrusted. Not only could it be wrong, your use case might expose you to malicious actors and exploits through prompt injection.

Unix Utility Maker​

Suppose you want to create a bot that generates unix text processing utilities. You ask ChatGPT through the API to generate some bash, then you pipe some standard input to that bash. The interface might look like this:

util-bot "Make a tool that counts the number of vowels in the input text" \ 
   > count_vowels.sh
echo "Hello World" | sh count_vowels.sh

# => 3

The bash it generates could do anything on your computer. So this feels like a scary idea. Here the worst case scenario is likely that your script is wrong and destroys some state on your machine, but you could see if you had a use case that allowed untrusted input into this code generator, someone could easily craft some malicious actions.

Prompt injection is an ongoing research topic and some believe it will be impossible to prevent. This is perhaps where Wasm can provide some relief. Wasm can create a sandbox for the generated code and you can give it only the capabilities it needs to get the job done.

Utility Maker in Extism​

Let’s build this using Extism. We’ll use LangChain to generate the code based on a description, and we’ll use the Extism Python SDK to orchestrate the Wasm module for executing the code. We are going to ask it to generate JavaScript code and run it in a sandbox we create with the JavaScript PDK.

Let’s start with our plug-in. We’ll create a simple JavaScript module that has 2 functions: one to store the code for a function in a plug-in variable and one to read it and invoke it.

// sandbox.js

function register_function() {
  let code = Host.inputString();
  Var.set("code", code);
}

function invoke() {
  let input = Host.inputString();
  let code = Var.getString("code");
  let func = eval(code);
  Host.outputString(func(input).toString());
}

module.exports = {register_function, invoke}

Let’s start the host side now. First we need some imports:

# host.py

import re
import sys
import json
import pathlib
from extism import Context
from langchain.chat_models import ChatOpenAI
from langchain.schema import (
    AIMessage,
    HumanMessage,
    SystemMessage
)

In order to generate runnable code, we need to coerce ChatGPT into generating a CJS module with a single function as an export. As an example, here is the code we expect from the prompt:

Write a function that takes a string of comma separated integers and
returns the sum as an integer

This prompt should generate the following code:

module.exports = function sumCommaSeparatedIntegers(str) {
  return str.split(',').reduce((acc, curr) => acc + parseInt(curr), 0);
}

This allows us to easily eval it into a function and call it in the plug-in code:

let func = eval(code)
let result = func(inputStr)

This process of coercion is colloquially called prompt engineering. LangChain gives us some tools to do this. Let’s write the function that generates the code from a description:

MARKDOWN_RE = re.compile("\`\`\`javascript([^\`]*)\`\`\`")
chat = ChatOpenAI(temperature=0)

def generate_code(description):
    messages = [
        SystemMessage(content="Your goal is to write a Javascript function."),
        SystemMessage(content="You must export the function as the default export, using the CommonJS module syntax"),
        SystemMessage(content="You must not include any comments, explanations, or markdown. The response should be JavaScript only."),
        HumanMessage(content=description),
    ]

    response = chat(messages)
    code = response.content
    # sometimes the LLM wraps the code in markdown, this removes it
    # if it does
    m = MARKDOWN_RE.match(code)
    if m and m.group(1):
        code = m.group(1)
    return code

Here we use a series of SystemMessages to pre-prompt the model with its goals. We put the user’s description at the end, in the HumanMessage.

Now we’ll write a function to execute this code in the Wasm module.

def execute_code(code, input):
    wasm_file_path = pathlib.Path(__file__).parent / "sandbox.wasm"
    config = {"wasm": [{"path": str(wasm_file_path)}], "memory": {"max": 5}}

    with Context() as context:
        plugin = context.plugin(config, wasi=True)
        plugin.call("register_function", code)
        return plugin.call("invoke", input)

And now a main to orchestrate it:

if __name__ == "__main__":
    code = generate_code(sys.argv[1])
    print(execute_code(code, sys.stdin.read()))

Let’s try it out by asking it to generate the canonical Extism count-vowels example

$ pip3 install extism langchain
$ export OPENAI_API_KEY=sk-mykeyhere # needed to query openai
$ echo "hello world" | python3 host.py "write a function that counts the number of vowels in a string"
b'3'

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())
}

Compiling and Running​

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

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

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.