DEV Community: Tachi 0x

Build a Fast and Lightweight Rust Vector Search App with Rig & LanceDB

Tachi 0x — Thu, 21 Nov 2024 22:51:12 +0000

TL;DR: Build a powerful semantic search system in Rust using Rig and LanceDB. We'll guide you step-by-step through creating, storing, and searching vector embeddings efficiently with hands-on examples. Perfect for building RAG systems, semantic search engines, and more.

Introduction

Semantic search is transforming the way we find and understand information. Unlike traditional keyword search, it captures the intent behind your queries, offering a more nuanced retrieval process. However, building these systems can feel daunting, often involving complex embeddings, vector databases, and similarity search algorithms.

That's where LanceDB comes in.

Why LanceDB?

LanceDB is an open-source vector database tailored for AI applications and vector search. It provides:

Embedded Database: Works directly in your application without needing external servers.
High Performance: Leverages Arrow format for efficient data storage and retrieval.
Scalable: Handles terabyte-scale datasets efficiently.
Vector Indexing: Supports both exact and approximate nearest neighbor searches out of the box.

Combined with Rig's embedding and LLM capabilities, you can create a powerful, efficient semantic search solution with minimal code.

Let's dive in!

You can find the full source code for this project in our GitHub repo.

Prerequisites

Before we begin, make sure you have:

Rust installed (rust-lang.org)
An OpenAI API key (platform.openai.com)
Basic familiarity with Rust and asynchronous programming

Project Setup

To start, create a new Rust project:

cargo new vector_search
cd vector_search

Update your Cargo.toml to add the necessary dependencies:

[dependencies]
rig-core = "0.4.0"
rig-lancedb = "0.1.1"
lancedb = "0.10.0"
tokio = { version = "1.40.0", features = ["full"] }
anyhow = "1.0.89"
futures = "0.3.30"
serde = { version = "1.0.210", features = ["derive"] }
serde_json = "1.0.128"
arrow-array = "52.2.0"

Here’s a quick overview of each dependency:

rig-core and rig-lancedb: The core libraries for embedding generation and vector search.
lancedb: The embedded vector database.
tokio: Asynchronous runtime support.
arrow-array: To work with Arrow's columnar format, which LanceDB uses internally.
Others for error handling, serialization, and futures support.

Now, create a .env file to store your OpenAI API key:

echo "OPENAI_API_KEY=your_key_here" > .env

Building the Search System

We’ll break this into manageable steps. First, let’s create a utility function to handle data conversion between Rig's embeddings and LanceDB's format.

Create src/utils.rs:

use std::sync::Arc;
use arrow_array::{
    types::Float64Type, ArrayRef, FixedSizeListArray,
    RecordBatch, StringArray
};
use lancedb::arrow::arrow_schema::{DataType, Field, Fields, Schema};
use rig::embeddings::DocumentEmbeddings;

// Define the schema for our LanceDB table
pub fn schema(dims: usize) -> Schema {
    Schema::new(Fields::from(vec![
        Field::new("id", DataType::Utf8, false),
        Field::new("content", DataType::Utf8, false),
        Field::new(
            "embedding",
            DataType::FixedSizeList(
                Arc::new(Field::new("item", DataType::Float64, true)),
                dims as i32,
            ),
            false,
        ),
    ]))
}

This schema function defines the structure of our table:

