DEV Community: Shuttle

A Comprehensive Guide to the llm-chain Rust crate

Shuttle — Thu, 06 Jun 2024 17:06:37 +0000

LLM orchestration is an important part of creating AI powered applications. Particularly in business use cases, AI agents and RAG pipelines are commonly utilised for refined LLM responses. Although Langchain is currently the most popular LLM orchestration suite, there are also similar crates in Rust that we can use. In this article, we’ll be diving into one of them - llm-chain.

What is llm-chain?

llm-chain is a collection of crates that describes itself as “the ultimate toolbox” for working with Large Language Models (LLMs) to create AI-powered applications and other tooling.

You can find the crate’s GitHub repository here.

Comparison to other LLM orchestration crates

If you’ve looked around for different crates, you probably noticed that there are a few crates for LLM orchestration in Rust:

llm-chain (this one!)
langchain-rust
anchor-chain

In comparison to the others, llm-chain is somewhat macro heavy. However, it is also the most developed in terms of having extra data processing utilities. If you’re looking for an all-in-one package, llm-chain is the crate that’s most likely to help you get over the line. They also have their own docs.

Getting Started

Pre-requisites

Before you add llm-chain to your project, make sure you have access to the prompting model you want to use. For example, if you want to use OpenAI, make sure you have an OpenAI API key (set as OPENAI_API_KEY in environment variables).

Setup

To get started, all you need to do is to add the crate to your project (as well as Tokio for async):

cargo add llm-chain
cargo add tokio -F full

Next, we'll want to add a provider for whatever method of model prompting you want to use. Here we'll add the OpenAI integration by adding the crate to our application:

cargo add llm-chain-openai

Note that a full list of integrations can be found here, split by package.

Basic usage

Prompting

To get started with llm-chain, we can use their basic example as a way to quickly get something working. In this code snippet, we will:

Initialise an executor using executor!()
Use prompt!() with the system message and prompt to store both in a struct that will get used when the prompt (or chain) gets ran.
Runs the prompt and returns the results, using a reference to the executor.
Prints the results.

use std::error::Error;
use llm_chain::{executor, parameters, prompt};

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let exec = executor!()?;

    let res = prompt!(
    "You are a robot assistant for making personalized greetings",
    "Make a personalized greeting for Joe"
    )
    .run(&parameters!(), &exec)
    .await?;

    println!("{}", res);

    Ok(())
}

Running this prompt should yield a result that looks like this:

Assistant: Hello Joe! I hope you're having a fantastic day filled with joy and success. Remember to keep shining bright and making a positive impact wherever you go. Have a great day!

The default model for the llm_chain_openai executor is gpt-3.5-turbo. The executor parameters can be defined in the macro - you can also find more about this here.

Using Templates

However, if we want to move onto more advanced pipelines the easiest way for us to do this would be to use a prompt template with parameters. You can see below that much like in the previous code snippet, we generate an executor and return the results. However, instead of using prompt!() by itself we use it in Step::for_prompt_template - which you can find more about here.

use llm_chain::step::Step;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let exec = executor!()?;
    let step = Step::for_prompt_template(prompt!(
        "You are a bot for making personalised greetings",
        "Make a personalized greeting tweet for {{text}}"
    ));

    let step_results = step.run(&parameters!("Emil"), &exec).await?;
    println!("step_results: {step_results}");

    let immediate_results = step_results.to_immediate().await?.as_content();
    println!("immediate results: {immediate_results}");

    Ok(())
}

The results of the output should look like this:

step_results: Assistant: "Hey @Emil! Wishing you a fantastic day filled with joy, success, and lots of smiles! Keep shining bright and making a positive impact in the world. Cheers to you! 🌟 #YouGotThis"

immediate results: Assistant: "Hey @Emil! Wishing you a fantastic day filled with joy, success, and lots of smiles! Keep shining bright and making a positive impact in the world. Cheers to you! 🌟 #YouGotThis"

Chaining prompts

Of course, one of the main reasons why we're using Langchain (or Langchain-like libraries) in the first place is to be able to orchestrate our LLM usage. The llm-chain Rust crate assists us with this by letting us create chains of LLM prompts using the Chain struct.

There are three types of chains that we can use with llm-chain:

Sequential chains, which apply steps sequentially
Map-reduce chains, which use a "map" step to apply to each chunk from a loaded file and then reduce the text. This is quite useful for text summarization.
Conversational chains, which keep track of the conversation history and manage context. Conversational chains are great for chatbot applications, multi-step interactions and other places where context is essential.

Sequential chaining

The easiest to use type of chaining is sequential chaining, which simply pipes the output from each step into the next step. When creating our steps, we will use the Chain struct instead of creating each step individually:

use llm_chain::step::Step;
use llm_chain::chains::sequential::Chain;
use llm_chain::prompt;

// Create a chain of steps with two prompts
let first_step = Step::for_prompt_template(
            prompt!("You are a bot for making personalized greetings", "Make personalized birthday e-mail to the whole company for {{name}} who has their birthday on {{date}}. Include their name")
    );

// Second step: summarize the email into a tweet. Importantly, the text parameter is the result of the previous prompt.
let second_step = Step::for_prompt_template(
    prompt!( "You are an assistant for managing social media accounts for a company", "Summarize this email into a tweet to be sent by the company, use emojis if you can. \\n--\\n{{text}}")
        );

let chain: Chain = Chain::new( vec![first_step, second_step] );

Next, we'll then use the parameters! macro to inject parameters into the prompt pipeline:

use llm_chain::parameters;
use llm_chain::traits::Executor as ExecutorTrait;
use llm_chain_openai::chatgpt::Executor;

    // Create a new ChatGPT executor with the default settings
    let exec = Executor::new()?;

    // Run the chain with the provided parameters
    let res = chain
        .run(
            // Create a Parameters object with key-value pairs for the placeholders
            parameters!("name" => "Emil", "date" => "February 30th 2023"),
            &exec,
        )
        .await?;

    // Print the result to the console
    println!("{}", res.to_immediate().await?.as_content());
    Ok(())
}

Running the code should yield a result that looks like this:

Assistant: 🎉🎂 Join us in celebrating Emil's birthday on February 30th! 🎈🎁 Emil, your dedication and hard work are truly commendable. Wishing you happiness and success on your special day! 🥳🎉 #HappyBirthdayEmil #TeamAppreciation 🎂

Map-reduce chains

Map-reduce chains typically consist of two steps:

A "Map" step that takes a document and applies an LLM chain to it, treating the output as a new document
The new documents are then passed to a new chain that combines the separate documents to get a single output. At the end of a Map-Reduce chain, the output can be taken for further processing by sending it to another prompting model (for instance) or as part of a sequential pipeline.

To use this pattern, we need to create a prompt template:

// note that we import Chain from a different module here!
use llm_chain::chains::map_reduce::Chain;

// Create the "map" step to summarize an article into bullet points
let map_prompt = Step::for_prompt_template(prompt!(
"You are a bot for summarizing wikipedia articles, you are terse and focus on accuracy",
"Summarize this article into bullet points:\\n{{text}}"
));

// Create the "reduce" step to combine multiple summaries into one
let reduce_prompt = Step::for_prompt_template(prompt!(
"You are a diligent bot that summarizes text",
"Please combine the articles below into one summary as bullet points:\\n{{text}}"
));

// Create a map-reduce chain with the map and reduce steps
let chain = Chain::new(map_prompt, reduce_prompt);

Next, we need to take some text from a file and add it as a parameter - the {{text}} parameter in the Map prompt will automatically take in the file content:

// Load the content of the article to be summarized
let article = include_str!("article_to_summarize.md");

// Create a vector with the Parameters object containing the text of the article
let docs = vec![parameters!(article)];

let exec = executor!()?;
// Run the chain with the provided documents and an empty Parameters object for the "reduce" step
// Note that there are multiple modules with a Chain struct
// This one takes two different modules for
let res = chain.run(docs, Parameters::new(), &exec).await?;

// Print the result to the console
println!("{}", res.to_immediate().await?.as_content());

Note here that because we have two steps, the chain.run() function will take two different vectors - one for each step, in order. This means that we are passing the article content to the first prompt, but no parameters to the second document.

Conversational Chains

Of course, the last chain we need to talk about is conversational chains. In a nutshell, conversational chains allow you to load context from memory by using saved chat history. In situations where the platform or model cannot access saved chat history, you might store the response and then use it as extra context in the next message.

To use conversational chains, like before, we need to create a Chain (now imported from the conversation module) and define the steps for it:

use llm_chain::{
    chains::conversation::Chain, executor, output::Output, parameters, prompt, step::Step,
};

let exec = executor!()?;

let mut chain = Chain::new(
    prompt!(system: "You are a robot assistant for making personalized greetings."),
)?;

// Define the conversation steps.
let step1 = Step::for_prompt_template(prompt!(user: "Make a personalized greeting for Joe."));
let step2 =
    Step::for_prompt_template(prompt!(user: "Now, create a personalized greeting for Jane."));
let step3 = Step::for_prompt_template(
    prompt!(user: "Finally, create a personalized greeting for Alice."),
);

let step4 = Step::for_prompt_template(prompt!(user: "Remind me who did we just greet."));

Next, we will individually send each prompt to the Chain in turn, printing out the response from each one. Note that at step 4, we should receive an answer that includes the names of the previous three people we just made a personalized greeting for (Joe, Jane and Alice).

// Execute the conversation steps.
let res1 = chain.send_message(step1, &parameters!(), &exec).await?;
println!("Step 1: {}", res1.to_immediate().await?.primary_textual_output().unwrap());

let res2 = chain.send_message(step2, &parameters!(), &exec).await?;
println!("Step 2: {}", res2.to_immediate().await?.primary_textual_output().unwrap());

let res3 = chain.send_message(step3, &parameters!(), &exec).await?;
println!("Step 3: {}", res3.to_immediate().await?.primary_textual_output().unwrap());

let res4 = chain.send_message(step4, &parameters!(), &exec).await?;
println!("Step 4: {}", res4.to_immediate().await?.primary_textual_output().unwrap());

Running this should get an output that looks something like this:

Step 1: Hello, Joe! I hope you are having a fantastic day filled with positivity and joy. Keep shining bright and making a difference in the world with your unique presence. Wishing you continued success and happiness in all that you do!
Step 2: Hello, Jane! Sending you warm greetings and positive vibes today. May your day be as wonderful and vibrant as you are. Remember to keep being your amazing self and always believe in the incredible things you are capable of achieving. Wishing you endless happiness and success in all your endeavors!
Step 3: Hello, Alice! I hope this message finds you well and thriving. You are such a remarkable individual with a heart full of kindness and a spirit full of strength. Keep inspiring those around you with your grace and resilience. May your day be filled with love, laughter, and countless blessings. Stay amazing, Alice!
Step 4: We just created personalized greetings for Joe, Jane, and Alice.

Using embeddings with llm-chain

In terms of using embeddings with llm-chain, it provides a helper method for using Qdrant as a vector store. It abstracts over the qdrant_client crate, providing an easy way to embed documents and carry out similarity search. Note that the Qdrant struct will assume your collection(s) that you want to use have already been created!

Basic usage

While we can use qdrant_client to manually create our own embeddings, llm-chain also has an integration for easy access. We will be required to create our own client through qdrant_client - which we can then use with the Qdrant struct to be able to parse stuff.

First, let's define a couple of passages that we want to insert into our Qdrant collection:

    const DOC_DOG_DEF: &str = r#"The dog (Canis familiaris[4][5] or Canis lupus familiaris[5]) is a domesticated descendant of the wolf. Also called the domestic dog, it is derived from the extinct Pleistocene wolf,[6][7] and the modern wolf is the dog's nearest living relative.[8] Dogs were the first species to be domesticated[9][8] by hunter-gatherers over 15,000 years ago[7] before the development of agriculture.[1] Due to their long association with humans, dogs have expanded to a large number of domestic individuals[10] and gained the ability to thrive on a starch-rich diet that would be inadequate for other canids.[11]

    The dog has been selectively bred over millennia for various behaviors, sensory capabilities, and physical attributes.[12] Dog breeds vary widely in shape, size, and color. They perform many roles for humans, such as hunting, herding, pulling loads, protection, assisting police and the military, companionship, therapy, and aiding disabled people. Over the millennia, dogs became uniquely adapted to human behavior, and the human–canine bond has been a topic of frequent study.[13] This influence on human society has given them the sobriquet of "man's best friend"."#;

    const DOC_WOODSTOCK_SOUND: &str = r#"Sound for the concert was engineered by sound engineer Bill Hanley. "It worked very well", he says of the event. "I built special speaker columns on the hills and had 16 loudspeaker arrays in a square platform going up to the hill on 70-foot [21 m] towers. We set it up for 150,000 to 200,000 people. Of course, 500,000 showed up."[48] ALTEC designed marine plywood cabinets that weighed half a ton apiece and stood 6 feet (1.8 m) tall, almost 4 feet (1.2 m) deep, and 3 feet (0.91 m) wide. Each of these enclosures carried four 15-inch (380 mm) JBL D140 loudspeakers. The tweeters consisted of 4×2-Cell & 2×10-Cell Altec Horns. Behind the stage were three transformers providing 2,000 amperes of current to power the amplification setup.[49][page needed] For many years this system was collectively referred to as the Woodstock Bins.[50] The live performances were captured on two 8-track Scully recorders in a tractor trailer back stage by Edwin Kramer and Lee Osbourne on 1-inch Scotch recording tape at 15 ips, then mixed at the Record Plant studio in New York.[51]"#;

The Qdrant struct will automatically assume you have your collection set up and have a QdrantClient that already exists, along with the collection name. We'll pass these as arguments into a new function that does the following:

Create embeddings using llm-chain-openai
Insert the embeddings into Qdrant
Conduct a similarity search using the prompt

Firstly, we’ll want to define a method for creating our Qdrant struct so that we can re-use it later on:

use llm_chain_openai::embeddings::Embeddings;
use llm_chain_qdrant::Qdrant;
use llm_chain::schema::EmptyMetadata;
use qdrant_client::prelude::{QdrantClient, QdrantClientConfig};

// note that the URL must connect to port 6334 - qdrant_client uses GRPC!
// feel free to replace this wit
async fn create_qdrant_client(url: String) -> QdrantClient {
    let mut config = QdrantClientConfig::from_url(url);
    // this part is only required if you're running Qdrant on the cloud
    // if running locally, no api key is required
    config.api_key = std::env::var("QDRANT_API_KEY").ok();
    QdrantClient::new(Some(config))
}

async fn create_qdrant_struct(qdrant_client: QdrantClient, collection_name: String) -> Qdrant {
    let embeddings = Embeddings::default();

    // Storing documents
    Qdrant::new(
        qdrant_client,
        collection_name,
        embeddings,
        None,
        None,
        None,
    )
}

Next, we can use the Qdrant struct to carry out a similarity search! We’ll add our documents to our collection, then conduct a similarity search and print out the stored documents:

async fn similarity_search_qdrant(qdrant: Qdrant) -> Result<(), Box<dyn Error>> {
    // embed and upsert the documents into Qdrant
    let doc_ids = qdrant
        .add_documents(
            vec![
                DOC_DOG_DEF.to_owned(),
                DOC_WOODSTOCK_SOUND.to_owned(),
            ]
            .into_iter()
            .map(Document::new)
            .collect(),
        )
        .await?;

    println!("Documents stored under IDs: {:?}", doc_ids);

    // conduct similarity search and find similar vectors
    let response = qdrant
        .similarity_search(
            "Sound engineering is involved with concerts and music events".to_string(),
            1,
        )
        .await?;

    // print out the stored documents with payload, embeddings etc
    println!("Retrieved stored documents: {:?}", response);
}

After this, we can then send it into a Chain or whatever else we need.

Usage within a prompt template

In isolation, the Qdrant struct is not particularly helpful and mainly provides convenience methods for embedding things. However, we can also add it as part of a ToolCollection which lets the pipeline know that it is able to use embeddings.

let qdrant = create_qdrant_struct( ,
    "mycollection".to_string()).await;

let exec = executor!().unwrap();

let mut tool_collection = ToolCollection::<Multitool>::new();

tool_collection.add_tool(
    QdrantTool::new(
        qdrant,
        "factual information and trivia",
        "facts, news, knowledgebuild_local_qdrant().await; or trivia",
    )
    .into(),
);

let task = "Tell me something about dogs";

let prompt = ChatMessageCollection::new()
    .with_system(StringTemplate::tera(
        "You are an automated agent for performing tasks. Your output must always be YAML.",
    ))
    .with_user(StringTemplate::combine(vec![
        tool_collection.to_prompt_template().unwrap(),
        StringTemplate::tera("Please perform the following task: {{task}}."),
    ]));

let result = Step::for_prompt_template(prompt.into())
    .run(&parameters!("task" => task), &exec)
    .await
    .unwrap();

println!("{}", res.to_immediate().await?.as_content());

Processing data using llm-chain

While llm-chain provides tooling for creating LLM pipelines, another important part of Langchain and libraries like it is being able to process and transform data. Prompts (and prompt engineering) are important to get right. However if we're also feeding data into our pipeline, we'll also want to make sure it's as easy as possible to find the most relevant context.

Below are a couple of useful use cases that you may want to check out.

Scraping search results

llm-chain provides a convenience struct for scraping Google Results using the GoogleSerper struct - using the Serper.dev service.

use llm_chain::tools::{tools::GoogleSerper, Tool};

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let serper_api_key = std::env::var("SERPER_API_KEY").unwrap();
    let serper = GoogleSerper::new(serper_api_key);
    let result = serper
        .invoke_typed(&"Who was the inventor of Catan?".into())
        .await
        .unwrap();
    println!("Best answer from Google Serper: {}", result.result);
}

As well as this, there is also support for Bing Search API which provides 1000 free searches a month - you can find more about the pricing here. Below is a code snippet of how you can use the API:

