Forem: Josh Mo

Rewriting Claude Code in Rust

Josh Mo — Mon, 02 Mar 2026 16:31:02 +0000

It's late. You're getting the builder itch - the kind of itch that never truly leaves you until you've actually built something. Whether you're actually going to use it in a week or so is another matter entirely, but that's besides the point. Your computer is turned on, you have Claude Code open in the terminal tab for your editor. You've been using it for a little while now.

And this gets you thinking. Surely it can't be that hard to build a coding agent, right? It's basically just an LLM with some tools and access to your filesystem?

In this post, we'll build a small demo that will be able to act as a minimal demo version of Claude Code written in Rust using the Rig AI framework. It will be able to use bash commands, read and write files, as well as stream output back to you.

Interested in just trying out the demo for yourself? Here it is.

Getting started

Before we get started, you'll need the Rust programming language installed. You will also need an API key for the model provider of your choice, although for the purposes of simplicity we'll be using an Anthropic API key.

To get started, we'll create our project and call it rig-cc:

cargo init rig-cc
cd rig-cc

Next, we'll add our dependencies:

cargo add rig-core@0.30 tokio@1 anyhow \
serde serde_json thiserror futures -F tokio/full,serde/derive

What did we add?

rig-core: The core Rig library.
serde/serde_json: A library for converting to and from Rust structs and various data formats, with serde_json for JSON interop.
tokio: the most popular async Rust runtime.
thiserror: Derive error handling
futures: A crate for dealing with streams and other async-related functionality. We need to use the Stream/StreamExt traits to be able to utilise the item stream returned from the agent.

Building our coding agent

Before we start, let's break down the components of a coding agent. It should be able to:

Hold a conversation and use context from previous messages
Read from and write to files
Use Bash commands, and additionally allow the user to manually interrupt the bash command if it takes too long
Allow users to exit the conversation

First, we'll build the ReadFile tool, for reading files. Tools in Rig are simply defined as types that implement the rig::tool::Tool type.

#[derive(Deserialize, Serialize)]
pub struct ReadFile;

#[derive(serde::Deserialize)]
pub struct ReadFileArgs {
    path: String,
}

You can see below that we get the following:

The tool definition (through fn definition())
The relevant types (arguments, outputs, the tool name and error type)
What should actually happen when the tool is called (through fn call())

impl rig::tool::Tool for ReadFile {
    const NAME: &'static str = "read_file";
    type Error = rig::tool::ToolError;
    type Args = ReadFileArgs;
    type Output = String;

    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: "read_file".to_string(),
            description: "Read the contents of a file at the specified path. Returns the file contents as a string.".to_string(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "The path to the file to read"
                    }
                },
                "required": ["path"]
            }),
        }
    }

    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        std::fs::read_to_string(&args.path).map_err(|e| {
            ToolError(format!(
                "Failed to read file '{path}': {e}",
                path = args.path
            ))
        })
    }
}

The WriteFile tool is largely the same, as it's primarily just writing to a given file (and a check to see if the directory exists):

#[derive(Deserialize)]
pub struct WriteFileArgs {
    path: String,
    content: String,
}

#[derive(Deserialize, Serialize)]
pub struct WriteFile;

impl Tool for WriteFile {
    const NAME: &'static str = "write_file";
    type Error = ToolError;
    type Args = WriteFileArgs;
    type Output = String;

    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: "write_file".to_string(),
            description: "Write content to a file at the specified path. Creates parent directories if they don't exist. Overwrites the file if it already exists.".to_string(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "The path to the file to write"
                    },
                    "content": {
                        "type": "string",
                        "description": "The content to write to the file"
                    }
                },
                "required": ["path", "content"]
            }),
        }
    }

    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        if let Some(parent) = std::path::Path::new(&args.path).parent()
            && !parent.as_os_str().is_empty()
        {
            std::fs::create_dir_all(parent)
                .map_err(|e| ToolError(format!("Failed to create directories: {}", e)))?;
        }

        std::fs::write(&args.path, &args.content)
            .map_err(|e| ToolError(format!("Failed to write file '{}': {}", args.path, e)))?;

        Ok(format!(
            "Successfully wrote {} bytes to '{}'",
            args.content.len(),
            args.path
        ))
    }
}

Finally, we're going to write our bash tool. This has some somewhat complicated logic in it, so we'll split this into two parts: the function to call the bash command, and then implementing the tool.

First, we need to call the Bash command using std::process::Command - and pipe both stderr and stdout. This creates a child process that we can programmatically control (and more specifically, waiting for the process to finish so we can collect the results).

use tokio::process::Command;

let mut child = Command::new("bash")
    .arg("-c")
    .arg(&args.command)
    .stdout(std::process::Stdio::piped())
    .stderr(std::process::Stdio::piped())
    .spawn()
    .map_err(|e| ToolError(format!("Failed to spawn command: {}", e)))?;

Next, we need to use the tokio timeout function which will take a std::time::Duration and a future. If the timeout is reached, it will return an error; if not, the result will be given as usual.

/// Timeout. If the execution time for a command exceeds the timeout, users should be able to manually interrupt it.
const WARNING_TIMEOUT_SECS: u64 = 60;

let warning_duration = Duration::from_secs(WARNING_TIMEOUT_SECS);

let status = match tokio::time::timeout(warning_duration, child.wait()).await {
    // If OK, automatically try to get the result
    Ok(result) => result.map_err(|e| ToolError(format!("Command failed: {}", e)))?,
    // An error here represents that the timeout was reached
    Err(_) => {
        eprintln!(
            "\n[Command running for >{WARNING_TIMEOUT_SECS}s. Press Ctrl+C to interrupt]");
        io::stderr().flush().ok();
        child.wait()
             .await
             .map_err(|e| ToolError(format!("Command failed: {}", e)))?
        }
};

After that, we need to take the stdout/stderr pipes and get their entire contents. Because the take() function results in an Option<T>, we'll use the if let-some pattern which allows for simpler matching without needing to deal with a None value (if the value is None we can probably assume there's nothing there anyway).

let stdout = child.stdout.take();
let stderr = child.stderr.take();

let mut stdout_content = String::new();
let mut stderr_content = String::new();

if let Some(mut stdout) = stdout {
    use tokio::io::AsyncReadExt;
    let mut buf = Vec::new();
    stdout.read_to_end(&mut buf).await.ok();
    stdout_content = String::from_utf8_lossy(&buf).to_string();
}

if let Some(mut stderr) = stderr {
    use tokio::io::AsyncReadExt;
    let mut buf = Vec::new();
    stderr.read_to_end(&mut buf).await.ok();
    stderr_content = String::from_utf8_lossy(&buf).to_string();
}

Finally, we'll create our tool output. If the command status is a success, we'll re-assign our contents of stdout as mutable and then push whatever's in stderr to the message. This will let the model know if there are any issues while running, since even on a successful run there may be errors causing issues that may hamper things like performance, latency and other things. If it fails, then we provide the exit code and the stdout/stderr contents following the previous method (appending a newline, then the contents of each).

Regardless of the status of the command, the final text will get truncated according to the maximum output (to avoid flooding the LLM with context and therefore increasing the context window usage by accident) - and then return the output.

/// Max output bytes. Prevents infinite context being stuffed into the response.
const MAX_OUTPUT_BYTES: usize = 50 * 1024;

let mut output = if status.success() {
    let mut out = stdout_content;
    if !stderr_content.is_empty() {
        if !out.is_empty() {
            out.push('\n');
        }
        out.push_str("stderr:\n");
        out.push_str(&stderr_content);
        }
    out
    } else {
        let mut out = format!("Exit code: {}\n", status.code().unwrap_or(-1));
        if !stdout_content.is_empty() {
            out.push_str("stdout:\n");
            out.push_str(&stdout_content);
            out.push('\n');
        }
        if !stderr_content.is_empty() {
            out.push_str("stderr:\n");
            out.push_str(&stderr_content);
        }
        out
};

let total_bytes = output.len();
if total_bytes > MAX_OUTPUT_BYTES {
    output.truncate(MAX_OUTPUT_BYTES);
    while !output.is_char_boundary(output.len()) {
        output.pop();
    }
    output.push_str(&format!(
        "\n... [output truncated, {total_bytes} bytes total]",
            ));
    }

Ok(output)

OK, so now if we set up our Bash tool and arguments structs it should look like this (combine all the previous snippets into a single function like in the next snippet!):

#[derive(Deserialize)]
pub struct BashArgs {
    command: String,
}

#[derive(Deserialize, Serialize)]
pub struct Bash;

impl Bash {
    async fn run(&self, args: BashArgs) -> Result<String, ToolError> {
        use tokio::process::Command;

        let mut child = Command::new("bash")
            .arg("-c")
            .arg(&args.command)
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .spawn()
            .map_err(|e| ToolError(format!("Failed to spawn command: {}", e)))?;

        let warning_duration = Duration::from_secs(WARNING_TIMEOUT_SECS);
        let status = match timeout(warning_duration, child.wait()).await {
            Ok(result) => result.map_err(|e| ToolError(format!("Command failed: {}", e)))?,
            Err(_) => {
                eprintln!(
                    "\n[Command running for >{}s. Press Ctrl+C to interrupt]",
                    WARNING_TIMEOUT_SECS
                );
                io::stderr().flush().ok();
                child
                    .wait()
                    .await
                    .map_err(|e| ToolError(format!("Command failed: {}", e)))?
            }
        };

        let stdout = child.stdout.take();
        let stderr = child.stderr.take();

        let mut stdout_content = String::new();
        let mut stderr_content = String::new();

        if let Some(mut stdout) = stdout {
            use tokio::io::AsyncReadExt;
            let mut buf = Vec::new();
            stdout.read_to_end(&mut buf).await.ok();
            stdout_content = String::from_utf8_lossy(&buf).to_string();
        }

        if let Some(mut stderr) = stderr {
            use tokio::io::AsyncReadExt;
            let mut buf = Vec::new();
            stderr.read_to_end(&mut buf).await.ok();
            stderr_content = String::from_utf8_lossy(&buf).to_string();
        }

        let mut output = if status.success() {
            let mut out = stdout_content;
            if !stderr_content.is_empty() {
                if !out.is_empty() {
                    out.push('\n');
                }
                out.push_str("stderr:\n");
                out.push_str(&stderr_content);
            }
            out
        } else {
            let mut out = format!("Exit code: {}\n", status.code().unwrap_or(-1));
            if !stdout_content.is_empty() {
                out.push_str("stdout:\n");
                out.push_str(&stdout_content);
                out.push('\n');
            }
            if !stderr_content.is_empty() {
                out.push_str("stderr:\n");
                out.push_str(&stderr_content);
            }
            out
        };

        let total_bytes = output.len();
        if total_bytes > MAX_OUTPUT_BYTES {
            output.truncate(MAX_OUTPUT_BYTES);
            while !output.is_char_boundary(output.len()) {
                output.pop();
            }
            output.push_str(&format!(
                "\n... [output truncated, {} bytes total]",
                total_bytes
            ));
        }

        Ok(output)
    }
}

To finish this section off, we'll write the actual rig::tool::Tool implementation. Because we've already written the required function, instead of writing huge amounts of code in the async fn call() section we can write a nice little one liner.

impl Tool for Bash {
    const NAME: &'static str = "bash";
    type Error = ToolError;
    type Args = BashArgs;
    type Output = String;

    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: "bash".to_string(),
            description: "Execute a bash command and return its output. Use this for running shell commands, git operations, running tests, installing packages, etc. The command runs in the current working directory.".to_string(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The bash command to execute"
                    }
                },
                "required": ["command"]
            }),
        }
    }

    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        self.run(args).await
    }
}

Writing the coding agent loop

With all the tools written, now we simply need to write our agent loop. First though we'll need to write our system prompt - for now we'll call our program Rig Code, and tell it what tools it has access to as well as some basic guidelines.

const SYSTEM_PROMPT: &str = r#"You are Rig Code, an interactive AI coding assistant running in the terminal.

You have access to these tools:
- bash: Execute shell commands (runs in current working directory)
- read_file: Read file contents
- write_file: Create or modify files

Guidelines:
- Use bash to explore projects, run tests, git operations, etc.
- Read files before modifying them to understand context
- Be concise and focused on solving the user's problem
- When making changes, explain what you're doing briefly
"#;

Finally, we need to write our main function. The agent part is quite simple - we need to create an Anthropic client, then create an agent using the agent builder in Rig. Then we'll instantiate a message history, as well as an input/output buffer.

#[tokio::main]
async fn main() -> Result<()> {
    let client = Client::from_env();

    let agent = client
        .agent("claude-sonnet-4-5")
        .preamble(SYSTEM_PROMPT)
        .tool(ReadFile)
        .tool(WriteFile)
        .tool(Bash)
        .max_tokens(8192)
        .build();

    println!("Rig Code v0.1.0");
    println!("Type 'exit' or 'quit' to exit.\n");

    let stdin = io::stdin();
    let mut stdout = io::stdout();
    let mut history: Vec<Message> = Vec::new();

    loop {
        // .. loop goes here
    }

    Ok(())
}

Now for the loop! This should print a marker that makes it obvious it's the user's turn to type, then try to read the input whenever the user finishes inputting (ie, they press Return).

Once done, we then check for the following in the trimmed input:

If the user writes "/exit" or "/quit", break the loop and exit
If nothing's been inputted, just continue
If input exists, attempt to send a streaming prompt request to the model provider (we'll deal with this specifically later).

After the stream has finished, we should also show how many tokens were used.

print!("> ");
stdout.flush()?;

let mut input = String::new();

match stdin.read_line(&mut input) {
    Ok(0) => {
        println!("\nGoodbye!");
        break;
    }

    Ok(_) => {
        let input = input.trim();

        if input.eq_ignore_ascii_case("exit") || input.eq_ignore_ascii_case("quit") {
            println!("Goodbye!");
            break;
        }

        if input.is_empty() {
            continue;
        }

        println!();

        let mut stream = agent
            .stream_prompt(input)
            .with_history(history.clone())
            .multi_turn(100)
            .await;

        let mut response_text = String::new();
        let mut input_tokens = 0u64;
        let mut output_tokens = 0u64;

        while let Some(chunk) = stream.next().await {
            // .. handle stream (see below)
        }

        println!("\n");
        println!(
            "[Tokens: {} in / {} out]",
            // .. see below for format_number function
            format_number(input_tokens),
            format_number(output_tokens)
        );
        println!();

        history.push(Message::user(input));

        if !response_text.is_empty() {
            history.push(Message::assistant(response_text));
        }
    }

    Err(e) => {
        eprintln!("Error reading input: {}", e);
    }
}

One small detail in the implementation is that if you're burning through a large number of tokens very quickly, it can be difficult to immediately tell at a glance what the number is. We'll implement a function below that makes it look slightly nicer and easier to look at.

fn format_number(n: u64) -> String {
    let s = n.to_string();
    let mut result = String::new();
    for (i, c) in s.chars().rev().enumerate() {
        if i > 0 && i % 3 == 0 {
            result.push(',');
        }
        result.push(c);
    }
    result.chars().rev().collect()
}

The stream itself isn't difficult to deal with per se as you're just matching on item variants. You can see that we get a number of different variants - tool calls, text, tool results, and the final stream item (ie the aggregated text and token usage).

while let Some(chunk) = stream.next().await {
    match chunk {
        Ok(MultiTurnStreamItem::StreamAssistantItem(
            StreamedAssistantContent::Text(text),
        )) => {
            print!("{}", text.text);
            stdout.flush()?;
            response_text.push_str(&text.text);
        }

        Ok(MultiTurnStreamItem::StreamAssistantItem(
            StreamedAssistantContent::ToolCall { tool_call, .. },
        )) => {
            println!("\n[Calling tool: {}]", tool_call.function.name);
            stdout.flush()?;
        }

        Ok(MultiTurnStreamItem::StreamUserItem(
            rig::streaming::StreamedUserContent::ToolResult { tool_result, .. },
        )) => {
            println!("[Tool result received for: {}]", tool_result.id);
            stdout.flush()?;
        }

        Ok(MultiTurnStreamItem::FinalResponse(final_response)) => {
            let usage = final_response.usage();
            input_tokens = usage.input_tokens;
            output_tokens = usage.output_tokens;
        }

        Err(e) => {
            eprintln!("\nError: {}", e);
            break;
        }

        _ => {}
    }
}

Beyond this tutorial

OK, OK, so maybe we didn't quite build out Claude Code. But that's okay! Now you know what the bare minimum is when it comes to building a coding agent in Rust.

Beyond this though if you want to productionise this, there are some additional features you may wish to consider:

Cost calculations (derived from token usage)
Auto compaction of context
Skills/MCP integrations
Improving the system prompt

If you're interested in learning more about Rig, check us out:

How I ported a Rust framework to TS with WASM

Josh Mo — Wed, 30 Jul 2025 21:06:18 +0000

During one of our internal weekly meetings at work, we decided that we wanted to get Rig to JavaScript so that we can run it in the browser and in Node applications. Of course, there was a potential tonne of code rewriting to do if we were to write it entirely from TypeScript. Therefore, we ended up using WASM to port the inner core to TypeScript and then using manual glue code from TypeScript to be able to finish off the parts that couldn't be done directly through the Rust-WASM bridge.

What needs to get ported?

Here is a short list of things we were able to port from Rust:

All model providers
Functionality for users to create their own vector store and tool integrations from JS/TS
Wrapper types and conversion impls to/from non-WASM Rust types

Thankfully, a large portion of the WASM implementation is mostly just wrappers around the original structs. We just need to move them all over to WASM-friendly bindings by attaching the #[wasm_bindgen] macro onto newtype-tuple pattern structs. Since we didn't need to export all of the model behaviour from Rust, the implementation was made significantly more simple.

However, there is a large portion of the library that uses generics and traits. Unfortunately as JavaScript doesn't understand lifetimes, this means we have to create a struct that represents each generic variant. For example, we have the Agent struct (which represents, in short, an LLM wrapper):

struct Agent<M> {
    model: M,
    // .. some other fields here
}

impl Agent<M> where M: CompletionModel {
    // methods, etc
}

The following would have to be created in our new rig-wasm crate:

struct OpenAIAgent {
    // .. fields here
}

struct AnthropicAgent {
    // .. fields here
}

There is also the need for users to be able to generate their own types and use it with the library. For example, when you're using an agent you can attach some tools to it which then get called by the language model. You might also want to attach a vector store onto it for RAG (eg Qdrant or Pinecone).

Finally, we should ensure the final resulting library uses idiomatic TypeScript. This is technically optional. However to ensure that users don't go through hell and back trying to use our library, we'll be doing it anyway. This means that we will have to write some TypeScript code ourselves. No worries - at least it's not raw JavaScript. It could always be worse.

Regenerating provider modules

One of the issues that we had originally was that we support about 20 different providers to date - that's a lot of code to re-write. One of the solutions that I proposed was using a build script to regenerate all of the original provider modules based on whatever traits they implement and then use a Jinja template.

It worked for the most part - the primary concern is that the build script can be quite brittle because of the Jinja templating. You can see the build file, which is pretty much just a load of string parsing (and then conditionally adding more code to the final file that gets produced).

One thing you will notice if you check the provider modules for the relevant Agent struct(s) is that we actually take a JavaScript value as the parameter for fn new(), then check it for the required shape and ensure it satisfies all the requirements where required - and then (and only then!) do we create the relevant struct. Fortunately with types that are JSON-compatible, we can simply deserialize it into the relevant Rust struct. However, this is not possible with types in JS/TS that have function signatures for obvious reasons - so we parse each field manually.

User integrations

Alright, so here's where most of the technical complexity starts. As mentioned before we have a couple of integration types that hook directly into our Agent type, and they are as follows:

Vector stores (used in dynamic tool RAG-ing and dynamic document context)
Tools (called by the LLM)

Both traits are mostly all-async functions and require Send + Sync futures to be returned. While this has been useful in Rust to ensure runtime properties, it also can be a huge nuisance sometimes. This is one of those times. Fortunately, we have an escape hatch in the form of wasm_bindgen_futures::spawn_blocking - which allows us to run said JavaScript (and Rust!) functions. Unfortunately, it doesn't allow us to return a raw value.

A solution to this that I ended up using was to essentially use oneshot channels to return the error (or the raw result), then use a tokio::select! macro:

fn call(
        &self,
        args: Self::Args,
    ) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send + Sync {
        let (result_tx, result_rx) = futures::channel::oneshot::channel();
        let (error_tx, error_rx) = futures::channel::oneshot::channel();
        let js_args: JsValue = serde_wasm_bindgen::to_value(&args).expect("This should be a JSON object!");

        let func = self.inner.clone();

        spawn_local(async move {
            // .. do some work here

            let res = match do_some_work() {
                Ok(res) => res,
                Err(e) => {
                    error_tx
                        .send(e)
                        .expect("shouldn't fail while sending to a oneshot channel");
                return;
                }

            }

            result_tx
                .send(res)
                .expect("sending a message to a oneshot channel shouldn't fail");
        });

        async {
            tokio::select! {
                res = result_rx => {
                    Ok(res.unwrap())
                },
                err = error_rx => {
                    Err(ToolError::ToolCallError(err.inspect_err(|x| println!("Future was cancelled: {x}")).unwrap().to_string().into()))
                }
            }
        }
    }

This essentially means that whichever channel receives a value first is the channel that a returned value will be pulled from (ie, if I have an error channel & a result channel, if the error channel receives an error first, it'll return the error from the function - if no error gets sent, ie the function fully completed, it returns the result). It's not what I would consider perfect considering that sending to a oneshot channel can technically fail and I had to include an initPanicHook() function so I can track panics from the JS/TS side, but it works and does so well enough that I haven't had to worry about it yet.

Here is an example of using a tool in rig-wasm:

import { Agent } from "@0xplaygrounds/rig-wasm/openai"

const counter = {
  counter: 5324,
  name: function () {
    return "counter";
  },
  definition: function (_prompt: string) {
    return {
      name: "counter",
      description: "a counter that can only be incremented",
      parameters: {
        $schema: "https://json-schema.org/draft/2020-12/schema",
        title: "ToolDefinition",
        type: "object",
        properties: {},
        required: [],
      },
    };
  },
  call: async function (args: any) {
    this.counter += 1;
    return { result: this.counter };
  },
};

let prompt = `Please increment the counter by 1 and let me know what the resulting number is, as well as the original number.`;
let key = process.env.OPENAI_API_KEY;

try {
  const agent = new Agent({
    apiKey: key,
    model: "gpt-4o",
    tools: [counter],
  });

  console.log(`Prompt: ${prompt}`);
  let res = await agent.prompt(prompt);
  console.log(`GPT-4o: ${res}`);
} catch (e) {
  if (e instanceof Error) {
    console.log(`Error while prompting: ${e.message}`);
  }
}

As you can see, we don't really need to do anything out of the ordinary here. Calling the agent should just tell us that the counter has increased by 1, and that the new total is 5325. Internally Rig already supports multi-turn calling.

Unfortunately while we have a plethora of vector database integrations, we basically have to write them in TS manually. For this, we had to write an interface that looks like this:

export interface VectorStore {
  topN: (req: VectorSearchOpts) => Promise<TopNResult[]>;
  topNIds: (req: VectorSearchOpts) => Promise<TopNIdsResult[]>;
}

Then we had to write classes that basically had those function identifiers (and matched the typedefinitions). The Qdrant integration for example looks like this:

export class QdrantAdapter {
  private client: QdrantClient;
  private collectionName: string;
  private params: QdrantClientParams;
  private embeddingModel: CanEmbed;

  constructor(
    collectionName: string,
    embeddingModel: CanEmbed,
    params: QdrantClientParams,
  ) {
    this.params = params;
    this.embeddingModel = embeddingModel;
    this.collectionName = collectionName;
  }

  async loadClient() {
    if (!this.client) {
      try {
        this.client = new QdrantClient(this.params);
      } catch (err) {
        throw new Error("Failed to load Qdrant client: " + err);
      }
    }
  }

  async init(dimensions: number) {
    await this.loadClient();

    const collections = await this.client.getCollections();
    const exists = collections.collections.some(
      (c: any) => c.name === this.collectionName,
    );

    if (!exists) {
      await this.client.createCollection(this.collectionName, {
        vectors: {
          size: dimensions,
          distance: "Cosine",
        },
      });
    }
  }

  async topN(opts: VectorSearchOpts): Promise<SearchResult[]> {
    await this.loadClient();

    const embedding = await this.embeddingModel.embedText(opts.query);

    const result = await this.client.search(this.collectionName, {
      vector: embedding.vec,
      limit: opts.samples,
    });

    return result.map((res: any) => ({
      id: res.id,
      score: res.score,
      payload: res.payload,
    }));
  }

  async topNIds(opts: VectorSearchOpts): Promise<SearchResult[]> {
    await this.loadClient();

    const embedding = await this.embeddingModel.embedText(opts.query);

    const result = await this.client.search(this.collectionName, {
      vector: embedding.vec,
      limit: opts.samples,
    });

    return result.map((res: any) => ({
      id: res.id,
      score: res.score,
    }));
  }
}

Basically speaking, as long as it implements topN and topNIds, it should satisfy the criteria to be considered a VectorStore.

Build pipelines

Another aspect of actually preparing the library for shipping to production was figuring out how to get the build tooling working. Admittedly, this was my first time ever shipping my own JS/TS library at all, so I ended up having more friction than perhaps someone who's more experienced with the JS/TS build tooling ecosystem might be. I made judicious use of LLM prompting here to solve specific issues, and then checked the documentation to ensure that everything was totally correct and figure out why some things did work (or not) after the fact.

Fortunately, the Rust part was simple: build the binary for the target then use the wasm-bindgen CLI. At the time of writing the build pipeline (roughly a couple of weeks ago) I wanted to experiment with going much more level and basically just using wasm-bindgen-cli itself. However, if you wanted to you could definitely do this with something like wasm-pack. Perhaps at some point, it would be ideal to revisit this aspect of the build pipeline.

In terms of JS build tooling, I chose Rollup as it seemed quite minimal, is good for making smaller bundles and has a lot of different plugins you can play around with. I also didn't really feel like I needed all of the features that Webpack has. However, if your requirements differ and you aren't just building a library, you may want to use Webpack. It's much more popular.

There were quite a few gotchas that I encountered along the way while trying to get the thing to build (mostly skill issues, but perhaps some valuable advice for anyone who's bad at TypeScript trying to do the same thing):

The node target for wasm-bindgen-cli actually compiles to CommonJS, but most modern projects actually use ESM. We ended up using the experimental-nodejs-module target. Not very ideal to be honest, but it did work and I'm quite thankful for that. Hopefully, this target will stabilise at some point.
I tried to use .d.ts files in my library alongside my .ts files because I needed to declare some typedefs. It turns out that if you import anything, they are now not legitimate .d.ts files anymore. I realised after several Google searches that this is in fact an extremely bad idea, not to mention that this actually caused the compilation to fail. I ended up just writing TS files.
Setting up tsconfig.json was quite a task in itself (as someone who previously interacted very rarely with the tsconfig file, even when I primarily used TS as a hobbyist).
rollup does not automatically add the wasm files to your dist directory (or wherever you have the final compiled output), so you've got to copy them manually. Fortunately this is probably the least of your issues here - you can just run cp <foo> <bar> on each of the required WASM files and it should do the job.

The other major issue was trying to add all of the exports/inputs to the package.json file and the rollup.config.ts file. Fortunately, the rollup file issue was simple. Iterate through a directory, check the file extensions and then do some string operations:

const readDir = (inputDir: string): string[] => {
  const files = fs
    .readdirSync(inputDir)
    .filter((file: string) => file.endsWith(".ts")) 
    .map((file: string) => path.join(inputDir, file));

  return files;
};

const input: Record<string, string> = {};

const dirs = [providersDir, vectorStoresDir, coreDir];

for (const dir of dirs) {
  for (const file of readDir(dir)) {
    const name = path.basename(file, ".ts");
    input[name] = file;
  }
}

Unfortunately, the package.json file issue was not quite so simple. I ended up basically creating a Bash script that iterated through every file/directory in src, then operated on the file/folder names and created raw JSON. At the end, it then used jq to edit the package.json file so that it would set the exports as the newly created list.

A short excerpt of the Bash script:

echo '{' > "$EXPORTS_FILE"

# Add the root export manually
echo '  ".": {' >> "$EXPORTS_FILE"
echo '    "import": "./dist/esm/index.js",' >> "$EXPORTS_FILE"
echo '    "require": "./dist/cjs/index.cjs",' >> "$EXPORTS_FILE"
echo '    "types": "./dist/esm/index.d.ts"' >> "$EXPORTS_FILE"
echo '  },' >> "$EXPORTS_FILE"

# Gather all .js files (excluding index.js)
mapfile -t core_files < <(find "$CORE_MODULE" -maxdepth 1 -type f -name "*.ts" | sort)
echo "Found ${#core_files[@]} core modules. Building exports..."

for i in "${!core_files[@]}"; do
  file="${core_files[$i]}"
  base=$(basename "$file" .js)

  echo "  \"./${base%.ts}\": {" >> "$EXPORTS_FILE"
  echo "    \"import\": \"./$OUT_DIR/${base%.ts}.js\"," >> "$EXPORTS_FILE"
  echo "    \"types\": \"./$OUT_DIR/${base%.ts}.d.ts\"" >> "$EXPORTS_FILE"
  if [[ $i -lt $((${#providers_files[@]} - 1)) ]]; then
    echo "  }," >> "$EXPORTS_FILE"
  else
    echo "  }" >> "$EXPORTS_FILE"
  fi
done

echo "}" >> "$EXPORTS_FILE"

echo "✅ Wrote updated exports to $EXPORTS_FILE"

jq --slurpfile exports exports.json '.exports = $exports[0]' package.json > package.new.json && mv package.new.json package.json

echo "✅ Updated exports in package.json"

rm $EXPORTS_FILE

I will admit that I'm not a Bash expert and had an LLM spin up the skeleton of the script, then I finished off what parts I couldn't get the LLM to figure out in a reasonable amount of time (~5 minutes). It was mostly quite quick and painless. I am sure there is probably a more reasonable (and maintainable!) way to write this, but it works for now.

Now that I'd written the script, the build pipeline was actually mostly done. I just had to set up the final build script:

#!/usr/bin/env bash
cargo build -p rig-wasm --release --target wasm32-unknown-unknown
wasm-bindgen \
    --target experimental-nodejs-module \
    --out-dir rig-wasm/pkg/src/generated \
    target/wasm32-unknown-unknown/release/rig_wasm.wasm

./scripts/exports.sh 
npm ci
rollup -c 
cp src/generated/rig_wasm_bg.wasm dist/esm/generated 
cp src/generated/rig_wasm_bg.wasm.d.ts dist/esm/generated

Then all that was left to do was to run the script and then use npm publish.

Stuff we couldn't port

Of course, now that we've ported pretty much everything we could, there are a number of things that we couldn't port:

the DynClientBuilder that basically allows you to create whatever client you want based on the provider/model input (as strings). Unfortunately, JS/TS doesn't understand lifetimes, so this was impossible.
the PromptRequest<'a> struct that we use internally to control multi-turn usage also, unfortunately, can't be turned into JS because of lifetimes
The vector database integrations, as they're not WASM friendly because they use libraries that are not compatible with WASM. Thankfully, basically every vector database that exists has a TypeScript SDK. Those integrations will have to be re-created on the TS side. Thankfully as SDKs are often (mostly!) well documented, it is not difficult to do so.

Qualifying Sales Leads with Rig & Google Sheets MCP

Josh Mo — Wed, 23 Apr 2025 18:04:19 +0000

In today’s fast-paced sales environments, separating high-quality leads from noise is more critical than ever. Manual qualification is time-consuming and error-prone — but with the right automation stack, it doesn’t have to be. In this article, we’ll walk through how to build a lean, intelligent lead qualification workflow using Rig, as well as a Google Sheets MCP server. Whether you're a solo founder or scaling a sales team, this setup offers a flexible, AI-powered pipeline to help you prioritize the leads that matter most — all while keeping your favorite tools in play.

A very common example where this can come into play might be analysing the results of Google Forms submissions (or, actually, any website that provides form submissions and can be exported to Google Sheets!). By creating an LLM-assisted data pipeline, we can significantly cut down the time required to be able to do the admin work and you can continue doing the parts of your work that can't be easily automated away.

Interested in just getting straight ahead to trying it out? Here's the GitHub repo.

Running Google Sheets with MCP

While Google Sheets doesn't officially have their own MCP server, there are quite a few unofficial servers you can use. Here's one I used which supports both OAuth as well as service account credentials. I've also additionally checked the code out and ensured it is 100% safe to use - it's written using FastMCP.

To create a service account (that the MCP server can use) you'll need to do the following:

Create a project in Google Cloud Console (if you don't have one already or need to make a new one)
Enable the Google Drive API and Google Sheets API for your project
Go to the Credentials screen, create a service account and download the credentials file (you'll need this later!). Additionally ensure you remember the email.
Create a new Google Drive folder and give Editor permissions to your service account (through using the email provided from the service account menu).
You should now be ready to go!

Note that you'll need to store the ID of the Google Drive folder you're using. You can find it from the URL like so:

https://drive.google.com/drive/u/0/folders/{ID_HERE}

Getting Started

Talking to our MCP server

Before we do anything else, we need to get the following values as environment variables specifically for our MCP server:

The path to our credentials file (for usage with Google service accounts)
The ID of the folder where we've given our service account access to as environment variables

Next, we need to git clone the repository:

git clone https://github.com/xing5/mcp-google-sheets

Ensure that you put your downloaded service account credentials file into the newly cloned git folder! This will allow the MCP server to find your credentials. You'll also need to set your service account path (wherever your credentials .json file is) and Google Drive folder ID as environment variables:

export SERVICE_ACCOUNT_PATH=./credentials.json
export DRIVE_FOLDER_ID=my-folder-id-here

Once done, simply use cargo run to start the server in SSE (Server-Sent Events) mode. This means it essentially spins up a web server which emits messages from a given SSE path. Server Sent Events are a great way to pass messages from the server to the client while not requiring a constant 2-way connection as the client doesn't need to send anything back to the server.

Project setup

To get started with creating our project, you'll need to run the following to initialise your Rust application:

cargo init gsheets-mcp-rig

Next, you'll want to install the project dependencies - which you can do with the one liner below:

cargo add anyhow mcp-core rig-core tokio -F rig-core/mcp,tokio/macros,\
tokio/rt-multi-thread

What did we just add?

mcp-core: A crate for working with MCP.
rig-core: The Rig core crate. Uses the mcp feature for integration with mcp-core.
tokio: An async Rust runtime. Uses the mcaros and rt-multi-thread features for convenient async runtime setup.

Additionally, you'll also want to set an OpenAI API key as an environment variable.

You can do this with the one-liner below:

export OPENAI_API_KEY=my-key-here

Building our project

Now for the fun part: building our project!

To start with, we will create a simple function that will attempt to connect to our MCP server by its URL and port (in this case, localhost:8000 assuming you have it running locally). We then feed the command into the ClientSseTransport::new function to start the program up, then initiate a connection and return the client itself.

async fn connect_to_gsheets_mcp()
-> Result<mcp_core::client::Client<ClientSseTransport>, Box<dyn std::error::Error>> {
    println!("Loading Google Sheets MCP server...");

    // Create the MCP client
    let client_transport = ClientSseTransport::new(
        "http://localhost:8000"
    )
    .unwrap();

    let mcp_client = ClientBuilder::new(client_transport).build();

    // Start the MCP client
    mcp_client.open().await?;

    Ok(mcp_client)
}

Next, we need to fetch our list of tools from the MCP server and create a ToolSet as well as a Vec<ToolDefinition> - both of which will be used to invoke functions as well as help the LLM select which tool to use. To break this function down:

We use the fold() function, which takes an accumulator (ToolSet and Vec which are both contained in a tuple)
Within the closure, the tuple (of accumulators) is on the left hand side with the iterator return value being on the right hand side
We then create a rig::tools::McpTool and and it to the ToolSet while also creating a ToolDefinition and pushing it to the Vec<ToolDefinition>

fn get_tools_from_mcp_tool_response(
    tools_list_res: ToolsListResponse,
    mcp_client: mcp_core::client::Client<ClientSseTransport>,
) -> (ToolSet, Vec<ToolDefinition>) {
    let (tools, tooldefs) = tools_list_res.tools.into_iter().fold(
        (ToolSet::builder().build(), Vec::new()),
        |(mut tools, mut tooldefs), tool| {
            let mcp_tool = McpTool::from_mcp_server(tool.clone(), mcp_client.clone());
            tools.add_tool(mcp_tool);

            let tooldef = ToolDefinition {
                description: tool.description.unwrap_or(String::new()),
                name: tool.name,
                parameters: tool.input_schema.clone(),
            };
            tooldefs.push(tooldef);

            (tools, tooldefs)
        },
    );

    (tools, tooldefs)
}

Creating our reasoning loop

Now that we've fetched our tools from the MCP server, let's add them to our completion request.

We also want to make our application competent: it should be fully capable of incorporating tool responses, as well as returning an answer from the LLM without the user needing to do anything.

Before we start, let's create a couple of helper methods. The first one will simply be a message that converts an ID and a content field into a UserContent::tool_result.

fn tool_result_from_text_response(id: impl Into<String>, content: String) -> UserContent {
    let content = OneOrMany::one(ToolResultContent::Text(content.into()));
    UserContent::tool_result(id, content)
}

Now we can implement our response loop! This is a relatively long function, so buckle up. Here's what the function will do:

Create a CompletioRequestBuilder and fill it out with all required items including tool definitions, prompts and chat history.
Send the completion request.
Match the latest message:
- If it's a text message, add prompt and assistant message to chat history then just return the text
- If it's a tool call, attempt to call the function (using the toolset we initialised)
- If the tool call results in an error, set the tool call error as the the prompt (explained below)
- If it was successful, push both the original prompt and assistant message to the chat history, then set the prompt as the tool call result (explained below)

When we're sending a completion request to the model, the tool prompt resolves to a message with the tool role. This allows us to continue the conversation without us explicitly having to send a prompt ourselves.

Now for the code! Let's begin by initialising our toolset and filling out the completion request, then sending it. Note that although the function does use generics, we don't need to specify what these are when we're actually calling this function as there is only one model we're using (the OpenAI one):

async fn call_until_response<M: CompletionModel>(
    mut prompt: Message,
    model: &M,
    preamble: &str,
    chat_history: &mut Vec<Message>,
    toolset: &ToolSet,
    tooldefs: Vec<ToolDefinition>,
) -> Result<String, anyhow::Error> {

    loop {
        let request = CompletionRequestBuilder::new(model.clone(), prompt.to_owned())
            .preamble(preamble.to_owned())
            .messages(chat_history.clone())
            .temperature(0.0)
            .max_tokens(1024)
            .tools(tooldefs.clone())
            .build();
        // call model
        let resp = model
            .completion(request)
            .await
            .map_err(|x| anyhow::anyhow!("Error when prompting: {x}"))?;

        // rest of code goes down here!
    }
}

Next, we need to match the response and see what type the response is. If it's text, return it; otherwise, try to execute the function. If the function fails, set the prompt to be the tool call result and recursively call the function again to get a string response and return the string. Otherwise, the variable assignment works as normal and we can continue.

// keep calling tools until we get human readable answer from the model
match resp.choice.first() {
    AssistantContent::Text(text) => {
        let text = text.text;
        chat_history.push(prompt.clone().into());
        chat_history.push(Message::assistant(&text));
        return Ok(text);
    }
    AssistantContent::ToolCall(tool_call) => {
        // Call the tool
        let tool_response = toolset
            .call(
                &tool_call.function.name,
                tool_call.function.arguments.to_string(),
            )
            .await;

        let tool_response = match tool_response {
            Ok(res) => res,
            Err(e) => {
                chat_history.push(prompt.clone());
                chat_history.push(Message::Assistant {
                            content: OneOrMany::one(AssistantContent::ToolCall(tool_call.clone())),
                        });

                prompt = Message::User {
                    content: OneOrMany::one(UserContent::ToolResult(ToolResult {
                        id: tool_call.id.to_string(),
                        content: OneOrMany::one(ToolResultContent::Text(rig::message::Text {
                            text: e.to_string(),
                        })),
                    })),
                };
                continue;
                }
            };
        }
    }
}

Finally, we convert the text returned from the tool to a tool response message, then push the original prompt & assistant message to the chat history. We then (optionally) print the tool response message, as well as setting the "user" prompt to be the tool response message and then call the call_until_response function recursively to get our String result (and then return the String).

let tool_response_message =                tool_result_from_text_response(tool_call.id.clone(), tool_response);

let tool_call = OneOrMany::one(AssistantContent::ToolCall(tool_call));

// add tool call and response into chat history and continue the loop
chat_history.push(prompt.clone().into());
chat_history.push(Message::Assistant { content: tool_call });

println!("{tool_response_message:?}");

let tool_result_message = Message::User {
    content: OneOrMany::one(tool_response_message),
};

prompt = tool_result_message;

After this, the function will loop back around.

Hooking everything back up

Now that we've actually wrote all the difficult parts of the code, now is the easy part: taking the input and writing our main function!

To take input from the terminal, we simply use stdio().read_line in a function that returns a String:

fn take_input() -> String {
    let mut str = String::new();
    print!("> ");
    stdout()
        .flush()
        .expect("Failed to flush stdout - this should normally never happen");

    stdin().read_line(&mut str).expect("to read line");

    str
}

Next is the main function. We need to create our preamble (system message), then we initialise everything and make a loop, taking input from the user and using our LLM reasoning loop to get a response (then printing the response). We will also add a quit command to be able to break the loop.

const PREAMBLE: &str = r###"You are an agent designed to qualify sales leads from Google Sheets.

Users will typically ask you to qualify leads from Google Forms submissions
(or imported spreadsheets from results of other form submission-type applications).

Your job is to qualify sales leads based on the user's criteria.
If they don't give you a criteria for qualification,
ask what demographic the user is trying to capture with the form and qualify leads based off of that.

When creating the results, use a new sheet in the spreadsheet file the user has provided you with.
When done, specify the location of the sheet so that the user can inspect the result for themselves."###;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mcp_client = connect_to_stripe_mcp().await?;

    let tools_list_res = mcp_client.list_tools(None, None).await?;

    let (tools, tooldefs) = get_tools_from_mcp_tool_response(tools_list_res, mcp_client);

    let openai_client = providers::openai::Client::from_env();
    let model = openai_client.completion_model("gpt-4o");

    println!("Hi! How can I help you today? (write \"quit\" to exit)");
    println!("------------");

    loop {
        let prompt = take_input();
        println!("------------");

        if prompt == *"quit" {
            println!("Thanks for using me! I am quitting now.");
            break;
        }

        let mut chat_history = Vec::new();

        let res = call_until_response(
            prompt.into(),
            &model,
            PREAMBLE,
            &mut chat_history,
            &tools,
            tooldefs.clone(),
        )
        .await
        .unwrap();

        println!("{res}");
        println!("------------");
    }

    Ok(())
}

Extending this demo

So, now that we're done: how can we extend this demo? Below are a couple of ideas:

Switch to a webpage interface: While this example does work, if you plan on mass adoption you may instead want to try hosting this as a webpage.
Add some personality to your chatbot: Depending on your target audience or customer base, adapting the chatbot's personality by prompt engineering the preamble (system message) can make it much easier for your audience to approach using an AI-assisted interface. Don't forget that you can use guardrails, semantic routing and other techniques to improve the effectiveness of your chatbot as well as preventing any potential disasters! ## Finishing up Thanks for reading! Hopefully this has helped you learn a little bit more about how you can use MCP servers in production, as well as how Rig can help you achieve that goal.

For additional Rig resources and community engagement:

Check out more examples in our gallery.
Contribute or report issues on our GitHub.
Join discussions in our Discord community!

Taking Payments with AI Agents and Rig using Stripe's MCP server

Josh Mo — Wed, 09 Apr 2025 19:07:11 +0000

Introduction

Hello world! In today's post we're going to take a look at how you can build AI agent systems that can take payments with Rig, a Rust-based AI agent framework, and the official Stripe MCP server.

By the end of this you'll have a CLI chatbot that can read products and prices that belong to your Stripe account, as well as creating payment links (using Stripe's MCP server). The chatbot will be modelled on a shop assistant.

Below is an example of a possible chat log with this chatbot:

And of course, here is the respective payment link output:

Interested in just jumping ahead to trying it out? Check out the repo.

Getting started

Pre-requisites

Before we get started, you'll need the latest version of Node as well as a Stripe API key. Note that if you don't actually own a company to use Stripe officially (or you want to ensure everything works), you can still use a test account with Stripe.

The Stripe MCP server itself requires Node to run but does not need an additional web server as it communicates through stdio (standard input/output). The program we'll be building will be instantiating the MCP server by itself, so outside of Node you shouldn't need anything else.

You'll also need an OpenAI API key.

Using Stripe

To start using Stripe, sign into their dashboard then create a few products and attach some prices to them - you can find out more about this on their documentation pages. Once done, we'll be creating payment links programatically through the MCP server so you won't need to manually handle this yourself.

Project setup

To get started with creating our project, you'll need to run the following to initialise your Rust application:

cargo init stripe-mcp-rig

Next, you'll want to install the project dependencies - which you can do with the one liner below:

cargo add anyhow mcp-core rig-core tokio -F rig-core/mcp,tokio/macros,\
tokio/rt-multi-thread

What did we just add?

mcp-core: A crate for working with MCP.
rig-core: The Rig core crate. Uses the mcp feature for integration with mcp-core.
tokio: An async Rust runtime. Uses the mcaros and rt-multi-thread features for convenient async runtime setup. ## Building our project Now for the fun part: building our project! ### Talking to our MCP server Before we do anything else, we need to get our Stripe API key as an environment variable as well as our OpenAI API key. You can do this with export command as below:

export STRIPE_API_KEY=my-key-here
export OPENAI_API_KEY=my-key-here

As mentioned before, currently the official Stripe MCP server supports stdio only - and we can start it up with the following one-line command (although later on, the program itself will instantiate the MCP server):

npx -y @stripe/mcp --tools=all --api-key=$STRIPE_API_KEY

The mcp_core crate allows us to run this in much the same way that you'd run something using the std::command::Command struct - providing the program name, then the arguments.

Next, a (relatively!) simple function that will attempt to get the Stripe API key from environment variables and automatically panics if it doesn't exist (as it's required to allow the program to run). We then feed the command into the ClientStdioTransport::new function to start the program up, then initiate a connection and return the client itself.

async fn connect_to_stripe_mcp()
-> Result<mcp_core::client::Client<ClientStdioTransport>, Box<dyn std::error::Error>> {
    println!("Loading Stripe MCP server...");

    let stripe_api_key =
        std::env::var("STRIPE_API_KEY").expect("STRIPE_API_KEY to exist as env var");
    let stripe_api_key = format!("--api-key={stripe_api_key}");
    // Create the MCP client

    let client_transport = ClientStdioTransport::new(
        "npx",
        &[
            "-y",
            "@stripe/mcp",
            "--tools=products.read,prices.read,paymentLinks.create",
            &stripe_api_key,
        ],
    )
    .unwrap();

    let mcp_client = ClientBuilder::new(client_transport).build();

    // Start the MCP client
    mcp_client.open().await?;

    Ok(mcp_client)
}

Note that in the command above, we have only allowed the following tools to be spawned with the MCP server:

Reading products
Reading prices
Creating payment links

Although the MCP server contains many, many more tools, in the wild this this would be a user-facing application. In this case, we should ensure that we grant the least power possible to the user that is still required for the application to work. Otherwise, they might try to create new products for example, or add fake customers to our Stripe account!

We use the fold() function, which takes an accumulator (ToolSet and Vec which are both contained in a tuple)
Within the closure, the tuple (of accumulators) is on the left hand side with the iterator return value being on the right hand side
We then create a rig::tools::McpTool and and it to the ToolSet while also creating a ToolDefinition and pushing it to the Vec<ToolDefinition>

fn get_tools_from_mcp_tool_response(
    tools_list_res: ToolsListResponse,
    mcp_client: mcp_core::client::Client<ClientStdioTransport>,
) -> (ToolSet, Vec<ToolDefinition>) {
    let (tools, tooldefs) = tools_list_res.tools.into_iter().fold(
        (ToolSet::builder().build(), Vec::new()),
        |(mut tools, mut tooldefs), tool| {
            let mcp_tool = McpTool::from_mcp_server(tool.clone(), mcp_client.clone());
            tools.add_tool(mcp_tool);

            let tooldef = ToolDefinition {
                description: tool.description.unwrap_or(String::new()),
                name: tool.name,
                parameters: tool.input_schema.clone(),
            };
            tooldefs.push(tooldef);

            (tools, tooldefs)
        },
    );

    (tools, tooldefs)
}

Creating our reasoning loop

Now that we've fetched our tools from the MCP server, let's add them to our completion request.

We also want to make our application competent: it should be fully capable of incorporating tool responses, as well as returning an answer from the LLM without the user needing to do anything.

Before we start, let's create a couple of helper methods. The first one will simply be a message that converts an ID and a content field into a UserContent::tool_result.

fn tool_result_from_text_response(id: impl Into<String>, content: String) -> UserContent {
    let content = OneOrMany::one(ToolResultContent::Text(content.into()));
    UserContent::tool_result(id, content)
}

Now we can implement our response loop! This is a relatively long function, so buckle up. Here's what the function will do:

Create a CompletioRequestBuilder and fill it out with all required items including tool definitions, prompts and chat history.
Send the completion request.
Match the latest message:
- If it's a text message, add prompt and assistant message to chat history then just return the text
- If it's a tool call, attempt to call the function (using the toolset we initialised)
- If the tool call results in an error, set the tool call error as the the prompt (explained below)
- If it was successful, push both the original prompt and assistant message to the chat history, then set the prompt as the tool call result (explained below)

async fn call_until_response<M: CompletionModel>(
    mut prompt: Message,
    model: &M,
    preamble: &str,
    chat_history: &mut Vec<Message>,
    toolset: &ToolSet,
    tooldefs: Vec<ToolDefinition>,
) -> Result<String, anyhow::Error> {

    loop {
        let request = CompletionRequestBuilder::new(model.clone(), prompt.to_owned())
            .preamble(preamble.to_owned())
            .messages(chat_history.clone())
            .temperature(0.0)
            .max_tokens(1024)
            .tools(tooldefs.clone())
            .build();
        // call model
        let resp = model
            .completion(request)
            .await
            .map_err(|x| anyhow::anyhow!("Error when prompting: {x}"))?;

        // rest of code goes down here!
    }
}

// keep calling tools until we get human readable answer from the model
match resp.choice.first() {
    AssistantContent::Text(text) => {
        let text = text.text;
        chat_history.push(prompt.clone().into());
        chat_history.push(Message::assistant(&text));
        return Ok(text);
    }
    AssistantContent::ToolCall(tool_call) => {
        // Call the tool
        let tool_response = toolset
            .call(
                &tool_call.function.name,
                tool_call.function.arguments.to_string(),
            )
            .await;

        let tool_response = match tool_response {
            Ok(res) => res,
            Err(e) => {
                prompt = Message::User {
                    content: OneOrMany::one(UserContent::ToolResult(ToolResult {
                        id: tool_call.id.to_string(),
                        content: OneOrMany::one(ToolResultContent::Text(rig::message::Text {
                            text: e.to_string(),
                        })),
                    })),
                };
                continue;
                }
            };
        }
    }
}

let tool_response_message =                tool_result_from_text_response(tool_call.id.clone(), tool_response);

let tool_call = OneOrMany::one(AssistantContent::ToolCall(tool_call));

// add tool call and response into chat history and continue the loop
chat_history.push(prompt.clone().into());
chat_history.push(Message::Assistant { content: tool_call });

println!("{tool_response_message:?}");

let tool_result_message = Message::User {
    content: OneOrMany::one(tool_response_message),
};

prompt = tool_result_message;

After this, the function will loop back around.

Hooking everything back up

Now that we've actually wrote all the difficult parts of the code, now is the easy part: taking the input and writing our main function!

To take input from the terminal, we simply use stdio().read_line in a function that returns a String:

fn take_input() -> String {
    let mut str = String::new();
    print!("> ");
    stdout()
        .flush()
        .expect("Failed to flush stdout - this should normally never happen");

    stdin().read_line(&mut str).expect("to read line");

    str
}

const PREAMBLE: &str = r###"You are a shop assistant.
Your job is to help customers find the product they are looking for using Stripe.
If they want to buy it, send a payment link where they can buy it.
If they are rude to you in any way, don't respond to them."###;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mcp_client = connect_to_stripe_mcp().await?;

    let tools_list_res = mcp_client.list_tools(None, None).await?;

    let (tools, tooldefs) = get_tools_from_mcp_tool_response(tools_list_res, mcp_client);

    let openai_client = providers::openai::Client::from_env();
    let model = openai_client.completion_model("gpt-4o");

    println!("Hi! How can I help you today? (write \"quit\" to exit)");
    println!("------------");

    loop {
        let prompt = take_input();
        println!("------------");

        if prompt == *"quit" {
            println!("Thanks for using me! I am quitting now.");
            break;
        }

        let mut chat_history = Vec::new();

        let res = call_until_response(
            prompt.into(),
            &model,
            PREAMBLE,
            &mut chat_history,
            &tools,
            tooldefs.clone(),
        )
        .await
        .unwrap();

        println!("{res}");
        println!("------------");
    }

    Ok(())
}

Extending this demo

So, now that we're done: how can we extend this demo? Below are a couple of ideas:

Switch to a webpage interface: While this example does work, if you plan on mass adoption you may instead want to try hosting this as a webpage.
Add some personality to your chatbot: Depending on your target audience or customer base, adapting the chatbot's personality by prompt engineering the preamble (system message) can make it much easier for your audience to approach using an AI-assisted interface. Don't forget that you can use guardrails, semantic routing and other techniques to improve the effectiveness of your chatbot as well as preventing any potential disasters! ## Finishing up Thanks for reading! Hopefully this has helped you learn a little bit more about how you can use MCP servers in production, as well as how Rig can help you achieve that goal.

For additional Rig resources and community engagement:

Check out more examples in our gallery.
Contribute or report issues on our GitHub.
Join discussions in our Discord community!

Using Model Context Protocol with Rig

Josh Mo — Wed, 02 Apr 2025 20:15:29 +0000

Introduction

The landscape of AI development is rapidly evolving, with large language models (LLMs) increasingly being called upon to interact with external tools and APIs to accomplish complex tasks. Anthropic's Model Context Protocol represents a significant step forward in this domain, providing a standardized approach for models like Claude to seamlessly interact with various tools. Recently, many companies have started looking to Model Context Protocol as a way for their users to be able to access their APIs through LLMs.

When combined with the Rig framework's robust agent architecture in Rust, developers can create powerful, reliable AI agents capable of executing sophisticated workflows with enhanced reasoning capabilities. In this article, we'll explore how to implement Anthropic's Model Context Protocol within the Rig framework, offering Rust developers a pathway to building state-of-the-art agentic AI systems. We'll walk through practical implementation steps, best practices, and example use cases that demonstrate the potential of this powerful combination.

Below is an example of the result of running the MCP server built in this article:

What is Model Context Protocol?

Model Context Protocol (MCP for short) is Anthropic's standardized interface for enabling AI models like Claude to interact with external tools and APIs. At its core, MCP defines a structured way for models to:

Request tool use: Models can explicitly indicate when they need to use a specific tool
Receive tool responses: Models can process and incorporate results from tool executions
Maintain context awareness: Models can track the state of ongoing interactions across multiple tool calls

Unlike traditional prompt engineering approaches, MCP provides a more reliable and consistent mechanism for tool use, reducing hallucinations and improving the model's ability to reason about when and how to use tools. The protocol uses a JSON-based format that clearly delineates between model reasoning, tool requests, and tool responses, making the interaction flow more predictable and traceable.

Why is MCP useful?

MCP is great because essentially, it allows you to create a proxy from your LLM to pretty much whatever you want. You can create an MCP to replicate certain services, or to create suites of utilities for your LLM. For example: many developers have already started creating and using MCP servers for various specific services like Stripe, Qdrant, and even things like Postgres - granting you full usage (at least insofar as the MCP server is capable) to the service from your LLM.

You can also additionally create MCP servers for all the functionality that a certain role might find useful: for example, an MCP server for accountants. Or even an MCP server that calls other LLMs, which then talk to MCP servers... and so on and so forth.

How do I use MCP with Rig?

Fortunately for us, Model Context Protocol has an official Rust SDK which you can use today to be able to spin up your own MCP server and use it with Rig.

There's two components to this:

Creating your own MCP server
Using an MCP server with Rig

We'll be covering both of these components in this article. The example we will be using in this article is an MCP server which has a tool for adding numbers together.

Building your own MCP server

Building your own MCP server requires you to write tools with names, descriptions and parameters (as well as descriptions of those parameters, ideally!) and then adding them all to a server.

Fortunately, the mcp-core crate provides most of this for us, leaving us to do primarily only the tool implementations. You can see what a tool looks like below:

use mcp_core_macros::tool;
use mcp_core::ToolResponseContent;
use anyhow::Result;

#[tool(
    name = "Add",
    description = "Adds two numbers together.",
    params(a = "The first number to add", b = "The second number to add")
)]
async fn add_tool(a: f64, b: f64) -> Result<ToolResponseContent> {
    Ok(tool_text_content!((a + b).to_string()))
}

The #[tool] macro generates a trait implementation which makes the tool compatible with the server (from the Rust side).

Next, we need to create a mcp_core::Server that lists all the tools that we can do. Note that we have a single tool in the ServerCapabilities struct called listChanged - this simply tells the server whether or not we want to receive notifications on our MCP client when the list of tools has changed.

let mcp_server_protocol = Server::builder("add".to_string(), "1.0".to_string())
    .capabilities(ServerCapabilities {
        tools: Some(json!({
            "listChanged": false,
        })),
        ..Default::default()
    })
    .register_tool(AddTool::tool(), AddTool::call())
    .build();

Note that you can basically register whatever tools you want to the server by using .register_tool() which takes the tool definition (from the add_tool function we created earlier), as well as a function pointer (which corresponds to the tool).

Next, we need to create an MCP transport server by simply creating a ServerSseTransport. We then plug this into Server::start() which starts up a web server and runs indefinitely (until it's shut down).

let mcp_server_transport =
    ServerSseTransport::new("127.0.0.1".to_string(), 3000, mcp_server_protocol);

Server::start(mcp_server_transport).await;

Note that if we also want to simultaneously do other things while the server is running, we can spawn a Tokio task which defers it to a separate concurrent task:

tokio::spawn( async move { Server::start(mcp_server_transport).await });

This can be quite useful for local testing purposes, although in production running the MCP server as a separate binary will be much more reliable as there won't be any other tasks which might hinder performance.

Using MCP servers with Rig

Using Rig, we can create an MCP client that connects to basically any MCP-compliant server we want. To do this, you'll need to enable the mcp feature and radd the mcp-core crate:

cargo add rig-core mcp-core -F rig-core/mcp

Next, we will need to create our client transport and use it in our builder. As a general rule, you can connect to MCP servers through one of two ways:

stdio ("Standard IO") transport.
sse transport (connecting to an MCP-compliant server with an SSE route)

The example below will use sse transport.

// Create the MCP client
let client_transport = ClientSseTransportBuilder::new("http://127.0.0.1:3000/sse".to_string()).build();
let mcp_client = ClientBuilder::new(client_transport).build();

// Start the MCP client
mcp_client.open().await?;

let init_res = mcp_client
    .initialize(
        Implementation {
            name: "mcp-client".to_string(),
            version: "0.1.0".to_string(),
        },
        ClientCapabilities::default(),
    )
    .await?;
println!("Initialized: {:?}", init_res);

let tools_list_res = mcp_client.list_tools(None, None).await?;
println!("Tools: {:?}", tools_list_res);

Once we've connected to the server and got our list of tools, we can then create our AI agent, add the MCP tools to our agent and then build & prompt the agent.

tracing::info!("Building RIG agent");
let completion_model = providers::openai::Client::from_env();
let mut agent_builder = completion_model.agent("gpt-4o");

// Add MCP tools to the agent
agent_builder = tools_list_res
    .tools
    .into_iter()
    .fold(agent_builder, |builder, tool| {
        builder.mcp_tool(tool, mcp_client.clone().into())
    });
let agent = agent_builder.build();

tracing::info!("Prompting RIG agent");
let response = agent.prompt("Add 10 + 10").await?;
tracing::info!("Agent response: {:?}", response);

Finishing up

Thanks for reading! I hope this article has enlightened you on using MCP with different tools and how to implement MCP servers using Rig, as well as the future possibilities that can be achieved by using MCP. Stay tuned for more!

For additional Rig resources and community engagement:

Check out more examples in our gallery.
Contribute or report issues on our GitHub.
Join discussions in our Discord community!

Automating onchain actions with AI using rig-onchain-kit

Josh Mo — Fri, 21 Mar 2025 15:57:49 +0000

Hello world! In today's article, we're going to go into how you can automate onchain actions with rig-onchain-kit, Rig's companion crate in collaboration with listen-rs that lets you execute actions onchain. By the end of this article, we'll have a CLI tool that you can use to trade onchain with Kraken using natural language.

In the screenshot below, we simply ask the bot to buy 0.02 SOL from Kraken at current market price:

Of course, that's not the only thing we can showcase - here is the real transaction that we've done on Kraken!

Note that although rig-onchain-kit supports both Solana and EVM, we will be focusing on Solana as our onchain network.

Interested in just running the project? Check out the GitHub repo.

Getting started

Pre-requisites

To get started, you'll need a Solana wallet private key as well as a Solana RPC URL if you want to execute code related to transferring Solana to and from your wallet. You will also need to create your own Kraken API key and secret if you don't already have one.

You will also need an OpenAI API key.

Note that we will need to export all of these to environment variables:

export OPENAI_API_KEY=<openai-api-key>
export SOLANA_RPC_URL=<solana-rpc-url>
export SOLANA_PRIVATE_KEY=<solana-private-key>
export KRAKEN_API_KEY=<kraken-api-key>
export KRAKEN_API_SECRET=<kraken-api-secret>

Project creation

To create your project, let's start by using cargo init:

cargo init kraken-trading-cli
cd kraken-trading-cli

Next, let's add our dependencies. You can use the one-liner below to add everything:

cargo add rig-onchain-kit tokio rig-core@0.9 anyhow \
kraken_rest_client rig-tool-macro serde serde-json thiserror --features \
rig-onchain-kit/solana,tokio/macros,serde/derive
cargo add rig-onchain-kit --no-default-features

What did we just add?

rig-onchain-kit: The companion crate to rig for executing actions onchain. We enable the solana feature as well as enabling no default features, which gets rid of some unnecessary features (for our application).
rig-core: The core crate for Rig.
tokio: An async runtime for Rust. The macros feature is enabled so we can use the #[tokio::main] macro.
kraken_rest_client: The unofficial community Kraken Rust SDK, which interacts with Kraken's REST API.
anyhow: A convenient error handling crate that can convert any error type into anyhow::Error.
rig-tool-macro: A macro crate for automatically generating tools from functions. Requires serde-json, serde with the derive feature and thiserror added as dependencies to work.
serde: A de/serialization crate, with the derive feature enabled for easy implementation. We don't use this crate directly but it is required by rig-tool-macro to work.
serde-json: A crate for converting serde::Deserialize and serde::Serialize types to and from raw JSON. Required by rig-tool-macro to work.
thiserror: A convenient error handling crate. Required by rig-tool-macro to work.

Building our project

So now that we're set up, let's dive into how you can build with the rig-onchain-kit crate. The crate itself already has a lot of tools which we'll be covering, but we'll also talk about how you can create your own tools and use them with the kit to make even more powerful AI agents!

Executing onchain actions with rig-onchain-kit

Executing actions onchain is quite simple. Using create_solana_agent creates an agent which has access to mostly every single tool that is already pre-existing in the crate (for the sake of simplicity). The signer context will then use an async closure to allow for scoped thread-local execution - see below for an example of swapping SOL to mocKJUP on Jupiter:

use rig_onchain_kit::agent::create_solana_agent;
use rig_onchain_kit::signer::SignerContext;
use rig_onchain_kit::signer::solana::LocalSolanaSigner;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let private_key = std::env::var("SOLANA_PRIVATE_KEY")?;
    let signer = LocalSolanaSigner::new(private_key);

    SignerContext::with_signer(Arc::new(signer), async {
        let agent = create_solana_agent();
        let response = agent.prompt("Please swap $5 worth of SOL from my wallet to mockJUP on Jupiter")?);
        println!("{}", response);
    });

    Ok(())
}

Simply by adding your Solana private key as well as the URL of the RPC you want to use as environment variables, you can now swap $5 worth of Solana from your wallet to $mockJUP (Jupiter Swap's testing token) by simply using cargo run. That's pretty much it - nothing else required.

Authenticating on Kraken's REST API

Kraken's REST API by default requires a number of different variables to authenticate manually. Thankfully, the kraken_rest_client crate does this for us already and we just need to provide our API key and secret! Note that creating an API key and secret requires you to sign in to Kraken via Kraken Pro.

Before we start work, let's make a function for quickly spinning up our HTTP client which will be used in all of our related tools that we'll be creating:

fn create_kraken_client() -> kraken_rest_client::Client {
    let kraken_api_key = env::var("KRAKEN_API_KEY").expect("KRAKEN_API_KEY env var to exist");

    let kraken_api_secret = env::var("KRAKEN_API_SECRET").expect("KRAKEN_API_SECRET env var to exist");

    Client::new(
        &kraken_api_key, 
        &kraken_api_secret
    )
}

Building your own tools for trading

By leveraging the #[tool] macro (from rig-tool-macro), we can generate tools that implement the rig::tool::Tool trait - which is important for us as we'll require that trait to be implemented for our types later on.

The first tool we'll make is a tool for price checking ticker pairs. The logic itself is relatively simple: just spin up a new Kraken client, then get a ticker and return the price.

#[tool(description = "
Attempts to check the price of a ticker pair on Kraken.

The ticker MUST be correct as this can lead to other tool calls you make being wrong if this is not right.

Input parameters:
- pair: The liquidity pair (for example: BTC to USD would be BTCUSD).
")]
async fn fetch_ticker_price_from_kraken(pair: String) -> Result<String, anyhow::Error> {
    let client = create_kraken_client();
    let res = client
        .get_ticker(&pair)
        // expire the order after an hour to prevent stale orders
        .send()
        .await
        .unwrap();

    // the response comes back as a HashMap<String, Ticker>
    // in this case we should simply iterate through the values
    // find the first value, then iterate on the inner array and find the first value - which is the current price
    let price = res.values().next().unwrap().a.iter().next().unwrap();

    Ok(price.to_string())
}

Next is the fun part: writing a tool to place buy orders! We need to explicitly tell the LLM that this function is dangerous and can lead to loss of funds if incorrect parameters have been set.

#[tool(description = "
Attempts to place a buy order for a token using the Kraken API given a token ticker, a volume and a price. The order will expire after one hour to ensure that the order doesn't go stale.

This function is dangerous, as it can lead to loss of funds if the price is incorrect. The price and volume MUST be correct as this can lead to a loss of funds.

Ensure you double check the address, the price & volume as well as the token ticker before you execute this function.

Input parameters:
- token_ticker: The token ticker.
- purchase_currency_ticker: The ticker of the token the user wants to make the purchase with.
- volume: How much of the token to buy.
- price: The price to buy the token at.
")]
async fn buy_token_from_kraken(
    token_ticker_to_buy: String,
    purchase_currency_ticker: String,
    volume: f64,
    price: f64,
) -> Result<String, anyhow::Error> {
    // .. some code here
}

Now we just need to implement the rest of our logic which means adding a limit order and then sending it and returning the . That's pretty much it.

use kraken_rest_client::{Client, PairName, OrderSide};
use std::env;

#[tool(description = "
Attempts to place a buy order for a token using the Kraken API given a token ticker, a volume and a price. The order will expire after one hour to ensure that the order doesn't go stale.

This function is dangerous, as it can lead to loss of funds if the price is incorrect. The public Solana key must be correct or the API call won't go through. The price and volume MUST be correct as this can lead to a loss of funds.

Ensure you double check the address, the price & volume as well as the token ticker before you execute this function.

Input parameters:
- token_ticker: The token ticker.
- purchase_currency_ticker: The ticker of the token the user wants to make the purchase with.
- volume: How much of the token to buy.
- price: The price to buy the token at.
")]
async fn buy_token_from_kraken(
    token_ticker_to_buy: String,
    purchase_currency_ticker: String,
    volume: f64,
    price: f64,
) -> Result<String, anyhow::Error> {
    let client = create_kraken_client();

    let pair = format!("{token_ticker_to_buy}{purchase_currency_ticker}");
    let resp = client
        .add_limit_order(
            &pair,
            OrderSide::Buy,
            volume.to_string().as_ref(),
            price.to_string().as_ref(),
        )
        // expire the order after an hour to prevent stale orders
        .expire_after(60 * 60)
        .userref(231421)
        .validate_only()
        .send()
        .await
        .unwrap();

    println!("{resp:?}");

    let txn_ids = if let Some(transactions) = resp.txid {
        transactions.join("\n")
    } else {
        "none".to_string()
    };

    Ok(format!(
        "Order successfully placed! Transactions IDs: {txn_ids}"
    ))
}

Note that the AddOrderRequest struct which gets generated from the add_limit_order request offers a validate_only() function. If you don't want to generate real buy orders, feel free to turn this on! An example of this can be found below:

    let resp = client
        .add_limit_order(
            &pair,
            OrderSide::Buy,
            volume.to_string().as_ref(),
            price.to_string().as_ref(),
        )
        .expire_after(60 * 60)
        .userref(231421)
        // we add it here before sending
        .validate_only()
        .send()
        .await
        .unwrap();

Note that the use of the Kraken API key means that the funds will be going into your Kraken account rather than your original Solana account. Thankfully, Kraken already has a method in their REST API to be able to withdraw funds to a pre-authorized wallet (so that you don't have to do this manually!) - let's have a look at the code:

use kraken_rest_client::{Client, PairName, OrderSide};
use std::env;

#[tool(description = "
Attempts to withdraw an asset from Kraken to a wallet that has been whitelisted on the given Kraken account. This method MUST be used if you are withdrawing from Kraken as due to being a centralized exchange, Kraken does not provide a private key for users to use.

Ensure you double check the key, the price & volume as well as the token ticker before you execute this function as the transaction won't go through (or will lead to loss of funds due to gas fees on the user's part).
")]
async fn withdraw_from_kraken(
    token_ticker: &str,
    wallet_key: &str,
    amount: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let client = create_kraken_client();

    let resp = client
    .withdraw_request(token_ticker, wallet_key, amount)
    // expire the order after an hour to prevent stale orders
    .send()
    .await?;
}

As you can see, pretty easy. This does however require that we have a pre-authorized wallet to withdraw to (with an associated key).

Integrating custom tools with your AI agent

Now that we've created some tools, let's add them to our agent. While create_solana_agent() does provide an agent with a variety of tools, we should ideally include as few tools as possible with our agent to prevent the LLM getting confused.

We also want to make our application competent: it should be fully capable of incorporating tool responses, as well as returning an answer from the LLM without the user needing to do anything.

Before we start, let's create a couple of helper methods. The first one will simply be a message that converts an ID and a content field into a UserContent::tool_result.

fn tool_result_from_text_response(id: impl Into<String>, content: String) -> UserContent {
    let content = OneOrMany::one(ToolResultContent::Text(content.into()));
    UserContent::tool_result(id, content)
}

The next function we'll implement is a function for converting a rig::message::Message into a String (ie, the text content that's within the message).

fn get_content_from_user_message(message: &Message) -> Result<String, anyhow::Error> {
    match message {
        Message::User { content } => match content.first() {
            UserContent::Text(text) => Ok(text.text),
            UserContent::ToolResult(ToolResult { content, .. }) => match content.first() {
                ToolResultContent::Text(text) => Ok(text.text),
                _ => Err(anyhow!(
                    "Expected text from tool result, got something else"
                )),
            },
            _ => {
                return Err(anyhow!(
                    "Expected text or a tool result, received something else"
                ));
            }
        },
        _ => Err(anyhow!(
            "Expected a user message, got an assistant message in prompt"
        )),
    }
}

The last helper function we'll write is a short function that simply returns tool definitions, using the prompt (and then returns a Vec<ToolDefinition> which we can use as part of our completion request):

async fn tool_definitions(message: &Message) -> Result<Vec<ToolDefinition>, anyhow::Error> {
    let prompt_as_string = get_content_from_user_message(&message)?;
    let buy_or_sell_token_on_kraken_tool = BuyOrSellTokenOnKraken
        .definition(prompt_as_string.clone())
        .await;
    let withdraw_from_kraken_tool = WithdrawFromKraken
        .definition(prompt_as_string.clone())
        .await;
    let transfer_sol_tool = TransferSol.definition(prompt_as_string.clone()).await;
    let fetch_ticker_price_from_kraken_tool = FetchTickerPriceFromKraken
        .definition(prompt_as_string.clone())
        .await;

    Ok(vec![
        buy_or_sell_token_on_kraken_tool,
        withdraw_from_kraken_tool,
        transfer_sol_tool,
        fetch_ticker_price_from_kraken_tool,
    ])
}

Now we can implement our response loop! This is a relatively long function, so buckle up. Here's what the function will do:

Initialise a toolset
Create a CompletioRequestBuilder and fill it out with all required items including tool definitions, prompts and chat history.
Send the completion request.
Match the latest message:
- If it's a text message, add prompt and assistant message to chat history then just return the text
- If it's a tool call, attempt to call the function (using the toolset we initialised)
- If the tool call results in an error, set the tool call error as the the prompt (explained below)
- If it was successful, push both the original prompt and assistant message to the chat history, then set the prompt as the tool call result (explained below)

async fn call_until_response<M: CompletionModel>(
    mut prompt: Message,
    model: &M,
    preamble: &str,
    chat_history: &mut Vec<Message>,
) -> Result<String, anyhow::Error> {
    let toolset = ToolSet::builder()
        .static_tool(BuyOrSellTokenOnKraken)
        .static_tool(WithdrawFromKraken)
        .static_tool(FetchTickerPriceFromKraken)
        .static_tool(TransferSol)
        .build();

    let tool_defs = tool_definitions(&prompt).await?;
    let request = CompletionRequestBuilder::new(model.clone(), prompt.to_owned())
        .preamble(preamble.to_owned())
        .messages(chat_history.clone())
        .temperature(0.0)
        .max_tokens(1024)
        .tools(tool_defs)
        .build();

    // call model
    let resp = model.completion(request).await.unwrap();

    // rest of code goes down here!
}

// keep calling tools until we get human readable answer from the model
match resp.choice.first() {
    AssistantContent::Text(text) => {
        let text = text.text;
        chat_history.push(prompt.clone().into());
        chat_history.push(Message::assistant(&text));
        return Ok(text);
    }
    AssistantContent::ToolCall(tool_call) => {
        // Call the tool
        let tool_response = toolset
            .call(
                &tool_call.function.name,
                tool_call.function.arguments.to_string(),
            )
            .await;

        let tool_response = match tool_response {
            Ok(res) => res,
            Err(e) => {
                prompt = Message::User {
                    content: OneOrMany::one(UserContent::ToolResult(ToolResult {
                        id: tool_call.id.to_string(),
                        content: OneOrMany::one(ToolResultContent::Text(rig::message::Text {
                            text: e.to_string(),
                        })),
                    })),
                };
                let res =
                    call_until_response(prompt.clone(), model, preamble, chat_history).await?;
                return Ok(res);
                }
            };
        }
    }
}