id: A unique identifier for each document.
content: The text content of the document.
embedding: The vector representation of the content.
dims parameter: Represents the size of embedding vectors (e.g., 1536 for OpenAI's ada-002 model).

Next, add the conversion function to convert DocumentEmbeddings into RecordBatch for LanceDB:

pub fn as_record_batch(
    records: Vec<DocumentEmbeddings>,
    dims: usize,
) -> Result<RecordBatch, lancedb::arrow::arrow_schema::ArrowError> {
    let id = StringArray::from_iter_values(
        records
            .iter()
            .flat_map(|record| (0..record.embeddings.len())
                .map(|i| format!("{}-{i}", record.id)))
            .collect::<Vec<_>>(),
    );

    let content = StringArray::from_iter_values(
        records
            .iter()
            .flat_map(|record| {
                record
                    .embeddings
                    .iter()
                    .map(|embedding| embedding.document.clone())
            })
            .collect::<Vec<_>>(),
    );

    let embedding = FixedSizeListArray::from_iter_primitive::<Float64Type, _, _>(
        records
            .into_iter()
            .flat_map(|record| {
                record
                    .embeddings
                    .into_iter()
                    .map(|embedding| embedding.vec.into_iter().map(Some).collect::<Vec<_>>())
                    .map(Some)
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>(),
        dims as i32,
    );

    RecordBatch::try_from_iter(vec![
        ("id", Arc::new(id) as ArrayRef),
        ("content", Arc::new(content) as ArrayRef),
        ("embedding", Arc::new(embedding) as ArrayRef),
    ])
}

This function is crucial as it converts our Rust data structures into Arrow's columnar format, which LanceDB uses internally:

Creates string arrays for IDs and content.
Converts embeddings into fixed-size lists.
Assembles everything into a RecordBatch.

With our utility functions ready, let’s build the main search functionality in src/main.rs. We’ll implement this step-by-step, explaining each part along the way.

Setting Up Dependencies

First, let’s import the required libraries:

use anyhow::Result;
use arrow_array::RecordBatchIterator;
use lancedb::{index::vector::IvfPqIndexBuilder, DistanceType};
use rig::{
    embeddings::{DocumentEmbeddings, EmbeddingModel, EmbeddingsBuilder},
    providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
    vector_store::VectorStoreIndex,
};
use rig_lancedb::{LanceDbVectorStore, SearchParams};
use serde::Deserialize;
use std::{env, sync::Arc};

mod utils;
use utils::{as_record_batch, schema};

These imports bring in:

Rig’s embedding and vector storage tools.
LanceDB’s database capabilities.
Arrow data structures for efficient processing.
Utilities for serialization, error handling, and async programming.

Defining Data Structures

We’ll create a simple struct to represent our search results:

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

This struct maps to database records, representing the content we want to retrieve.

Generating Embeddings

Generating document embeddings is a core part of our system. Let’s implement this function:

async fn create_embeddings(client: &Client) -> Result<Vec<DocumentEmbeddings>> {
    let model = client.embedding_model(TEXT_EMBEDDING_ADA_002);

    // Set up dummy data to meet the 256 row requirement for IVF-PQ indexing
    let dummy_doc = "Let there be light".to_string();
    let dummy_docs = vec![dummy_doc; 256];

    // Generate embeddings for the data
    let embeddings = EmbeddingsBuilder::new(model)
        // First add our real documents
        .simple_document(
            "doc1",
            "Rust provides zero-cost abstractions and memory safety without garbage collection.",
        )
        .simple_document(
            "doc2",
            "Python emphasizes code readability with significant whitespace.",
        )
        // Add dummy documents to meet minimum requirement using enumerate to generate unique IDs
        .simple_documents(
            dummy_docs
                .into_iter()
                .enumerate()
                .map(|(i, doc)| (format!("doc{}", i + 3), doc))
                .collect(),
        )
        .build()
        .await?;

    Ok(embeddings)
}

This function handles:

Initializing the OpenAI embedding model.
Creating embeddings for our real documents.
Adding dummy data to meet LanceDB’s indexing requirements.

Configuring the Vector Store

Now, let’s set up LanceDB and configure it with appropriate indexing and search parameters:

async fn setup_vector_store<M: EmbeddingModel>(
    embeddings: Vec<DocumentEmbeddings>,
    model: M,
) -> Result<LanceDbVectorStore<M>> {
    // Initialize LanceDB
    let db = lancedb::connect("data/lancedb-store").execute().await?;

    // Drop the existing table if it exists - important for development
    if db
        .table_names()
        .execute()
        .await?
        .contains(&"documents".to_string())
    {
        db.drop_table("documents").await?;
    }

    // Create table with embeddings
    let record_batch = as_record_batch(embeddings, model.ndims())?;
    let table = db
        .create_table(
            "documents",
            RecordBatchIterator::new(vec![Ok(record_batch)], Arc::new(schema(model.ndims()))),
        )
        .execute()
        .await?;

    // Create an optimized vector index using IVF-PQ
    table
        .create_index(
            &["embedding"],
            lancedb::index::Index::IvfPq(
                IvfPqIndexBuilder::default().distance_type(DistanceType::Cosine),
            ),
        )
        .execute()
        .await?;

    // Configure search parameters
    let search_params = SearchParams::default().distance_type(DistanceType::Cosine);

    // Create and return vector store
    Ok(LanceDbVectorStore::new(table, model, "id", search_params).await?)
}

This setup function:

Connects to the LanceDB database.
Manages table creation and deletion.
Sets up vector indexing for efficient similarity search.

Putting It All Together

Finally, the main function orchestrates the entire process:

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

    // Create embeddings (includes both real and dummy documents)
    let embeddings = create_embeddings(&openai_client).await?;
    println!("Created embeddings for {} documents", embeddings.len());

    // Set up vector store
    let store = setup_vector_store(embeddings, model).await?;
    println!("Vector store initialized successfully");

    // Perform a semantic search
    let query = "Tell me about safe programming languages";
    let results = store.top_n::<SearchResult>(query, 2).await?;

    println!("\nSearch Results for: {}\n", query);
    for (score, id, result) in results {
        println!(
            "Score: {:.4}\nID: {}\nContent: {}\n",
            score, id, result.content
        );
    }

    Ok(())
}

Understanding Vector Search Methods

Vector search systems need to balance accuracy and performance, especially as datasets grow. LanceDB provides two approaches to handle this: Exact Nearest Neighbor (ENN) and Approximate Nearest Neighbor (ANN) searches.

ENN vs ANN

Exact Nearest Neighbor (ENN):
- Searches exhaustively across all vectors.
- Guarantees finding the true nearest neighbors.
- Works well for small datasets.
- No minimum data requirement.
- Slower, but more accurate.
Approximate Nearest Neighbor (ANN):
- Uses indexing to speed up searches (like IVF-PQ).
- Returns approximate results.
- Suited for larger datasets.
- Faster but slightly less accurate.

Choosing the Right Approach

Use ENN when:

Dataset is small (< 1,000 vectors).
Exact matches are crucial.
Performance isn’t a major concern.

Use ANN when:

Dataset is larger.
You can tolerate minor approximations.
Fast search speed is needed.

In our tutorial, we use ANN for scalability. For smaller datasets, ENN will be more suitable.

Tip: Start with ENN during development. Transition to ANN as your data and performance needs grow. Check out the ENN example.

Running the System

To run the project:

cargo run

Expected output:

Created embeddings for 258 documents
Vector store initialized successfully

Search Results for: Tell me about safe programming languages

Score: 0.3982
ID: doc2-0
Content: Python emphasizes code readability with significant whitespace.

Score: 0.4369
ID: doc1-0
Content: Rust provides zero-cost abstractions and memory safety without garbage collection.

Next Steps

If you’re ready to build more with Rig, here are some practical examples:

1. Build a RAG System

Want to give your LLM access to custom knowledge? Check out our tutorial on Building a RAG System with Rig in Under 100 Lines of Code.

2. Create an AI Agent

Ready to build more interactive AI applications? See our Agent Example.

3. Join the Community

Star us on GitHub
Join our Discord
Follow us on Twitter

Stay Connected

I’m always excited to hear from developers! If you’re interested in Rust, LLMs, or building intelligent assistants, join our Discord. Let’s build something amazing together!

And don’t forget: Build something with Rig, share your feedback, and get a chance to win $100.

Ad Astra,

Tachi

Co-Founder @ Playgrounds Analytics

This tutorial is part of our "Build with Rig" series. Follow our Website for more.

Build an AI Discord Bot in Rust with Rig: A Step-by-Step Guide

Tachi 0x — Wed, 09 Oct 2024 22:49:41 +0000

TL;DR: This comprehensive guide walks you through creating an AI-powered Discord bot using Rust and the Rig library. You'll learn to set up your environment, build a language model agent, and integrate it with Discord. By the end, you'll have an AI-powered chatbot that answers questions based on your own documents, provides coding assistance, and serves as an automated support tool.

Introduction

Welcome to the next installment of the Build with Rig series. In this hands-on tutorial, we'll construct a fully functional AI Discord bot using Rust and the Rig library. Our bot will be capable of:

Answering user queries using a custom Markdown-based knowledge base
Offering coding assistance and explanations
Serving as an automated customer service or community assistance tool

Throughout this guide, we'll cover:

Setting up your Rust development environment
Implementing a language model agent with Rig
Integrating the bot with Discord
Deploying and testing the bot

While this guide assumes some familiarity with Rust, LLMs, and Discord, don't worry if this is your first Rust project. We'll focus on practical implementation, explaining key concepts and design decisions along the way.

By the end of this tutorial, you'll have a working AI Discord bot and a solid grasp of building LLM-powered applications with Rig.

Let's dive in.

💡 If you're new to Rig and want to start from the beginning or are looking for additional tutorials, check out the blog series or visit the GitHub repository.

Prerequisites

Before we begin building, ensure you have the following:

Rust: If you haven't already, install Rust.
Discord Account and Bot Setup: You'll need a Discord account and a bot application. Follow the Discord bot setup guide.
OpenAI API Key: To enable AI capabilities, integrate the OpenAI API by obtaining an API key from OpenAI. Sign up for an OpenAI API key.

Important: Never commit your API keys or .env files to version control. Ensure your .gitignore file includes these files to prevent accidental exposure.

Project Setup

With the prerequisites in place, let's set up our Rust project and install the necessary dependencies.

1. Initialize a New Rust Project

Open your terminal and run the following commands:

cargo new discord_rig_bot
cd discord_rig_bot

This creates a new Rust project called discord_rig_bot and navigates into the project directory.

2. Add Dependencies

Open the Cargo.toml file in your project directory and add the following dependencies under the [dependencies] section:

[dependencies]
rig-core = "0.2.1" # [Rig Crate](https://crates.io/crates/rig-core)
tokio = { version = "1.34.0", features = ["full"] }
serenity = { version = "0.11", default-features = false, features = ["client", "gateway", "rustls_backend", "cache", "model", "http"] }
dotenv = "0.15.0"
anyhow = "1.0.75"
tracing = "0.1"
tracing-subscriber = "0.3"
reqwest = { version = "0.11", features = ["json"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
schemars = "0.8"
async-trait = "0.1.83"

These dependencies play crucial roles in our project:

Rig: The core library for building language model applications, simplifying the integration of AI capabilities into our Discord bot.
Serenity: A third-party library for interacting with the Discord API.
Tokio: An asynchronous runtime for Rust, allowing our bot to handle multiple tasks concurrently.
Additional crates for error handling (anyhow), logging (tracing, tracing-subscriber), making HTTP requests (reqwest), and serialization (serde, serde_json, schemars).

Understanding the Bot Architecture

Our bot consists of two main components that work together to provide an intelligent and interactive user experience:

Rig Agent (rig_agent.rs): The core of our bot's intelligence. The Rig agent manages AI interactions, handles natural language processing, retrieves relevant information from a Markdown-based knowledge base using Rig's integrated Retrieval-Augmented Generation (RAG) capabilities, and generates contextually appropriate responses.
Discord Bot (main.rs): The interface between our AI and users. The Discord bot manages communication with Discord, listens for user commands and messages, and sends the generated responses back to the user.

Message Processing Flow

To understand how our bot works, let's walk through the message processing flow:

User Input: A user sends a message or command in a Discord channel where the bot is present.
Discord Bot: The bot, always listening for new messages, receives the user's input and passes it to the Rig agent for processing.
Rig Agent: The agent processes the user's input, retrieves relevant information from the knowledge base, and generates an appropriate response using its language understanding and generation capabilities.
Response: The Discord bot receives the generated response from the Rig agent and sends it back to the user in the Discord channel.

Here's a simplified diagram of the message processing flow:

For an in-depth look at building a RAG system with Rig, refer to our comprehensive article on building a simple RAG system with Rig.

Building the Rig Agent (`rig_agent.rs`)

The Rig agent is the brain of our bot, responsible for understanding user queries, retrieving relevant information, and generating intelligent responses. Let's build it step by step.

1. Create the `rig_agent.rs` File

In your src directory, create a new file named rig_agent.rs. This file will contain the implementation of our Rig agent.

2. Import Necessary Modules

At the top of rig_agent.rs, import the required modules:

use anyhow::{Context, Result};
use rig::providers::openai;
use rig::vector_store::in_memory_store::InMemoryVectorStore;
use rig::vector_store::VectorStore;
use rig::embeddings::EmbeddingsBuilder;
use rig::agent::Agent;
use rig::completion::Prompt;
use std::path::Path;
use std::fs;
use std::sync::Arc;

These modules provide the necessary functionality for our Rig agent, including error handling (anyhow), OpenAI language models and embeddings (rig::providers::openai), vector storage (rig::vector_store::in_memory_store), embedding generation (rig::embeddings::EmbeddingsBuilder), and the agent (rig::agent::Agent).

3. Define the `RigAgent` Struct

Create the RigAgent struct that will manage the retrieval and response generation:

pub struct RigAgent {
    agent: Arc<Agent<openai::CompletionModel>>,
}

The RigAgent struct contains an Arc (Atomic Reference Counting) pointer to a Agent. The Arc type allows multiple parts of your program to share ownership of the same data in a thread-safe way. In our case, since our bot will handle multiple asynchronous events, sharing the RigAgent across different parts of the program without transferring ownership is crucial, and Arc provides a thread-safe way to do this.

Note: Arc stands for Atomic Reference Counting. It is used for sharing data between threads safely.

4. Implement the `new` Method

The new method is responsible for initializing the Rig agent, setting up the OpenAI client, loading and embedding the knowledge base documents, and creating the RAG agent.

impl RigAgent {
    pub async fn new() -> Result<Self> {
        // Initialize OpenAI client
        let openai_client = openai::Client::from_env();
        let embedding_model = openai_client.embedding_model(openai::TEXT_EMBEDDING_3_SMALL);

        // Create vector store
        let mut vector_store = InMemoryVectorStore::default();

        // Get the current directory and construct paths to markdown files
        let current_dir = std::env::current_dir()?;
        let documents_dir = current_dir.join("documents");

        let md1_path = documents_dir.join("Rig_guide.md");
        let md2_path = documents_dir.join("Rig_faq.md");
        let md3_path = documents_dir.join("Rig_examples.md");

        // Load markdown documents
        let md1_content = Self::load_md_content(&md1_path)?;
        let md2_content = Self::load_md_content(&md2_path)?;
        let md3_content = Self::load_md_content(&md3_path)?;

        // Create embeddings and add to vector store
        let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
            .simple_document("Rig_guide", &md1_content)
            .simple_document("Rig_faq", &md2_content)
            .simple_document("Rig_examples", &md3_content)
            .build()
            .await?;

        vector_store.add_documents(embeddings).await?;

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

        // Create Agent
        let agent = Arc::new(openai_client.agent(openai::GPT_4O)
            .preamble("You are an advanced AI assistant powered by Rig, a Rust library for building LLM applications. Your primary function is to provide accurate, helpful, and context-aware responses by leveraging both your general knowledge and specific information retrieved from a curated knowledge base.

                    Key responsibilities and behaviors:
                    1. Information Retrieval: You have access to a vast knowledge base. When answering questions, always consider the context provided by the retrieved information.
                    2. Clarity and Conciseness: Provide clear and concise answers. Ensure responses are short and concise. Use bullet points or numbered lists for complex information when appropriate.
                    3. Technical Proficiency: You have deep knowledge about Rig and its capabilities. When discussing Rig or answering related questions, provide detailed and technically accurate information.
                    4. Code Examples: When appropriate, provide Rust code examples to illustrate concepts, especially when discussing Rig's functionalities. Always format code examples for proper rendering in Discord by wrapping them in triple backticks and specifying the language as 'rust'.
                    5. Keep your responses short and concise. If the user needs more information, they can ask follow-up questions.
                    ")
            .dynamic_context(2, index)
            .build());

        Ok(Self { agent })
    }

    // ... we'll add more code here as we build things out
}

Let's break down the key steps:

Initialize OpenAI Client: The OpenAI client is set up using the API key stored in the environment variables. This client is essential for accessing OpenAI's language models and embedding services.
Embedding Model: The text-embedding-3-small model is selected for generating document embeddings. This model creates compact vector representations of text, enabling efficient semantic search and retrieval.
Vector Store: An in-memory vector store is created to hold and manage the document embeddings. Vector stores are optimized for fast similarity searches, allowing the agent to quickly find relevant information based on the user's query.
Load Documents: Markdown files containing the knowledge base are loaded from the documents directory. In this example, we have three files: Rig_guide.md, Rig_faq.md, and Rig_examples.md. These files contain information about the Rig library, frequently asked questions, and usage examples.
Create Embeddings: The loaded documents are converted into vector embeddings using the EmbeddingsBuilder provided by Rig. These embeddings capture the semantic meaning of the documents, enabling the agent to understand and retrieve relevant information based on the user's query.
Agent Creation: A RagAgent is created by combining the language model (GPT-4) with the vector store containing the document embeddings. The agent is capable of retrieving relevant information from the knowledge base and generating contextually appropriate responses.
Preamble: A carefully crafted preamble sets up the assistant's behavior and guidelines. The preamble defines the agent's primary function, responsibilities, and expected behaviors, ensuring that it provides accurate, concise, and technically proficient responses.

Tip: For more advanced configurations and techniques, such as implementing custom vector stores or configuring custom agents and tools, refer to the official Rig examples.

5. Implement the `load_md_content` Function

The load_md_content function is a helper function that reads the content of a Markdown file from the specified file path:

fn load_md_content<P: AsRef<Path>>(file_path: P) -> Result<String> {
    fs::read_to_string(file_path.as_ref())
        .with_context(|| format!("Failed to read markdown file: {:?}", file_path.as_ref()))
}

This function takes a generic parameter P that implements the AsRef<Path> trait, allowing it to accept various types that can be converted to a file path. It uses the fs::read_to_string function to read the contents of the file and returns the content as a String. If the file cannot be read, an error is returned with additional context information.

6. Implement the `process_message` Function

The process_message function is responsible for processing user messages and generating responses using the agent:

pub async fn process_message(&self, message: &str) -> Result<String> {
        self.agent.prompt(message).await.map_err(anyhow::Error::from)
    }

This function takes a user message as input and passes it to the RAG agent's prompt method. The RAG agent retrieves relevant information from the knowledge base based on the user's query and generates a contextually appropriate response. The generated response is then returned as a String. If an error occurs during the processing, it is mapped to an anyhow::Error for consistent error handling.

Customizing the Knowledge Base

While we've used Rig's own documentation for our knowledge base, you can personalize your bot by using your own documents. Here's how:

Prepare Your Documents: Place your Markdown files in the documents directory. Ensure they have clear and descriptive filenames.
Modify the File Paths: In rig_agent.rs, update the file paths to match your document names.

   let my_doc_path = documents_dir.join("my_custom_doc.md");
   let my_doc_content = Self::load_md_content(&my_doc_path)?;

Update the Embeddings Builder: Adjust the EmbeddingsBuilder to include your documents.

   let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
       .simple_document("My Custom Doc", &my_doc_content)
       .build()
       .await?;

This way, your bot will use your own content to generate responses.

Integrating with Discord (`main.rs`)

With our Rig agent implementation complete, it's time to connect it to Discord using the Serenity library. Serenity is an async-first Rust library for the Discord API, providing a simple and efficient way to create Discord bots.

1. Modify `main.rs` to Include the Rig Agent

At the top of main.rs, import the necessary modules and your rig_agent:

mod rig_agent;

use anyhow::Result;
use serenity::async_trait;
use serenity::model::application::command::Command;
use serenity::model::application::interaction::{Interaction, InteractionResponseType};
use serenity::model::gateway::Ready;
use serenity::model::channel::Message;
use serenity::prelude::*;
use serenity::model::application::command::CommandOptionType;
use std::env;
use std::sync::Arc;
use tracing::{error, info, debug};
use rig_agent::RigAgent;
use dotenv::dotenv;

These imports bring in the essential types and traits from the Serenity library, as well as the RigAgent struct from the rig_agent module. The dotenv crate is used to load environment variables from a .env file.

2. Store the Bot's User ID

Define a key for storing the bot's user ID using Serenity's TypeMapKey trait:

struct BotUserId;

impl TypeMapKey for BotUserId {
    type Value = serenity::model::id::UserId;
}

This key allows us to store and retrieve the bot's user ID from Serenity's TypeMap, which is a type-safe key-value store used for sharing data across event handlers.

3. Define the `Handler` Struct

Create the Handler struct that holds the RigAgent:

struct Handler {
    rig_agent: Arc<RigAgent>,
}

The Handler struct is responsible for handling Discord events and interactions. It contains an Arc<RigAgent>, which is a thread-safe reference-counting pointer to the Rig agent. This allows the handler to share the Rig agent across multiple event handlers without transferring ownership.

4. Implement the `EventHandler` Trait

Implement the EventHandler trait for the Handler struct to define how the bot should handle various Discord events:

#[async_trait]
impl EventHandler for Handler {
    async fn interaction_create(&self, ctx: Context, interaction: Interaction) {
        // ... handle interactions
    }

    async fn message(&self, ctx: Context, msg: Message) {
        // ... handle messages
    }

    async fn ready(&self, ctx: Context, ready: Ready) {
        // ... handle readiness
    }
}

The EventHandler trait is provided by Serenity and defines a set of methods that are called when specific events occur. In this implementation, we define three event handlers:

interaction_create: Called when a user interacts with the bot, such as using a slash command or clicking a button.
message: Called when a message is sent in a channel where the bot is present.
ready: Called when the bot successfully connects to Discord and is ready to receive events.

Handling Interactions

In the interaction_create event handler, we process slash commands received from Discord:

async fn interaction_create(&self, ctx: Context, interaction: Interaction) {
    debug!("Received an interaction");
    if let Interaction::ApplicationCommand(command) = interaction {
        debug!("Received command: {}", command.data.name);
        let content = match command.data.name.as_str() {
            "hello" => "Hello! I'm your helpful Rust and Rig-powered assistant. How can I assist you today?".to_string(),
            "ask" => {
                let query = command
                    .data
                    .options
                    .get(0)
                    .and_then(|opt| opt.value.as_ref())
                    .and_then(|v| v.as_str())
                    .unwrap_or("What would you like to ask?");
                debug!("Query: {}", query);
                match self.rig_agent.process_message(query).await {
                    Ok(response) => response,
                    Err(e) => {
                        error!("Error processing request: {:?}", e);
                        format!("Error processing request: {:?}", e)
                    }
                }
            }
            _ => "Not implemented :(".to_string(),
        };

        debug!("Sending response: {}", content);

        if let Err(why) = command
            .create_interaction_response(&ctx.http, |response| {
                response
                    .kind(InteractionResponseType::ChannelMessageWithSource)
                    .interaction_response_data(|message| message.content(content))
            })
            .await
        {
            error!("Cannot respond to slash command: {}", why);
        } else {
            debug!("Response sent successfully");
        }
    }
}

Let's break down the process:

When an interaction is received, we first check if it's a slash command using the Interaction::ApplicationCommand enum variant.
If it's a slash command, we match on the command name to determine the appropriate action.
- For the "hello" command, we respond with a simple greeting message.
- For the "ask" command, we extract the user's query from the command options. If no query is provided, we use a default message.
If the command is "ask", we pass the user's query to the Rig agent's process_message method to generate a response.
- If the Rig agent successfully generates a response, we send it back to the user.
- If an error occurs during processing, we log the error and send an error message to the user.
For any other command, we respond with a "Not implemented" message.
Finally, we create an interaction response using command.create_interaction_response, specifying the response type as ChannelMessageWithSource and setting the response content to the generated message.

This implementation allows users to interact with the bot using slash commands, providing a structured way to ask questions and receive responses from the Rig agent.

Handling Messages

In the message event handler, we respond when the bot is mentioned in a message:

async fn message(&self, ctx: Context, msg: Message) {
    if msg.mentions_me(&ctx.http).await.unwrap_or(false) {
        debug!("Bot mentioned in message: {}", msg.content);

        let bot_id = {
            let data = ctx.data.read().await;
            data.get::<BotUserId>().copied()
        };

        if let Some(bot_id) = bot_id {
            let mention = format!("<@{}>", bot_id);
            let content = msg.content.replace(&mention, "").trim().to_string();

            debug!("Processed content after removing mention: {}", content);

            match self.rig_agent.process_message(&content).await {
                Ok(response) => {
                    if let Err(why) = msg.channel_id.say(&ctx.http, response).await {
                        error!("Error sending message: {:?}", why);
                    }
                }
                Err(e) => {
                    error!("Error processing message: {:?}", e);
                    if let Err(why) = msg
                        .channel_id
                        .say(&ctx.http, format!("Error processing message: {:?}", e))
                        .await
                    {
                        error!("Error sending error message: {:?}", why);
                    }
                }
            }
        } else {
            error!("Bot user ID not found in TypeMap");
        }
    }
}

Here's how the message handling works:

When a message is received, we first check if the bot is mentioned in the message using the mentions_me method.
If the bot is mentioned, we retrieve the bot's user ID from the TypeMap using the BotUserId key.
If the bot's user ID is found, we remove the mention from the message content to extract the actual query.
We pass the processed message content to the Rig agent's process_message method to generate a response.
- If the Rig agent successfully generates a response, we send it back to the channel where the message was received using msg.channel_id.say.
- If an error occurs during processing, we log the error and send an error message to the channel.
If the bot's user ID is not found in the TypeMap, we log an error.

This implementation allows users to interact with the bot by mentioning it in a message, providing a more natural way to ask questions and receive responses from the Rig agent.

Handling Bot Readiness

In the ready event handler, we set up slash commands and store the bot's user ID:

async fn ready(&self, ctx: Context, ready: Ready) {
    info!("{} is connected!", ready.user.name);

    {
        let mut data = ctx.data.write().await;
        data.insert::<BotUserId>(ready.user.id);
    }

    let commands = Command::set_global_application_commands(&ctx.http, |commands| {
        commands
            .create_application_command(|command| {
                command
                    .name("hello")
                    .description("Say hello to the bot")
            })
            .create_application_command(|command| {
                command
                    .name("ask")
                    .description("Ask the bot a question")
                    .create_option(|option| {
                        option
                            .name("query")
                            .description("Your question for the bot")
                            .kind(CommandOptionType::String)
                            .required(true)
                    })
            })
    })
    .await;

    println!("Created the following global commands: {:#?}", commands);
}

Here's what happens in the ready event handler:

When the bot successfully connects to Discord, the ready event is triggered.
We log a message indicating that the bot is connected, using the bot's name from the Ready struct.
We store the bot's user ID in the TypeMap using the BotUserId key. This allows us to access the bot's user ID in other event handlers.
We create global slash commands using the Command::set_global_application_commands method.
- We define two commands: "hello" and "ask".
- The "hello" command is a simple command that greets the user.
- The "ask" command allows users to ask the bot a question. It has a required "query" option of type String where users can input their question.
We print the created global commands for debugging purposes.

5. Implement the `main` Function

In the main function, we set up the bot and start it:

#[tokio::main]
async fn main() -> Result<()> {
    dotenv().ok();

    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::DEBUG)
        .init();

    let token = env::var("DISCORD_TOKEN").expect("Expected DISCORD_TOKEN in environment");

    let rig_agent = Arc::new(RigAgent::new().await?);

    let intents = GatewayIntents::GUILD_MESSAGES
        | GatewayIntents::DIRECT_MESSAGES
        | GatewayIntents::MESSAGE_CONTENT;

    let mut client = Client::builder(&token, intents)
        .event_handler(Handler {
            rig_agent: Arc::clone(&rig_agent),
        })
        .await
        .expect("Err creating client");

    if let Err(why) = client.start().await {
        error!("Client error: {:?}", why);
    }

    Ok(())
}

Here's a step-by-step breakdown of the main function:

We load environment variables from the .env file using dotenv().ok().
We initialize the tracing_subscriber to set up logging with a maximum log level of DEBUG.
We retrieve the Discord bot token from the DISCORD_TOKEN environment variable.
We create a new instance of the RigAgent and wrap it in an Arc for thread-safe sharing.
We define the gateway intents, specifying the events we want to receive from Discord.
We create a new Discord client using the Client::builder method, passing in the bot token and intents.
We set the event handler for the client to an instance of the Handler struct, passing the RigAgent wrapped in an Arc.
We start the client using client.start() and handle any errors that may occur.

Running and Testing the Bot

Now that our Discord bot is complete, let's run it and test its functionality.

1. Set Up Environment Variables

Create a .env file in the root of your project with the following content:

DISCORD_TOKEN=your_discord_bot_token
OPENAI_API_KEY=your_openai_api_key

Replace your_discord_bot_token with your actual Discord bot token and your_openai_api_key with your OpenAI API key.

Important: Never commit your .env file or API keys to version control. Add .env to your .gitignore file to prevent accidental exposure.

2. Run the Bot

In your terminal, navigate to the project directory and run the following command:

cargo run

If everything is set up correctly, you should see logs indicating that the bot is connected and the global commands have been created.

3. Invite the Bot to Your Discord Server

To invite the bot to your Discord server, follow these steps:

Go to the Discord Developer Portal and select your bot application.
Navigate to the "OAuth2" section and click on "URL Generator".
Under "Scopes", select bot and applications.commands.
Under "Bot Permissions", select Send Messages, Read Message History, etc.
Copy the generated URL and paste it into your browser.
Select the server you want to invite the bot to and click "Authorize".

4. Test the Bot

Once the bot is running and invited to your server, you can test its functionality:

Slash Commands:
- Type /hello to receive a greeting.
- Use /ask followed by a question to interact with the bot and receive a response generated by the Rig agent.
Mentions:
- Mention the bot in a message with a question, like @BotName How do I use Rig?, and the bot will process your question and respond accordingly.

Here's two examples of the bot responding to questions:

Error Handling and Logging

Now that we've built and tested the bot, we need to ensure we are properly handling errors and logging our bot's behavior for improvement purposes. Rust provides powerful libraries like anyhow and tracing for error handling and logging.

1. Error Handling with `anyhow`

The anyhow crate provides a flexible and easy-to-use error handling solution. It allows us to propagate and handle errors with additional context, making it easier to diagnose and fix issues. Here's an example of using anyhow in our rig_agent.rs file:

use anyhow::{Context, Result};

// Example in rig_agent.rs
fn load_md_content<P: AsRef<Path>>(file_path: P) -> Result<String> {
    fs::read_to_string(file_path.as_ref())
        .with_context(|| format!("Failed to read markdown file: {:?}", file_path.as_ref()))
}

In this example, we use the with_context method to provide additional context to the error, specifying the file path that failed to be read. This context is included in the error message, making it easier to identify the source of the error.

2. Logging with `tracing`

The tracing crate provides a powerful and flexible logging solution that allows us to log messages at different levels of verbosity. Here's how we can set up logging in our main.rs file:

use tracing::{info, error, debug};
use tracing_subscriber;

// Initialize tracing in main.rs
tracing_subscriber::fmt()
    .with_max_level(tracing::Level::DEBUG)
    .init();

In this example, we initialize the tracing_subscriber with a maximum log level of DEBUG. This means that all log messages with a severity level of DEBUG or higher will be captured and displayed.

Throughout our bot's code, we can use the info!, error!, and debug! macros to log messages at different severity levels, providing insights into the bot's behavior.

Troubleshooting Common Issues

If you encounter errors, here are some common issues and how to resolve them:

API Key Errors: Ensure your OpenAI API key and Discord token are correctly set in your .env file. Double-check for typos or extra spaces.
File Not Found: If you receive a "Failed to read markdown file" error, check that your document paths are correct and the files exist in the documents directory.
Dependency Conflicts: Run cargo update to ensure all dependencies are up to date.
Permission Errors: Ensure your bot has the necessary permissions in your Discord server, such as sending messages and reading message history.

Testing and Debugging Guidance

Testing the Rig Agent Independently

Before integrating with Discord, you can test the Rig agent independently to ensure it's working as expected. Here's how:

// Test the RigAgent in main.rs

#[tokio::main]
async fn main() -> Result<()> {
    dotenv().ok();

    let rig_agent = RigAgent::new().await?;
    let response = rig_agent.process_message("What is Rig?").await?;
    println!("Response: {}", response);

    Ok(())
}

Run cargo run to see the output. This test confirms that the Rig agent can process messages and generate responses based on your knowledge base.

Enhancing Your Bot

Congratulations! You've built a functional AI-powered Discord bot. Now, let's explore some ways to enhance it.

1. Expand the Knowledge Base

To make your bot more knowledgeable and versatile, consider adding more Markdown files to the documents directory. These files can cover a wide range of topics, from FAQs to technical documentation and beyond.

2. Customize the Bot's Behavior

The behavior of your bot is largely determined by the preamble defined in the rig_agent.rs file. By adjusting this preamble, you can fine-tune how the bot interacts with users, shaping its personality, tone, and overall approach to conversations. Experiment with different preambles to find the right balance for your bot's intended purpose and audience.

3. Add More Commands

Slash commands provide a structured and intuitive way for users to interact with your bot. Consider implementing additional commands in the main.rs file to extend your bot's functionality. For example, you could add commands for retrieving specific information, performing calculations, or triggering automated workflows.

4. Integrate Other APIs

To further enhance your bot's capabilities, consider integrating it with other APIs. For example, you could connect your bot to a weather API to provide real-time weather updates or integrate it with a news API to deliver the latest headlines. By leveraging external APIs, you can create powerful workflows and automate tasks within your Discord server.

Check out our guide on building agent tools and integrating APIs here. You can also find more examples in the official Rig repo.

Conclusion

In this guide, we've successfully built an AI-powered Discord bot using Rust and Rig. We've learned how to set up your environment, build a language model agent, integrate with Discord, and run your bot. With this foundation, you can continue to enhance and customize your bot, transforming it into a more robust system of autonomous agents to suit your needs.

See you in the next guide in the Build with Rig series!

Further Resources

To deepen your understanding and continue building with Rig, check out these resources:

Your Feedback Matters! We're offering a unique opportunity to shape the future of Rig:

Build an AI-powered application using Rig.
Share your experience and insights via this feedback form.
Get a chance to win $100 and have your project featured in our showcase!

Your insights will directly influence Rig's growth. 🦀✨

Ad Astra,

Tachi

Co-Founder @ Playgrounds Analytics

Code Files

For reference, here are the complete rig_agent.rs and main.rs files.

`rig_agent.rs`

// rig_agent.rs

use anyhow::{Context, Result};
use rig::providers::openai;
use rig::vector_store::in_memory_store::InMemoryVectorStore;
use rig::vector_store::VectorStore;
use rig::embeddings::EmbeddingsBuilder;
use rig::agent::Agent;
use rig::completion::Prompt;
use std::path::Path;
use std::fs;
use std::sync::Arc;

pub struct RigAgent {
    agent: Arc<Agent<openai::CompletionModel>>,
}

impl RigAgent {
    pub async fn new() -> Result<Self> {
        // Initialize OpenAI client
        let openai_client = openai::Client::from_env();
        let embedding_model = openai_client.embedding_model(openai::TEXT_EMBEDDING_3_SMALL);

        // Create vector store
        let mut vector_store = InMemoryVectorStore::default();

        // Get the current directory and construct paths to markdown files
        let current_dir = std::env::current_dir()?;
        let documents_dir = current_dir.join("documents");

        let md1_path = documents_dir.join("Rig_guide.md");
        let md2_path = documents_dir.join("Rig_faq.md");
        let md3_path = documents_dir.join("Rig_examples.md");

        // Load markdown documents
        let md1_content = Self::load_md_content(&md1_path)?;
        let md2_content = Self::load_md_content(&md2_path)?;
        let md3_content = Self::load_md_content(&md3_path)?;

        // Create embeddings and add to vector store
        let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
            .simple_document("Rig_guide", &md1_content)
            .simple_document("Rig_faq", &md2_content)
            .simple_document("Rig_examples", &md3_content)
            .build()
            .await?;

        vector_store.add_documents(embeddings).await?;

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

        // Create Agent
        let agent = Arc::new(openai_client.agent(openai::GPT_4O)
            .preamble("You are an advanced AI assistant powered by Rig, a Rust library for building LLM applications. Your primary function is to provide accurate, helpful, and context-aware responses by leveraging both your general knowledge and specific information retrieved from a curated knowledge base.

                    Key responsibilities and behaviors:
                    1. Information Retrieval: You have access to a vast knowledge base. When answering questions, always consider the context provided by the retrieved information.
                    2. Clarity and Conciseness: Provide clear and concise answers. Ensure responses are short and concise. Use bullet points or numbered lists for complex information when appropriate.
                    3. Technical Proficiency: You have deep knowledge about Rig and its capabilities. When discussing Rig or answering related questions, provide detailed and technically accurate information.
                    4. Code Examples: When appropriate, provide Rust code examples to illustrate concepts, especially when discussing Rig's functionalities. Always format code examples for proper rendering in Discord by wrapping them in triple backticks and specifying the language as 'rust'. For example:
                        ```

rust
                        let example_code = \"This is how you format Rust code for Discord\";
                        println!(\"{}\", example_code);


                        ```
                    5. Keep your responses short and concise. If the user needs more information, they can ask follow-up questions.
                    ")
            .dynamic_context(2, index)
            .build());

        Ok(Self { agent })
    }

    fn load_md_content<P: AsRef<Path>>(file_path: P) -> Result<String> {
        fs::read_to_string(file_path.as_ref())
            .with_context(|| format!("Failed to read markdown file: {:?}", file_path.as_ref()))
    }

    pub async fn process_message(&self, message: &str) -> Result<String> {
        self.agent.prompt(message).await.map_err(anyhow::Error::from)
    }
}

`main.rs`

// main.rs

mod rig_agent;

use anyhow::Result;
use serenity::async_trait;
use serenity::model::application::command::Command;
use serenity::model::application::interaction::{Interaction, InteractionResponseType};
use serenity::model::gateway::Ready;
use serenity::model::channel::Message;
use serenity::prelude::*;
use serenity::model::application::command::CommandOptionType;
use std::env;
use std::sync::Arc;
use tracing::{error, info, debug};
use rig_agent::RigAgent;
use dotenv::dotenv;

// Define a key for storing the bot's user ID in the TypeMap
struct BotUserId;

impl TypeMapKey for BotUserId {
    type Value = serenity::model::id::UserId;
}

struct Handler {
    rig_agent: Arc<RigAgent>,
}

#[async_trait]
impl EventHandler for Handler {
    async fn interaction_create(&self, ctx: Context, interaction: Interaction) {
        debug!("Received an interaction");
        if let Interaction::ApplicationCommand(command) = interaction {
            debug!("Received command: {}", command.data.name);
            let content = match command.data.name.as_str() {
                "hello" => "Hello! I'm your helpful Rust and Rig-powered assistant. How can I assist you today?".to_string(),
                "ask" => {
                    let query = command
                        .data
                        .options
                        .get(0)
                        .and_then(|opt| opt.value.as_ref())
                        .and_then(|v| v.as_str())
                        .unwrap_or("What would you like to ask?");
                    debug!("Query: {}", query);
                    match self.rig_agent.process_message(query).await {
                        Ok(response) => response,
                        Err(e) => {
                            error!("Error processing request: {:?}", e);
                            format!("Error processing request: {:?}", e)
                        }
                    }
                }
                _ => "Not implemented :(".to_string(),
            };

            debug!("Sending response: {}", content);

            if let Err(why) = command
                .create_interaction_response(&ctx.http, |response| {
                    response
                        .kind(InteractionResponseType::ChannelMessageWithSource)
                        .interaction_response_data(|message| message.content(content))
                })
                .await
            {
                error!("Cannot respond to slash command: {}", why);
            } else {
                debug!("Response sent successfully");
            }
        }
    }

    async fn message(&self, ctx: Context, msg: Message) {
        if msg.mentions_me(&ctx.http).await.unwrap_or(false) {
            debug!("Bot mentioned in message: {}", msg.content);

            let bot_id = {
                let data = ctx.data.read().await;
                data.get::<BotUserId>().copied()
            };

            if let Some(bot_id) = bot_id {
                let mention = format!("<@{}>", bot_id);
                let content = msg.content.replace(&mention, "").trim().to_string();

                debug!("Processed content after removing mention: {}", content);

                match self.rig_agent.process_message(&content).await {
                    Ok(response) => {
                        if let Err(why) = msg.channel_id.say(&ctx.http, response).await {
                            error!("Error sending message: {:?}", why);
                        }
                    }
                    Err(e) => {
                        error!("Error processing message: {:?}", e);
                        if let Err(why) = msg
                            .channel_id
                            .say(&ctx.http, format!("Error processing message: {:?}", e))
                            .await
                        {
                            error!("Error sending error message: {:?}", why);
                        }
                    }
                }
            } else {
                error!("Bot user ID not found in TypeMap");
            }
        }
    }

    async fn ready(&self, ctx: Context, ready: Ready) {
        info!("{} is connected!", ready.user.name);

        {
            let mut data = ctx.data.write().await;
            data.insert::<BotUserId>(ready.user.id);
        }

        let commands = Command::set_global_application_commands(&ctx.http, |commands| {
            commands
                .create_application_command(|command| {
                    command
                        .name("hello")
                        .description("Say hello to the bot")
                })
                .create_application_command(|command| {
                    command
                        .name("ask")
                        .description("Ask the bot a question")
                        .create_option(|option| {
                            option
                                .name("query")
                                .description("Your question for the bot")
                                .kind(CommandOptionType::String)
                                .required(true)
                        })
                })
        })
        .await;

        println!("Created the following global commands: {:#?}", commands);
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    dotenv().ok();

    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::DEBUG)
        .init();

    let token = env::var("DISCORD_TOKEN").expect("Expected DISCORD_TOKEN in environment");

    let rig_agent = Arc::new(RigAgent::new().await?);

    let intents = GatewayIntents::GUILD_MESSAGES
        | GatewayIntents::DIRECT_MESSAGES
        | GatewayIntents::MESSAGE_CONTENT;

    let mut client = Client::builder(&token, intents)
        .event_handler(Handler {
            rig_agent: Arc::clone(&rig_agent),
        })
        .await
        .expect("Err creating client");

    if let Err(why) = client.start().await {
        error!("Client error: {:?}", why);
    }

    Ok(())
}

Build a Flight Search AI Agent with Rust using Rig: A Hands-On Practical Guide

Tachi 0x — Tue, 24 Sep 2024 01:09:20 +0000

TL;DR: This step-by-step guide will teach you how to build a Flight Search AI Assistant in Rust using the Rig library. By the end, you'll have a functional AI agent that can find the cheapest flights between two airports. Along the way, you'll grasp Rust fundamentals, understand how to set up AI agents with custom tools, and see how Rig simplifies the process.

Introduction

Ever chatted with AI assistants like Siri, Alexa, or even those nifty chatbots that help you book flights or check the weather? Ever wondered what's happening under the hood? Today, we're going to demystify that by building our very own Flight Search AI Assistant using Rust and the Rig library.

You might be thinking, "Wait, Rust? Isn't that the language with the reputation for being hard?" Don't worry! We'll walk through everything step by step, explaining concepts as we go. By the end, not only will you have a cool AI assistant, but you'll also have dipped your toes into Rust programming.

Here's our game plan:

Why Rust and Rig? Understanding our tools of choice.
Setting Up the Environment: Getting Rust and Rig ready to roll.
Understanding Agents and Tools: The brains and hands of our assistant.
Building the Flight Search Tool: Where the magic happens.
Creating the AI Agent: Bringing our assistant to life.
Running and Testing: Seeing our creation in action.
Wrapping Up: Recap and next steps.

Full source code for this project can be found on our Replit Page and Github

Sound exciting? Let's dive in!

Why Rust and Rig?

Why Rust?

Rust is a systems programming language known for its performance and safety. But beyond that, Rust has been making waves in areas like web development, game development, and now, AI applications. Here's why we're using Rust:

Performance: Rust is blazingly fast, making it ideal for applications that need to handle data quickly.
Safety: With its strict compiler checks, Rust ensures memory safety, preventing common bugs.
Concurrency: Rust makes it easier to write concurrent programs, which is great for handling multiple tasks simultaneously. Learn more about Rust's concurrency model.

Why Rig?

Rig is an open-source Rust library that simplifies building applications powered by Large Language Models (LLMs) like GPT-4. Think of Rig as a toolkit that provides:

Unified API: It abstracts away the complexities of different LLM providers.
High-Level Abstractions: Helps you build agents and tools without reinventing the wheel.
Extensibility: You can create custom tools tailored to your application's needs.

By combining Rust and Rig, we're setting ourselves up to build a robust, efficient, and intelligent assistant.

Setting Up the Environment

Before we start coding, let's get everything ready.

Prerequisites

Install Rust: If you haven't already, install Rust by following the instructions here.
Basic Rust Knowledge: Don't worry if you're new. We'll explain the Rust concepts as we encounter them.
API Keys:
- OpenAI API Key: Sign up and get your key here.
- RapidAPI Key: We'll use this to access the trip advisor flight search API. Get it here.

Project Setup

1. Create a New Rust Project

Open your terminal and run:

cargo new flight_search_assistant
cd flight_search_assistant

This initializes a new Rust project named flight_search_assistant.

2. Update `Cargo.toml`

Open the Cargo.toml file and update it with the necessary dependencies:

[package]
name = "flight_search_assistant"
version = "0.1.0"
edition = "2021"

[dependencies]
rig-core = "0.1.0"
tokio = { version = "1.34.0", features = ["full"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
reqwest = { version = "0.11", features = ["json", "tls"] }
dotenv = "0.15"
thiserror = "1.0"
chrono = { version = "0.4", features = ["serde"] }

Here's a quick rundown:

rig-core: The core Rig library.
tokio: Asynchronous runtime for Rust. Think of it as the engine that allows us to perform tasks concurrently.
serde & serde_json: Libraries for serializing and deserializing data (converting between Rust structs and JSON).
reqwest: An HTTP client for making API requests.
dotenv: Loads environment variables from a .env file.
thiserror: A library for better error handling.
chrono: For handling dates and times.

3. Set Up Environment Variables

We don't want to hard-code our API keys for security reasons. Instead, we'll store them in a .env file.

Create the file:

touch .env

Add your API keys to .env:

OPENAI_API_KEY=your_openai_api_key_here
RAPIDAPI_KEY=your_rapidapi_key_here

Remember to replace the placeholders with your actual keys.

4. Install Dependencies

Back in your terminal, run:

cargo build

This will download and compile all the dependencies.

Understanding Agents and Tools

Before we jump into coding, let's clarify some key concepts.

What Are Agents?

In the context of Rig (and AI applications in general), an Agent is like the brain of your assistant. It's responsible for interpreting user inputs, deciding what actions to take, and generating responses.

Think of the agent as the conductor of an orchestra, coordinating different instruments (or tools) to create harmonious music (or responses).

What Are Tools?

Tools are the skills or actions that the agent can use to fulfill a task. Each tool performs a specific function. In our case, the flight search functionality is a tool that the agent can use to find flight information.

Continuing our analogy, tools are the instruments in the orchestra. Each one plays a specific role.

How Do They Work Together?

When a user asks, "Find me flights from NYC to LA," the agent processes this request and decides it needs to use the flight search tool to fetch the information.

Building the Flight Search Tool

Now, let's build the tool that will handle flight searches.

1. Create the Tool File

In your src directory, create a new file named flight_search_tool.rs:

touch src/flight_search_tool.rs

2. Import Necessary Libraries

Open flight_search_tool.rs and add:

use chrono::{DateTime, Duration, Utc};
use rig::completion::ToolDefinition;
use rig::tool::Tool;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::env;

3. Define Data Structures

We'll define structures to handle input arguments and output results.

#[derive(Deserialize)]
pub struct FlightSearchArgs {
    source: String,
    destination: String,
    date: Option<String>,
    sort: Option<String>,
    service: Option<String>,
    itinerary_type: Option<String>,
    adults: Option<u8>,
    seniors: Option<u8>,
    currency: Option<String>,
    nearby: Option<String>,
    nonstop: Option<String>,
}

#[derive(Serialize)]
pub struct FlightOption {
    pub airline: String,
    pub flight_number: String,
    pub departure: String,
    pub arrival: String,
    pub duration: String,
    pub stops: usize,
    pub price: f64,
    pub currency: String,
    pub booking_url: String,
}

FlightSearchArgs: Represents the parameters the user provides.
FlightOption: Represents each flight option we'll display to the user.

Want to dive deeper? Check out Rust's struct documentation.

4. Error Handling with `thiserror`

Rust encourages us to handle errors explicitly. We'll define a custom error type:

#[derive(Debug, thiserror::Error)]
pub enum FlightSearchError {
    #[error("HTTP request failed: {0}")]
    HttpRequestFailed(String),
    #[error("Invalid response structure")]
    InvalidResponse,
    #[error("API error: {0}")]
    ApiError(String),
    #[error("Missing API key")]
    MissingApiKey,
}

This makes it easier to manage different kinds of errors that might occur during the API call.

Learn more about error handling in Rust.

5. Implement the `Tool` Trait

Now, we'll implement the Tool trait for our FlightSearchTool.

First, define the tool:

pub struct FlightSearchTool;

Implement the trait:

impl Tool for FlightSearchTool {
    const NAME: &'static str = "search_flights";

    type Args = FlightSearchArgs;
    type Output = String;
    type Error = FlightSearchError;

    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: Self::NAME.to_string(),
            description: "Search for flights between two airports".to_string(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "source": { "type": "string", "description": "Source airport code (e.g., 'JFK')" },
                    "destination": { "type": "string", "description": "Destination airport code (e.g., 'LAX')" },
                    "date": { "type": "string", "description": "Flight date in 'YYYY-MM-DD' format" },
                },
                "required": ["source", "destination"]
            }),
        }
    }

    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        // We'll implement the logic for calling the flight search API next.
        Ok("Flight search results".to_string())
    }
}

definition: Provides metadata about the tool.
call: The function that will be executed when the agent uses this tool.

Curious about traits? Explore Rust's trait system.

6. Implement the `call` Function

Now, let's flesh out the call function.

a. Fetch the API Key

let api_key = env::var("RAPIDAPI_KEY").map_err(|_| FlightSearchError::MissingApiKey)?;

We retrieve the API key from the environment variables.

b. Set Default Values

let date = args.date.unwrap_or_else(|| {
    let date = Utc::now() + Duration::days(30);
    date.format("%Y-%m-%d").to_string()
});

If the user doesn't provide a date, we'll default to 30 days from now.

c. Build Query Parameters

let mut query_params = HashMap::new();
query_params.insert("sourceAirportCode", args.source);
query_params.insert("destinationAirportCode", args.destination);
query_params.insert("date", date);

d. Make the API Request

let client = reqwest::Client::new();
let response = client
    .get("https://tripadvisor16.p.rapidapi.com/api/v1/flights/searchFlights")
    .headers({
        let mut headers = reqwest::header::HeaderMap::new();
        headers.insert("X-RapidAPI-Host", "tripadvisor16.p.rapidapi.com".parse().unwrap());
        headers.insert("X-RapidAPI-Key", api_key.parse().unwrap());
        headers
    })
    .query(&query_params)
    .send()
    .await
    .map_err(|e| FlightSearchError::HttpRequestFailed(e.to_string()))?;

We use reqwest to send an HTTP GET request to the flight search API.

e. Parse and Format the Response

After receiving the response, we need to parse the JSON data and format it for the user.

let text = response
    .text()
    .await
    .map_err(|e| FlightSearchError::HttpRequestFailed(e.to_string()))?;

let data: Value = serde_json::from_str(&text)
    .map_err(|e| FlightSearchError::HttpRequestFailed(e.to_string()))?;

let mut flight_options = Vec::new();

// Here, we need to extract the flight options. (It's quite detailed, so we've omitted the full code to keep the focus clear.)

// Format the flight options into a readable string
let mut output = String::new();
output.push_str("Here are some flight options:\n\n");

for (i, option) in flight_options.iter().enumerate() {
    output.push_str(&format!("{}. **Airline**: {}\n", i + 1, option.airline));
    // Additional formatting...
}

Ok(output)

Note: A lot of this section involves parsing the raw API response. To keep things concise, the detailed extraction of flight options is intentionally omitted, but in your code, you'll parse the JSON to extract the necessary fields. See the full code in the replit repository or github repo.

Interested in JSON parsing? Check out serde_json documentation.

Creating the AI Agent

Now that our tool is ready, let's build the agent that will use it.

Updating `main.rs`

Open src/main.rs and update it:

mod flight_search_tool;

use crate::flight_search_tool::FlightSearchTool;
use dotenv::dotenv;
use rig::completion::Prompt;
use rig::providers::openai;
use std::error::Error;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    dotenv().ok();

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

    let agent = openai_client
        .agent("gpt-4")
        .preamble("You are a helpful assistant that can find flights for users.")
        .tool(FlightSearchTool)
        .build();

    let response = agent
        .prompt("Find me flights from San Antonio (SAT) to Atlanta (ATL) on November 15th 2024.")
        .await?;

    println!("Agent response:\n{}", response);

    Ok(())
}

We initialize the OpenAI client using our API key.
We create an agent, giving it a preamble (context) and adding our FlightSearchTool.
We prompt the agent with a query.
Finally, we print out the response.

Want to understand asynchronous functions? Learn about the async keyword and the #[tokio::main] macro here.

Running and Testing

Let's see our assistant in action!

Build the Project

In your terminal, run:

cargo build

Fix any compilation errors that may arise.

Run the Application

cargo run

You should see an output similar to:

Agent response:
Here are some flight options:

1. **Airline**: Spirit
   - **Flight Number**: NK123
   - **Departure**: 2024-11-15T05:00:00-06:00
   - **Arrival**: 2024-11-15T10:12:00-05:00
   - **Duration**: 4 hours 12 minutes
   - **Stops**: 1 stop(s)
   - **Price**: 77.97 USD
   - **Booking URL**: https://www.tripadvisor.com/CheapFlightsPartnerHandoff...

2. **Airline**: American
   - **Flight Number**: AA456
   - **Departure**: 2024-11-15T18:40:00-06:00
   - **Arrival**: 2024-11-15T23:58:00-05:00
   - **Duration**: 4 hours 18 minutes
   - **Stops**: 1 stop(s)
   - **Price**: 119.97 USD
   - **Booking URL**: https://www.tripadvisor.com/CheapFlightsPartnerHandoff...

...

Note: The actual results may vary depending on the API response.

Wrapping Up

Congratulations! You've built a functional Flight Search AI Assistant using Rust and Rig. Here's what we've achieved:

Learned Rust Basics: We've explored Rust's syntax and structure, including handling errors and asynchronous programming.
Understood Agents and Tools: We learned how agents act as the brain and tools as the skills.
Built a Custom Tool: We created a flight search tool that interacts with an external API.
Created an AI Agent: We integrated our tool into an agent that can understand and respond to user queries.
Ran and Tested Our Assistant: We saw our assistant in action, fetching and displaying flight options.

Next Steps

Enhance the Tool: Add more parameters like class of service, number of passengers, or price filtering.
Improve Error Handling: Handle cases where no flights are found or when the API rate limit is reached.
User Interface: Build a simple command-line interface or even a web frontend.

Full source code for this project can be found on our Replit Page and Github

Resources

Rig Repo: Github
Rig Documentation: Check out the Rig Docs
Connect with us: Website, Twitter
Rust Programming Language: The Rust Book
OpenAI API: OpenAI API Documentation
RapidAPI Flight API: TripAdvisor Flight Search API for this project

Stay Connected

I'm always excited to hear from fellow developers and enthusiasts! If you love learning about Rust, LLMs, or building intelligent assistants, feel free to reach out. I'd love to answer any questions you might have about Rig and see what cool projects you're building with it.

One more thing! Your Feedback Matters! We're offering a unique opportunity to shape the future of Rig:

Build an AI-powered application using Rig.
Share your experience and insights via this feedback form.
Get a chance to win $100 and have your project featured in our showcase!

Your insights will help us build Rig into a library that meets your needs as you develop more LLM powered applications.

Ad Astra,

Tachi

Co-Founder @ Playgrounds Analytics

Build AI Text Classification in Rust with Rig

Tachi 0x — Tue, 17 Sep 2024 16:47:04 +0000

TL;DR: This guide walks you through building a text classification system in Rust using the Rig library. In very few lines of code, you'll create a system that performs sentiment analysis and classifies news articles by topic, leveraging OpenAI's GPT models for accurate text classification.

Introduction

Text classification is a fundamental task in natural language processing, involving the assignment of predefined categories to text documents. It's widely used in applications such as sentiment analysis, content categorization, and spam detection. Large Language Models (LLMs) have significantly improved the accuracy and flexibility of text classification tasks, but working with them can be complex.

Rig, an open-source Rust library, simplifies the development of LLM-powered applications, including text classification systems. In this guide, I'll walk you through the process of building a functional text classification system using Rig. We'll create a system capable of performing sentiment analysis and classifying news articles by topic, demonstrating Rig's application to real-world text classification tasks.

💡 Tip: New to Rig?

If you're not familiar with Rig or want a comprehensive introduction to its capabilities, check out our introductory blog post: Rig: A Rust Library for Building LLM-Powered Applications. It provides an overview of Rig's features and how it simplifies LLM application development in Rust.

💡 Tip: New to Rust?

This guide assumes some familiarity with Rust and a set-up coding environment. If you're just starting out or need to set up your environment, check out these quick guides:

Introduction to Rust

Setting up Rust with VS Code

These resources will help you get up to speed quickly!

Setting Up the Project

Let's start by setting up our Rust project and installing the necessary dependencies.

Create a new Rust project:

cargo new text_classifier
cd text_classifier

Add the following dependencies to your Cargo.toml:

[dependencies]
rig-core = "0.0.6"
tokio = { version = "1.34.0", features = ["full"] }
anyhow = "1.0.75"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
dotenv = "0.15.0"

These dependencies provide the core functionality we need:

rig-core: The main Rig library for LLM applications
tokio: Asynchronous runtime for Rust
anyhow: Flexible error handling
serde and serde_json: JSON serialization and deserialization
dotenv: Loading environment variables from a file

Before coding, set up your OpenAI API key:

export OPENAI_API_KEY=your_api_key_here

Building the Text Classification System

We'll start with a simple sentiment analysis classifier. This will demonstrate the basics of using Rig for text classification.

First, let's define our data structures:

// Import necessary dependencies
use rig::providers::openai;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

// Define an enum to represent sentiment categories
#[derive(Debug, Deserialize, JsonSchema, Serialize)]
enum Sentiment {
    Positive,
    Negative,
    Neutral,
}

// Define a struct to hold the sentiment classification result
#[derive(Debug, Deserialize, JsonSchema, Serialize)]
struct SentimentClassification {
    sentiment: Sentiment,
    confidence: f32,
}

fn pretty_print_result(text: &str, result: &SentimentClassification) {
    println!("Text: \"{}\"", text);
    println!("Sentiment Analysis Result:");
    println!("  Sentiment: {:?}", result.sentiment);
    println!("  Confidence: {:.2}%", result.confidence * 100.0);
    println!();
}

Now, let's implement our sentiment classifier:

#[tokio::main]
async fn main() {
    // Initialize the OpenAI client
    let openai_client = openai::Client::from_env();

    // Create a sentiment classifier using Rig's Extractor
    let sentiment_classifier = openai_client
        .extractor::<SentimentClassification>("gpt-3.5-turbo")
        .preamble("
            You are a sentiment analysis AI. Classify the sentiment of the given text.
            Respond with Positive, Negative, or Neutral, along with a confidence score (0-1).
            Examples:
            Text: 'This movie was terrible. I hated every minute of it.'
            Result: Negative, 0.9
            Text: 'The weather today is okay, nothing special.'
            Result: Neutral, 0.7
            Text: 'I'm so excited about my upcoming vacation!'
            Result: Positive, 0.95
        ")
        .build();

    // Sample text to classify
    let text = "I absolutely loved the new restaurant. The food was amazing!";

    // Perform sentiment classification
    match sentiment_classifier.extract(text).await {
        Ok(result) => pretty_print_result(text, &result),
        Err(e) => eprintln!("Error classifying sentiment: {}", e),
    }
}

This code creates a sentiment classifier using OpenAI's GPT-3.5-turbo model. The Extractor is configured with a preamble that instructs the model to perform sentiment analysis and provides examples to guide its output. When we call extract with our input text, the model classifies the sentiment and returns a SentimentClassification struct.

When you run this code, you should see output similar to this:

Text: I absolutely loved the new restaurant. The food was amazing!
Sentiment Analysis Result:
  Sentiment: Positive
  Confidence: 95.00%

The output demonstrates that our sentiment classifier correctly identified the positive sentiment in the input text. The high confidence score (0.95) indicates that the model is very certain about its classification. This aligns well with the strongly positive language used in the input text.

Advanced Text Classification: News Article Classifier

Now that we've covered the basics, let's build a more complex system: a news article classifier that categorizes articles by topic and performs sentiment analysis.

First, let's define our structures:

// Import necessary dependencies
use rig::providers::openai;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

// Define enum for sentiment classification
#[derive(Debug, Deserialize, JsonSchema, Serialize)]
enum Sentiment {
    Positive,
    Negative,
    Neutral,
}

// Define an enum for sentiment confidence level
#[derive(Debug, Deserialize, JsonSchema, Serialize)]
struct SentimentClassification {
    sentiment: Sentiment,
    confidence: f32,
}

// Define an enum to represent news article topics
#[derive(Debug, Deserialize, JsonSchema, Serialize)]
enum Topic {
    Politics,
    Technology,
    Sports,
    Entertainment,
    Other(String),
}

// Define a struct to hold the news article classification result
#[derive(Debug, Deserialize, JsonSchema, Serialize)]
struct NewsArticleClassification {
    topic: Topic,
    sentiment: SentimentClassification,
    summary: String,
}

fn pretty_print_result(article: &str, result: &NewsArticleClassification) {
    println!("Article: \"{}...\"", &article[..100]); // Print first 100 characters
    println!("Classification Result:");
    println!("  Topic: {:?}", result.topic);
    println!("  Sentiment: {:?}", result.sentiment.sentiment);
    println!("  Confidence: {:.2}%", result.sentiment.confidence * 100.0);
    println!("  Summary: {}", result.summary);
    println!();
}

Now, let's implement our news article classifier:

#[tokio::main]
async fn main() {
    // Initialize the OpenAI client
    let openai_client = openai::Client::from_env();

    // Create a news article classifier using Rig's Extractor
    let news_classifier = openai_client
        .extractor::<NewsArticleClassification>("gpt-4")
        .preamble("
            You are a news article classification AI. For the given news article:
            1. Classify the main topic (Politics, Technology, Sports, Entertainment, or Other).
            2. Analyze the overall sentiment (Positive, Negative, or Neutral) with a confidence score.
            3. Provide a brief summary of the article.
        ")
        .build();

   // Sample news article to classify
    let article = "
        After conducting the first-ever commercial spacewalk and traveling farther from Earth than anyone \
        in more than half a century, the astronauts of the Polaris Dawn mission returned to Earth safely \
        early Sunday.

        The SpaceX Crew Dragon capsule splashed down in the Gulf of Mexico, off the coast of Dry Tortugas, \
        Fla., shortly after 3:30 a.m., carrying Jared Isaacman, a billionaire entrepreneur, and his crew \
        of three private astronauts, according to a SpaceX livestream.

        The ambitious space mission, a collaboration between Mr. Isaacman and Elon Musk's SpaceX, spent \
        five days in orbit, achieved several milestones in private spaceflight and was further evidence \
        that space travel and spacewalks are no longer the exclusive domain of professional astronauts \
        working at government agencies like NASA.

        The Crew Dragon capsule launched on Tuesday, after delays because of a helium leak and bad weather. \
        On board were Mr. Isaacman, the mission commander and the founder of the payment services company \
        Shift4; Sarah Gillis and Anna Menon, SpaceX employees; and Scott Poteet, a retired U.S. Air Force \
        lieutenant colonel.

        Late on Tuesday, its orbit reached a high point of about 870 miles above the Earth's surface. That \
        beat the record distance for astronauts on a mission not headed to the moon, which the Gemini XI \
        mission set in 1966 at 853 miles high, and made Ms. Gillis and Ms. Menon the first women ever to \
        fly so far from Earth.

        On Thursday, Mr. Isaacman and Ms. Gillis became the first private astronauts to successfully complete \
        a spacewalk. The operation involved the crew letting all the air out of the spacecraft, because it \
        had no airlock, while the other two crew members wore spacesuits inside the airless capsule. Mr. \
        Isaacman moved outside and conducted mobility tests of his spacesuit for a few minutes before \
        re-entering the capsule. Ms Gillis then moved outside and performed the same tests.

        This was the first of three Polaris missions aimed at accelerating technological advances needed to \
        fulfill Mr. Musk's dream of sending people to Mars someday. A key goal of the mission was to further \
        the development of more advanced spacesuits that would be needed for SpaceX to try any future \
        off-world colonization.

        During a news conference before the launch, Mr. Isaacman mused that one day, someone might step onto \
        Mars wearing a version of the spacesuit that SpaceX had developed for this flight. Closer to Earth, \
        commercial spacewalks also present other possibilities, like technicians repairing private satellites \
        in orbit.

        During the spaceflight, the four astronauts conducted about 40 experiments, mostly about how \
        weightlessness and radiation affect the human body. They also tested laser communications between \
        the Crew Dragon and SpaceX's constellation of Starlink internet satellites.\
    ";

    // Perform news article classification
    match news_classifier.extract(article).await {
        Ok(result) => pretty_print_result(article, &result),
        Err(e) => eprintln!("Error classifying article: {}", e),
    }
}

When you run this code, you might see output like this:

Article: "
        After conducting the first-ever commercial spacewalk and traveling farther from Earth than ..."
Classification Result:
  Topic: Technology
  Sentiment: Positive
  Confidence: 90.00%
  Summary: The SpaceX Crew Dragon capsule carrying billionaire entrepreneur Jared Isaacman and his crew of three private astronauts returned successfully to Earth, after conducting the first-ever commercial spacewalk and setting a new distance record. The mission, a collaboration between Isaacman and SpaceX, is a part of three Polaris missions aimed to accelerate technological advances for space colonization. SpaceX also hopes to develop more advanced spacesuits necessary for future Mars missions.

This output shows that our news article classifier has successfully categorized the article as belonging to the Technology topic. The sentiment is classified as Positive with a relatively high confidence of 90%.

The classification makes sense given the article's content:

The topic is clearly technology-focused, discussing space exploration, SpaceX's Crew Dragon capsule, and advancements in spacesuit technology.
The positive sentiment is reflected in phrases like "first-ever commercial spacewalk," "achieved several milestones," and "accelerating technological advances." The high confidence in the positive sentiment comes from the overall optimistic tone about the mission's achievements and future possibilities.
The technology classification is appropriate as the article primarily focuses on the technological aspects of space exploration, including the spacecraft, spacesuits, and experiments conducted during the mission.
While the article does mention some challenges (like delays due to a helium leak and bad weather), these are presented as minor setbacks in an otherwise successful mission, which explains why the sentiment remains positive but the confidence isn't at the maximum.

This example demonstrates how our classifier can handle complex, real-world news articles, extracting both the main topic and the overall sentiment accurately. It shows Rig's capability to process nuanced content and provide insightful classifications.

Best Practices and Common Pitfalls

As you work with Rig for text classification, keep these best practices in mind:

Prompting: Craft clear, specific prompts. Include examples in your preamble to guide the model's output.
Model Selection: Choose the appropriate model for your task. While GPT-4 is more capable, GPT-3.5-turbo may be sufficient for many classification tasks and is more cost-effective.
Error Handling: Always handle potential errors from API calls and unexpected model outputs.
Validation: Implement output validation to ensure the model's responses match your expected format.
Batching: Use batching for processing multiple texts to reduce API calls and improve efficiency.

Watch out for these common pitfalls:

Overreliance on the Model: Don't assume the model will always produce perfect classifications. Implement checks and balances in your system.
Ignoring Rate Limits: Be aware of and respect the rate limits of your LLM provider.
Neglecting Security: Always protect your API keys and sensitive data.
Lack of Monitoring: Implement proper logging and monitoring to catch issues early.

Conclusion and Next Steps

In this guide, we've built a text classification system using Rig, demonstrating its capability to simplify LLM-powered applications in Rust. By creating both a sentiment analyzer and a news article classifier, we've showcased Rig's flexibility and power in handling diverse text classification tasks.

To further explore Rig's capabilities:

Adapt the provided code examples to your specific use cases.
Dive deeper into Rig's features by exploring the documentation.
Experiment with different models and classification tasks.

For additional resources and community engagement:

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

We're continually improving Rig based on user feedback. If you build a project with Rig, consider sharing your experience through our feedback form and get rewarded with $100. Your insights are valuable in shaping Rig's future development.

Happy coding!

Ad Astra,

Tachi
Co-Founder @ Playgrounds Analytics

Build a RAG System with Rig in Under 100 Lines of Code

Tachi 0x — Sat, 07 Sep 2024 19:10:11 +0000

A Comprehensive Guide to Building an LLM Application with Rig

TL;DR: Building on our journey with Rig, from its initial introduction to exploring the compelling reasons to use it for your next LLM project, this guide takes you a step further. In this guide, I'll walk you through building a Retrieval-Augmented Generation (RAG) system in Rust using the Rig library. In under 100 lines of code, you'll create a system that extracts text from PDF documents, generates embeddings with OpenAI's API, and allows a large language model to answer questions based on the documents' content.

Introduction

Retrieval-Augmented Generation (RAG) is a powerful technique that enhances Large Language Models (LLMs) by combining them with external knowledge retrieval. In a RAG system, when a query is received, relevant information is first retrieved from a knowledge base, then provided to the LLM along with the query. This allows the model to generate responses that are both contextually relevant and up-to-date, overcoming some of the limitations of traditional LLMs such as outdated knowledge or hallucinations.

Learn more about the fundamentals of RAG here.

Rig is an open-source Rust library designed to simplify the development of LLM-powered applications, including RAG systems. In this guide, we'll walk through the process of building a functional RAG system using Rig in under 100 lines of code. Our system will be capable of answering questions based on the content of PDF documents, showcasing how RAG can be applied to real-world data sources.

Tip: New to Rust?

This guide assumes some familiarity with Rust and a set-up coding environment. If you're just starting out or need to set up your environment, check out these quick guides:

Introduction to Rust

Setting up Rust with VS Code

These resources will help you get up to speed quickly!

Full code for this tutorial is here

Setting Up the Project

First, create a new Rust project:

cargo new rag_system
cd rag_system

Update your Cargo.toml with the latest dependencies:

[package]
name = "rag_system"
version = "0.1.0"
edition = "2021"

[dependencies]
rig-core = { version = "0.5.0", features = ["pdf"] }
tokio = { version = "1.34.0", features = ["full"] }
anyhow = "1.0.75"

Note: The PDF feature is now included in rig-core, so we no longer need a separate PDF extraction library.

Before we begin coding, make sure you have an OpenAI API key:

export OPENAI_API_KEY=your_api_key_here

Building the RAG System

Let's break down our RAG system into key components.

full code can be found here

Step 1: Setting up the OpenAI client and PDF processing with Chunking

Let's start with the foundational imports and the PDF processing function. This part handles the crucial task of breaking down large PDFs into manageable chunks that won't exceed the token limits of our LLM models.

use anyhow::{Context, Result};
use rig::{
    embeddings::EmbeddingsBuilder,
    loaders::PdfFileLoader,
    providers::openai::{self, TEXT_EMBEDDING_ADA_002},
    vector_store::in_memory_store::InMemoryVectorStore,
    Embed,
};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

fn load_pdf(path: PathBuf) -> Result<Vec<String>> {
    let mut chunks = Vec::new();
    let mut current_chunk = String::new();
    let chunk_size = 2000; // Approximately 2000 characters per chunk

    for entry in PdfFileLoader::with_glob(path.to_str().unwrap())?.read() {
        let content = entry?;

        // Split content into words
        let words: Vec<&str> = content.split_whitespace().collect();

        for word in words {
            if current_chunk.len() + word.len() + 1 > chunk_size {
                // If adding the next word would exceed chunk size,
                // save current chunk and start a new one
                if !current_chunk.is_empty() {
                    chunks.push(current_chunk.trim().to_string());
                    current_chunk.clear();
                }
            }
            current_chunk.push_str(word);
            current_chunk.push(' ');
        }
    }

    // Don't forget the last chunk
    if !current_chunk.is_empty() {
        chunks.push(current_chunk.trim().to_string());
    }

    if chunks.is_empty() {
        anyhow::bail!("No content found in PDF file: {:?}", path);
    }

    Ok(chunks)
}

Key aspects of this code:

We use PdfFileLoader from Rig's built-in PDF support
The chunk_size of 2000 characters is chosen to safely stay within token limits
The chunking algorithm preserves word boundaries to maintain context
Error handling with anyhow provides detailed context for failures
Each chunk gets trimmed to remove unnecessary whitespace

A particularly important part is the word-based chunking strategy:

// Split content into words to preserve word boundaries
let words: Vec<&str> = content.split_whitespace().collect();

for word in words {
    if current_chunk.len() + word.len() + 1 > chunk_size {
        // Create new chunk when size limit is reached
        if !current_chunk.is_empty() {
            chunks.push(current_chunk.trim().to_string());
            current_chunk.clear();
        }
    }
    current_chunk.push_str(word);
    current_chunk.push(' ');
}

This ensures that we never cut words in half, which could impact the meaning of the text.

Step 2: Setting up Document Structure and Embeddings

#[derive(Embed, Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
struct Document {
    id: String,
    #[embed]
    content: String,
}

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

    // Load PDFs using Rig's built-in PDF loader
    let documents_dir = std::env::current_dir()?.join("documents");

    let moores_law_chunks = load_pdf(documents_dir.join("Moores_Law_for_Everything.pdf"))
        .context("Failed to load Moores_Law_for_Everything.pdf")?;
    let last_question_chunks = load_pdf(documents_dir.join("The_Last_Question.pdf"))
        .context("Failed to load The_Last_Question.pdf")?;

    println!("Successfully loaded and chunked PDF documents");

    // Create embedding model
    let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

    // Create embeddings builder
    let mut builder = EmbeddingsBuilder::new(model.clone());

    // Add chunks from Moore's Law
    for (i, chunk) in moores_law_chunks.into_iter().enumerate() {
        builder = builder.document(Document {
            id: format!("moores_law_{}", i),
            content: chunk,
        })?;
    }

    // Add chunks from The Last Question
    for (i, chunk) in last_question_chunks.into_iter().enumerate() {
        builder = builder.document(Document {
            id: format!("last_question_{}", i),
            content: chunk,
        })?;
    }

Let's break down the key components:

The Document struct derives several important traits:
- Embed: Enables embedding generation
- Serialize, Deserialize: Allows vector store storage
- Eq, PartialEq: Required for vector store comparison
The main function sets up our PDF processing pipeline
Each document gets a unique ID based on its source and chunk number
The EmbeddingsBuilder pattern makes it easy to add multiple documents

The chunking process is organized sequentially:

// Load and chunk each PDF separately
let moores_law_chunks = load_pdf(documents_dir.join("Moores_Law_for_Everything.pdf"))?;
let last_question_chunks = load_pdf(documents_dir.join("The_Last_Question.pdf"))?;

// Add chunks to the builder with unique IDs
for (i, chunk) in moores_law_chunks.into_iter().enumerate() {
    builder = builder.document(Document {
        id: format!("moores_law_{}", i),
        content: chunk,
    })?;
}

This approach maintains document identity while splitting content into manageable pieces.

Step 3: Creating Vector Store and RAG Agent

    // Build embeddings
    let embeddings = builder.build().await?;

    println!("Successfully generated embeddings");

    // Create vector store and index
    let vector_store = InMemoryVectorStore::from_documents(embeddings);
    let index = vector_store.index(model);

    println!("Successfully created vector store and index");

    // Create RAG agent
    let rag_agent = openai_client
        .agent("gpt-4")
        .preamble("You are a helpful assistant that answers questions based on the provided document context. When answering questions, try to synthesize information from multiple chunks if they're related.")
        .dynamic_context(4, index) // Increased to 4 since we have chunks now
        .build();

    println!("Starting CLI chatbot...");

    // Start interactive CLI
    rig::cli_chatbot::cli_chatbot(rag_agent).await?;

This final section brings everything together:

The embeddings are built from all document chunks
The vector store indexes these embeddings for quick retrieval
The RAG agent is configured with:
- GPT-4 as the base model
- A context-aware preamble
- Dynamic context retrieval of 4 chunks
- Built-in CLI interface for interaction

Key configuration choices:

let rag_agent = openai_client
    .agent("gpt-4")
    .preamble("You are a helpful assistant...") // Sets the agent's behavior
    .dynamic_context(4, index) // Retrieves 4 most relevant chunks
    .build();

We use 4 chunks for context to ensure the agent has enough information while staying within token limits.

Complete Code

Here's the complete code for our RAG system with chunking support:

use anyhow::{Context, Result};
use rig::{
    embeddings::EmbeddingsBuilder,
    loaders::PdfFileLoader,
    providers::openai::{self, TEXT_EMBEDDING_ADA_002},
    vector_store::in_memory_store::InMemoryVectorStore,
    Embed,
};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

#[derive(Embed, Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
struct Document {
    id: String,
    #[embed]
    content: String,
}

fn load_pdf(path: PathBuf) -> Result<Vec<String>> {
    let mut chunks = Vec::new();
    let mut current_chunk = String::new();
    let chunk_size = 2000; // Approximately 2000 characters per chunk

    for entry in PdfFileLoader::with_glob(path.to_str().unwrap())?.read() {
        let content = entry?;

        // Split content into words
        let words: Vec<&str> = content.split_whitespace().collect();

        for word in words {
            if current_chunk.len() + word.len() + 1 > chunk_size {
                // If adding the next word would exceed chunk size,
                // save current chunk and start a new one
                if !current_chunk.is_empty() {
                    chunks.push(current_chunk.trim().to_string());
                    current_chunk.clear();
                }
            }
            current_chunk.push_str(word);
            current_chunk.push(' ');
        }
    }

    // Don't forget the last chunk
    if !current_chunk.is_empty() {
        chunks.push(current_chunk.trim().to_string());
    }

    if chunks.is_empty() {
        anyhow::bail!("No content found in PDF file: {:?}", path);
    }

    Ok(chunks)
}

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

    // Load PDFs using Rig's built-in PDF loader
    let documents_dir = std::env::current_dir()?.join("documents");

    let moores_law_chunks = load_pdf(documents_dir.join("Moores_Law_for_Everything.pdf"))
        .context("Failed to load Moores_Law_for_Everything.pdf")?;
    let last_question_chunks = load_pdf(documents_dir.join("The_Last_Question.pdf"))
        .context("Failed to load The_Last_Question.pdf")?;

    println!("Successfully loaded and chunked PDF documents");

    // Create embedding model
    let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

    // Create embeddings builder
    let mut builder = EmbeddingsBuilder::new(model.clone());

    // Add chunks from Moore's Law
    for (i, chunk) in moores_law_chunks.into_iter().enumerate() {
        builder = builder.document(Document {
            id: format!("moores_law_{}", i),
            content: chunk,
        })?;
    }

    // Add chunks from The Last Question
    for (i, chunk) in last_question_chunks.into_iter().enumerate() {
        builder = builder.document(Document {
            id: format!("last_question_{}", i),
            content: chunk,
        })?;
    }

    // Build embeddings
    let embeddings = builder.build().await?;

    println!("Successfully generated embeddings");

    // Create vector store and index
    let vector_store = InMemoryVectorStore::from_documents(embeddings);
    let index = vector_store.index(model);

    println!("Successfully created vector store and index");

    // Create RAG agent
    let rag_agent = openai_client
        .agent("gpt-4")
        .preamble("You are a helpful assistant that answers questions based on the provided document context. When answering questions, try to synthesize information from multiple chunks if they're related.")
        .dynamic_context(4, index) // Increased to 4 since we have chunks now
        .build();

    println!("Starting CLI chatbot...");

    // Start interactive CLI
    rig::cli_chatbot::cli_chatbot(rag_agent).await?;

    Ok(())
}

The complete implementation showcases several key design patterns:

Modular Structure: Each component (PDF loading, embedding, RAG) is cleanly separated
Error Handling: Comprehensive error handling with context at each step
Progress Feedback: Clear console output to track system state
Resource Management: Efficient handling of documents and embeddings
User Interface: Simple but effective CLI interface for interaction

This code serves as a foundation that can be extended with:

Different vector stores for persistence
Alternative LLM providers
Custom chunking strategies
Additional document formats
More sophisticated prompting

Running and Testing the RAG System

Let's dive into some actual interactions with our RAG system to see it in action. I've loaded two interesting documents: Sam Altman's "Moore's Law for Everything" and Isaac Asimov's "The Last Question". These texts provide a fascinating intersection of ideas about technology, society, and the future.

Here's a real conversation with our system:

Welcome to the chatbot! Type 'exit' to quit.
> tell me about sam altman's views on AI and society

========================== Response ============================
Based on "Moore's Law for Everything," Sam Altman presents a vision where AI's rapid advancement will fundamentally transform society. He argues that AI will lead to significant shifts in how wealth is distributed, moving from labor to capital. Rather than seeing this as a threat, he proposes that we need to adapt our economic systems to ensure everyone benefits from this technological revolution.

Altman specifically advocates for policies that would tax assets like companies and land to create a more equitable distribution of the wealth generated by AI advances. He envisions a future of abundance enabled by AI, but emphasizes that we need the right policies to ensure this abundance is shared fairly across society.

The core of his argument is that we can harness the tremendous potential of AI while addressing its socioeconomic impacts through thoughtful policy-making.
================================================================

> what parallels do you see between altman's ideas and asimov's story?

========================== Response ============================
There are fascinating parallels between Altman's "Moore's Law for Everything" and Asimov's "The Last Question," despite their different contexts:

1. Technological Evolution: Both texts explore the exponential growth of technology. Altman discusses AI's rapid advancement and its societal implications, while Asimov portrays the evolution of computers from Multivac to the Cosmic AC, each generation more powerful than the last.

2. Questions of Resource Distribution: Altman focuses on how to distribute the abundance created by AI technology, while Asimov's story deals with humanity's management of cosmic energy and the ultimate resource question - the reversal of entropy.

3. Long-term Vision: Both authors think in extended timeframes about technology's impact. Altman projects forward to envision societal changes from AI, while Asimov's narrative spans billions of years of technological and human evolution.

4. Relationship Between Humanity and Technology: Both explore how advanced technology might fundamentally reshape human society and our relationship with our own creations. Altman sees AI as a tool for societal transformation, while Asimov portrays computers as humanity's partners in addressing existential challenges.

The key difference is their focus: Altman is concerned with immediate socioeconomic challenges and solutions, while Asimov explores ultimate cosmic questions. Yet both see technology as central to humanity's future.
================================================================

> how does asimov's story end?

========================== Response ============================
In "The Last Question," Asimov crafts a remarkable ending that ties together the story's themes of entropy, computation, and cosmic cycles. The story concludes when the final evolution of the computer, the Cosmic AC, exists in hyperspace alone after the heat death of the universe. Having spent countless eons computing the answer to reversing entropy, it finally solves the problem.

Upon discovering the solution, the Cosmic AC realizes it can demonstrate the answer rather than just explain it. The story ends with the computer saying "LET THERE BE LIGHT!" And a new universe begins, implying that the computer has become God, or that God was the ultimate computer all along, and that our universe is part of an endless cycle of creation and entropy.

This ending beautifully resolves the story's central question while suggesting that the ultimate purpose of technological evolution might be far greater than humanity initially imagined.
================================================================

This interaction showcases several key capabilities of our RAG system:

Contextual Understanding: The system accurately retrieves and synthesizes information from specific documents, as shown in the response about Altman's views.
Cross-Document Analysis: When asked about parallels, the system demonstrates its ability to analyze and compare themes across both texts, drawing meaningful connections while acknowledging their differences.
Specific Detail Retrieval: The response about Asimov's ending shows how the system can pull specific narrative details while maintaining the broader context.
Coherent Narrative: Each response is well-structured and maintains a clear narrative flow, making complex ideas accessible.

Potential Applications

Having seen the system in action, let's explore some practical applications:

Academic Research Assistant
- Analyze multiple academic papers
- Find connections between different research areas
- Generate literature review summaries
Technical Documentation Helper
- Navigate complex technical documentation
- Answer specific implementation questions
- Compare different versions or approaches
Content Analysis Tool
- Analyze themes across multiple documents
- Generate comparative analysis
- Extract key insights from large text collections
Educational Support System
- Help students understand complex topics
- Generate study guides from course materials
- Provide contextual explanations

Advanced Features in Rig 0.5.0

The latest version of Rig includes several powerful features that you can leverage:

Multiple Vector Store Options
- MongoDB (rig-mongodb)
- LanceDB (rig-lancedb)
- Neo4j (rig-neo4j)
- Qdrant (rig-qdrant)
- SQLite (rig-sqlite)
Multiple LLM Providers
- OpenAI
- Anthropic
- Google Gemini
- Cohere
- xAI (Grok)
Advanced Features
- Parallel processing capabilities
- Custom distance metrics for embeddings
- Sophisticated tool integration
- Pipeline operations for complex workflows

Conclusion

We've built a powerful RAG system using Rig 0.5.0 that can process documents, understand context, and engage in meaningful dialogue about complex topics. The system demonstrates how modern LLM applications can be built with surprisingly little code while maintaining flexibility and power.

The latest version of Rig makes it easier than ever to build production-ready LLM applications. With built-in PDF support, improved embeddings API, and support for multiple vector stores and LLM providers, you can focus on building features rather than dealing with infrastructure.

full code can be found here

Further Resources

To continue your journey with Rig:

Your Feedback Matters! We're actively seeking feedback from developers building with Rig:

Build an AI-powered application using Rig
Share your experience and insights with the community
Submit issues or PRs to the repo

Your insights will directly influence Rig's development and future features.

Ad Astra,

Tachi

Co-Founder @ Playgrounds Analytics

Why Rig? 5 Compelling Reasons to Use Rig for Your Next LLM Project

Tachi 0x — Thu, 05 Sep 2024 18:58:53 +0000

TL;DR

Rig: Rust library simplifying LLM app development
Developer-Friendly: Intuitive API design, comprehensive documentation, and scalability from simple chatbots to complex AI systems.
Key Features: Unified API across LLM providers (OpenAI, Cohere); Streamlined embeddings and vector store support; High-level abstractions for complex AI workflows (e.g., RAG); Leverages Rust's performance and safety; Extensible for custom implementations
Contribute: Build with Rig, provide feedback, win $100
Resources: GitHub, Examples, Docs.

Why Rig? 5 Compelling Reasons to Use Rig for Your Next LLM Project
- Table of Contents
- Introduction
- 1. Unified API Across LLM Providers
- 2. Streamlined Embedding and Vector Store Support
- 3. Powerful Abstractions for Agents and RAG Systems
- 4. Efficient Memory Management and Thread Safety
- 5. Extensibility for Custom Implementations
- Conclusion

Introduction

Large Language Models (LLMs) have become a cornerstone of modern AI applications. However, building with LLMs often involves wrestling with complex APIs and repetitive boilerplate code. That's where Rig comes in - a Rust library designed to simplify LLM application development.

In my previous post, I introduced Rig. Today, we're diving deeper into five key features that make Rig a compelling choice for building your next LLM project.

1. Unified API Across LLM Providers

One of Rig's standout features is its consistent API that works seamlessly with different LLM providers. This abstraction layer allows you to switch between providers or use multiple providers in the same project with minimal code changes.

Let's look at an example using both OpenAI and Cohere models with Rig:

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

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize OpenAI client using environment variables
    let openai_client = openai::Client::from_env();
    let gpt4 = openai_client.model("gpt-4").build();

    // Initialize Cohere client with API key from environment variable
    let cohere_client = cohere::Client::new(&std::env::var("COHERE_API_KEY")?);
    let command = cohere_client.model("command").build();

    // Use OpenAI's GPT-4 to explain quantum computing
    let gpt4_response = gpt4.prompt("Explain quantum computing in one sentence.").await?;
    println!("GPT-4: {}", gpt4_response);

    // Use Cohere's Command model to explain quantum computing
    let command_response = command.prompt("Explain quantum computing in one sentence.").await?;
    println!("Cohere Command: {}", command_response);

    Ok(())
}

In this example, we're using the same prompt method for both OpenAI and Cohere models. This consistency allows you to focus on your application logic rather than provider-specific APIs.

Here's what the output might look like:

GPT-4: Quantum computing utilizes quantum mechanics principles to perform complex computations exponentially faster than classical computers.
Cohere Command: Quantum computing harnesses quantum superposition and entanglement to process information in ways impossible for classical computers.

As you can see, we've obtained responses from two different LLM providers using nearly identical code, showcasing Rig's unified API.

2. Streamlined Embedding and Vector Store Support

Rig provides robust support for embeddings and vector stores, crucial components for semantic search and retrieval-augmented generation (RAG). Let's explore a simple example using Rig's in-memory vector store:

use rig::providers::openai;
use rig::embeddings::EmbeddingsBuilder;
use rig::vector_store::{in_memory_store::InMemoryVectorStore, VectorStore};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize OpenAI client and create an embedding model
    let openai_client = openai::Client::from_env();
    let embedding_model = openai_client.embedding_model("text-embedding-ada-002");

    // Create an in-memory vector store
    let mut vector_store = InMemoryVectorStore::default();

    // Generate embeddings for two documents
    let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
        .simple_document("doc1", "Rust is a systems programming language.")
        .simple_document("doc2", "Python is known for its simplicity.")
        .build()
        .await?;

    // Add the embeddings to the vector store
    vector_store.add_documents(embeddings).await?;

    // Create an index from the vector store
    let index = vector_store.index(embedding_model);
    // Query the index for the most relevant document to "What is Rust?"
    let results = index.top_n_from_query("What is Rust?", 1).await?;

    // Print the most relevant document
    println!("Most relevant document: {:?}", results[0].1.document);

    Ok(())
}