use llm_chain::tools::{tools::BingSearch, Tool};

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let bing_api_key = std::env::var("BING_API_KEY").unwrap();
    let bing = BingSearch::new(bing_api_key);
    let result = bing
        .invoke_typed(&"Who was the inventor of Catan?".into())
        .await
        .unwrap();
    println!("Best answer from bing: {}", result.result);
}

Should you need to change between one or the other, both are quite easy to use.

Extracting labelled text

llm-chain also has some convenience methods for extracting labelled text. If you have a string of bullet points for instance, you can use extract_labeled_text() to be able to extract the text.

use llm_chain::parsing::extract_labeled_text;

fn main() {
    let text = r"
- Title: The Matrix
- Actor: Keanu Reeves
- Director: The Wachowskis
";

    let result = extract_labeled_text(text);
    println!("{:?}", result);
}

Running this code should result in an output that looks like this:

[("Title", "The Matrix"), ("Actor", "Keanu Reeves"), ("Director", "The Wachowskis")]

You can find out more about the parsing module for llm-chain here as well as some of the examples.

Conclusion

Thanks for reading! With the power of llm-chain, you can easily leverage AI for your applications.

Implementing Semantic Caching using Qdrant and Rust

Shuttle — Thu, 30 May 2024 15:49:45 +0000

Hello world! Today we're going to learn about semantic caching with Qdrant, in Rust. By the end of this tutorial, you'll have a Rust application that can do the following:

Ingest a CSV file, turn it into an embedding with the help of an LLM and insert it into Qdrant
Create two collections in Qdrant - one for regular usage and one for caching
Utilize semantic caching for quicker access

Interested in deploying or got lost and want to find a repository with the code? You can find that here

What is semantic caching, and why use it?

In a regular data cache, we store information to enable faster retrieval later on. For example, you might have in a web service that's served behind Nginx. We can have Nginx cache either all responses, or only the most accessed endpoints. This improves performance and reduces load on your web server.

Semantic caching in this regard is quite similar. Using vector databases, we can create database collections that store the queries themselves. For example, these two questions semantically carry the same meaning:

What are some best practices for writing the Rust programming language?
What are some best practices for writing Rustlang?

We can store a copy of the query in a cache collection, with the answer as a JSON payload. If users then ask a similar question, we can retrieve the embedding and fetch the answer from the payload. This avoids us having to use an LLM to get our answer.

There are a couple of benefits to semantic caching:

Prompts that require long responses can see serious cost savings.
It's pretty easy to implement and fairly cheap - the only cost is in storage and using the embedding model
You can use a cheaper model than your regular embedding

Semantic caching is normally used with RAG - Retrieval Augmented Generation. RAG is a framework to allow context retrieval from pre-embedded materials. For example, CSV files or documents can be turned into embeddings using models and stored in a database. Whenever a user wants to find similar documents to a given prompt, they embed the prompt and search against it in a given database.

Of course, there are good reasons not to use semantic caching. Prompts that need differing, varied answers won't find any use for semantic caching. This is particularly relevant in generative AI usage. Fetching a stored query will reduce the creativity of the response. Regardless, if part of your pipeline is able to capitalise on semantic caching, it's a good idea to do so.

Project setup

Getting started

To get started, don't forget to use cargo shuttle init, with the Axum framework. We'll install our dependencies using the shell snippet below:

cargo add qdrant-client@1.7.0 anyhow async-openai serde serde-json \
shuttle-qdrant uuid -F uuid/v4,serde/derive

You can find our quickstart docs here.

Setting up secrets

To set up our secrets, we'll use a Secrets.toml file located in our project root (you will need to create this manually). You can then add whatever secrets you need using the format below:

OPENAI_API_KEY = ""
QDRANT_URL = ""
QDRANT_API_KEY = ""

Setting up Qdrant

Creating collections

Now that we can get started, we will add some more general methods for creating a regular collection as well as a cache collection, to simulate a real-world scenario (as well as a new() function to make creating the RAGSystem struct itself). We’ll create the struct first: Note here that although we're using vectors with 1536 dimensions, the number of dimensions you'll need may depend on the model you are using.

use qdrant_client::prelude::QdrantClient;
use async_openai::{config::OpenAIConfig, Client};

struct RagSystem {
    qdrant_client: QdrantClient,
    openai_client: Client<OpenAIConfig>
}

static REGULAR_COLLECTION_NAME: &str = "my-collection";
static CACHE_COLLECTION_NAME: &str = "my-collection-cached";

impl RAGSystem {
    fn new(qdrant_client: QdrantClient) -> Self {
        let openai_api_key = env::var("OPENAI_API_KEY").unwrap();

        let openai_config = OpenAIConfig::new()
            .with_api_key(openai_api_key)
            .with_org_id("qdrant-shuttle-semantic-cache");

        let openai_client = Client::with_config(openai_config);

        Self {
            openai_client,
            qdrant_client,
        }
    }
}

Now we’ll create the methods for initialising our regular collection. Note that we’ll only need to use these once. After the collections have already been created, if we try to initialise them again we’ll get an error.

use qdrant_client::prelude::CreateCollection;
use qdrant_client::qdrant::{
    vectors_config::Config,VectorParams,
    VectorsConfig, WithPayloadSelector,
};

impl RagSystem {
    async fn create_regular_collection(&self) -> Result<()> {
        self.qdrant_client
            .create_collection(&CreateCollection {
                collection_name: REGULAR_COLLECTION_NAME.to_string(),
                vectors_config: Some(VectorsConfig {
                    config: Some(Config::Params(VectorParams {
                        size: 1536,
                        distance: Distance::Cosine.into(),
                        ..Default::default()
                    })),
                }),
                ..Default::default()
            })
            .await?;

        Ok(())
    }
}

Next, we’ll create our cache collection. When creating this collection, note that we use Distance::Euclid instead of Distance::Cosine. Both of these can be defined as follows:

Distance::Cosine (or “cosine similarity”) measures how closely two vectors are pointing in the same direction. If we plot two vectors on a graph, for example, a vector located at [2,1] would be much closer to [1,1] than it would be [-1, -2]. Cosine similarity is overwhelmingly used in measuring document similarity in text analysis.
Distance::Euclid (or “Euclidean distance”) measures how closely two vectors are from each other - i.e., the distance from A to B where A and B are two points on a graph. Rather than trying to determine similarity, here we want to determine whether something is mostly or exactly the same.

impl RagSystem {
    async fn create_cache_collection(&self) -> Result<()> {
        self.qdrant_client
            .create_collection(&CreateCollection {
                collection_name: CACHE_COLLECTION_NAME.to_string(),
                vectors_config: Some(VectorsConfig {
                    config: Some(Config::Params(VectorParams {
                        size: 1536,
                        distance: Distance::Euclid.into(),
                        hnsw_config: None,
                        quantization_config: None,
                        on_disk: None,
                        ..Default::default()
                    })),
                }),
                ..Default::default()
            })
            .await?;

        Ok(())
    }
}

Creating embeddings

Next, we need to create an embedding from a file input - using a CSV file as an example. To do so, we'll need to do the following:

Read the file inputs and parse it to a string (std::fs::read_to_string() parses to a String type automatically)
Chunk the file contents into appropriate amounts (here we'll do it per-row naively, for illustration)
Bulk embed the embeddings and add them to Qdrant

Here we're using the async-openai library to be able to create the embedding - but if you don't want to use OpenAI, you can always use fastembed-rs as an alternative or another crate of your choice that allows embedding creation.

use std::path::PathBuf;
use async_openai::types::{CreateEmbeddingRequest, EmbeddingInput};
use anyhow::Result;

impl RagSystem {
    async fn embed_and_upsert_csv_file(&self, file_path: PathBuf) -> Result<()> {
        let file_contents = std::fs::read_to_string(&file_path)?;

        // note here that we skip 1 because CSV files typically have headers
        // if you don't have any headers, you can remove it
        let chunked_file_contents: Vec<String> =
            file_contents.lines().skip(1).map(|x| x.to_owned()).collect();

        let embedding_request = CreateEmbeddingRequest {
            model: "text-embedding-ada-002".to_string(),
            input: EmbeddingInput::StringArray(chunked_file_contents.to_owned()),
            encoding_format: None, // defaults to f32
            user: None,
            dimensions: Some(1536),
        };

        let embeddings = Embeddings::new(&self.openai_client)
            .create(embedding_request)
            .await?;

        if embeddings.data.is_empty() {
            return Err(anyhow::anyhow!(
                "There were no embeddings returned by OpenAI!"
            ));
        }

        let embeddings_vec: Vec<Vec<f32>> =
            embeddings.data.into_iter().map(|x| x.embedding).collect();

        // note that we create the upsert_embedding function later on
        for embedding in embeddings_vec {
            self.upsert_embedding(embedding, file_contents.clone())
                .await?;
        }

        Ok(())
    }
}

We'll then need to embed any further inputs to search for any matching embeddings. The embed_prompt function will look quite similar to the embedding part of our embed_and_upsert_csv_file function. However, it will instead return a Vec<f32> as we'll want to use this later to search our collection.

impl RagSystem {
    pub async fn embed_prompt(&self, prompt: &str) -> Result<Vec<f32>> {
        let embedding_request = CreateEmbeddingRequest {
            model: "text-embedding-ada-002".to_string(),
            input: EmbeddingInput::String(prompt.to_owned()),
            encoding_format: None, // defaults to f32
            user: None,
            dimensions: Some(1536),
        };

        let embeddings = Embeddings::new(&self.openai_client)
            .create(embedding_request)
            .await?;

        if embeddings.data.is_empty() {
            return Err(anyhow::anyhow!(
                "There were no embeddings returned by OpenAI!"
            ));
        }

        Ok(embeddings.data.into_iter().next().unwrap().embedding)
    }
}

Upserting embeddings

Once we've created our embeddings, we’ll create a method for adding embeddings into Qdrant called upsert_embedding. This will deal with creating the payload for our embedding and insert it into the database. Once added to the collection, we can search our collection later on and get the associated JSON payload alongside the embedding!

The function will look like this:

use qdrant_client::prelude::PointStruct;

impl RAGSystem {
    async fn upsert_embedding(&self, embedding: Vec<f32>, file_contents: String) -> Result<()> {
        let payload = serde_json::json!({
           "document": file_contents
        })
        .try_into()
        .map_err(|x| anyhow::anyhow!("Ran into an error when converting the payload: {x}"))?;

        let points = vec![PointStruct::new(
            uuid::Uuid::new_v4().to_string(),
            embedding,
            payload,
        )];

        self.qdrant_client
            .upsert_points(REGULAR_COLLECTION_NAME.to_owned(), None, points, None)
            .await?;

        Ok(())
    }
}

Here, we use a uuid::Uuid as a unique identifier for our embedding(s). You can also do the same thing by having a u64 counter that increases with every embedding. However, you'll want to make sure you don't accidentally overwrite your own embeddings! Inserting a new embedding with the same ID as a currently existing embedding in the collection will overwrite the embedding.

Of course, we also need to create a method for adding things to our cache. Note that our payload is different here. Instead of using the document payload field, we use answer since the payload will hold a pre-generated answer to the question.

impl RagSystem {
    pub async fn add_to_cache(&self, embedding: Vec<f32>, answer: String) -> Result<()> {
        let payload = serde_json::json!({
           "answer": answer
        })
        .try_into()
        .map_err(|x| anyhow::anyhow!("Ran into an error when converting the payload: {x}"))?;

        let points = vec![PointStruct::new(
            uuid::Uuid::new_v4().to_string(),
            embedding,
            payload,
        )];

        self.qdrant_client
            .upsert_points(CACHE_COLLECTION_NAME.to_owned(), None, points, None)
            .await?;

        Ok(())
    }
}

Searching Qdrant collections

Having made something we can search against in Qdrant, we'll need to implement some methods for our VectorDB. We'll split this up into two methods:

search_regular_collection
search_cache_collection

When searching for an embedding, we should attempt to search our semantic cache using search_cache_collection - if it doesn't find anything, we should then use the regular search_regular_collection method to get the document, prompt OpenAI with it and then return the result as

To make our methods a little bit more error-resistant, we have used .into_iter().next() on the results. This tries to find the first item in the vector by only going to the first item in the vector. This works because we're only looking for one single embedding, but you can increase or decrease the limit as you'd like.

Once we find a match, we need to get the document key from our JSON payload associated with the embedding match and return it. We'll be using this as context in our RAG prompt later on!

use qdrant_client::qdrant::{
    with_payload_selector::SelectorOptions, SearchPoints, WithPayloadSelector
};

impl RagSystem {
        pub async fn search(&self, embedding: Vec<f32>) -> Result<String> {
        let payload_selector = WithPayloadSelector {
            selector_options: Some(SelectorOptions::Enable(true)),
        };

        let search_points = SearchPoints {
            collection_name: REGULAR_COLLECTION_NAME.to_owned(),
            vector: embedding,
            limit: 1,
            with_payload: Some(payload_selector),
            ..Default::default()
        };

        let search_result = self
            .qdrant_client
            .search_points(&search_points)
            .await
            .inspect_err(|x| println!("An error occurred while searching for points: {x}"))
            .unwrap();

        let result = search_result.result.into_iter().next();

        let Some(result) = result else {
            return Err(anyhow::anyhow!("There's nothing matching."));
        };

        Ok(result.payload.get("document").unwrap().to_string())
    }
}

Of course, you'll also want to implement a function for searching your cache collection. Note that although the functions are mostly the same, we get the answer field from the payload instead of document for semantics.

impl RagSystem {
    pub async fn search_cache(&self, embedding: Vec<f32>) -> Result<String> {
        let payload_selector = WithPayloadSelector {
            selector_options: Some(SelectorOptions::Enable(true)),
        };

        let search_points = SearchPoints {
            collection_name: CACHE_COLLECTION_NAME.to_owned(),
            vector: embedding,
            limit: 1,
            with_payload: Some(payload_selector),
            ..Default::default()
        };

        let search_result = self
            .qdrant_client
            .search_points(&search_points)
            .await
            .inspect_err(|x| println!("An error occurred while searching for points: {x}"))?;

        let result = search_result.result.into_iter().next();

        let Some(result) = result else {
            return Err(anyhow::anyhow!("There's nothing matching."));
        };

        Ok(result.payload.get("answer").unwrap().to_string())
    }
}

Prompting

Of course, now that everything else is done, the last thing to do is prompting! Here, you can see below that we generate a prompt that basically consists of the prompt we want, as well as the provided context. We then grab the first result from OpenAI and return the message content.

use async_openai::types::{
    ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs,
    ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequestArgs
};

impl RagSystem {
    pub async fn prompt(&self, prompt: &str, context: &str) -> Result<String> {
        let input = format!(
            "{prompt}

            Provided context:
            {context}
            "
        );

        let res = self
            .openai_client
            .chat()
            .create(
                CreateChatCompletionRequestArgs::default()
                    .model("gpt-4o")
                    .messages(vec![
                        ChatCompletionRequestMessage::User(
                            ChatCompletionRequestUserMessageArgs::default()
                                .content(input)
                                .build()?,
                        ),
                    ])
                    .build()?,
            )
            .await
            .map(|res| {
                // We extract the first result
                match res.choices[0].message.content.clone() {
                    Some(res) => Ok(res),
                    None => Err(anyhow::anyhow!("There was no result from OpenAI")),
                }
            })??;

        println!("Retrieved result from prompt: {res}");

        Ok(res)
    }
}

Using Qdrant in a Rust web service

Let's have a quick look at a real world example. Below is a HTTP endpoint for the Axum framework that takes our RAGSystem as application state. It’ll embed the prompt and attempt to search the cache. If there’s no result, it searches in the regular collection for a match. The resulting document payload is added to an augmented prompt, and the question and answer are added to the cache. Finally, a response is returned from the endpoint.

use axum::{Json, extract::State, response::IntoResponse, http::StatusCode};
use serde::Deserialize;

#[derive(Deserialize)]
struct Prompt {
    prompt: String,
}

async fn prompt(
    State(state): State<RAGSystem>,
    Json(prompt): Json<Prompt>,
) -> Result<impl IntoResponse, impl IntoResponse> {
    let embedding = match state.embed_prompt(&prompt.prompt).await {
        Ok(embedding) => embedding,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("An error occurred while embedding the prompt: {e}"),
            ))
        }
    };

    if let Ok(answer) = state.search_cache(embedding.clone()).await {
        return Ok(answer);
    }

    let search_result = match state.search(embedding.clone()).await {
        Ok(res) => res,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("An error occurred while prompting: {e}"),
            ))
        }
    };

    let llm_response = match state.prompt(&prompt.prompt, &search_result).await {
        Ok(prompt_result) => prompt_result,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Something went wrong while prompting: {e}"),
            ))
        }
    };

    if let Err(e) = state.add_to_cache(embedding, &llm_response).await {
        return Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Something went wrong while adding item to the cache: {e}"),
        ));
    };

    Ok(llm_response)
}

The last thing to do is setting up our main function. Note that we add the shuttle_qdrant::Qdrant annotation to our main function, allowing us to provision a Qdrant instance locally with Docker automatically on a local run. In production though, we'll need the cloud_url and api_key keys filled out.

#[shuttle_runtime::main]
async fn main(
    #[shuttle_qdrant::Qdrant(
        cloud_url = "{secrets.QDRANT_URL}",
        api_key = "{secrets.QDRANT_API_KEY}"
    )]
    qdrant: QdrantClient,
    #[shuttle_runtime::Secrets] secrets: SecretStore,
) -> shuttle_axum::ShuttleAxum {
    secrets.into_iter().for_each(|x| env::set_var(x.0, x.1));

    let rag = RAGSystem::new(qdrant);

    let setup_required = true;

    if setup_required {
        rag.create_regular_collection().await?;
        rag.create_cache_collection().await?;

        rag.embed_csv_file("test.csv".into()).await?;
    }

    let rtr = Router::new().route("/prompt", post(prompt)).with_state(rag);

    Ok(rtr.into())
}

Deploying