let tool_response_message =                tool_result_from_text_response(tool_call.id.clone(), tool_response);

let tool_call = OneOrMany::one(AssistantContent::ToolCall(tool_call));

// add tool call and response into chat history and continue the loop
chat_history.push(prompt.clone().into());
chat_history.push(Message::Assistant { content: tool_call });

println!("{tool_response_message:?}");

let tool_result_message = Message::User {
    content: OneOrMany::one(tool_response_message),
};

prompt = tool_result_message;

let res = call_until_response(prompt.clone(), model, preamble, chat_history).await?;

Ok(res)

Completing the loop

Finally, we need to adjust our main function so that we have a full application. Because we already created all of the previous code, we can keep our main function relatively short:

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    let preamble = format!(
        "You are a bot designed to assist users with trading on Kraken and withdrawing them to a pre-whitelisted wallet.
        Users may also ask you to send coins from their own wallet to a given Kraken deposit address.

        Note that when placing an order using the buy_or_sell_token_on_kraken tool call, no transaction IDs may be returned.
        If this is the case, remind the user to turn off 'validation only' as this is primarily for the cases of debugging."
    );

    let openai_client = rig::providers::openai::Client::from_env();
    let model = openai_client.completion_model("gpt-4o");

    let mut chat_history = Vec::new();

    let private_key =
        std::env::var("SOLANA_PRIVATE_KEY").expect("SOLANA_PRIVATE_KEY env var to exist");

    println!("Hi! How can I help you today?");
    loop {
        let prompt = take_user_input();

        let signer = LocalSolanaSigner::new(private_key.clone());
        let res = SignerContext::with_signer(Arc::new(signer), async {
            call_until_response(prompt.into(), &model, &preamble, &mut chat_history).await
        })
        .await?;

        println!("{res}");
    }
}