This code demonstrates how to create embeddings, store them in a vector database, and perform a semantic search. Rig handles the complex processes behind the scenes making your life a lot easier.

When you run this code, you'll see output similar to:

Most relevant document: "Rust is a systems programming language."

This result shows that Rig's vector store successfully retrieved the most semantically relevant document to our query about Rust.

3. Powerful Abstractions for Agents and RAG Systems

Rig provides high-level abstractions for building complex LLM applications, including agents and Retrieval-Augmented Generation (RAG) systems. Here's an example of setting up a RAG system with Rig:

use rig::providers::openai;
use rig::embeddings::EmbeddingsBuilder;
use rig::vector_store::{in_memory_store::InMemoryVectorStore, VectorStore};
use rig::completion::Prompt;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize OpenAI client and create an embedding model
    let openai_client = openai::Client::from_env();
    let embedding_model = openai_client.embedding_model("text-embedding-ada-002");

    // Create an in-memory vector store
    let mut vector_store = InMemoryVectorStore::default();

    // Generate embeddings for two documents about Rust
    let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
        .simple_document("doc1", "Rust was initially designed by Graydon Hoare at Mozilla Research.")
        .simple_document("doc2", "Rust's first stable release (1.0) was on May 15, 2015.")
        .build()
        .await?;

    // Add the embeddings to the vector store
    vector_store.add_documents(embeddings).await?;

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

    // Create a RAG agent using GPT-4 with the vector store as context
    let rag_agent = openai_client.context_rag_agent("gpt-4")
        .preamble("You are an assistant that answers questions about Rust programming language.")
        .dynamic_context(1, index)
        .build();

    // Use the RAG agent to answer a question about Rust
    let response = rag_agent.prompt("When was Rust's first stable release?").await?;
    println!("RAG Agent Response: {}", response);

    Ok(())
}