To deploy, all you need to do is use cargo shuttle deploy (with the --ad flag if on a Git branch with uncommitted changes) and wait for it to deploy! Once you've deployed, any further deploys needed will only need to re-compile your application (and any extra dependencies if added) then it'll be done much, much faster.

Extending this example

Want to extend this example? Here's a couple ways you can do that.

Use a cheaper model for semantic caching

While using a high-performance model is great and all, one thing that we want to save on in particular is costs. One thing that we can do here to save tokens is by using a cheaper model and asking the model if one question is semantically the same as another. Here's a prompt you can use:

Are these two questions semantically the same? Answer either 'Yes' or 'No'. Do not answer with anything else. If you don't know the answer, say 'I don't know'.

Question 1: <question 1 goes here>
Question 2: <question 2 goes here>

Smaller payload indexes

It should be noted of course that while our example does work, one thing you might need to think about is payload indexes or associated data connected to a particular embedding. If you're inserting the whole file contents as the payload for every single embedding in a large file, chances are you are going to increase your resource usage quite rapidly. You can mitigate this by only inserting a relevant slice of the file per embedding (so for example in this case, it might be the row).

Finishing up

Thanks for reading! By using semantic caching, we can create a much more performant RAG system that saves on both time and costs.

Building Agentic RAG with Rust, OpenAI & Qdrant

Shuttle — Fri, 24 May 2024 08:58:09 +0000

Hey there! In this article, we're gonna talk about building an agentic RAG workflow with Rust! We'll be building an agent that can take a CSV file, parse it and embed it into Qdrant, as well as retrieving the relevant embeddings from Qdrant to answer questions from users about the contents of the CSV file.

Interested in deploying or just want to see what the final code looks like? You can find the repository here.

What is Agentic RAG?

Agentic RAG, or Agentic Retrieval Augmented Generation, is the concept of mixing AI agents with RAG to be able to produce a workflow that is even better at being tailored to a specific use case than an agent workflow normally would be.

Essentially, the difference between this workflow and a regular agent workflow would be that each agent can individually access embeddings from a vector database to be able to retrieve contextually relevant data - resulting in more accurate answers across the board in an AI agent workflow!

Getting Started

To get started, use cargo shuttle init to create a new project.

Next, we'll add the dependencies we need using a shell snippet:

cargo add anyhow
cargo add async-openai
cargo add qdrant-client
cargo add serde -F derive
cargo add serde-json
cargo add shuttle-qdrant
cargo add uuid -F v4

We'll also need to make sure to have a Qdrant URL and an API key, as well as an OpenAI API key. Shuttle uses environment variables via a SecretStore macro in the main function, and can be stored in the Secrets.toml file:

OPENAI_API_KEY = ""

Next, we'll update our main function to have our Qdrant macro and our secrets macro. We'll iterate through each secret and set it as an environment variable - this allows us to use our secrets globally, without having to reference the SecretStore variable at all:

#[shuttle_runtime::main]
async fn main(
    #[shuttle_qdrant::Qdrant] qdrant_client: QdrantClient,
    #[shuttle_runtime::Secrets] secrets: SecretStore,
) -> shuttle_axum::ShuttleAxum {
    secrets.into_iter().for_each(|x| {
        set_var(x.0, x.1);
    });

    let router = Router::new()
        .route("/", get(hello_world));

    Ok(router.into())
}

Building an agentic RAG workflow

Setting up our agent

The agent itself is quite simple: it holds an OpenAI client, as well as a Qdrant client to be able to search for relevant document embeddings. Other fields can also be added here, depending on what capabilities your agent requires.

use async_openai::{config::OpenAIConfig, Client as OpenAIClient};
use qdrant_client::prelude::QdrantClient;

pub struct MyAgent {
    openai_client: OpenAIClient<OpenAIConfig>,
    qdrant_client: QdrantClient,
}

Next we'll want to create a helper method for creating the agent, as well as a system message which we'll feed into the model prompt later.

static SYSTEM_MESSAGE: &str = "
        You are a world-class data analyst, specialising in analysing comma-delimited CSV files.

        Your job is to analyse some CSV snippets and determine what the results are for the question that the user is asking.

        You should aim to be concise. If you don't know something, don't make it up but say 'I don't know.'.
"

impl MyAgent {
    pub fn new(qdrant_client: QdrantClient) -> Self {
        let api_key = std::env::var("OPENAI_API_KEY").unwrap();
        let config = OpenAIConfig::new().with_api_key(api_key);

        let openai_client = OpenAIClient::with_config(config);

        Self {
            openai_client,
            qdrant_client,
        }
    }
}

File parsing and embedding into Qdrant

Next, we will implement a File struct for CSV file parsing - it should be able to hold the file path, contents as well as the rows as a Vec<String> (string array, or more accurately a vector of strings). There's a few reasons why we store the rows as a Vec<String>:

Smaller chunks improve the retrieval accuracy, one of the biggest challenges that RAG has to deal with. Retrieving a wrong or otherwise inaccurate document can hamper accuracy significantly.
Improved retrieval accuracy leads to enhanced contextual relevance - which is quite important for complex queries that require specific question.
Processing and indexing smaller chunks

pub struct File {
    pub path: String,
    pub contents: String,
    pub rows: Vec<String>,
}

impl File {
    pub fn new(path: PathBuf) -> Result<Self> {
        let contents = std::fs::read_to_string(&path)?;

        let path_as_str = format!("{}", path.display());

        let rows = contents
            .lines()
            .map(|x| x.to_owned())
            .collect::<Vec<String>>();

        Ok(Self {
            path: path_as_str,
            contents,
            rows
        })
    }
}

While the above parsing method is serviceable (collecting all the lines into a Vec<String>), note that it is a naive implementation. Based on how your CSV files are delimited and/or if there is dirty data to clean up, you may want to either prepare your data so that it is already well-prepared, or include some form of data cleaning or validation. Some examples of this might be:

unicode-segmentation - a library crate for splitting sentences
csv_log_cleaner - a binary crate for cleaning CSVs
validator - a library crate for validating struct/enum fields

Next, we'll go back to our agent and implement a method for embedding documents into Qdrant that will take the File struct we defined.

To do this, we need to do the following:

Take the rows we created earlier and add them as the input for our embed request.
Create the embeddings (with openAI) and create a payload for storing alongside the embeddings in Qdrant. Note that although we use a uuid::Uuid for unique storage, you could just as easily use numbers by adding a number counter to your struct and incrementing it by 1 after you've inserted an embedding.
Assuming there are no errors, return Ok(())

use async_openai::types::{ CreateEmbeddingRequest, EmbeddingInput };
use async_openai::Embeddings;
use qdrant_client::prelude::{Payload, PointStruct};

static COLLECTION: &str = "my-collection";

// text-embedding-ada-002 is the model name from OpenAI that deals with embeddings
static EMBED_MODEL: &str = "text-embedding-ada-002";

impl MyAgent {
pub async fn embed_document(&self, file: File) -> Result<()> {
        if file.rows.is_empty() {
            return Err(anyhow::anyhow!("There's no rows to embed!"));
        }

        let request = CreateEmbeddingRequest {
            model: EMBED_MODEL.to_string(),
            input: EmbeddingInput::StringArray(file.rows.clone()),
            user: None,
            dimensions: Some(1536),
            ..Default::default()
        };

        let embeddings_result = Embeddings::new(&self.openai_client).create(request).await?;

        for embedding in embeddings_result.data {
            let payload: Payload = serde_json::json!({
                "id": file.path.clone(),
                "content": file.contents,
                "rows": file.rows
            })
            .try_into()
            .unwrap();

            println!("Embedded: {}", file.path);

            let vec = embedding.embedding;

            let points = vec![PointStruct::new(
                uuid::Uuid::new_v4().to_string(),
                vec,
                payload,
            )];
            self.qdrant_client
                .upsert_points(COLLECTION, None, points, None)
                .await?;
        }
        Ok(())
    }
}

Document searching

Now that we’ve embedded our document, we’ll want a way to check whether our embeddings are contextually relevant to whatever prompt the user gives us. For this, we’ll create a search_document function that does the following:

Embed the prompt using CreateEmbeddingRequest and get the embedding from the results. We’ll be using this embedding in our document search. Because we’ve only added one sentence to embed here (the prompt), it will only return one sentence - so we can create an iterator from the vector and attempt to find the first result.
Create a parameter list for our document search through the SearchPoints struct (see below). Here we need to set the collection name, the vector that we want to search against (ie the input), how many results we want to be returned if there are any matches, as well as the payload selector.
Search the database for results - if there are no results, return an an error; if there is a result, then return the result back.

use qdrant_client::qdrant::{
    with_payload_selector::SelectorOptions, SearchPoints, WithPayloadSelector,
};

impl MyAgent {
    async fn search_document(&self, prompt: String) -> Result<String> {
        let request = CreateEmbeddingRequest {
            model: EMBED_MODEL.to_string(),
            input: EmbeddingInput::String(prompt),
            user: None,
            dimensions: Some(1536),
            ..Default::default()
        };

        let embeddings_result = Embeddings::new(&self.openai_client).create(request).await?;

        let embedding = &embeddings_result.data.first().unwrap().embedding;

        let payload_selector = WithPayloadSelector {
            selector_options: Some(SelectorOptions::Enable(true)),
        };

        // set parameters for search
        let search_points = SearchPoints {
            collection_name: COLLECTION.to_string(),
            vector: embedding.to_owned(),
            limit: 1,
            with_payload: Some(payload_selector),
            ..Default::default()
        };

        // if the search is successful
        // attempt to iterate through the results vector and find a result
        let search_result = self.qdrant_client.search_points(&search_points).await?;
        let result = search_result.result.into_iter().next();

        match result {
            Some(res) => Ok(res.payload.get("contents").unwrap().to_string()),
            None => Err(anyhow::anyhow!("There were no results that matched :(")),
        }
    }
}

Now that everything we need to use our agent effectively is set up, we can set up a prompt function!

use async_openai::types::{
    ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs,
    ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequestArgs,
};

static PROMPT_MODEL: &str = "gpt-4o";

impl MyAgent {
    pub async fn prompt(&self, prompt: &str) -> anyhow::Result<String> {
        let context = self.search_document(prompt.to_owned()).await?;
        let input = format!(
            "{prompt}

            Provided context:
            {}
            ",
            context // this is the payload from Qdrant
        );

        let res = self
            .openai_client
            .chat()
            .create(
                CreateChatCompletionRequestArgs::default()
                    .model(PROMPT_MODEL)
                    .messages(vec![
                        //First we add the system message to define what the Agent does
                        ChatCompletionRequestMessage::System(
                            ChatCompletionRequestSystemMessageArgs::default()
                                .content(SYSTEM_MESSAGE)
                                .build()?,
                        ),
                        //Then we add our prompt
                        ChatCompletionRequestMessage::User(
                            ChatCompletionRequestUserMessageArgs::default()
                                .content(input)
                                .build()?,
                        ),
                    ])
                    .build()?,
            )
            .await
            .map(|res| {
                //We extract the first one
                res.choices[0].message.content.clone().unwrap()
            })?;

        println!("Retrieved result from prompt: {res}");

        Ok(res)
    }
}

Hooking the agent up to our web service

Because we separated the agent logic from our web service logic, we just need to connect the bits together and we should be done!

Firstly, we'll create a couple of structs - the Prompt struct that will take a JSON prompt, and the AppState function that will act as shared application state in our Axum web server.

#[derive(Deserialize)]
pub struct Prompt {
    prompt: String,
}

#[derive(Clone)]
pub struct AppState {
    agent: MyAgent,
}

We'll also introduce our prompt handler endpoint here:

async fn prompt(
    State(state): State<AppState>,
    Json(json): Json<Prompt>,
) -> Result<impl IntoResponse> {
    let prompt_response = state.agent.prompt(&json.prompt).await?;

    Ok((StatusCode::OK, prompt_response))
}

Then we need to parse our CSV file in the main function, create our AppState and embed the CSV, as well as setting up our router:

#[shuttle_runtime::main]
async fn main(
    #[shuttle_qdrant::Qdrant] qdrant_client: QdrantClient,
    #[shuttle_runtime::Secrets] secrets: SecretStore,
) -> shuttle_axum::ShuttleAxum {
    secrets.into_iter().for_each(|x| {
        set_var(x.0, x.1);
    });

    // note that this already assumes you have a file called "test.csv"
    // in your project root
    let file = File::new("test.csv".into())?;

    let state = AppState {
        agent: MyAgent::new(qdrant_client),
    };

    state.agent.embed_document(file).await?;

    let router = Router::new()
        .route("/", get(hello_world))
        .route("/prompt", post(prompt))
        .with_state(state);

    Ok(router.into())
}

Deploying

To deploy, all you need to do is use cargo shuttle deploy (with the --ad flag if on a Git branch with uncommitted changes), sit back and watch the magic happen!

Finishing Up

Thanks for reading! With the power of combining AI agents and RAG, we can create powerful workflows to be able to satisfy many different use cases. With Rust, we can leverage performance benefits to be able to run our workflows safely and with a low memory footprint.

Building an AI Content Writer using Rust and GPT-4o

Shuttle — Thu, 16 May 2024 17:32:54 +0000

Hello world! In this guide, we're going to talk about how you can get started with using AI agents to create a content writer that will use the Serper.dev API to search Google for results on your query, then use the results together with GPT-4o to create a summary of the results and finally create an article about it.

Interested in just deploying or got stuck during the tutorial? Have a look at the repository.

Setting up

To get started, we'll create a new project using cargo-shuttle init, making sure to pick Axum as the framework. After that, we'll install the dependencies we need:

cargo add async-openai
cargo add reqwest -F json
cargo add serde -F derive
cargo add serde_json
cargo add thiserror

You will also additionally need API keys from Serper and OpenAI, which you will put in a new Secrets.toml file in your project root:

SERPER_API_KEY = ""
OPENAI_API_KEY = ""

Error handling

Before we get started, let's quickly define some error types. We can add this to an errors.rs file and use thiserror. The reason why we'll do this is for error propagation: instead of manually pattern matching, if we want the error to be propagated we can use this enum as the error return type then use error propagation:

use async_openai::error::OpenAIError;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use thiserror::Error;