Now all we need to do is use cargo run and it'll just work!

Don't forget to set your environment variables from the Prerequisites section if you haven't already.

Finishing up

Thanks for reading! Hopefully after reading this article, you have a better understanding of how we can leverage the Rig Onchain Kit to build some really cool stuff.

For additional Rig resources and community engagement:

Check out more examples in our gallery.
Contribute or report issues on our GitHub.
Join discussions in our Discord community!

Creating AI Memories using Rig & MongoDB

Josh Mo — Wed, 12 Mar 2025 20:12:43 +0000

Introduction

Artificial intelligence (AI) applications that engage in conversations often benefit from memory—storing past interactions to provide more contextually aware responses. Memory in AI can be categorized into short-term and long-term memory. Short-term memory involves tracking recent exchanges within a single session, while long-term memory enables AI to recall past interactions across different sessions.

In this article, we’ll build a CLI tool using the Rig AI framework and MongoDB for retrieval-augmented generation (RAG). This tool will store summarized conversations in a database and retrieve them when needed, enabling the AI to maintain context over time.

What is Memory in AI?

Memory in AI allows systems to retain and utilize past interactions to improve responses. There are two key types of memory:

Short-term memory: Captures recent exchanges in an ongoing conversation. Typically stored in-memory, it resets after a session ends. An example of this might be a chat history.
Long-term memory: Stores relevant conversation history and facts persistently, allowing recall across sessions. Usually, this is handled using vector databases or structured databases like MongoDB.

