Call WASM functions with Parakeet

Call WASM functions with Parakeet

Use Mistral Function Calling support and Ollama to run WASM Rust functions with Parakeet and Extism

Since version 0.0.6, Parakeet supports "Function Calling" and allows you to execute WebAssembly functions with the help of the Extism framework.

Function calling?

First, it's not a feature where a LLM can call and execute a function. "Function Calling" is the ability for certain LLMs to provide a specific output with the same format (we could say: "a predictable output format"). To make it short, given a list of tools, an LLM can guess, from a prompt, which tool to use and with which arguments:

  • User: Say "hello" to Bob

  • System: {"name":"hello","arguments":{"name":"Bob"}}

If you want more details about how to use the "Function Calling" of LLMs, you can read this blog post: https://k33g.hashnode.dev/function-calling-with-ollama-mistral-7b-bash-and-jq

WebAssembly?

It's a long story. WASM is a binary format executed thanks to a WASM runtime. WASM is polyglot: you can build WASM modules with several languages like Rust, Go, C/C++, AssemblyScript, and other languages. The Extism project allows the execution of WASM programs in various languages, such as Go. Then, it was easy to add its support to the Parakeet library.

Function Calling and Parakeet

Using the "Function Calling" feature with Parakeet is straightforward:

  1. We use mistral:7b because it supports natively "Function Calling."

  2. We define a list of tools ([]llm.Tool{}). You must be specific in the description of the function and the arguments; this will help the LLM choose the tool.

  3. We generate the content of the prompt from the tools list and the instructions of the users with tools.GenerateContent(toolsList) and tools.GenerateInstructions(`say "hello" to Bob`)

  4. Then, we can query the LLM. ✋ Important: set the Temperature to 0.0, the Format to "json" and Raw to true. The "Raw mode" allows the overriding of the default template of the LLM.

package main

import (
    "github.com/parakeet-nest/parakeet/completion"
    "github.com/parakeet-nest/parakeet/llm"
    "github.com/parakeet-nest/parakeet/tools"
    "github.com/parakeet-nest/parakeet/gear"
    "fmt"
    "log"
)

func main() {
    ollamaUrl := "http://localhost:11434"
    model := "mistral:7b"

    toolsList := []llm.Tool{
        {
            Type: "function",
            Function: llm.Function{
                Name:        "hello",
                Description: "Say hello to a given person with his name",
                Parameters: llm.Parameters{
                    Type: "object",
                    Properties: map[string]llm.Property{
                        "name": {
                            Type:        "string",
                            Description: "The name of the person",
                        },
                    },
                    Required: []string{"name"},
                },
            },
        },
        {
            Type: "function",
            Function: llm.Function{
                Name:        "hey",
                Description: "Say hey to a given person with his name",
                Parameters: llm.Parameters{
                    Type: "object",
                    Properties: map[string]llm.Property{
                        "name": {
                            Type:        "string",
                            Description: "The name of the person",
                        },
                    },
                    Required: []string{"name"},
                },
            },
        },
    }

    toolsContent, err := tools.GenerateContent(toolsList)
    if err != nil {
        log.Fatal("😡:", err)
    }

    userContent := tools.GenerateInstructions(`say "hello" to Bob`)

    messages := []llm.Message{
        {Role: "system", Content: toolsContent},
        {Role: "user", Content: userContent},
    }

    options := llm.Options{
        Temperature:   0.0, // no fantasy
        RepeatLastN:   2,
        RepeatPenalty: 2.0,
    }

    query := llm.Query{
        Model: model,
        Messages: messages,
        Options: options,
        Format:  "json",
        Raw:     true,
    }

    answer, err := completion.Chat(ollamaUrl, query)
    if err != nil {
        log.Fatal("😡:", err)
    }
    result, err := gear.PrettyString(answer.Message.Content)
    if err != nil {
        log.Fatal("😡:", err)
    }
    fmt.Println(result)
}

When you run the program, you should get:

{
  "name": "hello",
  "arguments": {
    "name": "Bob"
  }
}

If you try with: userContent := tools.GenerateInstructions(say "hey" to Sam) you should get:

{
  "name": "hey",
  "arguments": {
    "name": "Sam"
  }
}

Wouldn't it be nice to be then able to run functions for real?

Using WebAssembly to call the functions.

With Parakeet, it's simple. But first, let's create a Rust WASM plugin with two functions:

use extism_pdk::*;

#[plugin_fn]
pub fn hello(name: String) -> FnResult<String> {
    Ok(format!("👋 Hello, {}!", name))
}

#[plugin_fn]
pub fn hey(name: String) -> FnResult<String> {
    Ok(format!("🙂 Hey, {}! What's up?", name))
}

You can find the source of the plugin here: https://github.com/parakeet-nest/blog-post-samples/tree/main/2024-06-02/01-mistral-call-wasm/plugin

Before building the wasm plugin, don't forget to add the wasm target: rustup target add wasm32-unknown-unknown, then you can run the cargo build command.

Then, we need to create an instance of the WASM plugin in our Go code:

wasmPlugin, _ := wasm.NewPlugin(
  "./plugin/target/wasm32-unknown-unknown/debug/greetings.wasm", nil)

And use it to run the appropriate function from the output of the LLM; thanks to the Call method of the plugin, wasmPlugin.Call(functionName, []byte(name)) :

answer, err := completion.Chat(ollamaUrl, query)
if err != nil {
    log.Fatal("😡:", err)
}

jsonRes, err := gear.JSONParse(answer.Message.Content)

if err != nil {
    log.Fatal("😡:", err)
}

functionName := jsonRes["name"].(string)
name := jsonRes["arguments"].(map[string]interface{})["name"].(string)

fmt.Println("Calling", functionName, "with", name)

// call the function of the wasm plugin
res, _ := wasmPlugin.Call(functionName, []byte(name))

// display the result
fmt.Println(string(res))

And the result of the Rust WASM function will be:

Calling hey with Sam
🙂 Hey, Sam! What's up?

Now, you can add superpowers to your LLMs 🦸‍♂️.

You will find the source code of the examples here: https://github.com/parakeet-nest/blog-post-samples/tree/main/2024-06-02.

On the Parakeet repository, there is an example with a plugin developed with Go and built with Tinygo: https://github.com/parakeet-nest/parakeet/tree/main/examples/18-call-functions-for-real.

Happy experiments! Stay tuned for the next episode.