#[derive(Error, Debug)]
pub enum ApiError {
    #[error("OpenAI error: {0}")]
    OpenAI(#[from] OpenAIError),
    #[error("Reqwest error: {0}")]
    Reqwest(#[from] reqwest::Error),
    #[error("De/serialization error: {0}")]
    SerdeJson(#[from] serde_json::Error),
}

impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        let (status, body) = match self {
            Self::OpenAI(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
            Self::Reqwest(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
            Self::SerdeJson(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
        };

        println!("An error happened: {body}");

        (status, body).into_response()
    }
}

This is enabled by ApiError implementing From<T>, allowing easy propagation.

Building AI Agents

Our first step will be defining a generic interface that all of our autonomous agents will work with. We'll be creating two agents:

A researcher (takes some data from a Google search, feeds it into ChatGPT and asks it to summarize the information)
A writer (takes the summary and writes an article about it)

It will look something like this:

use async_openai::{config::OpenAIConfig, Client as OpenAIClient};

pub trait Agent {
    fn name(&self) -> String;
    fn client(&self) -> OpenAIClient<OpenAIConfig>;
    fn system_message(&self) -> String;

    // to be given a default implementation later
    async fn prompt(&self, input: &str, data: String) -> Result<String, ApiError>;
}

Why do we need the other 3 methods if prompt() already uses self? This is because as part of the Agent trait, the prompt function cannot reference types that it doesn't know about. If we have a struct that has the async_openai::Client type already, we need to create a method from the Agent trait to be able to access the client.

If you're only creating one agent, typically you don't need a specific trait. However, if you wanted to create more agents in the same library or application, it would be a good idea to have a generic interface to hold all of the relevant methods!

Let's define our Researcher struct, which will hold the data and methods for us to query the Serper API and then use the data to prompt a model:

#[derive(Clone)]
pub struct Researcher {
    http_client: reqwest::Client,
    system: Option<String>,
    openai_client: OpenAIClient<OpenAIConfig>,
}

Next, we can implement the Agent trait for our struct:

impl Agent for Researcher {
    fn name(&self) -> String {
        "Researcher".to_string()
    }

    fn client(&self) -> OpenAIClient<OpenAIConfig> {
        self.openai_client.clone()
    }

    fn system_message(&self) -> String {
        if let Some(message) = &self.system {
            message.to_owned()
        } else {
            "You are an agent.

        You will receive a question that may be quite short or does not have much context.
        Your job is to research the Internet and to return with a high-quality summary to the user, assisted by the provided context.
        The provided context will be in JSON format and contains data about the initial Google results for the website or query.

        Be concise.

        Question:
        ".to_string()
        }
    }
}

Note that here, the name() and client() functions are mostly boilerplate. If you wanted to extend this even further, you could use a macro to get rid of this totally.

The system message will be for pre-prompting the model. When we prompt the model, the system message will be passed in and then the bot will respond to the prompt according to the system message. In models where system messages don't exist, this would simply represent the text before you put your prompt.

We also implement two methods for impl Researcher: one to initialise the struct itself, and then one for preparing the data to send into our agent pipeline:

use std::env;
use reqwest::header::HeaderMap;

impl Researcher {
    pub fn new() -> Self {
        let api_key = env::var("OPENAI_API_KEY").unwrap();
        let config = OpenAIConfig::new().with_api_key(api_key);

        let openai_client = OpenAIClient::with_config(config);

        let mut headers = HeaderMap::new();
        headers.insert(
            "X-API-KEY",
            env::var("SERPER_API_KEY").unwrap().parse().unwrap(),
        );
        headers.insert("Content-Type", "application/json".parse().unwrap());

        let http_client = reqwest::Client::builder()
            .default_headers(headers)
            .build()
            .unwrap();

        Self {
            http_client,
            system: None,
            openai_client,
        }
    }

    pub async fn prepare_data(&self, prompt: &str) -> Result<String, ApiError> {
        let json = serde_json::json!({
            "q": prompt
        });

        let res = self
            .http_client
            .post("<https://google.serper.dev/search>")
            .json(&json)
            .send()
            .await
            .unwrap();

        let json = res.json::<Value>().await?;
        Ok(serde_json::to_string_pretty(&json)?)
    }
}

The Writer half of our AI agents will mostly be the same, save the reqwest::Client. We need to implement Agent alongside it, however,

#[derive(Clone)]
pub struct Writer {
    system: Option<String>,
    client: OpenAIClient<OpenAIConfig>,
}

impl Writer {
    pub fn new() -> Self {
        let api_key = env::var("OPENAI_API_KEY").unwrap();
        let config = OpenAIConfig::new().with_api_key(api_key);

        let client = OpenAIClient::with_config(config);
        Self {
            system: None,
            client,
        }
    }
}

impl Agent for Writer {
    fn name(&self) -> String {
        "Writer".to_string()
    }

    fn client(&self) -> OpenAIClient<OpenAIConfig> {
        self.client.clone()
    }

    fn system_message(&self) -> String {
        if let Some(message) = &self.system {
            message.to_owned()
        } else {
            "You are an agent.

        You will receive some context from another agent about some Google results that a user has searched.
        Your job is to research the Internet and to write a high-quality article that a user has written. The article must not appear to be AI written. The article should be SEO optimised without overly compromising the
        quality of the article.

        You are free to be as creative as you wish. However, each paragraph must have the following:
        - The point you are trying to make
        - If there is a follow up action point
        - Why the follow up action point exists (or why the user needs to carry it out)

        Search query:
".to_string()
        }
    }
}

Finally, we'll go back and fill the prompt method back in on the default Agent trait so that we have a default method implementation (and therefore don't need to keep re-implementing it):

use async_openai::types::{
    ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs,
    ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequestArgs,
};
    async fn prompt(&self, input: &str, data: String) -> Result<String, ApiError> {
        let input = format!(
            "{input}

            Provided context:
            {}
            ",
            serde_json::to_string_pretty(&data)?
        );
        let res = self
            .client()
            .chat()
            .create(
                CreateChatCompletionRequestArgs::default()
                    .model("gpt-4o")
                    .messages(vec![
                        //First we add the system message to define what the Agent does
                        ChatCompletionRequestMessage::System(
                            ChatCompletionRequestSystemMessageArgs::default()
                                .content(&self.system_message())
                                .build()?,
                        ),
                        //Then we add our prompt
                        ChatCompletionRequestMessage::User(
                            ChatCompletionRequestUserMessageArgs::default()
                                .content(input)
                                .build()?,
                        ),
                    ])
                    .build()?,
            )
            .await
            .map(|res| {
                //We extract the first one
                res.choices[0].message.content.clone().unwrap()
            })?;

        println!("Retrieved result from prompt: {res}");

        Ok(res)
    }

Writing our web service

Now that the hard work is over - we can implement the AI agents in our web application!

To get started, we'll create an AppState struct that implements Clone. Typically, this is a trait bound set by Axum or pretty much any Rust-based framework that you use. In it we'll have our Researcher and Writer struct:

use crate::agent::{Researcher, Writer};

#[derive(Clone)]
pub struct AppState {
    pub researcher: Researcher,
    pub writer: Writer,
}

impl AppState {
    pub fn new() -> Self {
        let researcher = Researcher::new();
        let writer = Writer::new();

        Self { researcher, writer }
    }
}

Next, we will write our handler endpoint that will take in a JSON input, run the agent pipeline and then return the end result:

#[derive(Deserialize, Serialize)]
pub struct Prompt {
    q: String,
}

#[axum::debug_handler]
async fn prompt(
    State(state): State<AppState>,
    Json(prompt): Json<Prompt>,
) -> Result<impl IntoResponse, ApiError> {
    let data = state.researcher.prepare_data(&prompt.q).await?;
    let resarcher_result = state.researcher.prompt(&prompt.q, data).await?;
    let writer_result = state.writer.prompt(&prompt.q, res).await?;

    Ok(writer_result)
}

And that's basically it!

We can then hook it all up by adding our endpoint to the router:

    let router = Router::new()
        .route("/", get(hello_world))
        .route("/prompt", post(prompt))
        .with_state(state);

Deployment

Now all we need to do is use cargo shuttle deploy (adding the --ad flag if on an uncommitted Git branch) and watch the magic happen!

Extending this project

Making a Pipeline struct for your agents

So let's say you've built this example, and want to go even further. What about pulling in another agent that generates a Twitter post or a LinkedIn post. At this point, you probably want to build a pipeline that holds all your agents, then you can just write .run_pipeline() and it'll do everything for you.

To do this, you could create a Pipeline trait does two things:

Initialise the agents set (and return it as a Vec<Box<dyn Agent>>)
Run the vector as a pipeline, where the results of the previous agent gets fed into the next one

An issue that you may come across with However, you may run into an issue with your Agents needing to implement Sized. This is because Clone requires an object to be a known size at compile-time - otherwise it won't work! To fix this, we can wrap the dyn type in a Box, allocating it on the heap. This works because the Clone trait requires the type to have a known, static size at compile-time.

Additionally, you might also receive an error with not being able to compile because of the prompt() method being async. You can fix this by adding the async_trait crate, then using the attribute macro above your code:

#[async_trait::async_trait]
trait Agent {
     // ..
}

Note that for every impl Agent for T, you’ll also need to remember to add the async trait macro. Otherwise, you’ll get an error about the lifetime annotations not matching!

Updating the prompt

If you’re not happy with the prompt results, don’t forgetyou can always update the message prompt that gets sent to your model!

Finishing Up

By leveraging Rust with the power of GPT-4o, you can develop a robust AI-powered content writer.

Additional Resources

For further learning and details, refer to:

Intro to advanced Rust traits & generics

Shuttle — Thu, 18 Apr 2024 16:09:26 +0000

Hello world! In this post we’re going to give a quick refresher course on Rust traits and generics, as well as implementing some more advanced trait bounds and type signatures.

A quick refresher on Rust traits

Writing a Rust trait is as simple as this:

pub trait MyTrait {
    fn some_method(&self) -> String;
}

Whenever a type implements MyTrait, you can guarantee that it will implement the some_method() function. To implement a trait simply requires that you implement the required methods (the ones with a semi-colon at the end).

struct MyStruct;

impl MyTrait for MyStruct {
    fn some_method(&self) -> String {
        "Hi from some_method!".to_string()
    }
}

You can also implement traits you don’t own on types you do own, or traits you do own on a type you don’t own - but not both! The reason you can’t do this is because of trait coherence. We want to make sure that we don’t accidentally have conflicting trait implementations:

// implementing Into<T>, a trait we don't own, on MyStruct
impl Into<String> for MyStruct {
    fn into(self) -> String {
        "Hello world!".to_string()
    }
}

// implementing MyStruct for a type we don't own
impl MyTrait for String {
    fn some_method(&self) -> String {
        self.to_owned()
    }
}

// You can't do this!
impl Into<String> for &str {
   fn into(self) -> String {
       self.to_owned()
   }
}

A common workaround for this is to create a newtype pattern - that is, a one-field tuple struct encapsulating the type we want to extend.

struct MyStr<'a>(&'a str);

impl<'a> Into<String> for MyStr<'a> {
    fn into(self) -> String {
        self.0.to_owned()
    }
}

fn main() {
    let my_str = MyStr("Hello world!");
    let my_string: String = my_str.into();

    println!("{my_string}");
}

If you have multiple traits that have the same method name, you need to manually declare what trait implementation you’re calling the type from:

pub trait MyTraitTwo {
    fn some_method(&self) -> i32;
}

impl MyTraitTwo for MyStruct {
    fn some_method(&self) -> i32 {
        42
    }
}

fn main() {
    let my_struct = MyStruct;
    println("{}", MyTraitTwo::some_method(&my_struct);
}

Sometimes, you might want the user to be able to have a default implementation as it may otherwise be quite tricky to do so. We can do this by simply defining the method within the trait.

trait MyTrait {
    fn some_method(&self) -> String {
        "Boo!".to_string()
    }
}

Traits can also require other traits! Take the std::error::Error trait for example:

trait Error: Debug + Display {
    // .. re-implement the provided methods here if you want
}

Here, we explicitly tell the compiler that our type must implement both Debug and Display traits before it can implement Error.

An introduction to marker traits

Marker traits are used as a “marker” for the compiler to understand that when a marker trait is implemented for a type, certain guarantees can be upheld. They have no methods or specific properties but are often used to ensure certain behaviors by the compiler.

There’s a couple of reasons why you would want marker traits:

The compiler needs to know if something can be guaranteed to do something
They’re an implementation-level detail that you can also implement manually

Two marker traits in particular, in conjunction with other lesser-used marker traits, are quite important to us: Send and Sync. Send and Sync are unsafe to implement manually - this is typically because you need to manually ensure that it is implemented safely. Unpin is also another example of this. You can find more about why it's unsafe to manually implement these traits here.

In addition to this, marker traits are also (generally speaking) auto traits. If a struct has fields that all implement an auto trait, the struct itself will also implement the auto trait. For example:

If all field types within a struct are Send, the struct is now automatically marked Send by the compiler with no input required from the user.
If all but one of your struct fields implement Clone but one doesn’t, your struct now cannot derive Clone anymore. You can get around this by wrapping the relevant type in an Arc or Rc - but this depends on your use case. In certain cases, this is not possible and you may need to think about an alternative solution.

Why do marker traits matter in Rust?

Marker traits in Rust form the core of the ecosystem and allows us to provide garuantees that may not be possible in other languages. For example, Java has marker interfaces which are analogous to Rust’s marker traits. However, marker traits in Rust are not just for behavior like Cloneable or Serializable; they also ensure that types can be sent across threads, for example. This is a subtle but far-reaching difference within the Rust ecosystem. By being able to declare a type as !Send for example, we can ensure that the user never attempts to try to send the type across a thread. This makes the problem of concurrency much easier to handle, as well as a few other things:

The Clone trait being required to clone things, but Clone also requires Copy (bitwise copying of a variable).
?Sized - Generics initially are given a static lifetime. You can allow them to be unsized (ie, use references) by adding ?Sized as a trait bound. Essentially what this means is when we implement ?Sized for a type, we are allowing it to be dynamically sized - so the compiler won't be able to allocate space for it at compile-time because it's not constant-size.

Object traits and dynamic dispatch

In addition to all of the above, traits can also make use of dynamic dispatch. Dynamic dispatch is essentially moving the process of selecting which implementation of a polymorphic function to use at runtime. While Rust does favour static dispatch for performance reasons, there are benefits to using dynamic dispatch through trait objects.

The most common pattern for using trait objects would be Box<dyn MyTrait>, where we are required to wrap the trait object in Box to make it implement the Sized trait. Because we’re moving the polymorphism process to runtime, the compiler can’t know what size the type is. Wrapping the type in a pointer (or "boxing" it) puts it on the heap instead of the stack.

// a struct with your object trait in it
struct MyStruct {
     my_field: Box<dyn MyTrait>
}

// an illegal struct that won't compile because of the Sized bound
struct MyStruct: Sized {
    my_field: Box<dyn MyTrait>
}

// this works!
fn my_function(my_item: Box<dyn MyTrait>) {
     // .. some code here
}

// this doesn't!
fn my_function(my_item: dyn MyTrait) {
     // .. some code here
}

The object type will then be computed during runtime, as opposed to generics which use compile-time.

The main advantages of dynamic dispatch are that your function doesn’t need to know the concrete type; as long as the type implements the trait, you can use it as a trait object (as long as it’s trait object safe). This is similar to the concept of duck-typing in other languages, where the functions and properties available to an object determine the typing. You also save some code bloat, which depending on your use can be a good thing.

Errors are also easier to understand from a library user perspective. From a library developer’s point of view this is not such an issue, but if you need to use a generics-heavy library, you can get some very confusing errors! Axum and Diesel are two libraries that can sometimes be guilty of this and have workarounds for this (Axum’s #[debug_handler] macro and Diesel’s documentation, respectively). Because you’re moving the dispatch process to runtime, you also save compilation time.

The downsides are that you need to ensure object trait safety. This basically means ensuring that your type doesn’t require Self: Sized. As mentioned earlier, this stems from the fact that by moving dispatch to runtime, the compiler can’t guess the size of your type - object traits don’t have a constant size at compile-time. This is also why we need to box dynamically dispatched objects and put them on the heap as mentioned earlier. Because of this, your application also takes a performance hit - although of course this depends on how many dynamically dispatched objects you’re using and how large they are!

To illustrate these points further, there are two HTML templating libraries that come to mind:

Askama, which uses macros and generics for compile-time checking
Tera, which uses dynamic dispatch for hot-reloading

Both of these libraries, while they can be used somewhat interchangeably for most use cases, have different trade-offs. Askama takes longer to compile and any errors will show in compile-time, but Tera only throws compilation errors at runtime and takes a performance hit due to dynamic dispatch. Zola, the Static Site Generator, uses Tera specifically because of certain design conditions that are unable to be satisfied by Askama.

Combining traits and generics

Getting started

Traits and generics synergise very well together and are easy to use. You can write a struct that implements generics like this without much trouble:

struct MyStruct<T> {
    my_field: T
}

However, to be able to use our struct with types from other crates, we will need to ensure that our struct can garuantee certain behavior. This is where we add trait bounds: conditions that a type must satisfy in order for the type to compile. A common trait bound you may find would be Send + Sync + Clone:

struct MyStruct<T: Send + Sync + Clone> {
    my_field: T
}

Now we can use any value we want for T as long as that type implements the Send, Sync and Clone traits!

As a more complex example of using traits with generics that you may occasionally need to re-implement for your own types, take the FromRequest trait from Axum for example (the below code snippet is a simplification of the original trait to illustrate the point):

trait FromRequest<S>
   where S: State
    {
    type Rejection: impl IntoResponse;

    fn from_request(r: Request, _state: S) -> Result<Self, Self::Rejection>;
}

Here we can also add trait bounds by using the where clause. This trait simply tells us that S implements State. However, State also requires the inner object to be Clone. By using complex trait bounds, we can create framework systems that make heavy use of traits to be able to do what some might refer to as “trait magic”. Take a look at this trait bound, for example:

use std::future::Future;

struct MyStruct<T, B> where
   B: Future<Output = String>,
   T: Fn() -> B
   {
    my_field: T
}

#[tokio::main]
async fn main() {
    let my_struct = MyStruct { my_field: hello_world };
    let my_future = (my_struct.my_field)();
    println!("{:?}", my_future.await);
}

async fn hello_world() -> String {
    "Hello world!".to_string()
}

The above one-field struct stores a function closure that returns impl Future<Output = String>, that we store hello_world in and then call it in the main function. We then wrap parenthesis around the field to be able to call it, then await the future. Note that we do not have brackets at the end of the field. This is because adding () at the end actually invokes the function! You can see where we invoke the function after declaring the struct, then awaiting it.

Usage in libraries

Combining traits and generics like this is extremely powerful. One use case where this is effectively leveraged is in HTTP frameworks. Actix Web for example, has a trait called Handler<Args> that takes a number of arguments, calls itself and then has a function called call that produces a Future:

pub trait Handler<Args>: Clone + 'static {
     type Output;
     type Future: Future<Output = Self::Output>;

     fn call(&self, args: Args) -> Self::Future;
}

This then allows us to extend this trait to a handler function. We can tell the web service that we have a function that has an inner function, some arguments and implements Responder (Actix Web’s HTTP response trait):

pub fn to<F, Args>(handler: F) -> Route where
    F: Handler<Args>,
    Args: FromRequest + 'static,
    F::Output: Responder + 'static {
         // .. the actual function  code here
    }

Note that other frameworks like Axum also follow this same methodology to provide an extremely ergonomic developer experience.

Finishing up

Thanks for reading! While traits and generics can be a confusing topic to understand, hopefully this guide to using Rust traits and generics has shed some light on the subject!

Working with OpenAPI using Rust

Shuttle — Thu, 04 Apr 2024 15:24:07 +0000

Hello world! In this article we’re going to talk about how you can make the most of OpenAPI with Rust, by learning all the different ways we can use OpenAPI in a Rust context. By the end of this article, you'll learn the following:

Adding OpenAPI to a Rust web service
How to generate API client libraries from OpenAPI specifications
Working with the OpenAPI spec directly

What is OpenAPI?

From Swagger.io:

The OpenAPI specification (OAS) defines a standard, language-agnostic interface to HTTP APIs which allows both humans and computers to discover and understand the capabilities of the service without access to source code.

No matter what language you’re using, OpenAPI allows you to easily read the specification and understand how to use an API without needing specific documentation.

Here is a simple example of what an OpenAPI specification file may look like:

swagger: "3.0"
info:
  version: "1.0"
  title: "Hello World API"
paths:
  /hello/{user}:
    get:
      description: Returns a greeting to the user!
      parameters:
        - name: user
          in: path
          type: string
          required: true
          description: The name of the user to greet.
      responses:
        200:
          description: Returns the greeting.
          schema:
            type: string
        400:
          description: Invalid characters in "user" were provided.

There are several benefits of using OpenAPI:

You can improve cross-team collaboration by allowing teammates to quickly experiment with endpoints by providing a frontend
You can quickly get an understanding of endpoints that your teammates have made
It’s widely used, so you can get an understanding of official APIs that utilise it much faster
We can generate code from it as it’s machine parseable!

Adding OpenAPI to a Rust API

utoipa

Adding an OpenAPI specification to a Rust API can be done with the utoipa family of crates. utoipa is a crate that primarily uses macros to set up the OpenAPI specification. There is also support for frontend GUIs like Swagger UI, Redoc and Rapidoc that allow you to visualise working with your API

A simple Axum example that shows a JSON representation of your OpenAPI specification looks like this:

use std::net::SocketAddr;

use axum::{routing::get, Json};
use utoipa::OpenApi;

#[derive(OpenApi)]
#[openapi(paths(openapi))]
struct ApiDoc;

/// Return JSON version of an OpenAPI schema
#[utoipa::path(
    get,
    path = "/api-docs/openapi.json",
    responses(
        (status = 200, description = "JSON file", body = ())
    )
)]
async fn openapi() -> Json<utoipa::openapi::OpenApi> {
    Json(ApiDoc::openapi())
}