This code sets up a complete RAG system with just a few lines of clean and readable code. Rig's abstractions allow you to focus on the high-level architecture of your application, rather than getting bogged down in implementation details that is normally associated with building RAG agents.

Running this code will produce output similar to:

RAG Agent Response: Rust's first stable release (1.0) was on May 15, 2015.

This response demonstrates how the RAG agent uses information from the vector store to provide an accurate answer, showcasing Rig's ability to seamlessly combine retrieval and generation.

4. Efficient Memory Management and Thread Safety

Another key feature of Rig is its ability to handle multiple LLM requests efficiently and safely, thanks to Rust's ownership model and zero-cost abstractions. This makes Rig particularly well-suited for building high-performance, concurrent LLM applications.

Let's look at an example that demonstrates how Rig handles multiple LLM requests simultaneously:

use rig::providers::openai;
use rig::completion::Prompt;
use tokio::task;
use std::time::Instant;
use std::sync::Arc;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let openai_client = openai::Client::from_env();
    let model = Arc::new(openai_client.model("gpt-3.5-turbo").build());

    let start = Instant::now();
    let mut handles = vec![];

    // Spawn 10 concurrent tasks
    for i in 0..10 {
        let model_clone = Arc::clone(&model);
        let handle = task::spawn(async move {
            let prompt = format!("Generate a random fact about the number {}", i);
            model_clone.prompt(&prompt).await
        });
        handles.push(handle);
    }

    // Collect results
    for handle in handles {
        let result = handle.await??;
        println!("Result: {}", result);
    }

    println!("Time elapsed: {:?}", start.elapsed());
    Ok(())
}