By combining Rig's AI processing capabilities with MongoDB’s persistence, we can build an AI assistant that not only responds intelligently in the moment but also remembers past interactions.

Getting Started

Pre-requisites

Before we get started, you'll need a MongoDB instance for storing long-term memory. You'll need to ensure that you sign up for
MongoDB Atlas (they provide one free cluster which we'll be using!), and create a database (called knowledgebase) in which we'll put a collection (called memories).

On that collection, we need an index that looks like the following (using the JSON editor):

{
    "fields": [
        {
            "numDimensions": 1536,
            "path": "embedding",
            "similarity": "cosine",
            "type": "vector"
        }
    ]
}

You'll also want an OpenAI API key as we'll be using their embedding and chat completion models.

Creating our new project

Before we begin, ensure you have Rust installed. Then, create a new Rust project:

cargo new ai-memory-cli
cd ai-memory-cli

This project will serve as the foundation for our AI memory system.

Adding Dependencies

To integrate Rig-Core and MongoDB, we need to add the necessary dependencies. You can copy the command below for easy installation:

cargo add rig-core rig-mongodb mongodb tokio nanoid serde serde-json chrono -F rig-core/derive,tokio/macros,serde/derive

What did we just add?

rig-core: The core rig library required to use Rig. We add the derive feature so we can easily implement the Embed trait for structs we want to add as documents into our vector store.
rig-mongodb: rig's MongoDB integration.
mongodb: The MongoDB Rust SDK.
tokio: The Tokio runtime, used for ergonomic access to async Rust. We use the macros feature for access to the #[tokio::main] macro, which allows for convenient async runtime setup.
nanoid: A quick, compact and easy way to get random IDs without using UUIDs (via nanoid!(<id-length>)).
serde: A de/serialization library. We attach the derive feature to make it easy to implement serialization and deserialization for our structs.
serde-json: serde's JSON crate.
chrono: A crate for working with time. We use this crate in this project to get UNIX timestamps.

Taking User Input

Taking user input is quite easy. We simply create a mutable empty string, prompt the user to input something (and flush the output so it actually prints instead of potentially staying in buffer) and return the user input from the function.

fn take_user_input() -> String {
    let mut string = String::new();

    print!("> ");
    let _ = stdout().flush();

    stdin().read_line(&mut string).unwrap();

    string
}

From this, we can create a simple async fn main that instantiates an OpenAI client, then loops infinitely (until it either crashes or the user presses ctrl-c) and answers user questions.

use rig::providers::openai::Client;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let openai_client = Client::from_env();

    loop {
        let query = take_user_input();

        let agent = openai_client
            .agent("gpt-4o")
            .preamble("You are a helpful agent.")
            .build();

        let response = agent.prompt(query.as_ref()).await?;

        println!("{response}");
    }
}

However, we can do much better than this. Let's add some short term memory to our AI so it can more competently answer user questions.

Adding short term memory

Relative to long term memory, short term memory is pretty easy to add. In terms of practical implementation, you can do something as simple as storing all the messages in a Vec<Message> and then sending that list of messages in your prompt. See below:

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut messages = Vec::new();

    let openai_client = Client::from_env();
    let embedding_model = openai_client.embedding_model("text-embedding-ada-002");

    loop {
        let agent = openai_client
            .agent("gpt-4o")
            .preamble("You are a helpful agent.")
            .build();

        let response = agent.chat(query.as_ref(), messages.clone()).await.unwrap();

        println!("{response}");

        messages.push(Message::user(query));
        messages.push(Message::assistant(response));
    }
}

Adding long term memory

Now for the hard part: adding long term memory! There's several different ways we can do this. A naive impementation might try to summarize a conversation on every turn: however, this is extremely resource intensive in terms of tokens (and potentially GPU or CPU usage if you're running your model locally). Typically, a system might have several different criterias for adding memory (note that this list is non-exhaustive):

A new topic
A given checkpoint in the conversation (ie, send the conversation to another model to check if there is anything significant to keep, if there is, store it)
A given conversation chunk size

We'll be implementing our long term memory system to try and summarize conversation chunks every 3 turns (so 6 messages).

Summarizing a Conversation and Adding It to the Vector Store

To efficiently store memory, we first summarize conversations before adding them to the database. This prevents excessive storage usage and allows the AI to retrieve meaningful insights instead of raw message logs. The system first extracts relevant information by identifying key points from the discussion. Once the essential details are identified, we generate a concise summary using an LLM that captures the core context. This structured summary is then stored in MongoDB, allowing for efficient retrieval and future use in maintaining AI memory.

To get started, we'll define a struct that holds the MongoDB object ID, the conversation ID (which will be randomly generated later on), the memory itself and the UNIX timestamp at which the memory was created. Note that because of the way ID is handled in Mongo, we create a custom function that can deserialize either as a string or an object.

use nanoid::nanoid;
use chrono::Utc;

use serde::{Deserialize, Deserializer};
use serde_json::Value;



// Shape of data that needs to be RAG'ed.
// The definition field will be used to generate embeddings.
#[derive(Embed, Clone, Deserialize, Debug)]
struct Memory {
    #[serde(rename = "_id", deserialize_with = "deserialize_object_id")]
    id: String,
    conversation_id: String,
    #[embed]
    memory: String,
    timestamp_created: usize,
}

impl Memory {
    fn new(conversation_id: &str, memory: String) -> Self {
        let id = nanoid!(10);

        let conversation_id = conversation_id.to_string();
        let timestamp_created = Utc::now().timestamp();

        Self {
            id,
            conversation_id,
            memory,
            timestamp_created,
        }
    }
}

fn deserialize_object_id<'de, D>(deserializer: D) -> Result<String, D::Error>
where
    D: Deserializer<'de>,
{
    let value = Value::deserialize(deserializer)?;
    match value {
        Value::String(s) => Ok(s),
        Value::Object(map) => {
            if let Some(Value::String(oid)) = map.get("$oid") {
                Ok(oid.to_string())
            } else {
                Err(serde::de::Error::custom(
                    "Expected $oid field with string value",
                ))
            }
        }
        _ => Err(serde::de::Error::custom(
            "Expected string or object with $oid field",
        )),
    }
}

We'll also create a function for connecting to MongoDB. Note that this requires us to set our MongoDB connection string as an environment variable using MONGODB_CONNECTION_STRING:

use mongodb::{Client as MongoClient, Collection};
use mongodb::bson;
use mongodb::options::ClientOptions;

async fn connect_to_mongodb() -> Collection<bson::Document> {
    // Initialize MongoDB client
    let mongodb_connection_string = std::env::var("MONGODB_CONNECTION_STRING").expect("MONGODB_CONNECTION_STRING not set");
    let options = ClientOptions::parse(mongodb_connection_string)
        .await
        .expect("MongoDB connection string should be valid");

    let mongodb_client =
        MongoClient::with_options(options).expect("MongoDB client options should be valid");

    // Initialize MongoDB vector store
    let collection: Collection<bson::Document> = mongodb_client
        .database("knowledgebase")
        .collection("memories");

    collection
}

Next, we'll implement our summarize_chunks function which will ensure there are enough new messages, send the messages to an LLM and ask it to summarize the messages, then store the messages in MongoDB.

At the start of the function, we try to get the last 6 messages by reversing the iterator for our Vec<Message> (so it starts at the last pushed message) then takes 6 messages. If there's less than 6, we return an error stating there are not enough messages to summarize.

use rig::providers::openai::EmbeddingModel;

async fn summarize_chunks(
    openai_client: &Client,
    mongo: &Collection<bson::Document>,
    messages: Vec<Message>,
    conversation_id: &str,
    embedding_model: &EmbeddingModel,
) -> Result<(), Box<dyn std::error::Error>> {
    let messages = messages.into_iter().rev().take(6).collect::<Vec<Message>>();

    // Ensure there are a reasonable number of turns that have passed before trying to summarize the message.
    // In this case, we want to summarize every 3 turns - so we need to ensure the remainder after dividing by 6 is 0
    // (because 3 turns = 6 messages)
    // If it doesn't meet the required criteria, just return early
    if messages.len() % 6 != 0 {
        return Err("Not enough new messages to summarize".into());
    }

    // rest of code goes here
    Ok(())
}

To be able to get the memories back in the format that we can use easily, we'll set up an extractor asking the LLM to summarize an inputted conversation chunk into a list of factoids, using the JSON schema provided. In this case it's a Vec<String>, but you could easily extend this if you wanted to.

For embedding our memories into Mongo, we will additionally enrich the data by using Memory::new() to create a Vec<Memory> out of the strings.

    let agent = openai_client.extractor::<Vec<String>>("gpt-4o").preamble("Please summarize the inputted conversation chunk by the user into a list of factoids, using the JSON schema provided.").build();

    let messages_as_string = serde_json::to_string_pretty(&messages).unwrap();

    let res = agent.extract(&messages_as_string).await.unwrap();

    let response_as_memory_entries = res
        .into_iter()
        .map(|x| Memory::new(&conversation_id, x))
        .collect::<Vec<Memory>>();

Finally, we need to use EmbeddingsBuilder to create an iterator that will output both embeddings and documents. We'll use this to create our final Vec<mongodb::bson::Document> (an array of MongoDB BSON documents) which we'll then finally insert into Mongo. Done!

use rig::embeddings::EmbeddingsBuilder;
use mongodb::bson::doc;

    let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
        .documents(response_as_memory_entries)?
        .build()
        .await?;

    let mongo_documents = embeddings
        .iter()
        .map(
            |(
                Memory {
                    id,
                    conversation_id,
                    memory,
                    timestamp_created,
                },
                embedding,
            )| {
                doc! {
                    "id": id.clone(),
                    "conversation_id": conversation_id.clone(),
                    "memory": memory.clone(),
                    "timestamp_created": timestamp_created.clone(),
                    "embedding": embedding.first().vec.clone(),
                }
            },
        )
        .collect::<Vec<_>>();

    mongo.insert_many(mongo_documents).await?;

The full code for the summarize_chunks function can found below.

async fn summarize_chunks(
    openai_client: &Client,
    mongo: &Collection<bson::Document>,
    messages: Vec<Message>,
    conversation_id: &str,
    embedding_model: &EmbeddingModel,
) -> Result<(), Box<dyn std::error::Error>> {
    let messages = messages.into_iter().rev().take(6).collect::<Vec<Message>>();

    // Ensure there are a reasonable number of turns that have passed before trying to summarize the message.
    // In this case, we want to summarize every 3 turns - so we need to ensure the remainder after dividing by 6 is 0
    // (because 3 turns = 6 messages)
    // If it doesn't meet the required criteria, just return early
    if messages.len() % 6 != 0 {
        return Err("Not enough new messages to summarize".into());
    }

    let agent = openai_client.extractor::<Vec<String>>("gpt-4o").preamble("Please summarize the inputted conversation chunk by the user into a list of factoids, using the JSON schema provided.").build();

    let messages_as_string = serde_json::to_string_pretty(&messages).unwrap();

    let res = agent.extract(&messages_as_string).await.unwrap();

    let response_as_memory_entries = res
        .into_iter()
        .map(|x| Memory::new(&conversation_id, x))
        .collect::<Vec<Memory>>();

    let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
        .documents(response_as_memory_entries)?
        .build()
        .await?;

    let mongo_documents = embeddings
        .iter()
        .map(
            |(
                Memory {
                    id,
                    conversation_id,
                    memory,
                    timestamp_created,
                },
                embedding,
            )| {
                doc! {
                    "id": id.clone(),
                    "conversation_id": conversation_id.clone(),
                    "memory": memory.clone(),
                    "timestamp_created": timestamp_created.clone(),
                    "embedding": embedding.first().vec.clone(),
                }
            },
        )
        .collect::<Vec<_>>();

    mongo.insert_many(mongo_documents).await?;

    Ok(())
}

To use this function, all we need to do is to use it at the end of the loop:

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // .. setup goes here

    loop {
        let query = take_user_input();

        let agent = openai_client
            .agent("gpt-4o")
            .preamble("You are a helpful agent.")
            .append_preamble(&additional_input)
            .build();

        let response = agent.chat(query.as_ref(), messages.clone()).await.unwrap();

        println!("{response}");

        messages.push(Message::user(query));
        messages.push(Message::assistant(response));

        if summarize_chunks(
            &openai_client,
            &mongo,
            messages.clone(),
            &conversation_id,
            &embedding_model,
        )
        .await
        .is_ok()
        {
            println!("Saved new memories.");
        };
    }
}

Retrieving and Using Memories

To enable AI memory retrieval we need to search the vector index for relevant past conversations and return relevant summarized conversation chunks.

Search the database for relevant past conversations based on the current query.
Retrieve stored summaries and use them as context for the next response.

To get started, we will create an in-memory vector index. We'll then query the index for data in the shape of the struct we want to use (in this case Memory), and then join them all up into a single string that splits the memories by newlines for LLM ingestion.

use rig::vector_store::VectorStoreIndex;
use rig_mongodb::{MongoDbVectorIndex, SearchParams};

async fn get_relevant_memories(
    mongo: &Collection<bson::Document>,
    query: &str,
    model: &EmbeddingModel,
) -> Result<Option<Vec<String>>, Box<dyn std::error::Error>> {
    // Create a vector index on our vector store.
    // Note: a vector index called "vector_index" must exist on the MongoDB collection you are querying.
    // IMPORTANT: Reuse the same model that was used to generate the embeddings
    let index = MongoDbVectorIndex::new(
        mongo.clone(),
        model.clone(),
        "vector_index",
        SearchParams::new(),
    )
    .await?;

    // Query the index
    let results = index.top_n::<Memory>(query, 5).await?;

    let res_as_vec_strings = results
        .into_iter()
        .map(|(_, embedded_string, _)| embedded_string)
        .collect::<Vec<String>>();

    if res_as_vec_strings.is_empty() {
        return Ok(None);
    }

    Ok(Some(res_as_vec_strings))
}