#[tokio::main]
async fn main() {
    let socket_address: SocketAddr = "127.0.0.1:8080".parse().unwrap();
    let listener = tokio::net::TcpListener::bind(socket_address).await.unwrap();

    let app = axum::Router::new().route("/api-docs/openapi.json", get(openapi));

    axum::serve(listener, app.into_make_service())
        .await
        .unwrap()
}

Let’s break this down. We have the following:

An ApiDoc struct that takes the OpenApi derive macro and sets all the attributes required to serve the OpenAPI specification from your API.
We have an attribute macro above our function handler. This macro will document the given information about a handler endpoint and show it in the OpenAPI spec when we open it.
We have a “list” of responses in the macro. Note that because we have no exact type to give to OpenAPI, we leave the body as ().

Interested in checking out what attributes the utoipa::path macro can take? Have a look here.

If you run this code and visit [localhost:8080/api-docs/openapi.json](http://localhost:8080/api-docs/openapi.json) you should see a JSON response of the API specification.

This is typically good enough for just providing a basic representation. However, for internal exploration you may want to add a GUI to your OpenAPI specification. You can do this by installing the utoipa_swagger_ui crate and changing your Router to the following:

let app = Router::new().merge(
    SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", ApiDoc::openapi())
);

If you run your code again and go to localhost:8080/swagger-ui, you’ll get a Swagger UI menu that should look something like this:

Two small things to note here are that utoipa-service is taken from the crate name and crate is the module where our route comes from. Here because our function is taken from the top level, it says crate - but we can fix this later on if we wanted by putting the route in a module.

If we then expand this section by clicking on the route, it will allow us to then execute the endpoint! Pretty helpful, huh?

The utoipa crate is quite comprehensive, so you can be assured that it will support mostly everything you want to do. If you’d like to check out their documentation, you can do so here.

When it comes to having an API with hundreds or thousands of endpoints though, this may not be entirely ergonomic. You’ll be spending quite a lot of time writing macros which can bloat your files. To that end, you can also use the utoipauto crate which lets you automate all of the work with only one macro. However, this adds additional compilation time. Whether you’ll want to use it depends on your use case.

poem-openapi

Should you be happening to use the Poem framework (which hit 3.0.0 recently!), you can also use the poem-openapi crate to add OpenAPI functionality to your Poem service. Similarly to the utoipa crate, poem-openapi also uses macros to get OpenAPI documentation.

use poem::{listener::TcpListener, Route};
use poem_openapi::{param::Query, payload::PlainText, OpenApi, OpenApiService};

struct Api;

#[OpenApi]
impl Api {
    #[oai(path = "/hello", method = "get")]
    async fn index(&self, name: Query<Option<String>>) -> PlainText<String> {
        match name.0 {
            Some(name) => PlainText(format!("hello, {}!", name)),
            None => PlainText("hello!".to_string()),
        }
    }
}

#[tokio::main]
async fn main() -> Result<(), std::io::Error> {
    let api_service =
        OpenApiService::new(Api, "Hello World", "1.0").server("http://localhost:3000/api");
    let ui = api_service.swagger_ui();
    let app = Route::new().nest("/api", api_service).nest("/", ui);

    poem::Server::new(TcpListener::bind("0.0.0.0:3000"))
        .run(app)
        .await
}

Running this code will also generate a Swagger UI GUI as above, but at a different endpoint (localhost:3000/api).

Generating Rust from OpenAPI specifications

Now let’s talk about generating Rust code from OpenAPI specifications. The OpenAPI collective have made a tool to generate a server client library - Rust support included! We’ll be using npm to install the OpenAPI generator. You can install it with the following shell snippet:

npm install @openapitools/openapi-generator-cli -g

There are also alternative ways to install the OpenAPI generator, which you can check out here.

Next, we’ll need a specification to generate a client library from. The generator takes YAML or JSON files as input. Thankfully for us, we already have a JSON filefrom the OpenAPI service we just built. If we head to localhost:8080/api-docs/openapi.json, we can select the Raw Text option, prettify it and then put it all into a JSON file for input. Our file name will be utoipa-client.json.

Next, we’ll actually generate the client code. This can be done with the following shell snippet:

npx @openapitools/openapi-generator-cli generate -i utoipa-client.json -g rust -o ./utoipa-client

This looks like quite a long command! What’s happening here?

The -i flag is our input file
The -g flag is for the generator we should use (in this case, the rust one). Generators are not necessarily one-per-language, hence the flag convention.
The -o flag is for the output directory. If the directory doesn’t exist, the generator will attempt to create it.

Once done, you should see a new Rust crate in the automatically generated utoipa-client folder. At this point in time, the generator is mostly correct. However, your generated code can have syntactical errors if your OpenAPI specification input is malformed. You can find out more about this here.

Interested in further customization? You can check out more OpenAPI generator configuration options here.

Working with OpenAPI specifications directly in Rust

Interested in building tools that use the OpenAPI specification? There’s a crate for that! Using the openapiv3 crate, you can also deserialize (and serialize) to and from the OpenAPI spec format. Below, we deserialize it from a JSON string - you would additionally need serde_json installed:

use serde_json;
use openapiv3::OpenAPI;

fn main() {
    let data = include_str!("openapi.json");
    let openapi: OpenAPI = serde_json::from_str(data)
        .expect("Could not deserialize input");
    println!("{:?}", openapi);
}

There are several crates for handling this; notably, the openapiv3 crate only supports the V3 specification. For V3.1 you want to use the oas3 crate, which can take both YAML and JSON:

fn main() {
    match oas3::from_path("path/to/openapi.yaml") {
      Ok(spec) => println!("spec: {:?}", spec),
      Err(err) => println!("error: {}", err)
    }
}

Finishing up

Thanks for reading! With this article, you should be able to tackle OpenAPI with Rust no problem.

Sending logs to Grafana Loki in a Rust application

Shuttle — Fri, 29 Mar 2024 16:49:44 +0000

Hello world! We will look at how you can leverage Grafana Loki for log storage and analysis. Application monitoring tools are a crucial part of monitoring and observability. They allow you to examine exactly how your application works, what is going in/out, and what is going wrong. Using application monitoring tools saves time and money by being able to fix production issues faster and retaining users.

At the end of this article, you’ll have a Rust web service deployed freely to Shuttle that logs traces to Grafana. Interested in checking out the final repository? You can find that here.

Why should I use Grafana Loki?

Loki is a logging service for Grafana designed with scalability and cost-effectiveness in mind. Rather than indexing the contents of your logs, it only stores the metadata and labels. A set of labels for each log stream is also used. Every unique set of labels represents each new stream - so if you add or remove any labels from a stream for example, you create an entirely new stream. With Grafana, you can also additionally visualize all of your data quite easily using either pre-created dashboards or making your own.

Getting started

Pre-requisites

To get the most out of this article, you need to either self-host Grafana services or sign up to Grafana Cloud. For this article we’ll mostly be referencing Grafana Cloud as this is the easiest way to use their services (without manual setup).

To get started, you will want an API token that has the write:logs permissions. This can be done from Grafana Cloud user management. Make sure you save your Grafana user and API token variables, as you’ll need them in just a little bit.

For initializing our service, we’ll use cargo shuttle init --template axum (requires cargo-shuttle installed) to create a new Shuttle project with the Axum template. We will then add the following dependencies with this shell snippet:

cargo add url
cargo add tracing-subscriber -F fmt,env-filter
cargo add tracing-loki
cargo add tracing
cargo add base64
cargo add shuttle-runtime --no-default-features

We’ll also additionally need to add the #[shuttle_runtime::Secrets] annotation macro to our main function. This allows us to automatically grab secrets from our Secrets.toml file when we use cargo shuttle run to run our Shuttle service. It should look like this:

[..]
use shuttle_runtime::SecretStore;

#[shuttle_runtime::main]
async fn main(
    #[shuttle_runtime::Secrets] secrets: SecretStore
) -> shuttle_axum::ShuttleAxum {
    let router = Router::new().route("/", get(hello_world));

    Ok(router.into())
}

The Secrets.toml file will be located at the project folder root and takes a key-value format.

GRAFANA_USER = "<your-grafana-user-here>"
GRAFANA_API_KEY = "<your-grafana-api-token-here>"

Building

To get started, we will be using the tracing_loki crate to build the base of our tracing subscriber. tracing_loki allows us to create a tracing_subscriber layer that collects and exports logs to a Grafana Loki instance. On creation of a Grafana Cloud account we are given a free data source by default for Loki, which tracing_loki is compatible with.

The default Grafana data source uses the basic authentication scheme, which you can find more about here. Later on when we create the tracing_loki tracing layer, we will add this as a HTTP header.

We’ll start by grabbing our secrets and making a Base64 string out of them:

fn init_grafana_subscriber(store: SecretStore) {
    let grafana_user = store.get("GRAFANA_USER").unwrap();
    let grafana_password = store.get("GRAFANA_API_KEY").unwrap();

    let basic_auth = format!("{grafana_user}:{grafana_password}");
    let encoded_basic_auth = BASE64_STANDARD.encode(basic_auth.as_bytes());

    // .. rest of code
}

Next, we’ll need to create the tracing_loki layer which sends the logs to a Grafana instance. A few things are going on in the below snippet:

We add a label (with a key and a value)
We add an extra field for the Process ID (PID). This also takes a key-value format, so if you wanted to add anything else, you would do it here.
We set the HTTP header using the .http_header() function, setting the Authorization header as required.

use url::Url;

let url = Url::parse("https://logs-prod-012.grafana.net").expect("Failed to parse Grafana URL");

let (layer, task) = tracing_loki
    ::builder()
    .label("application", "shuttle-grafana")
    .unwrap()
    .extra_field("pid", format!("{}", process::id()))
    .unwrap()
    .http_header("Authorization", format!("Basic {encoded_basic_auth}"))
    .unwrap()
    .build_url(url)
    .unwrap();

Additionally, we will want to create an EnvFilter. Having the tracing subscriber set at the default trace logging level can be useful. However, due to parsing headers and other similar actions that all use trace! spans, there are quite a lot of them. This can lead you to go over the free tier limits unintentionally. Here, we will set the default directive so that we only get errors where the logging level is DEBUG or above (i.e. warning or an error).

use tracing_subscriber::filter{EnvFilter, LevelFilter};

let filter = EnvFilter::builder()
    .with_default_directive(LevelFilter::DEBUG.into())
    .parse("").unwrap();

The last thing to do is to create a tracing subscriber with the layers that we’ve created, and then initialise it! You may have noticed earlier that the tracing_loki::builder() method also generates a task that deals with log aggregation. We will need to spawn a Tokio task to handle this.

// We need to register our layer with `tracing`.
tracing_subscriber::registry()
    .with(filter)
    .with(tracing_subscriber::fmt::Layer::new())
    .with(layer)
    // One could add more layers here, for example logging to stdout:
    // .with(tracing_subscriber::fmt::Layer::new())
    .init();

// The background task needs to be spawned so the logs actually get
// delivered.
tokio::spawn(task);

Now the function is done! We can add it to our fn main at the start of the function.

#[tracing::instrument]
async fn hello_world() -> &'static str {
    tracing::debug!("An event happened!");
    "Hello, world!"
}

#[shuttle_runtime::main]
async fn main(
    #[shuttle_runtime::Secrets] secrets: SecretStore
) -> shuttle_axum::ShuttleAxum {
    setup_tracing(&secrets);

    let router = Router::new().route("/", get(hello_world));

    Ok(router.into())
}

When using cargo shuttle run now, your terminal will become populated with traces. If you visit localhost:8000 in the browser, you should see a debug tracing event with the description An event happened! in your traces. Note that it may take Grafana 5-10 minutes to receive your traces.

Reading your logs

To read your logs, you need to create a dashboard for your data source. Head over to your Grafana Cloud instance, find the Grafana Cloud data source and create a dashboard for it. Then once you’re on the dashboard, you can query your logs! A best practices guide for building Grafana dashboards can be found here.

The label we added to our logs should show up under the Labels column. It will contain the extra label we put in (the application label), as well as the logging level of the trace. The trace will also contain exactly what was in the trace message. If you need to add fields in the #[tracing::instrument] macro, you can do so and it will show up in Grafana. You can also find a guide for understanding labels here.

Deploying

To deploy, simply use cargo shuttle deploy and watch the magic happen!

Finishing up

Thanks for reading! Application monitoring is just one step to ensuring our Rust web services are performing better than ever. By using instrumentation, we can reduce the need for manual debugging.

Making a notification service with Rust & AWS SNS

Shuttle — Wed, 20 Mar 2024 22:25:28 +0000

Notifications are a very helpful tool for notifying services. This is particularly relevant if you’re using microservice architecture. One notification mechanism we can use is AWS SNS (Simple Notification Service), which lets us create topics, subscribe to them and push messages to topics. By the end of this article, you’ll have two fully functioning web services with the following overall functionality:

One web service will have a single endpoint that can receive a message from AWS SNS
One web service will handle the job of sending messages

Interested in just deploying? You can find the repository here, complete with instructions on how to deploy.

Getting started

Pre-requisites

To get started, you will need two keys from AWS:

Your access key ID
Your secret access key

Receiving AWS SNS messages

To start with, let’s have a look at how we can receive SNS messages. We can do this with cargo shuttle init, picking Axum as our framework. This service will have a single endpoint that takes the SNS messages and prints the message out.

To make this idiomatic, we’ll want to create a struct called SnsMessage that implements axum::FromRequest. This will allow us to use it as an extractor, rather than trying to parse the POST requests from SNS manually.

use axum::{routing::{get, post}, Router, response::{IntoResponse, Response},
    http::StatusCode, extract::{Request, FromRequest},
    Json, RequestExt
};
use serde::Deserialize;

#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct SnsMessage {
    #[serde(rename = "Type")]
    kind: String,
    message_id: String,
    topic_arn: String,
    subject: String,
    message: String,
    timestamp: String,
    signature_version: String,
    signature: String,
    signing_cert_url: String,
    unsubscribe_url: String
}

#[axum::async_trait]
impl<S> FromRequest<S> for SnsMessage {
    type Rejection = Response;

    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
        let headers = req.headers();

        // if none of these headers are sent in the request
        // automatically send a BAD_REQUEST status code
        if !headers.contains_key("x-amz-sns-message-type")
            | !headers.contains_key("x-amz-message-id")
            | !headers.contains_key("x-amz-topic-arn")
            | !headers.contains_key("x-amz-subscription-arn") {
            return Err((StatusCode::BAD_REQUEST).into_response())
        }

        let Json(payload): axum::Json<SnsMessage> = req.extract()
            .await.map_err(|_| (StatusCode::BAD_REQUEST).into_response())?;

        Ok(payload)
    }
}

Note that this FromRequest<S> implementation is primarily to show the fundamentals. In production use cases, you would want to verify the signature sent from AWS to make sure it’s accurate. You can find out more about the basics of this here.

To use our new struct, you can use SnsMessage as an extractor for a handler function like this:

async fn receive_sns(
    sns_message: SnsMessage
) -> StatusCode {
    println!("{}", sns_message.message);
    StatusCode::OK
}

You can find out more about receiving POST requests from AWS SNS here. To extend this, you may want to think about verifying the signature sent by AWS here.

We’ll tie this all together by adding the handler function to our router:

use axum::routing::post;

#[shuttle_runtime::main]
async fn axum() -> shuttle_axum::ShuttleAxum {
    let router = Router::new().route("/", get(hello_world))
        .route("/sns", post(receive_sns));

    Ok(router.into())
}

Now it’s complete!

Deploying

To deploy, you can use cargo shuttle deploy (with --allow-dirty if on a Git branch and uncommitted changes). Take note of the deployment URL! We will need this later on.

Sending AWS SNS messages

Now for the second part of our notification service: sending messages to our receiver!

Setup

We’ll get started by using cargo shuttle init once again, picking Axum as our choice of framework. Make sure to follow the prompt until the end to finalise your project.

Once done, we can install our dependencies with this script:

cargo add aws-config@1.1.8 -F behavior-version-latest
cargo add aws-credential-types@1.1.8 -F hardcoded-credentials
cargo add aws-sdk-sns@1.18.0
cargo add anyhow@1.0.81
cargo add serde@1.0.197 -F derive
cargo add serde-json@1.0.114
cargo add thiserror@1.0.58

Remember your AWS keys from earlier? We’ll add them to Secrets.toml file in the root of our project folder. It should look like this:

AWS_ACCESS_KEY_ID = "your-access-key-id-here"
AWS_SECRET_ACCESS_KEY = "your-secret-access-key-here"

Setting up our AWS config

To get started, we’ll want to add our resource annotations to our main function - which should look like this:

#[shuttle_runtime::main]
async fn main(
    #[shuttle_runtime::Secrets] secrets: SecretStore,
) -> shuttle_axum::ShuttleAxum {
   // .. your code here
}

When we use cargo shuttle run, our secrets will now get provisioned to us without needing to do anything!

Before we start adding more code, let’s write our AppState struct to hold Axum state. This will allow our handler functions to access our AWS SNS client whenever we want.

use aws_sdk_sns::Client;

#[derive(Clone)]
pub struct AppState {
    sns: Client,
    topic_arn: String
}

Next, we’ll want to add code to our main function that gets the secrets, creates aws_credential_types::Credentials and creates an AWS config from the credentials, as well as adding a region. We’ll use eu-west-02 to reduce latency (because the Shuttle servers are hosted on eu-west-02). We then create a new client, initialize our app state and then add it to the axum::Router;