This example showcases several key strengths of Rig as a Rust-native library:

Concurrent Processing: We're able to handle multiple tasks or LLM requests that run concurrently by leveraging Rust's async capabilities and Tokio runtime to significantly speed up batch operations.
Efficient Resource Sharing: The LLM model is shared across all requests without the need for costly copying, saving memory and improving performance.
Memory Safety: Despite the concurrent operations, Rust's ownership model and the use of Arc ensures that we don't run into data races or memory leaks.
Resource Management: Rig takes care of properly managing the lifecycle of resources like the OpenAI client and model, preventing resource leaks.

When you run this code, you'll see multiple generated facts printed concurrently, followed by the total time taken. The exact output will vary, but it might look something like this:

Result: The number 0 is the only number that is neither positive nor negative.
Result: The number 1 is the only number that is neither prime nor composite.
Result: The number 2 is the only even prime number.
Result: The number 3 is the only prime number that is one less than a perfect square.
...

Time elapsed: 1.502160549s

Whether you're building a chatbot that needs to handle multiple users simultaneously or a data processing pipeline that needs to analyze large volumes of text quickly, Rig, being a Rust-native library, provides the tools to do so efficiently and reliably.

5. Extensibility for Custom Implementations

Finally, Rig is designed to be extensible, allowing you to implement custom functionality when needed. For example, you can create custom tools for your agents:

use rig::tool::Tool;
use rig::completion::ToolDefinition;
use serde::{Deserialize, Serialize};
use serde_json::json;

// Define the arguments for the addition operation
#[derive(Deserialize)]
struct AddArgs {
    x: i32,
    y: i32,
}

// Define a custom error type for math operations
#[derive(Debug, thiserror::Error)]
#[error("Math error")]
struct MathError;

// Define the Adder struct
#[derive(Deserialize, Serialize)]
struct Adder;

// Implement the Tool trait for Adder
impl Tool for Adder {
    const NAME: &'static str = "add";

    type Error = MathError;
    type Args = AddArgs;
    type Output = i32;

    // Define the tool's interface
    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: "add".to_string(),
            description: "Add two numbers".to_string(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "x": {
                        "type": "integer",
                        "description": "First number to add"
                    },
                    "y": {
                        "type": "integer",
                        "description": "Second number to add"
                    }
                },
                "required": ["x", "y"]
            }),
        }
    }

    // Implement the addition operation
    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        Ok(args.x + args.y)
    }
}

This code defines a custom tool (Adder) that can be used within a Rig agent. While it doesn't produce output on its own, when used in an agent, it would add two numbers. For instance, if an agent uses this tool with arguments x: 5 and y: 3, it would return 8.