Once we've ran the function, we declare a new variable that's either a newline-joined list of memories, or an empty string if there are no relevant memories (or the database just doesn't have any).

let additional_input = if let Some(messages) =
    get_relevant_memories(&mongo, &query, &embedding_model).await?
    {
        format!(
            "\n\nRelevant memorires from previous conversations: {}",
            messages.join("\n")
        )

    } else {
        String::new()
    };

let agent = openai_client
    .agent("gpt-4o")
    .preamble("You are a helpful agent.")
    .append_preamble(&additional_input)
    .build();

Trying it out

Your final async fn main should look like this:

use chrono::Utc;
use mongodb::bson::{self, doc};
use mongodb::options::ClientOptions;
use mongodb::{Client as MongoClient, Collection};
use nanoid::nanoid;
use rig::embeddings::EmbeddingsBuilder;
use rig::providers::openai::EmbeddingModel;
use rig::vector_store::VectorStoreIndex;
use rig_mongodb::{MongoDbVectorIndex, SearchParams};
use serde_json::Value;
use std::io::{Write, stdin, stdout};

use rig::{Embed, completion::Chat, message::Message, providers::openai::Client};
use serde::{Deserialize, Deserializer};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut messages = Vec::new();
    let mongo = connect_to_mongodb().await;

    let conversation_id = nanoid!(6);

    let openai_client = Client::from_env();
    let embedding_model = openai_client.embedding_model("text-embedding-ada-002");

    loop {
        let query = take_user_input();

        let additional_input = if let Some(messages) =
            get_relevant_memories(&mongo, &query, &embedding_model).await?
        {
            format!(
                "\n\nRelevant memorires from previous conversations: {}",
                messages.join("\n")
            )
        } else {
            String::new()
        };

        let agent = openai_client
            .agent("gpt-4o")
            .preamble("You are a helpful agent.")
            .append_preamble(&additional_input)
            .build();

        let response = agent.chat(query.as_ref(), messages.clone()).await.unwrap();

        println!("{response}");

        messages.push(Message::user(query));
        messages.push(Message::assistant(response));

        if summarize_chunks(
            &openai_client,
            &mongo,
            messages.clone(),
            &conversation_id,
            &embedding_model,
        )
        .await
        .is_ok()
        {
            println!("Saved new memories.");
        };
    }
}

To use it, make sure you have your OpenAI API key and MongoDB URLs added as OPENAI_API_KEY and MONGODB_CONNECTION_STRING as environment variables, then simply use cargo run and try it out!

Conclusion

By combining Rig’s AI capabilities with MongoDB’s robust storage, we’ve created a system that allows an AI assistant to retain and recall relevant past conversations. This setup enables both short-term and long-term memory, making interactions feel more natural and context-aware.

For additional Rig resources and community engagement:

Check out more examples in our gallery.
Contribute or report issues on our GitHub.
Join discussions in our Discord community!

Local Embeddings with Fastembed, Rig & Rust

Josh Mo — Wed, 05 Mar 2025 15:52:51 +0000

Introduction

Hello world! Today we're going to be talking about embeddings. Used in recommendation systems, semantic search and Retrieval Augmented Generation (RAG) pipelines, embeddings are the core of being able to easily retrieve contextual information for a given query. This result can then be plugged into an LLM to retrieve a result - or alternatively, the resulting embedding can be stored in a database.

In this article, we're going to get into why you should use local embeddings vs. embeddings from managed models, as well as the technical implementation for how to create and use your local embeddings.

Interested in checking out the code for this example? Have a look at the example on our GitHub repo.

Why local embeddings?

Embeddings, while an important part of the RAG ecosystem, can be relatively expensive to create, especially as you start storing embeddings from bigger and bigger amounts of documents. Some simple use cases where this can happen might be:

If you're running a directory that aggregates a list of businesses
If you have a large amount of products that you sell (a supermarket or large grocery store being a good example of this)
You're storing user-submitted uploads

For all the convenience that managed models provide, there are also data concerns. If you're working on a commercial project, you may not want to hand over your data to OpenAI - especially in industries where data privacy is absolutely paramount. In this case, a local model has exactly what you need: text embedding in the knowledge that your data will be staying totally on your infrastructure.

Building your project

Without further do, let's get to building! Before we start, don't forget to have the Rust programming language installed.

Getting started

To get started, we'll spin up a new project:

cargo init local-embeddings-rig
cd local-embeddings-rig

Next, we'll need to add our dependencies:

cargo add rig-core serde rig-fastembed -F rig-core/derive,serde/derive

Note that this one-line command adds the following dependencies:

rig-core: The rig crate which houses all the traits (interfaces) we need to create the embeddings. We've added the derive feature to get the Embed derive macro which will make it much easier to embed whatever we need to.
rig-fastembed: The fastembed-rs integration for rig.
serde: A de/serialization library. We've added the derive feature for the Serialize and Deserialize macros, both of which are extremely helpful for serializing and deserializing simple structs without hassle.

Let's get embedding with Rig

Next, we'll go into main.rs and define a struct which will serve as the payload we want to associate with our struct.

use rig::Embed;
use rig_fastembed::FastembedModel;
use serde::{Deserialize, Serialize};

// Shape of data that needs to be RAG'ed.
// The definition field will be used to generate embeddings.
#[derive(Embed, Clone, Deserialize, Debug, Serialize, Eq, PartialEq, Default)]
struct WordDefinition {
    id: String,
    word: String,
    #[embed]
    definitions: Vec<String>,
}

Once done, we'll define our main function (which will use the Tokio runtime macro to allow for easy async) as well as a new rig-fastembed::Client which we'll be using to create our embedding models from.

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // Create OpenAI client
    let fastembed_client = rig_fastembed::Client::new();

    let embedding_model = fastembed_client.embedding_model(&FastembedModel::AllMiniLML6V2Q);


    // .. some more code goes here

    Ok(())
}

The rig-core library exposes an EmbeddingsBuilder struct which we can use to store our documents in. Note that the field which we've annotated previously with #[embed] will be used as the text source for embedding, but if you want to change this, you can create custom behavior by implementing rig::Embed manually.

let documents = vec![
        WordDefinition {
            id: "doc0".to_string(),
            word: "flurbo".to_string(),
            definitions: vec![
                "A green alien that lives on cold planets.".to_string(),
                "A fictional digital currency that originated in the animated series Rick and Morty.".to_string()
            ]
        },
        WordDefinition {
            id: "doc1".to_string(),
            word: "glarb-glarb".to_string(),
            definitions: vec![
                "An ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.".to_string(),
                    "A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy.".to_string()
            ]
        },
        WordDefinition {
            id: "doc2".to_string(),
            word: "linglingdong".to_string(),
            definitions: vec![
                "A term used by inhabitants of the sombrero galaxy to describe humans.".to_string(),
                "A rare, mystical instrument crafted by the ancient monks of the Nebulon Mountain Ranges on the planet Quarm.".to_string()
            ]
        },
    ];

let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
    documents(documents)?
    .build()
    .await?;

And that's it! The embeddings variable will now be an iterator that outputs (embedding, T). You can use this however you want. For example:

for (embedding, document) in embeddings {
    println!("{embedding}");
    println!("{document:?}");
}

Using embeddings with vector stores in Rig

Next, we create our vector store and add the embeddings to it. Note that embeddings actually returns an (embedding, doc) tuple. The from_documents_with_id_f takes a tuple, as well as a closure (an anonymous function) for the document ID. Additionally we also have other integrations for vector stores if you'd like to use production-ready alternatives such as Qdrant, LanceDB and MongoDB. Check them out on our GitHub repo.

// Create vector store with the embeddings
let vector_store = InMemoryVectorStore::from_documents_with_id_f(embeddings, |doc| doc.id.clone());

// Create vector store index
let index = vector_store.index(embedding_model);

Finally, to use our vector store we use top_n to grab the results, annotating the function with the type we want to return from the function.

Note that here the returned type of the results variable is Vec<(f64, String, T)> (where T in this case is the WordDefinition struct). However, because we don't necessarily need to know what type it is (or it's obvious what type it is), we've left it out in the type annotation (Vec<_>).

let results = index
    .top_n::<WordDefinition>("I need to buy something in a fictional universe. What type of money can I use for this?", 1)
    .await?
    .into_iter()
    .map(|(score, id, doc)| (score, id, doc.word))
    .collect::<Vec<_>>();

    println!("Results: {:?}", results);

let id_results = index
    .top_n_ids("I need to buy something in a fictional universe. What type of money can I use for this?", 1)
    .await?
    .into_iter()
    .collect::<Vec<_>>();

println!("ID results: {:?}", id_results);

So what's next?

Thanks for reading! If you have any thoughts, feel free to let us know in the comments.

For additional Rig resources and community engagement:

Check out more examples in our gallery.
Contribute or report issues on our GitHub.
Join discussions in our Discord community!

Implementing Design Patterns for Agentic AI with Rig & Rust

Josh Mo — Wed, 26 Feb 2025 15:23:23 +0000

In this tutorial, we'll be re-implementing some of the common design patterns you might find in agentic AI, using Rig. These examples will primarily be following Anthropic's article on building effective agents - so you can expect to see these commonly being used whatever kind of agent you're building!

By the end of this article, you'll have a good understanding of how to create the following agentic AI design patterns in Rust using the rig framework:

Prompt chaining
Routing
Parallelization
Orchestrator-worker
Evaluator-optimizer
Autonomous agent

Note that if you want to find any of the full code examples, you can find them in the rig repo.

Prompt chaining

Prompt chaining simply decomposes a task into a list of smaller steps, where we make LLM calls one after the other, piping the result of the previous LLM call into the next (or the prompt, if it's the first LLM call). You can additionally add gates between LLM calls to exit promptly if the given answer from the LLM is either unsatisfactory or has escaped the bounds of what the answer should be.

To get started, we'll create our OpenAI client then create our agents:

use std::env;

use rig::providers::openai::Client;

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // Create OpenAI client
    let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
    let openai_client = Client::new(&openai_api_key);

    let rng_agent = openai_client.agent("gpt-4")
        .preamble("
            You are a random number generator designed to only either output a single whole integer that is 0 or 1. Only return the number.
        ")
        .build();

    let adder_agent = openai_client.agent("gpt-4")
        .preamble("
            You are a mathematician who adds 1000 to every number passed into the context, except if the number is 0 - in which case don't add anything. Only return the number.
        ")
        .build();

    // .. more code down here!
}

Next, we'll create our pipeline then simply put the agents in our pipeline sequentially after each other. Simple!

use rig::pipeline::{self, Op};

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // .. previoiusly entered code goes here!

    let chain = pipeline::new()
        // Generate a whole number that is either 0 and 1
        .prompt(rng_agent)
        .map(|x| x.unwrap())
        .prompt(adder_agent);

    // Prompt the agent and print the response
    let response = chain
        .call("Please generate a single whole integer that is 0 or 1".to_string())
        .await;

    println!("Pipeline result: {response:?}");

    Ok(())
}

Although this is a relatively basic technique, it's very common and can be applied to a wide variety of use cases.

Prompt routing

Prompt routing is a way to mitigate prompt injection by categorising the user statement into a number of topics. To illustrate, our example will categorise a user's statement by either sheep, cows or dogs and then changing the prompt based on the given topic.

As before, we'll create our agents:

use std::env;
use rig::providers::openai::Client;

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // Create OpenAI client
    let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
    let openai_client = Client::new(&openai_api_key);

    // Note that you can also create your own semantic router for this
    // that uses a vector store under the hood
    let animal_agent = openai_client.agent("gpt-4")
        .preamble("
            Your role is to categorise the user's statement using the following values: [sheep, cow, dog]

            Return only the value.
        ")
        .build();

    let default_agent = openai_client.agent("gpt-4").build();

    // .. more code to go down here!
}

Note that while we use a default agent as the final step to illustrate how to use this pattern, you can essentially put whatever you want there - be it an AI agent with tools, another mapping function or something else.

use rig::pipeline::{self, Op, TryOp};

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // previous code goes here

    let chain = pipeline::new()
        // Use our classifier agent to classify the agent under a number of fixed topics
        .prompt(animal_agent)
        // Change the prompt depending on the output from the prompt
        .map_ok(|x: String| match x.trim() {
            "cow" => Ok("Tell me a fact about the United States of America.".to_string()),
            "sheep" => Ok("Calculate 5+5 for me. Return only the number.".to_string()),
            "dog" => Ok("Write me a poem about cashews".to_string()),
            message => Err(format!("Could not process - received category: {message}")),
        })
        .map(|x| x.unwrap().unwrap())
        // Send the prompt back into another agent with no pre-amble
        .prompt(default_agent);

    // Prompt the agent and print the response
    let response = chain.try_call("Sheep can self-medicate").await?;

    println!("Pipeline result: {response:?}");

    Ok(())
}

This technique is quite helpful as it allows you to essentially categorise your prompt by topic then create a response or do something else based on the topic.

Parallelization

In some workloads, you may want to parallelize your API calls by sending them at exactly the same time, return the result from all of them and display the results.

Below is an example of how you can create a pipeline using Rig and parallelize your API calls extracting the information we want as a JSON object, then map the three results into a single result that we then present to the user.

First, let's declare our JSON output - we can do this by declaring a struct that derives schemars::JsonSchema which allows a JSON schema to be outputted for a struct.

use rust::{Deserialize, Serialize};
use schemars::JsonSchema;

#[derive(Deserialize, JsonSchema, Serialize)]
struct DocumentScore {
    /// The score of the document
    score: f32,
}

Next, let's define our agents. In this example, we define three different agents to get three different types of sentiment for a given statement (note these are arbitrary characteristics - you can use basically whatever you want):

How manipulative it sounds
How depressive it sounds
How intelligent it sounds

use std::env;
use rig::providers::openai::Client;

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // Create OpenAI client
    let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
    let openai_client = Client::new(&openai_api_key);

    let manipulation_agent = openai_client
        .extractor::<DocumentScore>("gpt-4")
        .preamble(
            "
            Your role is to score a user's statement on how manipulative it sounds between 0 and 1.
        ",
        )
        .build();

    let depression_agent = openai_client
        .extractor::<DocumentScore>("gpt-4")
        .preamble(
            "
            Your role is to score a user's statement on how depressive it sounds between 0 and 1.
        ",
        )
        .build();

    let intelligent_agent = openai_client
        .extractor::<DocumentScore>("gpt-4")
        .preamble(
            "
            Your role is to score a user's statement on how intelligent it sounds between 0 and 1.
        ",
        )
        .build();

    // .. more code down here
}

Once that's done, we'll create our pipeline and run it. We use the parallel!() macro, which allows us to be able to chain multiple function calls in a single operation and allows us to then map the results as a tuple of results in the .map() function. Note that we also use passthrough() which allows us to simply pass the original prompt through to the next step in the chain.

use rig::pipeline::agent_ops::extract;
use rig::{
    parallel,
    pipeline::{self, passthrough, Op},.
};

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // .. previous code goes here!

    let chain = pipeline::new()
        .chain(parallel!(
            passthrough(),
            extract(manipulation_agent),
            extract(depression_agent),
            extract(intelligent_agent)
        ))
        .map(|(statement, manip_score, dep_score, int_score)| {
            format!(
                "
                Original statement: {statement}
                Manipulation sentiment score: {}
                Depression sentiment score: {}
                Intelligence sentiment score: {}
                ",
                manip_score.unwrap().score,
                dep_score.unwrap().score,
                int_score.unwrap().score
            )
        });

    // Prompt the agent and print the response
    let response = chain
        .call("I hate swimming. The water always gets in my eyes.")
        .await;

    println!("Pipeline run: {response:?}");

    Ok(())
}

Orchestrator-worker

Many AI-assisted production workloads use a central orchestrator to manage the execution of agents across a pipeline. The orchestrator decides which agent to trigger based on factors like the current task’s progress, complexity, priority, and resource availability. For instance, it may start with an agent handling data preprocessing, then move on to model training, and finally route the data for post-processing or analysis, depending on the results.

The orchestrator can also monitor system performance, ensuring that tasks are executed efficiently and that delays or failures in one part of the pipeline don’t affect the rest of the workflow. In production, it ideally should dynamically adjust the pipeline based on real-time metrics, scaling or optimizing agent operations as needed to maintain performance and avoid bottlenecks.

To get started, we'll define our structs for extracting JSON output from a response:

use schemars::JsonSchema;

#[derive(serde::Deserialize, JsonSchema, serde::Serialize, Debug)]
struct Specification {
    tasks: Vec<Task>,
}

#[derive(serde::Deserialize, JsonSchema, serde::Serialize, Debug)]
struct Task {
    original_task: String,
    style: String,
    guidelines: String,
}

#[derive(serde::Deserialize, JsonSchema, serde::Serialize, Debug)]
struct TaskResults {
    style: String,
    response: String,
}

Next, we'll create our agents. The premable (system prompt) is relatively long for the classify agent, so we'll give it its own constant value to assist with readability.

use std::env;
use rig::providers::openai::Client;

const CLASSIFY_PREAMBLE: &str = "Analyze the given task and break it down into 2-3 distinct approaches.

            Provide an Analysis:
            Explain your understanding of the task and which variations would be valuable.
            Focus on how each approach serves different aspects of the task.

            Along with the analysis, provide 2-3 approaches to tackle the task, each with a brief description:

            Formal style: Write technically and precisely, focusing on detailed specifications
            Conversational style: Write in a friendly and engaging way that connects with the reader
            Hybrid style: Tell a story that includes technical details, combining emotional elements with specifications

            Return only JSON output.";

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // Create OpenAI client
    let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
    let openai_client = Client::new(&openai_api_key);

    // create an agent that gives a specification for a task
    // given a couple of differnet provided approaches
    let classify_agent = openai_client.extractor::<Specification>("gpt-4")
        .preamble(CLASSIFY_PREAMBLE)
        .build();

    // using the previous agent, we'll extract a specification from this prompt
    let specification = classify_agent.extract("
        Write a product description for a new eco-friendly water bottle.
        The target_audience is environmentally conscious millennials and key product features are: plastic-free, insulated, lifetime warranty
        ").await.unwrap();

    // finally, we'll generate content based on the generated response
    let content_agent = openai_client
        .extractor::<TaskResults>("gpt-4")
        .preamble(
            "
                Generate content based on the original task, style, and guidelines.

                Return only your response and the style you used as a JSON object.
                ",
        )
        .build();

    // .. more code goes down here
}

For each task result, we want to try and create some content using each of the style, the task given and the guidelines. We'll then collect this into an array of task results and use an LLM to figure out which one is the most appropriate for our task.

use std::env;
use rig::providers::openai::Client;

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // .. previous code goes here

    let mut vec: Vec<TaskResults> = Vec::new();

    for task in specification.tasks {
        let results = content_agent
            .extract(&format!(
                "
            Task: {},
            Style: {},
            Guidelines: {}
            ",
                task.original_task, task.style, task.guidelines
            ))
            .await
            .unwrap();

        vec.push(results);
    }

    // .. more code goes here
}

Finally now that we have the content results that we want, we use an LLM-as-judge to help us figure out what the best piece of writing was, and to give us its reasoning. It'll return the style it has chosen as well as the corresponding material.

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // .. previous code goes here
    let judge_agent = openai_client
        .extractor::<Specification>("gpt-4")
        .preamble(
            "
            Analyze the given written materials and decide the best one, giving your reasoning.

            Return the style as well as the corresponding material you have chosen as a JSON object.
            ",
        )
        .build();

    let task_results_raw_json = serde_json::to_string_pretty(&vec).unwrap();

    let results = judge_agent.extract(&task_results_raw_json).await.unwrap();

    println!("Results: {results:?}");

    Ok(())
}