let access_key_id = secrets.get("AWS_ACCESS_KEY_ID").expect("AWS_ACCESS_KEY_ID not set in Secrets.toml");
    let secret_access_key = secrets.get("AWS_SECRET_ACCESS_KEY").expect("AWS_ACCESS_KEY_ID not set in Secrets.toml");

        let creds = Credentials::from_keys(
               access_key_id,
               secret_access_key,
               None
            );

        let cfg = aws_config::from_env()
            .region(Region::new("eu-west-02"))
            .credentials_provider(creds)
            .load()
            .await;

       let sns = Client::new(&cfg);

    let state = AppState { sns };

    let router = Router::new().route("/", get(hello_world)).with_state(state);

Error handling

Error handling for the ShuttleAxum type

The shuttle_axum::ShuttleAxum return type defaults to anyhow::Error for user errors. This doesn’t affect handler functions. For functions used in the main function however, they will need to either be able to convert to anyhow::Error or be of the same type.

Error handling for API route functions

To create errors easily, we’ll use the thiserror crate we installed to generate From<T> implementations for our error type, as well as the error messages. Let’s have a look.

use thiserror::Error;

use aws_sdk_sns::error::SdkError;
use aws_sdk_sns::operation::publish::PublishError

#[derive(Debug, Error)]
pub enum ApiError {
    #[error("Error while publishing message: {0}")]
    PublishMessage(#[from] SdkError<PublishError>),
}

We use the thiserror::Error derive macro to enable the attribute macros. A few things are happening here:

The #[from] implementation automatically derives From<T> so that the given type will automatically turn into a given enum variant. Note however that if the error type is an enum, it will convert it regardless of the enum variant. If you want more custom error handling, you may want to implement From<T> manually.
#[error("..")] automatically generates the std::fmt::Display implementation for our error type. Note that we need to use this on every enum variant, otherwise we will receive compile-time errors.
Using thiserror::Error automatically implements std::error::Error for ApiError so there’s no need for us to re-implement it!

Next, we’ll want to implement axum::response::IntoResponse so that we can return ApiError from our Axum handler functions.

use axum::{
    http::StatusCode,
    response::{IntoResponse, Response},
};

impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        (StatusCode::INTERNAL_SERVER_ERROR, self.to_string().into_response())
    }
}

Currently we just have all of our error variants return internal server errors. To go further into this, you could create a response depending on the internal error enum variant or source (via pattern matching).

Topics

To create a topic, we can use the AWS SNS client to create a builder object and then add tags or attributes to it depending on what we want.

For the endpoint, it is mostly about creating the builder, appending the attributes and using send() to create it.

async fn create_topic(sns: &Client, name: &str) -> AnyhowResult<String> {
    let topic = sns.create_topic().name(name);

    let output = topic.send().await.unwrap();

    println!("Topic created: {output:?}");
    Ok(output.topic_arn.unwrap())
}

The output from sending the topic will contain the ARN. We need to store this to be able to retrieve a topic later on. Note that attempting to create two topics with the same name will return the current ARN instead of creating a new resource.

If you’re interested in customising this further, you can check out more about the CreateTopicFluentBuilder here.

Subscriptions

Receiving published messages requires a topic subscription. Creating a subscription requires declaring an endpoint, the topic ARN as well as the protocol used. We can refer to the SubscribeFluentBuilder for this here.

The function will look something like this:

async fn subscribe_to_topic(sns: &Client, url: &str, arn: &str) -> AnyhowResult<()> {
    let sub = sns
        .subscribe()
        .protocol("https".to_string())
        .endpoint(url)
        .topic_arn(arn);

    let output = sub
        .send()
        .await
        .map_err(|e| anyhow::anyhow!("error: {e}"))?;
    println!("New subscriber created: {output:?}");

    Ok(())
}

There is not much to talk about here as it’s primarily just setting it up and then sending it. If we wanted to, we could also create multiple subscriptions to the same topic ARN, giving us the benefit of being able to fan-out our notifications.

Because we’re setting up an HTTPS endpoint, we need to confirm the subscription. When we create the subscription, SNS will send a subscription confirmation message to our receiver. This requires our receiver to already be deployed. Because we already print out the confirmation message in our logs, we can easily access and visit the URL to confirm the subscription. You can find out more about this here.

To check that the subscription exists, we can do this:

async fn subscription_exists(sns: &Client, url: &str, arn: &str) -> AnyhowResult<bool> {
    let subscribers = sns
        .list_subscriptions()
        .send()
        .await
        .map_err(|e| anyhow::anyhow!("error: {e}"))?;

    if let Some(subs) = subscribers.subscriptions {
        if subs.iter().any(|sub| {
            sub.clone().endpoint.unwrap() == *url && sub.clone().topic_arn.unwrap() == *arn
        }) {
            return Ok(true);
        }

        return Ok(false);
    }

    Ok(false)
}

Publishing messages

To be able to publish messages, we’ll need the topic ARN we want to subscribe to, the message as well as the subject.

#[derive(Deserialize)]
struct PublishMessageParams {
    message: String,
    subject: String
}

async fn publish_message(
    State(state): State<AppState>,
    Json(json): Json<PublishMessageParams>
) -> Result<impl IntoResponse, ApiError> {
    let res = state.sns.publish()
        .topic_arn(state.arn)
        .message(json.message)
        .subject(json.subject)
        .send().await?;

    Ok(StatusCode::OK)
}

All of the endpoints who are subscribed to the topic will now automatically receive the message!

Interested in customising this further? You can check out more about the PublishFluentBuilder here.

Connecting it all together

To connect it all, let’s revisit our main function and fill it in with the new functions:

#[shuttle_runtime::main]
async fn main(#[shuttle_secrets::Secrets] secrets: SecretStore) -> shuttle_axum::ShuttleAxum {
    let access_key_id = secrets
        .get("AWS_ACCESS_KEY_ID")
        .expect("AWS_ACCESS_KEY_ID not set in Secrets.toml");
    let secret_access_key = secrets
        .get("AWS_SECRET_ACCESS_KEY")
        .expect("AWS_ACCESS_KEY_ID not set in Secrets.toml");

    let creds = Credentials::from_keys(access_key_id, secret_access_key, None);

    let cfg = aws_config::from_env()
        .region(Region::new("eu-west-02"))
        .credentials_provider(creds)
        .load()
        .await;

    let sns = aws_sdk_sns::Client::new(&cfg);
    let topic_arn = create_topic(&sns, "my_topic").await?;

    // NOTE: Change this to your deployment URL for your receiver service!
    // The SNS receiver route should be on `/sns` as instructed previously
    let url = "https://sns-receiver.shuttleapp.rs/sns";

    if !subscription_exists(&sns, url, &topic_arn).await? {
        subscribe_to_topic(&sns, url, &topic_arn).await?;
    }

    let state = AppState { sns, topic_arn };

    let router = Router::new().route("/", get(hello_world)).with_state(state);

    Ok(router.into())
}

Deploying

To deploy, we just need to use cargo shuttle deploy (with the --allow-dirty flag if on an uncommitted Git branch) and watch the magic happen!

Finishing up

Thanks for reading! AWS SNS is a powerful tool to be able to get going with notifying services, particularly if you need fan-out from something like AWS SQS or other message queues.

Learn more about how to competently write an API with Axum
Secure your API with JWT authentication
Learn about getting started with logging in Rust

Using PostHog with Rust

Shuttle — Thu, 14 Mar 2024 18:00:47 +0000

Hello world! In this article, we’re going to talk about how you can use Posthog with Rust to be able to track user behavior across an application. Analytics are more important than ever to be able to improve your product; Posthog makes it easy to do exactly that.

How does PostHog work?

Posthog works by allowing you to capture events and then convert those events into “insights” by aggregating the data. PostHog allows you to then query the data using HogQL (a DSL made by PostHog for querying data,, similar to SQL). Alhough there isn’t an officially supported Rust SDK for PostHog, they do have a community-supported one (posthog-rs) that can capture events in your application.

We’ll also look at using reqwest to set up our own custom client for sending requests to the PostHog API!

Capturing events

We can get started with the posthog-rs client by creating it:

use posthog::{ client, Event };
use std::env;

let posthog_api_key = env::var("POSTHOG_API_KEY").expect("POSTHOG_API_KEY env var missing");
let client = client(posthog_api_key);

The event essentially allows us to add anything we want - for example, we can add user IDs, user cookies or timestamps and whatever else we need.

Next, we’ll want to make a function for sending an event to the PostHog API:

fn send_posthog_event(client: &posthog::Client) {
    let mut event = Event::new("my_database_event", "1234");
    event.insert_prop("action", "insert").unwrap();
    client.capture(event).unwrap();
}

Now whenever we need to track an action, we can do so. Below is an example of using the client in an Axum handler function to be able to send an event:

use axum::{http::StatusCode, extract::State};

#[derive(Clone)]
struct AppState {
    db: sqlx::PgPool,
    posthog: posthog::Client
}

async fn add_to_database(
    State(state): State<AppState>
) -> StatusCode {
    let res = sqlx::query("SELECT * FROM USERS")
        .fetch_all(&state.db)
        .await
        .unwrap();

    send_posthog_event(&state.client);

    StatusCode::OK
}

Using the posthog-rs client in an async context

It should be noted that the posthog-rs client is not async-friendly due to using the blocking feature of the reqwest crate. This means that if you want to use it with async (for example in a web service), you will need to use tokio::task::block_in_place() like so:

use axum::{Extension, StatusCode, response::IntoResponse};

async fn my_example_axum_handler(
    Extension(client): Extension<posthog::Client>
) -> impl IntoResponse {
    tokio::task::block_in_place(|| move {
        let mut evt = Event::new("my_database_event", "Hello world!");
        evt.insert_prop("action", "insert").unwrap();

        client.capture(evt).unwrap();
    });

    StatusCode::OK
}

Alternatively, Shuttle has a fork of a fork of posthog-rs that contains an async client. You can find out more here. Due to being a fork it’s not an officially published crate, but you can add it to your Rust project with this shell snippet:

cargo add posthog-async --git https://github.com/shuttle-hq/posthog-rs.git --branch main

Now when you refer to it in your Rust files, you can use posthog_async as the dependency name.

Using a custom client

In this section, we’re going to create a custom client for working with Posthog. We’ll get started by installing reqwest, which is a library for writing HTTP requests and is async by default (requiring a feature to be blocking):

cargo add reqwest

Next, we will want to write a function that we can use to make our request client easily. We’ll need to add bearer authorization on the header, which we can do like so:

let posthog_api_key = std::env::var("POSTHOG_API_KEY")
    .expect("Could not find POSTHOG_API_KEY environment variable");

let client = reqwest::Client::builder()
    .header("Authorization", format!("Bearer {}"))
    .build().unwrap();

Now whenever we use client, it will always have the appropriate header attached.

Using PostHog insights

Creating insights

Insights in Posthog are the main building blocks of dashboards and allow you to visualise how users use your product. To create insights, we can use our reqwest::Client to send an API key manually. There’s quite a few endpoints - we’ll be making POST requests to /api/projects/:project_id/insights/ for this example.

To be able to set the JSON body, we will need to make use of the serde_json crate which we can install with this shell snippet:

cargo add serde_json

Next, we’ll want to make a JSON map - which we can do below with the serde_json::json!() macro:

use serde_json::json;

fn my_json_body() -> serde_json::Value {
    json!({
        "name": "my_insight",
    })
}

You can find a more extensive list of request body parameters here. Once you’re finished creating the request body you want to send to Posthog, you can then create the request using the client we created.

async fn create_insight(
    client: request::Client,
    project_id: String
) -> Result<(), Box<dyn std::error::Error>> {
    let url = format!(
        "https://app.posthog.com/api/projects/{project_id}/insights/"
    );

    let json_body = my_json_body();
    let response = client.post(&url).json(json_body).send.await?;
}

However, this is just the start of insight creation. To make our insight do anything, we need to add a funnel to it. Similarly to the last API endpoint where we attached some basic values, we’ll create a new JSON body using the json!() macro and then send it:

async fn create_insight_funnel(
    client: request::Client,
    project_id: String
) -> Result<(), Box<dyn std::error::Error>> {
    let url = format!(
        "https://app.posthog.com/api/projects/{project_id}/insights/"
    );

    let json_body = json!({
        events: [{"id":"my_database_event"}],
        date_from: "-1m"
    })
    let response = client.post(&url).json(json_body).send.await?;
}

Here, we’ve specified in the JSON body we want to get events from up to a month ago and the name of the event we want to get is my_database_event, which we created earlier. Note that the $pageview event signifies any time a user views a page that PostHog is added to.

Retrieving an insight

Once you have some events and insights built up, it’s time to retrieve them! You can do this by making a GET request to the insights URL. See below:

async fn get_insights(
    client: request::Client,
    project_id: String
) -> Result<(), Box<dyn std::error::Error>> {
    let url = format!(
        "https://app.posthog.com/api/projects/{project_id}/insights/"
    );

    let response = client.get(&url).send.await?;
}

When you get the response, you will need to turn it into a struct that implements Deserialize. Thankfully, reqwest provides a method for turning the response into JSON-compatible structs. Let’s make a struct that represents some of the properties of the JSON response from PostHog. serde completely ignores undeclared fields by default, allowing us to only take what we need.

#[derive(Deserialize, Debug)]
struct PosthogResponse {
    count: i32,
    next: String,
    previous: String,
    results: Vec<PosthogResult>
}

#[derive(Deserialize, Debug)]
struct PosthogResult {
    id: String,
    name: String,
    deleted: bool
}

Now we can expand our previous function to include the struct conversion:

async fn get_insights(
    client: request::Client,
    project_id: String
) -> Result<(), Box<dyn std::error::Error>> {
    let url = format!(
        "https://app.posthog.com/api/projects/{project_id}/insights/"
    );

    let response = client.get(&url).send.await?;
    let json: PosthogResponse = response.json().await?;

    println!("{json:?}");
    Ok(())
}

Finishing Up

Posthog makes it easy to track how users are using your application and find easy wins. While the Rust SDK is not complete, hopefully you’ve found some clarity in communicating with the Posthog API with Rust!

Get started with using Rust's most popular framework here
Learn about logging for your Rust application here
Check out our guide for using Stripe with Rust here

Building a simple web server in Rust

Shuttle — Wed, 13 Mar 2024 15:37:45 +0000

Building a web server with Rust doesn’t need to be complex. With frameworks like Axum, you can write a web server without hassle. Leveraging Rust allows you to easily take on the job of web services written in other languages and more. In this post we’re going to talk about how you can build and deploy a simple web server using Axum.

What Rust framework should I use?

While there’s a lot of options we can use, our personal choice is Axum for a few reasons:

Axum uses generics and traits. This allows you to leverage Rust language tooling in ways that other frameworks may not.
Axum has familiar syntax (using handler functions for routing).
It has an extremely high level of compatibility with tower crates. This allows you to go very low-level if required.

Additionally, we also have an article about what framework you should use here.

Getting started

To get started, you’ll want Rust installed on your system. Don’t have it installed? You can get it from this install page.

Next, you’ll want to use cargo shuttle init (requires cargo-shuttle installed). and then pick Axum for our framework.

Hello world!

When creating a project with cargo init, you will need to manually create (or copy!) the initial boilerplate. This may look something like this:

use axum::{Router, routing::get};
use std::net::SocketAddr;
use tokio::net::TcpListener;

async fn hello_world() -> &'static str {
    "Hello world!"
}

#[tokio::main]
async fn main() {
    let router = Router::new().route("/", get(hello_world));

    let addr = SocketAddr::from(([127,0,0,1], 8000));
    let tcp = TcpListener::bind(&addr).await.unwrap();

    axum::serve(tcp, router).await.unwrap();
}

As you can see from the above, we do a few things:

We set up a router with given routes and the functions that need to be called.
We define an address for our web server to receive requests at, bind it to a TCP listener.
The TCP listener then responds to requests and responds accordingly.

If you used cargo run to load this program up then go to [localhost:8000](http://localhost:8000) in the browser, it would return “Hello world!” as a raw text response.

When initialising a project with Shuttle, all of this gets set up for you. The basic project comes with a native integration so that you don’t need to set up the socket binding manually. The integration code is quite short and revolves around the usage of the shuttle_runtime::Service trait. If you have a non-standard service you want to run, you can run the service in a struct that implements the Service trait and you’ll be ready to go!

See below for an example of how a Shuttle “Hello World” project looks like:

use axum::{Router, routing::get};

async fn hello_world() -> &'static str {
    "Hello world!"
}

#[shuttle_runtime::main]
async fn main() -> shuttle_axum::ShuttleAxum {
    let router = Router::new().route("/", get(hello_world));

    Ok(router.into())
}

Routing

HTTP responses from routing within Rust frameworks can be done by anything that implements a trait that represents a HTTP response. In Axum, this would be the IntoResponse trait (or IntoResponseParts for headers and other non-response body parts). Web servers can only return things that are valid HTTP responses. Implementing IntoResponse (and IntoResponseParts respectively) ensures this!

It is possible to use impl IntoResponse as the return type for a function (for convenience). However, we would need to make sure all of the responses are of the same type. This can lead to confusion later down the line, particularly if you’re working in a team.

For JSON responses, Axum provides the handy Json<T> struct we can use as a response type by wrapping a type with it. For example, this snippet below shows how you can return some raw JSON:

use serde_json::{json, Value};
use axum::Json;

async fn return_some_json() -> Json<Value> {
    let json = json!({"hello":"world"})

    Json(json)
}

However, a more likely situation is that you’ll want to return data that follows a schema. We can use the Deserialize and Serialize traits from the serde crate to do this. We can easily apply these traits by adding the derive feature and then adding it as a derive macro to a struct:

use serde::{Deserialize, Serialize};

#[derive(Deserialize, Serialize)]
struct MyStruct {
    my_field: String
}

Now we can wrap it in axum::Json and return the struct in our HTTP response.