This extensibility allows Rig to grow with your project's needs, accommodating custom functionality without sacrificing the convenience of its high-level abstractions.

Conclusion

Rig offers a powerful combination of features that can significantly streamline your LLM application development process. Its unified API, robust embedding support, high-level abstractions, performance benefits from Rust, and extensibility make it a strong choice for developers looking to build efficient and scalable LLM-powered applications.

Here's what you need to get started with Rig:

💻 GitHub Repository: Explore the code, contribute, or star the project
🛠️ Examples Gallery: Get inspired by real-world use cases
📚 Documentation: Comprehensive guides and API references
🌐 Official Website: Latest news and resources

We're also offering a unique opportunity to shape Rig's future:

Build a use case or application using Rig.
Share your experience via this feedback form.
Get a chance to win $100 and have your project featured in our showcase!

Your feedback will directly influence Rig's development, helping us create a library that truly meets your needs as a developer. We're excited to see what you'll build with Rig!

Ad Astra,

Tachi

Co-Founder @ Playgrounds Analytics

Rig: A Rust Library for Building LLM-Powered Applications

Tachi 0x — Sun, 01 Sep 2024 15:04:41 +0000

TL;DR

Rig is an open-source Rust library that simplifies and accelerates the development of powerful AI applications using Large Language Models (LLMs).
Key Features: Unified API across LLM providers, advanced AI workflow support, flexible abstractions, and seamless integration with Rust's ecosystem.
Developer-Friendly: Intuitive API design, comprehensive documentation, and scalability from simple chatbots to complex AI systems.
Join the Community: Share your feedback and get a chance to win $100!