Although this pattern can be a little long-winded when it comes to writing all the code out, it's a very effective design for large systems where you don't really want a person to have to intervene in the loop.

Evaluator-optimizer

For coding use cases and other similar tasks where you need your agent to carry out some task without human guidance, you may want to use the evaluator-optimizer pattern which asks the LLM to complete the given task based on user input. If there is any feedback from previous message generation, then the AI agent should reflect and use them to improve the solution.

This workflow is quite similar to Chain of Thought prompting where you use multiple intermediate steps to be able to increase the LLM's reasoning capabilities.

To start with, we'll create our structs for extracting JSON from an LLM output:

use schemars::JsonSchema;

#[derive(serde::Deserialize, JsonSchema, serde::Serialize, Debug)]
struct Evaluation {
    evaluation_status: EvalStatus,
    feedback: String,
}

#[derive(serde::Deserialize, JsonSchema, serde::Serialize, Debug, PartialEq)]
enum EvalStatus {
    Pass,
    NeedsImprovement,
    Fail,
}

Once done, we'll need to create our agents as usual. The preamble for each agent has been extracted into its own variable for readability - note that we ask provide specific details on what the AI agent should output.

use std::env;

use rig::{completion::Prompt, providers::openai::Client};

const TASK: &str = "Implement a Stack with:
1. push(x)
2. pop()
3. getMin()
All operations should be O(1).
";

const GEN_AGENT_PREAMBLE: &str = "Your goal is to complete the task based on <user input>. If there are feedback
            from your previous generations, you should reflect on them to improve your solution

            Output your answer concisely in the following format:

            Thoughts:
            [Your understanding of the task and feedback and how you plan to improve]

            Response:
            [Your code implementation here]";

const EVAL_AGENT_PREAMBLE: &str = "Evaluate this following code implementation for:
            1. code correctness
            2. time complexity
            3. style and best practices

            You should be evaluating only and not attempting to solve the task.

            Only output \"PASS\" if all criteria are met and you have no further suggestions for improvements.

            Provide detailed feedback if there are areas that need improvement. You should specify what needs improvement and why.

            Only output JSON.";
#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // Create OpenAI client
    let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
    let openai_client = Client::new(&openai_api_key);

    let generator_agent = openai_client
        .agent("gpt-4")
        .preamble(GEN_AGENT_PREAMBLE)
        .build();

    let evaluator_agent = openai_client.extractor::<Evaluation>("gpt-4")
        .preamble(EVAL_AGENT_PREAMBLE)
        .build();
    // .. more code below
}

Next, we'll need to prompt our generate agent with the task. Initially there will be no chat history, so we simply just prompt it and add the response to a list of memories. We then loop, extracting an answer from our evaluation agent and checking if the struct (deserialized from JSON) has passed the evaluation. If it has then break the loop - if not, add the feedback and prompt the generator agent, adding the response from the generator agent back to the list of memories.

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // .. previous code goes here

    let mut memories: Vec<String> = Vec::new();

    let mut response = generator_agent.prompt(TASK).await.unwrap();
    memories.push(response.clone());

    loop {
        let eval_result = evaluator_agent
            .extract(&format!("{TASK}\n\n{response}"))
            .await
            .unwrap();

        if eval_result.evaluation_status == EvalStatus::Pass {
            break;
        } else {
            let context = format!("{TASK}\n\n{}", eval_result.feedback);

            response = generator_agent.prompt(context).await.unwrap();
            memories.push(response.clone());
        }
    }

    println!("Response: {response}");

    Ok(())

The evaluator-optimizer pattern is a relatively useful pattern because it's a good way to get your agents to do work autonomously and self-improve the output by itself without needing any kind of human interaction. In comparison to autonomous agents, this design pattern primarily revolves around using the LLM as a judge to figure out when it's considered "done" rather than an arbitrary condition in the environment.

Autonomous agent

The last workflow to showcase is the autonomous agent, which carries out a task until a goal has been achieved. This pattern is relatively simple, but requires no human intervention which can be pretty helpful.

In this example, we will simply have an LLM add a random number between 1 and 64 to a number that we give it, using only whole numbers. To start with, we'll define a Counter struct:

use schemars::JsonSchema;

#[derive(Debug, serde::Deserialize, JsonSchema, serde::Serialize)]
struct Counter {
    /// The score of the document
    number: u32,
}

Next, we'll create our Counter agent which will extract the number as the Counter struct. This is somewhat relatively simple as there isn't really much to this - we'll just tell the LLM to add a random number between 1 and 64 to the given number.

use rig::providers::openai::Client;
use std::env;

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // Create OpenAI client
    let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
    let openai_client = Client::new(&openai_api_key);

    let agent = openai_client.extractor::<Counter>("gpt-4")
        .preamble("
            Your role is to add a random number between 1 and 64 (using only integers) to the previous number.
        ")
        .build();

    // .. more code below
}

Finally, we'll create our original number, the amount of time we want to wait in between each loop iteration, then simply loop until we have the number we want (in this case, we've set it at 2000 but this can technically be whatever number you want).

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // previous code
    let mut number: u32 = 0;

    let mut interval = tokio::time::interval(std::time::Duration::from_secs(1));

    // Loop the agent and allow it to run autonomously. If it hits the target number (2000 or above)
    // we then terminate the loop and return the number
    // Note that the tokio interval is to avoid being rate limited
    loop {
        // Prompt the agent and print the response
        let response = agent.extract(&number.to_string()).await.unwrap();

        if response.number >= 2000 {
            break;
        } else {
            number += response.number
        }

        interval.tick().await;
    }

    println!("Finished with number: {number:?}");

    Ok(())
}

Autonomous agents, while a bit less complex than some of the other design patterns, are quite simple to implement (being that they are just LLM prompts on a loop) and so are useful for simple tasks.

Beyond this article

Thanks for reading! I hope this has helped you gain a better understanding of how to use design patterns with agentic AI to significantly improve the effectiveness of your AI-assisted applications.

For additional Rig resources and community engagement:

Check out more examples in our gallery.
Contribute or report issues on our GitHub.
Join discussions in our Discord community!

Building CSV RAG with Rig and Rust 🔥🔥🔥

Josh Mo — Wed, 19 Feb 2025 14:13:54 +0000

Introduction

In this article, we're going to look at how you can use the rig Rust AI framework to create an application that is able to load a CSV file, embed it into a vector store and have an LLM answer questions based on provided context from our vector store. We'll serve the application in the form of a command-line tool.

Interested in trying it out for yourself? Check out the Github repo.

What is RAG?

Retrieval-Augmented Generation (RAG) is a technique that combines the power of information retrieval with the generative capabilities of large language models (LLMs). The process typically involves two main components: retrieving relevant knowledge from a pre-existing database or knowledge base and generating contextually accurate answers based on that information. When a question or query is posed, the model first retrieves relevant snippets or pieces of information from the knowledge base. It then augments its generative model with this newly acquired knowledge, enabling it to provide a more accurate, informed response.

In contrast to traditional LLMs that rely purely on their trained parameters, RAG systems dynamically incorporate external data, ensuring that responses are grounded in factual and up-to-date knowledge. When a question is asked, the relevant knowledge from a pre-existing dataset is retrieved and fed into the model. The LLM then generates an answer that combines the retrieved information and its generative abilities, aiming to offer a detailed and accurate response. This hybrid approach allows RAG models to provide answers that are not only linguistically fluent but also well-supported by factual data.

Getting started

To get started, make sure you have the Rust programming language installed.

Next, you'll need to create a new project using cargo init:

cargo init csv-rag
cd csv-rag

Let's get building!

Before we start, let's add the relevant dependencies which we need:

cargo add serde rig-core tokio csv -F \
serde/derive,tokio/macros,tokio/rt-multi-thread

This adds the following dependencies:

rig-core - The rig library.
tokio - an asynchronous Rust runtime. We additionally add the macros and rt-multi-thread feature as we want to use the macro.
csv - The Rust crate for using CSVs.

Parsing CSVs

Before we do anything else, we'll need to declare our struct type as well as derive the correct macros for it. Note that the embedding trait for Rig also depends on Eq (which also depends on PartialEq), hence the derivation.

#[derive(serde::Deserialize, serde::Serialize, Debug, Eq, PartialEq)]
struct Record {
    first_name: String,
    last_name: String,
    email: String,
    role: String,
    salary: u32, //salary in dollars
}

Next to be able to use Rig's embeddings API, we need to implement the Embed trait. This trait simply defines a method for what to embed. Because we don't have a singular field we want to embed and will need the whole record, we can simply just implement fmt::Display for the struct then format the whole struct into a string (which we can then embed).

impl std::fmt::Display for Record {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let Self {
            first_name,
            last_name,
            email,
            role,
            salary,
        } = self;
        write!(
            f,
            "First name: {first_name}\nLast name: {last_name}\nEmail: {email}\nRole: {role}\nSalary: {salary}"
        )
    }
}

impl Embed for Record {
    fn embed(
        &self,
        embedder: &mut rig::embeddings::TextEmbedder,
    ) -> Result<(), rig::embeddings::EmbedError> {
        Ok(embedder.embed(self.to_string()))
    }
}

Embedding our CSV

Thankfully, embedding our CSV is actually quite easy now that we've done the hard part. We can create an OpenAI client and get the embedding model, then iterate through each record and build an embedding from it and store it in Rig's in-memory vector store.

You can see below the code to generate an OpenAI client is fairly simple:

let openai_client = rig::providers::openai::Client::from_env();

let agent = openai_client.agent("gpt-4o")
    .preamble("You are a helpful assistant. 
        Your job is to answer a user's questions based on 
        the context given.").build();

// Create embedding model
let embedding_model = openai_client.embedding_model("text-embedding-ada-002");

We'll also keep the code for loading the CSV in as a separate function as it's relatively long.

async fn load_csv(
    model: rig::providers::openai::EmbeddingModel,
) -> Result<
    InMemoryVectorIndex<rig::providers::openai::EmbeddingModel, Record>,
    Box<dyn std::error::Error>,
> {
    let reader = Reader::from_path("employees.csv")?;

    // this iterator attempts to deserialize all rows into the Record type
    // any failed rows will be skipped over rather than causing a panic
    // due to use of filter_map and x.ok() 
    let documents: Vec<Record> = reader
        .into_deserialize::<Record>()
        .filter_map(|x| x.ok()) 
        .collect();

    let documents = match EmbeddingsBuilder::new(model.clone())
        .documents(documents)?
        .build()
        .await
    {
        Ok(ok) => ok,
        Err(e) => return Err(format!("Got error while embedding: {e}").into()),
    };

    let vector_store = InMemoryVectorStore::from_documents(documents);

    Ok(vector_store.index(model))
}

Talking to our CSV file

The next step will be to actually take input from the user and prompt the LLM using not only the prompt, but additional (relevant!) snippets from the vector store.

We'll start by taking some input from standard input (i.e., you load up the application in the terminal, then you type in something and press Enter). We then check if the prompt is equal to quit - if it is then immediately break the loop and close the application, if not, carry on and retrieve a response from the LLM.

println!(
        "Hi! This is your PDF ragger. Write a prompt and press Enter or write \"quit\" to exit."
    );

loop {
    print!("> ");
    let stdin = std::io::stdin();
    let mut prompt = String::new();
    io::stdout().flush()?;
    stdin.read_line(&mut prompt)?;

    match prompt.trim() {
        "quit" => break,
        _ => {}
    }

    let docs = index.top_n::<Record>(prompt.trim(), 4).await?;

    let prompt = format!(
        "Relevant employees:\n{}\n\n{}",
        docs.into_iter()
            .map(|(_, _, doc)| doc.to_string())
            .collect::<Vec<String>>()
            .join("\n"),
        prompt,
    );

    let answer = agent.prompt(prompt.as_ref()).await?;

    println!("{answer}");
}

Creating a chat loop

In terms of storing your conversation history, it is actually quite simple to do so. All you need to do is to create a Vec<Message> and then feed it into the agent, using the chat function instead of prompt. You then add a user message and assistant message (in that order) at the end of the loop iteration. This allows you to build up a conversation with the LLM.

Additionally, we'll also add a reset command which allows the entire conversation to be wiped should the user want to ask the LLM about something else.

println!(
        "Hi! This is your PDF ragger. Write a prompt and press Enter or write \"quit\" to exit. Alternatively, use \"reset\" to reset the conversation."
    );

let mut chat_history: Vec<Message> = Vec::new();

loop {
    print!("> ");
    let stdin = std::io::stdin();
    let mut prompt = String::new();
    io::stdout().flush()?;
    stdin.read_line(&mut prompt)?;

    match prompt.trim() {
        "quit" => break,
        "reset" => {
            chat_history.clear();
            println!("Your conversation has been reset.");
            continue;
        }
        _ => {}
    }

    let docs = index.top_n::<Record>(prompt.trim(), 4).await?;

    let prompt = format!(
        "Relevant employees:\n{}\n\n{}",
        docs.into_iter()
            .map(|(_, _, doc)| doc.to_string())
            .collect::<Vec<String>>()
            .join("\n"),
        prompt,
    );

    let answer = agent.chat(prompt.as_ref(), chat_history.clone()).await?;

    println!("{answer}");
    chat_history.push(Message::user(prompt));
    chat_history.push(Message::assistant(answer));
}

Beyond this tutorial

Now that we've set up a base for your command-line tool, here are a few ways you can extend this demo to do much more than just simple ragging:

Try adding some tools to your agent!
What about semantic routing to make sure your LLM stays on track?
Try incorporating CSV RAG into a new, bigger pipeline!

Additionally, for additional Rig resources and community engagement:

Check out more examples in our gallery.
Contribute or report issues on our GitHub.
Join discussions in our Discord community!

Conclusion

Thanks for reading! Hopefully you will have a good example of how to improve how you work with data by creating a RAG pipeline from your CSV.

Building RAG using Rig, SurrealDB & DeepSeek on Cloudflare Workers

Josh Mo — Wed, 05 Feb 2025 16:13:56 +0000

Introduction

Retrieval-Augmented Generation (RAG) enhances AI responses by grounding them in relevant external knowledge. In this tutorial, you'll learn how to build a Cloudflare Worker that runs a RAG pipeline using SurrealDB, Rig AI, and DeepSeek via GroqCloud - all at blazing speeds of 275 tokens per second!

By the end of this guide, you'll have an understanding of how you can effectively use SurrealDB to create a Cloudflare worker that can execute Retrieval Augmented Generation (RAG) by creating a Cargo workspace with the following:

An application that can seed a SurrealDB instance
A Cloudflare worker that can carry out cosine similarity based vector search in a SurrealDB instance and use the found document to help answer a user query.

Interested in checking out the final code or just want to skip to the end result? Have a look at the code repository.

Getting Started

Pre-requisites

To get started with this tutorial, you will need the following:

A SurrealDB Cloud account
A Cloudflare account
An OpenAI account, with an API key
A GroqCloud account, with an API key

Before we go any further, you will need to ensure that your OpenAI and Groq API keys are stored as env variables locally to be able to run the SurrealDB instance locally. You can do this like so:

export OPENAI_API_KEY=<your-openai-api-key>
export GROQ_API_KEY=<your-groq-api-key>

Preparing our SurrealDB instance

To get started, register an account with SurrealDB Cloud. You'll then need to create a SurrealDB instance (they have a free instance available which you can use).

Next, open the instance in Surrealist (using the Connect drop-down) and you'll find a query box. Paste the following into the query and run it:

-- define and use namespace/database for queries
DEFINE NAMESPACE cfw_surrealdb_rig;
USE NAMESPACE cfw_surrealdb_rig;
DEFINE DATABASE cfw_surrealdb_rig;
USE DATABASE cfw_surrealdb_rig;

-- define root user
-- note that in prod, your roles should be properly locked down
DEFINE USER root ON ROOT PASSWORD 'root' ROLES OWNER;

-- define table & fields
DEFINE TABLE words SCHEMAFULL;
DEFINE field word on table words type string;
DEFINE field description on table words type string;
DEFINE field embedding on table words type array<float>;

-- define index on embedding field
DEFINE INDEX IF NOT EXISTS words_embedding_vector_index ON words
     FIELDS embedding
     MTREE DIMENSION 1536
     DIST COSINE;

This will do the following:

Create a new namespace & database, and then use it for the additional steps in the query
Define a new user with Owner permissions
Defines a table that has a schema with relevant fields for the word, a description and the embedding representation of the description
Defines an index on the embedding field to make vector search faster

Note that you will need the namespace & database names (cfw_surrealdb_rig) and the username & password for logging in later (root/root). Make sure you also store the SurrealDB hostname somewhere - you can find this by clicking the three vertical dots at the top right of your instance (from the Instances menu), then clicking "Copy hostname".

Project Setup

To get started, we'll create a workspace:

mkdir cfw-surrealdb-rig
cd cfw-surrealdb-rig

This will eventually hold the following:

A small program for filling the SurrealDB instance with data
Our Cloudflare Worker

For now, we will create a Cargo.toml file in our main file to signify that we're using Cargo workspaces:

# Cargo.toml
[workspace]
members = []
resolver = "2"

[workspace.dependencies]
rig-core = { version = "0.7.0", default-features = false, features = [
    "worker",
    "derive",
] }
surrealdb = { version = "2.1.4", features = ["protocol-ws"] }
serde = { version = "1.0.217", features = ["derive"] }
serde_json = "1.0.138"

See the list below for an explanation of dependencies:

rig-core: The AI framework we'll be using to connect with the various providers we need.
surrealdb: The Rust crate for using SurrealDB. We connect to it via Websocket, so the websocket protocol has been added.
serde: A (de)serialization library. We add the derive feature for extremely easy derive macro usage on structs.
serde-json: A library that allows (de)serialization to/from JSON.

Because most of our dependencies will be shared, we can store them in the workspace Cargo.toml file. When creating our new projects, we can then simply set workspace = true rather than having to manually add everything.

Seeding our SurrealDB instance

(note: If you want to skip writing all the code for this step, you can do this by visiting the repo and following the instructions)

To get started we'll create a new program in our workspace:

cargo init seeder
cd seeder

Next, we will add the following dependencies specific to this workspace crate with cargo add:

cargo add tokio

We'll also additionally add our workspace dependencies:

# cfw-surrealdb-rig/seeder/Cargo.toml
## .. rest of your .toml file up here

[dependencies]
tokio = "1.43.0" # or whatever the latest version is
surrealdb = { workspace = true }
rig-core = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }

Now it's time to start coding!

We'll firstly set up our main function to use the Tokio runtime and properly return an error instead of forcing a panic at runtime - this will allow us to use the ? operator which saves space and allows callstack propagation:

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // .. your code here
}

Next, we'll connect to our SurrealDB instance (put this code in your main function!). Note that these values have been filled in according to the assumed values, so if you're using different values you'll want to quickly double-check you have the right values here.