async fn return_a_struct_as_json() -> Json<MyStruct> {
     let my_struct = MyStruct { my_field: "Hello world!".to_string() };

     Json(my_struct)
}

Extractors

Extractors are handler function arguments. They extract parts of the HTTP request and turn it into simple variables that we can use with our application. We can use extractors for a lot of things:

Accessing shared mutable state (by adding state to our application then accessing it in the function)
Extracting a typed header for authorization purposes
Consuming the request body to extract a JSON or Form body, depending on what you want.
If there’s nothing readily available for our use case, we can just implement it ourselves!

Here is an example of how you can use extractors. Note here that we use de-structuring to automatically get the inner variable in Json<MyStruct> as it looks cleaner.


async fn function_with_extractors(
    Json(json): Json<MyStruct>,
) -> impl IntoResponse {
    format!("The contents of my_field is: {}", json.my_field)
}

Databases

Using databases generally isn’t much different in Rust than any other language. The main difference is Rust web frameworks generally use shared mutable state to pass around data. This means that you’ll want to first initialise your database connection (pool) and pass it around using state. In Axum, state is required to implement Clone. If your type cannot implement Clone because one or more types don’t implement Clone, you can wrap the state struct in a std::sync::Arc which does implement Clone.

Here’s an example of how you can do exactly that, using SQLx with Postgres as example. We initialise the PgPool, initialise our state struct and attach it to the router.

use sqlx::PgPool;
use axum::{extract::State, Router, routing::get, http::StatusCode};

#[derive(Clone)]
struct MyState {
    db: PgPool
}

#[tokio::main]
async fn main() {
    let db: PgPool = PgPool::connect("<your-db-connection-url-here>").await.unwrap();
    let state = MyState { db };
    let router = Router::new().route("/", get(hello_world)).with_state(state);

    // the rest of your code...
}

With Shuttle, we can provision a database by simply adding a database macro. This saves time both locally and in production! To get started, we’ll need to add the shuttle_shared_db dependency.

cargo add shuttle-shared-db -F postgres,sqlx

Then we simply add it to our main function and initialise the state struct again like before.

use sqlx::PgPool;
use axum::{Router, routing::get};

#[shuttle_runtime::main]
async fn main(
    #[shuttle_shared_db::Postgres] db: PgPool,
) -> shuttle_axum::ShuttleAxum {
    let state = MyState { db };
    let router = Router::new().route("/", get(hello_world)).with_state(state);

    Ok(router.into())
}

To use our state struct, we can use the axum::extract::State extractor:

use axum::{extract::State, http::StatusCode};

async fn hello_world(
    State(state): State<MyState>
) -> StatusCode {
    sqlx::query("SELECT 'Hello world!")
         .execute(&state.db)
         .await
         .unwrap();

    StatusCode::OK
}

Static files

To get started with static files on Axum, we’ll create a folder in the root of our project called static. We’ll then define it in a file named Shuttle.toml in the project root:

assets = ["static/*"]

Using the wildcard tag allows us to serve any file contained within the folder by using the file path.

We can write a HTML file in the static folder called index.html:

<h1>Hello world</h1>

To serve this static folder on Axum, we’ll need to import some things from tower-http. We will run the following shell snippet:

cargo add tower-http -F fs

Then we can add it like below:

use tower_http::services::{ServeDir};

let router = Router::new()
    .route_service("/", ServeDir::new("static"));

If you are using a SPA framework like React or Vue for your static files, you will want to additionally set up a .not_found_service() to be able to serve index.html:

let router = Router::new()
    .route_service("/", ServeDir::new("static")
        .not_found_service(ServeFile::new("static/index.html") 
     )
);

Deployment

To deploy Rust to Shuttle, you can use cargo shuttle deploy and watch the magic happen! Don’t forget to add the --allow-dirty flag if on a Git branch with uncommitted changes. Besides web server development, Shuttle is a lifesaver when it comes to easily deploying and making your web service public.

Interested? You can find our docs here.

Finishing up

Web development doesn’t have to be complicated. With Rust, we can achieve our goal of building a simple web server quickly and easily.

Get more in-depth with Axum here.
Learn more about tracing libraries and improve application logging here.

Writing a WASM module in Rust

Shuttle — Wed, 06 Mar 2024 22:35:15 +0000

Hello world! In today’s post we’re going to talk about how you can write a WebAssembly module in Rust. WebAssembly is a portable compilation target for programming languages to be able to conveniently inter-op with JavaScript on the web. Rust being able to take advantage of this has made it extremely useful for many use cases, such as:

CPU intensive workloads (encryption)
GPU intensive workloads (image/video processing, image recognition)

This article will focus on writing a WASM module for image processing that can be used on the backend, as well as exploring common ways to deploy WASM and its targets.

Getting started

To get started, you’ll need Rust installed. If you don’t, you can install it here.

We’ll be focusing on trying out writing a WASM module in three different ways:

Using the wasm-bindgen CLI
Using wasm-pack
Using napi-rs

We will initially use wasm-bindgen-cli to create our application, then look at using wasm-pack. The focus of this article will be creating a simple module for image processing. Byte array manipulation and data processing is an area where Rust can significantly speed up your application.

Before we start, make sure you have the wasm32-unknown-unknown target installed. If you don’t, you can add it like so:

rustup target add wasm32-unknown-unknown

Note that for trying out our module, you’ll also additionally want npm (or any alternatives) installed.

Writing a WASM module

The basics

To set up our project, we’ll get started with using cargo init --lib wasm-example to create a new library project named wasm-example. We will then install our dependencies with a small shell snippet:

cargo add wasm-bindgen@0.2.91
cargo add js-sys@0.3.68
cargo add image@0.24.9

We will also want to add the dynamic library flag to our Cargo.toml file. Normally, it lets Cargo know that we want to make a dynamic system library - but when using it with the WebAssembly target, it simply means “make a *.wasm file without a start function”. To do this, we can add this small snippet below:

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

JavaScript types in Rust

To be able to use JavaScript types in Rust, we need to use extern C in addition to using the wasm-bindgen macros. This allows us to import functions straight out of JavaScript and into Rust!

The Hello World application in WASM looks like this (from the book):

use wasm_bindgen::prelude::*;

#[wasm_bindgen]
extern "C" {
    fn alert(s: &str);
}

#[wasm_bindgen]
pub fn greet(name: &str) {
    alert(&format!("Hello, {}!", name));
}

Note that the alert function in the extern C is taken directly from JavaScript, and is what allows us to call it in our Rust function. If we were to compile this and execute it in a JavaScript file, it would be the same as if you had called alert() from regular JavaScript.

We can apply the same logic to be able to work with other types and functions - namely, buffers. Vec<u8> in JavaScript can either be represented in one of two ways:

The Uint8Array type (the direct JavaScript equivalent of Vec<u8>)
A Buffer type

Buffer is a subclass of Uint8Array. This is because when Node.js first released, there was no Uint8Array type - which is what led to the creation of the Buffer type. Later down the line when Uint8Arrays were introduced with ES6, both were eventually merged as it made sense to do so. Many JavaScript libraries still use Buffer. By using js-sys, we can get interoperability between JavaScript and Rust - which we can see below by defining the Buffer type and providing a method with the buffer() method:

use js_sys::ArrayBuffer;
// This defines the Node.js Buffer type
#[wasm_bindgen]
extern "C" {
    pub type Buffer;

    #[wasm_bindgen(method, getter)]
    fn buffer(this: &Buffer) -> ArrayBuffer;

    #[wasm_bindgen(method, getter, js_name = byteOffset)]
    fn byte_offset(this: &Buffer) -> u32;

    #[wasm_bindgen(method, getter)]
    fn length(this: &Buffer) -> u32;
}

Now when we write our WASM function, we can refer to the Buffer type directly!

Let’s write our Rust function for converting our image file format. We’ll make it require our Buffer and then have it return Vec<u8> - when we compile it through wasm-pack or another compiler, it will automatically get converted to a Uint8Array.

use js_sys::{ArrayBuffer, Uint8Array};
use wasm_bindgen::prelude::wasm_bindgen;
use image::ImageFormat;
use image::io::Reader;
use std::io::Cursor;

// .. extern C stuff goes here

#[wasm_bindgen]
pub fn convert_image(buffer: &Buffer) -> Vec<u8> {
    // This converts from a Node.js Buffer into a Vec<u8>
    let bytes: Vec<u8> = Uint8Array::new_with_byte_offset_and_length(
        &buffer.buffer(),
        buffer.byte_offset(),
        buffer.length()
    ).to_vec();

    let img2 = Reader::new(Cursor::new(bytes)).with_guessed_format().unwrap().decode().unwrap();

    let mut new_vec: Vec<u8> = Vec::new();
    img2.write_to(&mut Cursor::new(&mut new_vec), ImageFormat::Jpeg).unwrap();

    Ok(new_vec)
}

Building via wasm-bindgen-cli

Here, we need to compile from Rust to WASM by building our package for the wasm32-unknown-unknown target, which we can do like so:

cargo build --target=wasm32-unknown-unknown

Next, we need to use wasm-bindgen to generate the JS glue code to make it all work. We’ll use the nodejs target which will generate a CommonJS module and put it in the ./pkg folder which we can then implant anywhere we want.

wasm-bindgen --target nodejs --out-dir ./pkg \
./target/wasm32-unknown-unknown/release/wasm_example.wasm

Now we can either publish our WASM code as a package or implant it anywhere we want to use it!

I don’t want to use CommonJS!

If you don’t want to use CommonJS because you’re using ESM (EcmaScript modules, or ES6 modules), that’s cool! The CLI currently allows several targets:

bundler (produces code for usage with bundlers like Webpack)
web (directly loadable in a web browser)
nodejs (loadable via require as a CommonJS Node.js module)
deno (usable as a Deno module)
no-modules (like the web target but doesn’t use ES Modules).

There are specific docs for usage with ES which you can check out here. The easiest way to do it in terms of what compiler to use is typically with Webpack as it’s the most compatible. There is also an additional guide you can use for compiling to ES6 modules without a bundler here - though it involves initializing the WASM module manually before running which adds some overhead.

Test driving our new module

Now that we’ve written our code, let’s try it out! We will spin up a JavaScript backend server using Express.js. We will assume you’re running the following in the same folder as your Rust project (for convenience purposes). We’ll get started with the following shell snippet:

npm init -y
npm i express express-fileupload

Next, we’ll create a server.js file in our root directory and insert the following code:

const fileUpload = require('express-fileupload');
const express = require('express');
const { convert_image } = require('./pkg/wasmmeme.js');

const app = express();
const port = 3030;

app.get('/', (req, res) => {
  res.send(`
    <h2>With <code>"express"</code> npm package</h2>
    <form action="/api/upload" enctype="multipart/form-data" method="post">
      <div>Text field title: <input type="text" name="title" /></div>
      <div>File: <input type="file" name="file"/></div>
      <input type="submit" value="Upload" />
    </form>
  `);
});

app.post('/api/upload', (req, res, next) => {
        const image = convertImage(req.files.file.data)

  res.setHeader('Content-disposition', 'attachment; filename="meme.jpeg"');
  res.setHeader('Content-type', 'image/jpg');
    res.send(image);
  });

app.listen(port, () => {
  console.log(`Example app listening on port ${port}`)
})

This snippet does the following:

We set up an Express server at port 3030
We have a route at / that will give us a HTML form when we visit it in the browser
We have an API route that will grab the data from our file upload, convert it to a new format, set the correct headers and return the new image.