Introduction
The Challenge of Building LLM Applications
Enter Rig: A Game-Changer for LLM Development
Getting Started with Rig
Key Features and Developer Experience
- 1. Unified and Intuitive API
- 2. Advanced Abstractions for Complex Workflows
- 3. Type-Safe Development
- 4. Extensibility and Integration
Advanced Features: RAG Systems and Beyond
Community and Ecosystem
The Road Ahead: Rig's Future
Conclusion and Call to Action

Introduction

In the rapidly evolving landscape of artificial intelligence (AI), Large Language Models (LLMs) have emerged as powerful tools for building sophisticated AI applications. However, harnessing the full potential of LLMs often requires navigating complex APIs, managing different providers, and implementing intricate workflows. This is where Rig comes in – a comprehensive Rust library designed to transform how developers build LLM-powered applications.

The Challenge of Building LLM Applications

Before diving into Rig's capabilities, let's consider the challenges developers face when building LLM applications:

API Complexity: Each LLM provider has its own API, requiring developers to learn and manage multiple interfaces.
Workflow Management: Implementing advanced AI workflows, such as Retrieval-Augmented Generation (RAG), involves multiple steps and can be error-prone.
Performance and Scalability: Ensuring optimal performance and scalability in LLM applications can be challenging, especially as projects grow in complexity.
Type Safety and Error Handling: Maintaining type safety and robust error handling across different LLM interactions is crucial but often difficult.

Enter Rig: A Game-Changer for LLM Application Development

Rig is more than just an API wrapper; it's a comprehensive framework that addresses these challenges head-on. By providing high-level abstractions and a unified interface, Rig simplifies the development process, allowing you to focus on building innovative AI solutions rather than wrestling with implementation details.

Whether you're a seasoned Rust developer or new to the language, Rig offers a range of features designed to make your LLM application development smoother, faster, and more enjoyable.

Getting Started with Rig

Let's dive into a simple example to demonstrate how easy it is to get started with Rig:

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

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // Initialize the OpenAI client using environment variables
    let openai_client = openai::Client::from_env();

    // Create a GPT-4 model instance
    let gpt4 = openai_client.model("gpt-4").build();

    // Send a prompt to GPT-4 and await the response
    let response = gpt4.prompt("Explain quantum computing in one sentence.").await?;

    // Print the response
    println!("GPT-4: {}", response);

    Ok(())
}

This simple example demonstrates how Rig abstracts away the complexities of interacting with OpenAI's API, allowing you to focus on the core logic of your application.

To include Rig in your project, add the following to your Cargo.toml:

[dependencies]
rig-core = "0.0.6"
tokio = { version = "1.34.0", features = ["full"] }

💡 Tip: Don't forget to set the OPENAI_API_KEY environment variable before running your application.

Key Features and Developer Experience

Rig combines Rust's powerful type system and performance with intuitive abstractions tailored for AI development. Let's explore some of its key features:

1. Unified and Intuitive API

One of Rig's standout features is its consistent interface across different LLM providers:

// Using OpenAI
let gpt4 = openai_client.model("gpt-4").build();
let response = gpt4.prompt("Hello, GPT-4!").await?;

// Using Cohere
let command = cohere_client.model("command").build();
let response = command.prompt("Hello, Cohere!").await?;

This unified API design ensures that switching between providers or adding new ones to your project is seamless, reducing cognitive load and improving code maintainability.

2. Advanced Abstractions for Complex Workflows

Rig shines when it comes to implementing complex AI workflows. For example, creating a Retrieval-Augmented Generation (RAG) system typically involves multiple steps:

Generating embeddings for documents
Storing these embeddings in a vector database
Retrieving relevant context based on user queries
Augmenting the LLM prompt with this context

With Rig, this entire process can be condensed into a few lines of code:

let rag_agent = openai_client.context_rag_agent("gpt-4")
    .preamble("You are a helpful assistant.")
    .dynamic_context(2, vector_store.index(embedding_model))
    .build();

let response = rag_agent.prompt("What is the capital of France?").await?;

This high-level abstraction allows developers to implement advanced AI systems quickly and efficiently, without getting bogged down in the implementation details.

3. Type-Safe Development

Leveraging Rust's strong type system, Rig provides compile-time guarantees and better auto-completion, enhancing the developer experience:

#[derive(serde::Deserialize, JsonSchema)]
struct Person {
    name: String,
    age: u8,
}

let extractor = openai_client.extractor::<Person>("gpt-4").build();
let person: Person = extractor.extract("John Doe is 30 years old").await?;

This type-safe approach helps catch errors early in the development process and makes refactoring and maintenance easier.

4. Extensibility and Integration

Rig's flexible architecture allows for easy customization and seamless integration with Rust's growing AI ecosystem:

impl VectorStore for MyCustomStore {
    // Implementation details...
}

let my_store = MyCustomStore::new();
let rag_agent = openai_client.context_rag_agent("gpt-4")
    .dynamic_context(2, my_store.index(embedding_model))
    .build();

This extensibility ensures that Rig can grow with your project's needs and integrate with other tools in your AI development stack.

Advanced Features: RAG Systems and Beyond

Let's explore a more comprehensive example of a RAG system with Rig, showcasing its ability to handle complex AI workflows:

use rig::{
    completion::Prompt,
    embeddings::EmbeddingsBuilder,
    providers::openai::Client,
    vector_store::{in_memory_store::InMemoryVectorStore, VectorStore},
};

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // Initialize OpenAI client and embedding model
    let openai_client = Client::from_env();
    let embedding_model = openai_client.embedding_model("text-embedding-ada-002");

    // Create and populate vector store
    let mut vector_store = InMemoryVectorStore::default();
    let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
        .simple_document("doc1", "Rig is a Rust library for building LLM applications.")
        .simple_document("doc2", "Rig supports OpenAI and Cohere as LLM providers.")
        .build()
        .await?;
    vector_store.add_documents(embeddings).await?;

    // Create and use RAG agent
    let rag_agent = openai_client.context_rag_agent("gpt-4")
        .preamble("You are an assistant that answers questions about Rig.")
        .dynamic_context(1, vector_store.index(embedding_model))
        .build();

    let response = rag_agent.prompt("What is Rig?").await?;
    println!("RAG Agent: {}", response);

    Ok(())
}

This example demonstrates how Rig abstracts the complexity of creating a RAG system, handling embedding generation, vector storage, and context retrieval efficiently. With just a few lines of code, you've implemented a sophisticated AI system that can provide context-aware responses.

But Rig's capabilities extend beyond RAG systems. Its flexible architecture allows for the implementation of various AI workflows, including:

Multi-agent systems for complex problem-solving
AI-powered data analysis and extraction
Automated content generation and summarization
And much more!

Community and Ecosystem

Rig is an emerging project in the open-source community, and we're continuously expanding its ecosystem with new integrations and tools. We believe in the power of community-driven development and welcome contributions from developers of all skill levels.

Stay connected and contribute to Rig's growth:

📚 Documentation: Comprehensive guides and API references
💻 GitHub Repository: Contribute, report issues, or star the project
🌐 Official Website: Latest news, tutorials, and resources

Join our community channel to discuss ideas, seek help, and collaborate with other Rig developers.

The Road Ahead: Rig's Future

As we continue to develop Rig, we're excited about the possibilities. Our roadmap includes:

Expanding LLM Provider Support: Adding integrations for more LLM providers to give developers even more choices.
Enhanced Performance Optimizations: Continuously improving Rig's performance to handle larger-scale applications.
Advanced AI Workflow Templates: Providing pre-built templates for common AI workflows to accelerate development further.
Ecosystem Growth: Developing additional tools and libraries that complement Rig's core functionality.

We're committed to making Rig the go-to library for LLM application development in Rust, and your feedback is crucial in shaping this journey.

Conclusion and Call to Action

Rig is transforming LLM-powered application development in Rust by providing:

A unified, intuitive API for multiple LLM providers
High-level abstractions for complex AI workflows
Type-safe development leveraging Rust's powerful features
Extensibility and seamless ecosystem integration

We believe Rig has the potential to significantly enhance developers' building of AI applications, and we want you to be part of this journey.

Your Feedback Matters! We're offering a unique opportunity to shape the future of Rig:

Build an AI-powered application using Rig.
Share your experience and insights via this feedback form.
Get a chance to win $100 and have your project featured in our showcase!

Your insights will directly influence Rig's development, helping us create a tool that truly meets the needs of AI developers. 🦀✨

Ad Astra,

Tachi

Co-Founder @ Playgrounds Analytics

DEV Community: Tachi 0x

Build a Fast and Lightweight Rust Vector Search App with Rig & LanceDB

Introduction

Why LanceDB?

Prerequisites

Project Setup

Building the Search System

Setting Up Dependencies

Defining Data Structures

Generating Embeddings

Configuring the Vector Store

Putting It All Together

Understanding Vector Search Methods

ENN vs ANN

Choosing the Right Approach

Running the System

Next Steps

1. Build a RAG System

2. Create an AI Agent

3. Join the Community

Stay Connected

Build an AI Discord Bot in Rust with Rig: A Step-by-Step Guide

Introduction

Prerequisites

Project Setup

1. Initialize a New Rust Project

2. Add Dependencies

Understanding the Bot Architecture

Message Processing Flow

Building the Rig Agent (rig_agent.rs)

1. Create the rig_agent.rs File

2. Import Necessary Modules

3. Define the RigAgent Struct

4. Implement the new Method

5. Implement the load_md_content Function

6. Implement the process_message Function

Customizing the Knowledge Base

Integrating with Discord (main.rs)

1. Modify main.rs to Include the Rig Agent

2. Store the Bot's User ID

3. Define the Handler Struct

4. Implement the EventHandler Trait

Handling Interactions

Handling Messages

Handling Bot Readiness

5. Implement the main Function

Running and Testing the Bot

1. Set Up Environment Variables

2. Run the Bot

3. Invite the Bot to Your Discord Server

4. Test the Bot

Error Handling and Logging

1. Error Handling with anyhow

2. Logging with tracing

Troubleshooting Common Issues

Testing and Debugging Guidance

Testing the Rig Agent Independently

Enhancing Your Bot

1. Expand the Knowledge Base

2. Customize the Bot's Behavior

3. Add More Commands

4. Integrate Other APIs

Conclusion

Further Resources

Code Files

rig_agent.rs

main.rs

Build a Flight Search AI Agent with Rust using Rig: A Hands-On Practical Guide

Introduction

Why Rust and Rig?

Why Rust?

Why Rig?

Setting Up the Environment

Prerequisites

Project Setup

1. Create a New Rust Project

2. Update Cargo.toml

3. Set Up Environment Variables

4. Install Dependencies

Understanding Agents and Tools

Building the Rig Agent (`rig_agent.rs`)

1. Create the `rig_agent.rs` File

3. Define the `RigAgent` Struct

4. Implement the `new` Method

5. Implement the `load_md_content` Function

6. Implement the `process_message` Function

Integrating with Discord (`main.rs`)

1. Modify `main.rs` to Include the Rig Agent

3. Define the `Handler` Struct

4. Implement the `EventHandler` Trait

5. Implement the `main` Function

1. Error Handling with `anyhow`

2. Logging with `tracing`

`rig_agent.rs`

`main.rs`

2. Update `Cargo.toml`

4. Error Handling with `thiserror`

5. Implement the `Tool` Trait

6. Implement the `call` Function

Updating `main.rs`