rust-rag-toolchain is a Rust native library designed to empower developers with seamless access to common Retrieval Augmented Generation (RAG) workflows. It provides functionalities for embedding generation, storage, and retrieval, allowing developers to build AI applications with their own data.
Retrieval Augmented Generation (RAG): Enhance your AI applications with better prompt completions based on your own knowledge base. See BasicRagChain !.
PG Vector Support Able to store and retrieve from a postgres database with pg_vector enabled !
OpenAI Support Able to generate chat completions and embeddings via OpenAI to use an your AI workflows
Token Chunking Able to chunk your text based on a token size.
Add the latest version of rust-rag-toolchain to your Cargo.toml file.
The library is structured into several modules, each responsible for a specific functionality:
chains: This module contains the implementation of the RAG workflows.
chunkers: This module provides utilities for breaking down data into manageable chunks.
clients: This module contains client implementations for interacting with various services.
common: This module contains common utilities and helper functions used across the library.
loaders: This module provides functionalities for loading and processing data.
retrievers: This module contains implementations for retrieving data from various sources.
stores: This module provides functionalities for storing and managing data.
const SYSTEM_MESSAGE: &'static str =
"You are to give straight forward answers using the supporting information you are provided";
#[tokio::main]
async fn main() {
// Initialize the PostgresVectorStore
let store: PostgresVectorStore =
PostgresVectorStore::try_new("embeddings", TextEmbeddingAda002)
.await
.unwrap();
// Create a new embedding client
let embedding_client: OpenAIEmbeddingClient =
OpenAIEmbeddingClient::try_new(TextEmbeddingAda002).unwrap();
// Convert our store into a retriever
let retriever: PostgresVectorRetriever<OpenAIEmbeddingClient> =
store.as_retriever(embedding_client, DistanceFunction::Cosine);
// Create a new chat client
let chat_client: OpenAIChatCompletionClient =
OpenAIChatCompletionClient::try_new(Gpt3Point5Turbo).unwrap();
// Define our system prompt
let system_prompt: PromptMessage = PromptMessage::SystemMessage(SYSTEM_MESSAGE.into());
// Create a new BasicRAGChain with over our open ai chat client and postgres vector retriever
let chain: BasicRAGChain<OpenAIChatCompletionClient, PostgresVectorRetriever<_>> =
BasicRAGChain::builder()
.system_prompt(system_prompt)
.chat_client(chat_client)
.retriever(retriever)
.build();
// Define our user prompt
let user_message: PromptMessage =
PromptMessage::HumanMessage("what kind of alcohol does Morwenna drink".into());
// Invoke the chain. Under the hood this will retrieve some similar text from
// the retriever and then use the chat client to generate a response.
let response = chain
.invoke_chain(user_message, NonZeroU32::new(2).unwrap())
.await
.unwrap();
println!("{}", response.content());
}Contributions are welcome! If you have ideas for improvements or find any issues, please open an issue or submit a pull request.
For any questions or assistance, feel free to contact the maintainers:
Please raise an issue and I will get back to you quickly
Enjoy leveraging the power of retrieval augmented generation and embedding generation with rust-rag-toolchain!