If we use node server.js, head to [http://localhost:3030](http://localhost:3030) in our browser then fill the form out and attach an image, we should get an image download response back!

Note that depending on the settings you are using for your image file format conversion, your file size may increase post-conversion; this is because you may be using lossless conversion. If you want to use lossy conversion to decrease your file size, you want the new_with_quality method while instantiating the image encoder in your Rust code.

Building our app with alternative CLIs

While wasm-bindgen-cli is useful, it’s also the most low level CLI out of our options and you can run into spurious issues while using it such as wasm-bindgen version incompatibility issues. There are some additional quality of life changes that we could benefit from, such as automatic versioning and wasm-opt usage. Let’s take a quick look at some of these other options and see how they compare.

Wasm-pack

wasm-pack is a tool that aims to be a one-stop-shop for compiling Rust to WASM. It includes a CLI that you can use by installing it here. In comparison to using wasm-bindgen-cli, it has a number of quality of life upgrades:

Comes with wee_alloc, a WebAssembly allocator with a (pre-compression) 1kB code footprint.
Comes with a panic hook that allows you to debug Rust panic messages in the browser

To initialise our project, we can use wasm-pack new wasm-example which will do everything for us. Code-wise, our main function (and C/JS bindings) will remain the same as wasm-pack provides primarily tooling additions to make compilation easier and does not have any library code we can use.

napi-rs

napi-rs is a framework for building pre-compiled Node.js addons in Rust. If you find using wasm-bindgen too complicated to use and just want to write Node.js stuff, this is a good alternative. To use it, Node v0.10.0 or later is required. You can install it with the following shell snippet (requires npm or its alternatives):

npm install -g @napi-rs/cli

Once done, you can then use napi new wasm-example to build your new NAPI project!

napi-rs does come with some code changes, which you’ll be able to see below: we can finally get rid of the extern C block and instead use napi’s bindgen_prelude to include whatever we need.

use napi::bindgen_prelude::*;
use image::io::Reader;
use image::ImageFormat;
use image::ImageOutputFormat;
use std::io::Cursor

#[macro_use]
extern crate napi_derive;

#[napi]
pub fn convert_image(buffer: Buffer) -> Result<Buffer> {
    let bytes: Vec<u8> = buffer.into();

    let img2 = Reader::new(Cursor::new(bytes)).with_guessed_format().unwrap().decode().unwrap();

    let mut new_vec: Vec<u8> = Vec::new();
    img2.write_to(&mut Cursor::new(&mut new_vec), ImageFormat::Jpeg).unwrap();

    Ok(new_vec.into())
}

The advantages of this are clear:

We don’t need to manually import anything using extern C
We can easily use Node.js internals without trouble

Of course for all its advantages, napi-rs is only compatible with Node.js. If you want to write some WASM code for the browser, you would need to default to wasm-pack or wasm-bindgen. You additionally also need to use the Node ecosystem to keep your CLI updated, which from a Rust-first standpoint is a bit of a strange decision. Needless to say, napi-rs is a pretty easy way to start writing Node.js with Rust.

Finishing up

Thanks for reading! Rust has great inter-op with WASM, and there's no reason why we shouldn't be able to take advantage of this to help us when using other lanaugages.

We wrote a list of our top 8 Rust tools to speed up your productivity
Here’s our guide to getting started with Axum, Rust's most popular framework
We wrote about our top 8 tools to help you be more productive in Rust

A Full Stack SaaS Template with Loco

Shuttle — Thu, 29 Feb 2024 16:30:19 +0000

Loco is a Rust framework that aims to do it all - authentication, tasks, migrations and more. While initially being more complex to use than using pure Axum, it can save a lot of time setting up boilerplate and provides a great platform to build on.

In this guide, we’ll deploying Loco on Shuttle via a fullstack template that also additionally includes SaaS subscription payments. By using their SaaS starter, we can leverage their pre-built authentication features to build payments out quickly and easily. Interested in deploying or trying out the final repo? You can check it out here.

Steps to deploy from cloning:

Run cargo shuttle init --from joshua-mo-143/shuttle-stripe-ex and follow the prompt (requires cargo-shuttle installed)
Set up API keys (see below)
Use cargo shuttle deploy --allow-dirty and watch the magic happen!

Prerequisites

Before we get started, you will need a Stripe API key. It’s free to sign up with Stripe and you can turn on Test Mode before you do anything in production. Stripe also has docs on this if you need assistance here.

Once you’ve created your project, make sure you create a Secrets.toml file in the root of your project and add it like so:

STRIPE_API_KEY = "<YOUR-STRIPE-KEY-HERE>"

Getting started

We’re going to use the following command to initialise our project (requires cargo-shuttle installed), following the prompt to initialise a project with our name. The --from flag allows us to take the starter from

cargo shuttle init --from loco-rs/loco --subfolder starters/saas

We will then use cargo loco generate deployment to generate a Shuttle deployment for our Loco project!

While we’re here, make sure you’re using the latest version of cargo-shuttle and Shuttle dependencies in your project. We released v0.40.0 today!

We’ll also add the following dependencies with a shell snippet:

cargo add async-stripe@0.34.1 -F runtime-tokio-hyper-rustls
cargo add shuttle-shared-db -F postgres
cargo add shuttle-secrets

Note that you will get a blank Shuttle deployment with no database installed. To remedy this, we’ll adjust the main function to provide our database annotation and make sure that the Migrator (from the migrations folder) is added:

use migrations::Migrator;

#[shuttle_runtime::main]
async fn main(
    #[shuttle_shared_db::Postgres] conn_str: String,
    #[shuttle_metadata::ShuttleMetadata] meta: shuttle_metadata::Metadata,
    #[shuttle_secrets::Secrets] secrets: shuttle_secrets::SecretStore
) -> shuttle_axum::ShuttleAxum {
    std::env::set_var("DATABASE_URL", conn_str);

    let stripe_api_key = secrets
        .get("STRIPE_API_KEY")
        .expect("STRIPE_API_KEY not found in secrets");
    std::env::set_var("STRIPE_API_KEY", stripe_api_key);

    let environment = match meta.env {
        shuttle_metadata::Environment::Local => Environment::Development,
        shuttle_metadata::Environment::Deployment => Environment::Production,
    };
    let boot_result = create_app::<App, Migrator>(StartMode::ServerOnly, &environment)
        .await
        .unwrap();

    let router = boot_result.router.unwrap();
    Ok(router.into())
}

In the regular src/bin/main.rs, you may also need to adjust the app so that Migrator is also included.

Once done, you can use cargo shuttle run and it should just automatically work! You’ll get a database connection URL (save this for later!).

Migrations

To get started, we’ll make some migrations that we can then reference later on in the program. You can do this like so:

cargo loco generate model --migration-only subscription_tiers tier:string! stripe_item_id:string! stripe_price_id:string!
cargo loco generate model --migration-only user_subscriptions user:references stripe_customer_id:string! stripe_subscription_id:string! user_tier:string!

Then we’ll use the following to migrate your database and generate entities:

DATABASE_URL=<DB_URL_HERE> cargo loco db migrate
DATABASE_URL=<DB_URL_HERE> cargo loco db entities

Note that you will need a database URL for this. If you don’t have one yet, you can use cargo shuttle run to automatically spin up a Postgres container with a provided connection string or spin up your own Docker container.

These two commands will generate some files in the migrations folder as well as in src/models, which we’ll be making heavy use of as they are the main way to interface with the database when using Loco.

Frontend

We won’t be covering the frontend in this tutorial as there’s a lot of different ways you can do it - however, if you’d like to look at the way that we’ve done it, feel free to check out the repo here! We use React as provided by Loco with react-router-dom for routing and zustand for state management, with vanilla CSS.

The following pages in the repo have been provided:

A home page
Login and register pages
A dashboard page that allows users to downgrade/upgrade their subscription tier, cancel it and check what tier they are.
Pricing and payment checkout pages
A payment success/fail page

Error handling

Loco by default uses anyhow to be able to provide easy error handling. However, for our purposes, let’s create our own error type. This will allow us a couple of things:

We know exactly what error is happening and where
We can customise the behavior of our error handling

Let’s start by using the thiserror crate we added earlier to add macros for automatically implementing std::fmt::DIsplay and std::error::Error. The thiserror::Error derive macro also allows us to add attribute macros to our struct for automatic From<T> implementations, which saves a lot of time! That being said, you can also implement From<T> manually if you want to create more than one enum variant for an error based on what the reason of the error is.

use thiserror::Error;

#[derive(Error, Debug)]
pub enum ApiError {
    #[error("Stripe error: {0}")]
    Stripe(#[from] stripe::StripeError),
    #[error("User already has this subscription tier!")]
    UserTierAlreadyExists,
    #[error("SQL error: {0}")]
    SQL(#[from] sea_orm::DbErr),
    #[error("Model error: {0}")]
    Model(#[from] loco_rs::model::ModelError),
}

To be able to use this in our API, we will need to implement axum::response::IntoResponse. We can do this by simply matching each enum variant like below:

use axum::{http::StatusCode, response::{Response, IntoResponse}};

impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        let res = match self {
            Self::Stripe(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
            Self::SQL(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
            Self::UserTierAlreadyExists => (
                StatusCode::BAD_REQUEST,
                "User already has this tier!".to_string(),
            ),
            Self::Model(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
        };
        res.into_response()
    }
}

Using Stripe

To get started with using Stripe, we’ll want to create a new loco controller file that will hold all of our routes. We can do this with cargo loco generate controller stripe, which will generate a new controller route at src/controllers/stripe.rs and inject some code into src/app.rs to make sure the controller gets automatically included.

Looking inside src/controllers/stripe.rs should give you a function that returns Routes (a struct that builds on axum::Router) and a couple of routes for returning “Hello, world!” and the contents of a given request.

Let’s start first by defining what our user tiers are. Let’s say we have the Pro tier, and the Team tier. We can write an enum with relevant impls like so:

use std::fmt;
use serde::{Deserialize, Serialize};
use sea_orm::{EnumIter, DeriveActiveEnum};

#[derive(EnumIter, DeriveActiveEnum, Clone, Deserialize, Debug, Serialize, PartialEq, Eq)]
#[sea_orm(rs_type = "String", db_type = "String(StringLen::N(1))")]
pub enum UserTier {
    #[sea_orm(string_value = "P")]
    Pro,
    #[sea_orm(string_value = "T")]
    Team,
}

impl UserTier {
    fn get_price(&self) -> Option<i64> {
        match self {
            Self::Pro => Some(1000),
            Self::Team => Some(2500),
        }
    }

    fn from_str(str: &str) -> Self {
        match str {
            "Pro" => Self::Pro,
            "Team" => Self::Team,
            _ => panic!("There should only be the Pro and Team tier!")
        }
    }
}

impl fmt::Display for UserTier {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Pro => write!(f, "Pro"),
            Self::Team => write!(f, "Team"),
        }
    }
}

Note that the macros we are importing from sea_orm will let the enum be stored as a varchar type in Postgres.

Creating Stripe products and prices

Before we do anything else, we will want to create a Stripe product that has some prices attached to it. To do that, we can create two functions; one for creating the Product item, and then one for adding a price (as products on Stripe can have multiple prices). Creating a price requires a product. To keep it simple, we’ll keep one price attached to one item. Both functions will be used as part of other functions.

use stripe::{Client, Product, Price, CreatePrice, CreateProduct, IdOrCreate, 
CreatePriceRecurring, CreatePriceRecurringInterval, Currency};

async fn create_product_item(client: &Client, user_tier: &UserTier) -> Result<Product, ApiError> {
    let product = {
        let mut create_product = match user_tier {
            UserTier::Pro => CreateProduct::new("Pro User Subscription"),
            UserTier::Team => CreateProduct::new("Team Subscription"),
        };

        create_product.metadata = Some(std::collections::HashMap::from([(
            String::from("async-stripe"),
            String::from("true"),
        )]));

        Product::create(client, create_product).await?
    };

    Ok(product)
}

async fn create_product_price(
    client: &Client,
    user_tier: &UserTier,
    product: &Product,
) -> Result<Price, ApiError> {
    let price = {
        let mut create_price = CreatePrice::new(Currency::USD);
        create_price.product = Some(IdOrCreate::Id(&product.id));
        create_price.metadata = Some(std::collections::HashMap::from([(
            String::from("async-stripe"),
            String::from("true"),
        )]));
        create_price.unit_amount = user_tier.get_price();

        create_price.recurring = Some(CreatePriceRecurring {
            interval: CreatePriceRecurringInterval::Month,
            ..Default::default()
        });
        create_price.expand = &["product"];
        Price::create(client, create_price).await?
    };

    Ok(price)
}

This then allows us to build a more higher-level function that can either simply retrieve the Stripe product ID if it already exists in the database, or create a product with price then save the details in the database. For simplicity (and not wanting to deal with securing financial information in the database), we will only be storing the IDs and relevant information, such as what product tier a user is.

use sea_orm::DatabaseConnection;

async fn retrieve_product(
    client: &Client,
    user_tier: &UserTier,
    db: &DatabaseConnection,
) -> Result<Price, ApiError> {
    use crate::models::_entities::subscription_tiers;

    let product = subscription_tiers::Entity::find()
        .filter(subscription_tiers::Column::Tier.contains(user_tier.to_string()))
        .one(db)
        .await?;

    let price = match product {
        Some(product) => {
            Price::retrieve(
                client,
                &PriceId::from_str(&product.stripe_price_id.to_string()).unwrap(),
                &["product"],
            )
            .await?
        }
        None => {
            let product = create_product_item(client, &user_tier).await?;
            let price = create_product_price(client, &user_tier, &product).await?;

            let tier_model = subscription_tiers::ActiveModel {
                tier: ActiveValue::Set(user_tier.to_string()),
                stripe_item_id: ActiveValue::Set(product.id.to_string()),
                stripe_price_id: ActiveValue::Set(price.id.to_string()),
                ..Default::default()
            };

            subscription_tiers::Entity::insert(tier_model)
                .exec(db)
                .await?;

            price
        }
    };

    Ok(price)
}

Later on, when we run the web service, the API should automatically be able to know if it is required to remake the Stripe product or not.

Let’s write a function for creating a subscription. To get started, we will define what the JSON input should look like when the API receives a request:

#[derive(Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct UserSubscription {
    name: String,
    email: String,
    card_num: String,
    exp_year: i32,
    exp_month: i32,
    cvc: String,
    user_tier: UserTier,
}

Note that when sending the request, the variables must be in camel case and not snake case.

To make things a little bit easier for ourselves later, we will add an impl for UserSubscription which will allow us to automatically turn it into some structs that we’ll be using later on to create the customer and CardDetailsParams.

impl UserSubscription {
    fn as_create_customer(&self) -> CreateCustomer {
        CreateCustomer {
            name: Some(&self.name),
            email: Some(&self.email),
            description: Some(
                "A paying user.",
            ),
            metadata: Some(std::collections::HashMap::from([(
                String::from("async-stripe"),
                String::from("true"),
            )])),

            ..Default::default()
        }
    }

    fn as_card_details_params(&self) -> CardDetailsParams {
        CardDetailsParams {
            number: self.card_num.to_string(),
            exp_year: self.exp_year,
            exp_month: self.exp_month,
            cvc: Some(self.cvc.clone()),
            ..Default::default()
        }
    }
}

Creating subscriptions

Next, we’ll want to get started on writing an endpoint for creating user subscriptions! We’ll create a stripe::Client here from the API key that we stored earlier.

use crate::models::_entities::user_subscriptions;
use crate::models::_entities::users;
use loco_rs::controller::middleware::auth::JWTWithUser;

pub async fn create_subscription(
    State(ctx): State<AppContext>,
    auth: middleware::auth::JWTWithUser<users::Model>,
    Json(json): Json<UserSubscription>,
) -> Result<StatusCode, ApiError> {
    let secret_key = std::env::var("STRIPE_API_KEY").expect("Missing STRIPE_API_KEY in env");
    let client = Client::new(secret_key);

    // .. rest of your code
}

Here, note that we specifically use the JWTWithUser middleware to extract a Bearer JWT from the Authorization header and return a user model.

Next, we want to create a customer and add it to our Stripe account. We will also create the payment method and attach it to the customer. Note here that while we’re using card payments as it’s a very common form of payment, Stripe also has quite a few other types of payments you can try.

use stripe::{Customer, PaymentMethod, PaymentMethodTypeFilter, 
CreatePaymentMethodCardUnion, AttachPaymentMethod};

let customer = Customer::create(&client, json.as_create_customer()).await?;

let payment_method = {
    let pm = PaymentMethod::create(&client, CreatePaymentMethod {
        type_: Some(PaymentMethodTypeFilter::Card),
        card: Some(CreatePaymentMethodCardUnion::CardDetailsParams(json.as_card_details_params())),
        ..Default::default()
    }).await?;

    PaymentMethod::attach(&client, &pm.id, AttachPaymentMethod {
        customer: customer.id.clone(),
    }).await?;

    pm
};

Next, we’ll want to add the part that creates the subscription. This part is relatively simple as there’s only one item in our subscription list we want - which is the base price for our SaaS subscription (although if you want to extend it, there are options for adding more too!):

use stripe::{CreateSubscription, CreateSubscriptionItems, Subscription,

let mut params = CreateSubscription::new(customer.id.clone());
params.items = Some(
    vec![CreateSubscriptionItems {
        price: Some(price.id.to_string()),
        ..Default::default()
    }]
);

params.default_payment_method = Some(&payment_method.id);
params.expand = &["items", "items.data.price.product", "schedule"];

let subscription = Subscription::create(&client, params).await?;

let subscription_activemodel = user_subscriptions::ActiveModel {
    user_id: ActiveValue::Set(auth.user.id),
    stripe_customer_id: ActiveValue::Set(customer.id.to_string()),
    stripe_subscription_id: ActiveValue::Set(subscription.id.to_string()),
    user_tier: ActiveValue::Set(json.user_tier),
    ..Default::default()
};

user_subscriptions::Entity::insert(subscription_activemodel).exec(&ctx.db).await?;

Ok(StatusCode::OK);

If you get stuck on errors while writing this function, you can find the function in the repo here.

Cancelling Stripe Subscriptions

Okay, now let’s say you want your users to be able to use self-service to cancel the SaaS subscription. This primarily involves canceling the subscription and then making sure to update the database. In this case, we are choosing to outright delete the record from the database on successful cancellation.

Firstly, we’ll create the Stripe client again by grabbing our API key:

pub async fn cancel_subscription(
    State(ctx): State<AppContext>,
    auth: middleware::auth::JWTWithUser<UsersModel>,
) -> Result<StatusCode, ApiError> {
    let user = users::Model::find_by_pid(&ctx.db, &auth.claims.pid).await?;

    let secret_key = std::env::var("STRIPE_API_KEY").expect("Missing STRIPE_API_KEY in env");
    let client = Client::new(secret_key);

    // .. rest of your code
}

Next, we’ll find our user subscription from the user_subscriptions table based on the user ID foreign key. We will then use the Stripe subscription ID to cancel the subscription.

let subscription = user_subscriptions::Entity
    ::find()
    .filter(user_subscriptions::Column::UserId.eq(user.id))
    .one(&ctx.db).await?
    .unwrap();

let _ = Subscription::cancel(
    &client,
    &SubscriptionId::from_str(&subscription.stripe_subscription_id).unwrap(),
    CancelSubscription {
        cancellation_details: None,
        invoice_now: Some(true),
        prorate: Some(true),
    }
).await?;

Once the subscription has been successfully canceled and there’s nothing else we need to do with Stripe, we can go back and update our database to delete the record:

let subscription_to_delete = user_subscriptions::Entity
    ::find()
    .filter(user_subscriptions::Column::UserId.eq(user.id))
    .one(&ctx.db).await?
    .unwrap();

subscription_to_delete.delete(&ctx.db).await?;

Ok(StatusCode::OK);

Note that there are quite a few different ways to handle this. You may find that you want to explicitly mark a customer’s subscription as “expired” rather than outright deleting it from the database if you want users to be able to check their subscription history and other such details.

Upgrading/Downgrading Subscription Tiers

Finally, we will add the ability to upgrade and downgrade subscription tiers. In Stripe terms, we’re grabbing information about a user’s subscription and updating the price ID of an existing item on the subscription.

As before, we’ll start with creating the stripe::Client:

pub async fn update_subscription_tier(
    State(ctx): State<AppContext>,
    auth: middleware::auth::ApiToken<UsersModel>,
    Json(new_user_tier): Json<UpdateUserTier>,
) -> Result<StatusCode, ApiError> {
    use crate::models::_entities::subscription_tiers;
    use crate::models::_entities::user_subscriptions;

    let secret_key = std::env::var("STRIPE_API_KEY").expect("Missing STRIPE_API_KEY in env");
    let client = Client::new(secret_key);
    // .. rest of your code
}

After that, we will find the user_subscription based on the user ID. We will then retrieve the subscription data from Stripe and get the first item on the subscription list. Note that while we are using a vector index which can technically panic, there should always be at least one item so it is safe to use index 0. We will also return an error if the user’s current tier is the same as the requested tier.

let user_subscription = user_subscriptions::Entity
    ::find()
    .filter(user_subscriptions::Column::UserId.eq(auth.user.id))
    .one(&ctx.db).await?
    .unwrap();

let subscription_item = Subscription::retrieve(
    &client,
    &SubscriptionId::from_str(&user_subscription.stripe_subscription_id).unwrap(),
    &["items"]
).await?.items;

let subscription_item = &subscription_item.data[0];

if new_user_tier.user_tier == user_subscription.user_tier {
    return Err(ApiError::UserTierAlreadyExists);
}

Once done, we can then find the subscription tier data from our database according to the requested tier change and update the subscription using the new tier’s Stripe price object ID.

let new_subscription = subscription_tiers::Entity
    ::find()
    .filter(subscription_tiers::Column::Tier.contains(new_user_tier.user_tier.to_string()))
    .all(&ctx.db).await?;

let new_sub_tier: String = new_subscription
    .iter()
    .find(|x| x.tier == new_user_tier.user_tier.to_string())
    .map(|x| x.stripe_price_id.to_string())
    .unwrap();

let updated_subscription = Subscription::update(
    &client,
    &SubscriptionId::from_str(&user_subscription.stripe_subscription_id).unwrap(),
    UpdateSubscription {
        items: Some(
            vec![UpdateSubscriptionItems {
                id: Some(subscription_item.id.to_string()),
                price: Some(new_sub_tier),
                ..Default::default()
            }]
        ),
        ..Default::default()
    }
).await?;

Before we finish this function up, we will need to update the user tier in the user_subscriptions table.

use sea_orm::IntoActiveModel;

let mut updated_user_subscription = user_subscription.into_active_model();
updated_user_subscription.user_tier = ActiveValue::Set(new_user_tier.user_tier);

let _ = updated_user_subscription.update(&ctx.db).await?;

Ok(StatusCode::OK);

Grabbing a user’s product tier

Of course, for the frontend you’ll probably want a quick way to be able to grab the user’s product tier. You can do this like so:

pub async fn get_current_tier(
    State(ctx): State<AppContext>,
    auth: middleware::auth::JWTWithUser<UsersModel>,
) -> Result<Json<UpdateUserTier>, ApiError> {
    let user = users::Model::find_by_pid(&ctx.db, &auth.claims.pid).await?;

    let subscription = user_subscriptions::Entity::find()
        .filter(user_subscriptions::Column::UserId.eq(user.id))
        .one(&ctx.db)
        .await?
        .unwrap();

    Ok(Json(UpdateUserTier { user_tier: subscription.user_tier  }))

}

Hooking it all up

Now that we’re done, we can add it back to our router file for the controller. We can add it back in, like so:

pub fn routes() -> Routes {
    Routes::new()
        .prefix("stripe")
        .add("/get_current_tier", get(get_current_tier))
        .add("/create", post(create_subscription))
        .add("/update", post(update_subscription_tier))
        .add("/cancel", delete(cancel_subscription))
}

Deployment

To deploy, you just need to run cargo shuttle deploy --allow-dirty and let Shuttle make the magic happen! Once done, you’ll see information about your service.

Finishing up

Loco is a really strong framework to get started with, and by implementing subscription payments we’ve managed to slot in the final piece of the puzzle for a potential SaaS in the making.

Interested in reading more?

Try using Qdrant and OpenAI to add a RAG-based LLM to your SaaS here.
Add tracing to your project for better logging here.
Try implementing Oauth2-based authentication here.