use surrealdb::{
    engine::remote::ws::Wss,
    opt::auth::Root,
    Surreal,
};

let surreal =
        Surreal::new::<Wss>("<your-surrealdb-hostname-here>")
            .await?;

surreal.signin(Root {
            username: "root",
            password: "root",
        })
        .await?;

surreal.use_ns("cfw_surrealdb_rig")
       .use_db("cfw_surrealdb_rig")
       .await?;

Next, we will create an OpenAI client that will be used to embed documents. We'll use the OPENAI_API_KEY environment variable to set up our OpenAI client and we'll select our embedding model.

    // Create OpenAI client
    let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
    let openai_client = Client::new(&openai_api_key);

    let embedding_model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

Next we'll create our data to be embedded. Note that we implement a new trait, rig::Embed, which allows the struct to be automatically embedded. During the embedding process, our seeder will only embed the description, then when we iterate on the final Embeddings struct it'll map into embeddings and the documents themselves.

use rig::embeddings::EmbeddingModel;

#[derive(rig::Embed, Serialize, Deserialize, Clone, Debug, Eq, PartialEq, Default)]
struct WordDefinition {
    id: String,
    word: String,
    #[embed]
    description: String,
}

#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, Default)]
struct WordDefinitionEmbed {
    word: String,
    description: String,
    embedding: Vec<f64>,
}


let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
            .documents(vec![
                WordDefinition {
                    id: "doc0".to_string(),
                    word: "flurbo".to_string(),
                    description:
                        "1. *flurbo* (name): A flurbo is a green alien that lives on cold planets.".to_string(),
                },
                WordDefinition {
                    id: "doc1".to_string(),
                    word: "glarb-glarb".to_string(),
                    description: "2. *glarb-glarb* (noun): A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy.".to_string()

                },
                WordDefinition {
                    id: "doc2".to_string(),
                    word: "linglingdong".to_string(),
                    description: "1. *linglingdong* (noun): A rare, mystical instrument crafted by the ancient monks of the Nebulon Mountain Ranges on the planet Quarm.".to_string()

                },
            ])?
            .build()
            .await?;

let embeddings = embeddings
        .into_iter()
        .map(|(d, embeddings)| {
            let embedding: Vec<f64> = embeddings.first().vec.iter().map(|&x| x as f64).collect();
            WordDefinitionEmbed {
                word: d.word,
                description: d.description,
                embedding,
            }
        })
        .collect::<Vec<WordDefinitionEmbed>>();

for embedding in embeddings {
    let _: Option<WordDefinitionEmbed> = surreal.create("words").content(embedding).await?;
    }

Ok(())

If we run this with cargo run -p seeder, you should now have some records in your SurrealDB instance. You can verify this by going back to your SurrealDB instance and connecting with Surrealist (then using SELECT * FROM words as the query), or using the Explorer.

Setting up the worker

Next, we'll need to build our worker. Note that to do this, we'll need to be able to use only libraries that compile to WebAssembly - fortunately everything we're using in this post does by default!

If you haven't already, you'll need to install the wasm32-unknown-unknown compilation target, as well as cargo-generate which is used to initialise a lot of different types of boilerpalate:

rustup target add wasm32-unknown-unknown
cargo install cargo-generate

Next, we'll use cargo-generate to create our boilerplate:

cargo generate cloudflare/workers-rs

Make sure to select the hello world template, as that's what we'll be using. Note that upon generation, the package will be automatically added to your workspace - so you don't need to add it yourself.

We'll also additionally add our workspace dependencies:

# cfw-surrealdb-rig/seeder/Cargo.toml
## .. rest of your .toml file up here

[lib]
crate-type = ["cdylib"]

[dependencies]
worker = { version = "0.5.0" }
worker-macros = { version = "0.5.0" }
console_error_panic_hook = { version = "0.1.1" }
rig-core = { workspace = true }
surrealdb = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }

Like before, we need to connect to our SurrealDB instance - but this time we'll be using env.secret instead of std::env::var. This is because Cloudflare Workers use the Cloudflare deployment environment which supports environment variables in this manner rather than environment variables.

use surrealdb::{
    engine::remote::ws::Wss,
    opt::auth::Root,
    Surreal,
};

#[derive(Deserialize)]
struct UserQuery {
    contents: String,
}

#[event(fetch)]
async fn fetch(mut req: Request, env: Env, _ctx: Context) -> Result<Response> {
    // serialize raw json body to UserQuery struct type
    let query: UserQuery = req.json().await?;

    // environment secrets (we will talk about this later)
    let surrealdb_url = env.secret("SURREALDB_URL")?.to_string();
    let surrealdb_user = env.secret("SURREALDB_USER")?.to_string();
    let surrealdb_password = env.secret("SURREALDB_PASSWORD")?.to_string();
    let surrealdb_ns = env.secret("SURREALDB_NS")?.to_string();
    let surrealdb_db = env.secret("SURREALDB_DB")?.to_string();
    let openai_api_key = env.secret("OPENAI_API_KEY")?.to_string();
    let groq_api_key = env.secret("GROQ_API_KEY")?.to_string();

    // then we comnnect to SurrealDB as usual
    let surreal = Surreal::new::<Wss>(&surrealdb_url)
        .await
        .map_err(|x| worker::Error::RustError(x.to_string()))?; // Log into the database

    surreal
        .signin(Root {
            username: &surrealdb_user,
            password: &surrealdb_password,
        })
        .await
        .map_err(|x| worker::Error::RustError(x.to_string()))?;

    surreal
        .use_ns(&surrealdb_ns)
        .use_db(&surrealdb_db)
        .await
        .map_err(|x| worker::Error::RustError(x.to_string()))?;

    // .. rest of code goes here
}

Next, we'll create our OpenAI embedding model as usual and embed our text. We will query SurrealDB using the new embedding to find documents with the highest cosine similarity to our embedding. We then use .take(0) to retrieve a result.

Note that because we haven't set a threshold for the similarity score, this should always return a record as long as there is at least 1 record in the words table.

use rig::providers::openai::{Client, TEXT_EMBEDDING_ADA_002};
use rig::Embed;

// create openAI client
let openai_client = Client::new(&openai_api_key);

#[derive(Serialize, Deserialize, Clone, Debug, Default)]
struct EmbedResult {
    word: String,
    description: String,
    score: f32,
}

let embedding_model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

let text: Vec<f64> = embedding_model
        .embed_text(&query.contents)
        .await
        .map_err(|x| worker::Error::RustError(x.to_string()))?
        .vec
        .iter()
        .map(|&x| x as f64)
        .collect();

let mut result = surreal
        .query("SELECT word, description, vector::similarity::cosine($vec, embedding) as score FROM words order by score desc limit 1")
        .bind(("vec", text))
        .await
        .map_err(|x| worker::Error::RustError(x.to_string()))?;

let result: Vec<EmbedResult> = result
        .take(0)
        .map_err(|x| worker::Error::RustError(x.to_string()))?;

let result_as_string = result
        .into_iter()
        .map(|x| x.description)
        .collect::<Vec<String>>()
        .join("\n");

Finally, we just need to add our Groq client. We can leverage the openai module to be able to quickly and efficiently access the GroqCloud chat completion API as it has full compatibility with the OpenAI chat completion API.

use rig::completion::Prompt;

    let groq_client = Client::from_url(&groq_api_key, "https://api.groq.com/openai/v1");

    let chat_model = groq_client
        .agent("deepseek-r1-distill-llama-70b")
        .preamble("You are an agent designed to answer questions based on user input.")
        .build();

    let prompt = format!(
        "{}

            Context: {result_as_string}",
        query.contents
    );

    let answer = chat_model
        .prompt(&prompt)
        .await
        .map_err(|x| worker::Error::RustError(x.to_string()))?;

    Response::ok(answer)

Deployment

Now for the fun part: deploying our worker! For this you'll need to make sure to either have the wrangler CLI installed or npm (which you can then either use npm install to install wrangler, or use npx wrangler to use the CLI that way). Don't forget to use wrangler login to give wrangler access to your Cloudflare account if this is your first time!

Secrets

Although you can add secrets directly from the Workers web UI. That being said if you're already in your development environment it's often easier to just write your secrets in a git-ignored file and then use them that way. The wrangler CLI accepts uploading secrets in bulk by including them as a JSON file - which we'll do below.

To get started, create a file called secrets.json that looks like this (make sure to fill out the relevant variables):

{
  "SURREALDB_URL": "<surrealdb-instance-url-here>",
  "SURREALDB_USER": "root",
  "SURREALDB_PASSWORD": "root",
  "SURREALDB_NS": "cfw_surrealdb_rig",
  "SURREALDB_DB": "cfw_surrealdb_rig",
  "OPENAI_API_KEY": "<openai-api-key-here>",
  "GROQ_API_KEY": "<groq-api-key-here>"
}

Once done, use wrangler secret bulk --json secrets.json to upload your secrets in bulk. Nice and easy!

Observability

Enabling logging for your Worker, especially in production, can be hugely useful for debugging issues and getting your worker back up again.

To enable it, add the observability object to your wrangler.toml file then add the enabled = true property under it.

# wrangler.toml
[observability]
enabled = true

Now if there's any issues, you should be able to find what the issue is by simply looking at your logs for the worker from the web UI!

How to deploy

To deploy, all you need to do is use wrangler deploy. That's it! The wrangler CLI will then build your worker and upload it to Cloudflare to be used whenever someone triggers the endpoint.

Finishing Up

Thanks for reading! Hopefully this article has given you a better idea of how you can use SurrealDB with Rig to make effective AI agents.

Extending this further

Interested in extending this example? Here's a couple of ideas you can try:

Got a bunch of related documents you want to embed? You could try implementing Graph RAG to improve on the initial RAG method.
You could also try implementing semantic routing and other common RAG techniques that greatly improve the effectiveness of your AI agents.

Semantic Routing with Qdrant, Rig & Rust

Josh Mo — Thu, 16 Jan 2025 16:26:54 +0000

Introduction

In conversational AI and intelligent systems, decision-making is often the most critical layer. By themselves, LLMs can be mislead by users into making the wrong decisions. This has manifested itself many times - particularly to the detriment of companies who don't safeguard against it. However, what if I told you there's an easy way to protect your model against injection attacks and the like that's also cheap on compute?

In this article, we’ll delve into building an efficient semantic router using Qdrant, Rig, and Rust. By combining Qdrant’s vector search capabilities, the Rig LLM framework, and Rust’s performance, we’ll craft a system that empowers your agents and model prompts to make precise, context-aware decisions.

What even is a semantic router?

A semantic router, in short, is a decision making layer for your agents (and model prompts). It works as a lightweight abstraction by grouping utterances (short snippets of sentences) into topics, then comparing the original query against the router for semantic similarity. A topic is then returned based on the most similar sentence - or none, if there's no similar topic.

Semantic routing has several advantages:

Your agents won't deviate from the topics you want them to talk about
It's scalable to however many "routes" you want to implement
It's fast and doesn't require much compute

Getting Started

Before we get started, make sure the Rust programming language is installed! If you plan to run this, you'll also want an OpenAI API key. Make sure you grab one from the dashboard.

Before you run cargo run, make sure you export your API key:

export OPENAI_API_KEY=<key-goes-here>

Let's get building!

Ok so first of all let's create our project. Create your project like so with cargo init then cd into it:

cargo init semantic-router
cd semantic-router

The first thing we'll need to create is the thing that we want to embed. Additionally because we'll be grouping our embeddings by topic, we will want to declare a Topic struct that can be used to easily generate new instances of Utterance that can be easily embedded.

Note that although there is a Derive macro that rig provides, we embed manually here as there are other fields that we need to preserve which will act as metadata.

// topics.rs
use rig::{embeddings::{EmbedError, TextEmbedder}, Embed};

#[derive(serde::Deserialize, serde::Serialize, Debug, Clone)]
pub struct Utterance {
    pub id: String,
    pub topic: String,
    pub content: String,
}

impl Utterance {
    pub fn new(topic: &str, content: &str) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            topic: topic.to_string(),
            content: content.to_string(),
        }
    }
}
impl Embed for Utterance {
    fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError> {
        // Embeddings only need to be generated for `content` field.
        // Queries will be compared against the content        
        embedder.embed(self.content.to_owned());

       Ok(())
    }
}

Note that while Utterance::new() does take a topic name, it is a bit cumbersome to have to keep writing the topic name every time you need to instantiate a new Utterance. It might not matter for development workloads, but what about when you need to do a large amount of snippets for say, 50-60 topics? How can we make this a bit more ergonomic?

This is where we'll create our Topic struct. This struct will be used to make it easier to create instances of Utterance per topic. See below:

// topics.rs
pub struct Topic {
    name: String
}

impl Topic {
    pub fn new(name: &str) -> Self {
        let name = name.to_string();
        Self {
            name
        }
    }

    pub fn new_utterance(&self, content: &str) -> Utterance {
        Utterance::new(&self.name, content)
    }
}

Now to build our semantic router! For this part, we'll create a new() method that tries to initialise an OpenAI client using environment variables then attempts to connect to Qdrant locally, creating a collection if none exists (with the given name) and returns the struct.

use rig::{embeddings::EmbeddingsBuilder, providers::openai::{Client, EmbeddingModel as Model, TEXT_EMBEDDING_ADA_002}};
use rig_qdrant::QdrantVectorStore;
use qdrant_client::{qdrant::{CreateCollectionBuilder, PointStruct, QueryPointsBuilder, VectorParamsBuilder}, Payload, Qdrant};


pub struct SemanticRouter {
    model: Model,
    qdrant: Arc<Qdrant>,
    vector_store: QdrantVectorStore<Model>,
}

const COLLECTION_NAME: &str = "SEMANTIC_ROUTING";
const COLLECTION_SIZE: usize = 1536;

impl SemanticRouter {
    pub async fn new() -> Result<Self, Box<dyn std::error::Error>> {

        // Initialize OpenAI client.
        // Get your API key from https://platform.openai.com/api-keys
        let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
        let openai_client = Client::new(&openai_api_key);

        let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

        // note that this assumes you're hosting Qdrant locally on localhost:6334
        let qdrant = Arc::new(Qdrant::from_url("http://localhost:6334").build()?);

        // note that we use `Arc::into_inner() here` because
        // the Qdrant client doesn't actually implement Clone
        // so we need to create an `Arc::clone` then get the inner value
        let qdrant_inner = Arc::clone(&qdrant);
        let qdrant_inner = Arc::into_inner(qdrant_inner).unwrap();

            // Create a collection with 1536 dimensions if it doesn't exist
            // Note: Make sure the dimensions match the size of the embeddings returned by the
            // model you are using
            if !qdrant.collection_exists(COLLECTION_NAME).await? {
                qdrant
                    .create_collection(
                        CreateCollectionBuilder::new(COLLECTION_NAME)
                            .vectors_config(VectorParamsBuilder::new(COLLECTION_SIZE as u64, qdrant_client::qdrant::Distance::Cosine)),
                    )
                    .await?;
            }


        let query_params = QueryPointsBuilder::new(COLLECTION_NAME).with_payload(true);
        let vector_store = QdrantVectorStore::new(qdrant_inner, model.clone(), query_params.build());
            Ok(Self {
                model,
                qdrant,
                vector_store
            })

    }
}

Next, we need to implement methods for embedding utterances and adding them to our Qdrant collection. We iterate through the collection of documents, mapping them into a Vec<f32> and creating a PointStruct. Then all of our points are added.

// semantic_router.rs
use rig::vector_store::{VectorStoreIndex};

use crate::topics::Utterance;

impl SemanticRouter {
    // .. other method(s) above
    pub async fn embed_utterances(&self, utterances: Vec<Utterance>) -> Result<(), Box<dyn std::error::Error>> {
        let mut documents = EmbeddingsBuilder::new(self.model.clone())
            .documents(utterances)?
            .build()
            .await?;

        let points: Vec<PointStruct> = documents
            .into_iter()
            .map(|(d, embeddings)| {
                let vec: Vec<f32> = embeddings.first().vec.iter().map(|&x| x as f32).collect();
                PointStruct::new(
                    d.id.clone(),
                    vec,
                    Payload::try_from(serde_json::to_value(&d).unwrap()).unwrap(),
                )
            })
            .collect();

            self.qdrant
                .upsert_points(UpsertPointsBuilder::new(COLLECTION_NAME, points))
                .await?;

            Ok(())
    }
}

Finally, we need to add our query method for checking whether or not a prompt matches a conversation topic. Note that we only want results that are semantically similar to our prompt - so we need to make sure we check the similarity score is high and only return the metadata if it matches or passes the threshold. If not, we should return an error here - which will then get handled higher up the stack.

// semantic_router.rs
impl SemanticRouter {
    // .. other methods
    pub async fn query(&self, query: &str) -> Result<Utterance, Box<dyn std::error::Error>> {

        let results = self.vector_store
            .top_n::<Utterance>(query, 1)
            .await?;

        if results[0].0 <= 0.85 {
            return Err("No relevant snippet found.".into());
        }

        Ok(results[0].2.clone())
    }

}

Now to wrap this up, we need to add some code to our main function to show that this works. See the following code below:

// main.rs
pub mod semantic_router;
pub mod topics;

use semantic_router::SemanticRouter;
use topics::{Topic, Utterance};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let router = SemanticRouter::new().await.unwrap();

    let bees_topic = Topic::new("bees");

    // create a vector of strings then iterate through the vector
    // and map them all to `Utterance` instances
    let bee_facts = vec![
    "Bees communicate with their hive mates through intricate dances that convey the location of nectar-rich flowers.",
    "A single bee can visit up to 5,000 flowers in a day, tirelessly collecting nectar and pollen.",
    "The queen bee can lay up to 2,000 eggs in a single day during peak season.",
].into_iter().map(|x| bees_topic.new_utterance(x)).collect::<Vec<Utterance>>();

    // embed utterances into Qdrant
    router.embed_utterances(bee_facts).await.unwrap();

    let bee_answer = router.query("how many flowers does a bee visit in a day?").await?;
    println!("Topic: {}", bee_answer.topic);

    // note that this query *should* error out as it's unrelated
    // in which case, we simply tell the user we can't help them
     match router.query("what is skibidi toilet").await {
         Ok(res) => println!("Unexpectedly found a topic: {}", res.topic),
        Err(_) => println!("Sorry, I can't help you with that.")
    };

    Ok(())

}

Extending this example

So, now you know how to create your own semantic routing! Some ways to extend this can be found below:

Usage in a web application

For usage with web applications, you may find that you need to implement Clone for your shared application state - our router doesn't actually implement Clone. An easy way to fix this is to simply use the Arc<T> pattern, which will allow it to be shared. Note that this only works if the inner type is already Send + Sync - wrapping it in Arc doesn't automatically guarantee you can use it!

Finishing up

Thanks for reading! Hopefully this article has helped you learn a little bit more about how you can make your RAG and LLM-assisted applications even better than before.