diff --git a/Cargo.lock b/Cargo.lock index 128f75bc..2d82c8ca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1211,7 +1211,7 @@ dependencies = [ [[package]] name = "rig-core" -version = "0.0.6" +version = "0.0.7" dependencies = [ "anyhow", "futures", @@ -1228,7 +1228,7 @@ dependencies = [ [[package]] name = "rig-mongodb" -version = "0.0.6" +version = "0.0.7" dependencies = [ "anyhow", "futures", diff --git a/rig-core/Cargo.toml b/rig-core/Cargo.toml index 57b97d02..5a82b548 100644 --- a/rig-core/Cargo.toml +++ b/rig-core/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "rig-core" -version = "0.0.6" +version = "0.0.7" edition = "2021" license = "MIT" readme = "README.md" diff --git a/rig-core/src/agent.rs b/rig-core/src/agent.rs index 86b7b87e..cac1f5b8 100644 --- a/rig-core/src/agent.rs +++ b/rig-core/src/agent.rs @@ -10,12 +10,15 @@ //! //! # Example //! ```rust -//! use rig::{completion::Prompt, providers::openai}; +//! use rig::{ +//! completion::{Chat, Completion, Prompt}, +//! providers::openai, +//! }; //! -//! let openai_client = openai::Client::from_env(); +//! let openai = openai::Client::from_env(); //! //! // Configure the agent -//! let agent = client.agent("gpt-4o") +//! let agent = openai.agent("gpt-4o") //! .preamble("System prompt") //! .context("Context document 1") //! .context("Context document 2") @@ -26,9 +29,28 @@ //! .build(); //! //! // Use the agent for completions and prompts -//! let completion_req_builder = agent.completion("Prompt", chat_history).await; -//! let chat_response = agent.chat("Prompt", chat_history).await; -//! let chat_response = agent.prompt("Prompt").await; +//! // Generate a chat completion response from a prompt and chat history +//! let chat_response = agent.chat("Prompt", chat_history) +//! .await +//! .expect("Failed to chat with Agent"); +//! +//! // Generate a prompt completion response from a simple prompt +//! let chat_response = agent.prompt("Prompt") +//! .await +//! .expect("Failed to prompt the Agent"); +//! +//! // Generate a completion request builder from a prompt and chat history. The builder +//! // will contain the agent's configuration (i.e.: preamble, context documents, tools, +//! // model parameters, etc.), but these can be overwritten. +//! let completion_req_builder = agent.completion("Prompt", chat_history) +//! .await +//! .expect("Failed to create completion request builder"); +//! +//! let response = completion_req_builder +//! .temperature(0.9) // Overwrite the agent's temperature +//! .send() +//! .await +//! .expect("Failed to send completion request"); //! ``` use std::collections::HashMap; @@ -50,9 +72,9 @@ use crate::{ /// ``` /// use rig::{completion::Prompt, providers::openai}; /// -/// let openai_client = openai::Client::from_env(); +/// let openai = openai::Client::from_env(); /// -/// let comedian_agent = client +/// let comedian_agent = openai /// .agent("gpt-4o") /// .preamble("You are a comedian here to entertain the user using humour and jokes.") /// .temperature(0.9) @@ -60,7 +82,7 @@ use crate::{ /// /// let response = comedian_agent.prompt("Entertain me!") /// .await -/// .expect("Failed to prompt GPT-4"); +/// .expect("Failed to prompt the agent"); /// ``` pub struct Agent { /// Completion model (e.g.: OpenAI's `gpt-3.5-turbo-1106`, Cohere's `command-r`) @@ -168,9 +190,9 @@ impl Chat for Agent { /// ``` /// use rig::{providers::openai, agent::AgentBuilder}; /// -/// let openai_client = openai::Client::from_env(); +/// let openai = openai::Client::from_env(); /// -/// let gpt4 = openai_client.completion_model("gpt-4"); +/// let gpt4o = openai.completion_model("gpt-4o"); /// /// // Configure the agent /// let agent = AgentBuilder::new(model) diff --git a/rig-core/src/cli_chatbot.rs b/rig-core/src/cli_chatbot.rs index e46e1f23..362ba798 100644 --- a/rig-core/src/cli_chatbot.rs +++ b/rig-core/src/cli_chatbot.rs @@ -2,6 +2,8 @@ use std::io::{self, Write}; use crate::completion::{Chat, Message, PromptError}; +/// Utility function to create a simple REPL CLI chatbot from a type that implements the +/// `Chat` trait. pub async fn cli_chatbot(chatbot: impl Chat) -> Result<(), PromptError> { let stdin = io::stdin(); let mut stdout = io::stdout(); diff --git a/rig-core/src/completion.rs b/rig-core/src/completion.rs index fef5fedd..f642d292 100644 --- a/rig-core/src/completion.rs +++ b/rig-core/src/completion.rs @@ -24,7 +24,6 @@ //! responses, and errors. //! //! Example Usage: -//! //! ```rust //! use rig::providers::openai::{Client, self}; //! use rig::completion::*; @@ -32,14 +31,15 @@ //! // Initialize the OpenAI client and a completion model //! let openai = Client::new("your-openai-api-key"); //! -//! let gpt_4 = openai.completion_model(openai::GPT_4).build(); +//! let gpt_4 = openai.completion_model(openai::GPT_4); //! //! // Create the completion request -//! let builder = gpt_4.completion_request("Who are you?"); +//! let request = gpt_4.completion_request("Who are you?") //! .preamble("\ //! You are Marvin, an extremely smart but depressed robot who is \ //! nonetheless helpful towards humanity.\ //! ") +//! .temperature(0.5) //! .build(); //! //! // Send the completion request and get the completion response @@ -130,11 +130,16 @@ pub struct ToolDefinition { // ================================================================ // Implementations // ================================================================ -/// Trait defining a high-level LLM on-shot prompt interface (i.e.: prompt in, response out). +/// Trait defining a high-level LLM simple prompt interface (i.e.: prompt in, response out). pub trait Prompt: Send + Sync { - /// Send a one-shot prompt to the underlying completion model. - /// If the response is a message, then it is returned as a string. If the response - /// is a tool call, then the tool is called and the result is returned as a string. + /// Send a simple prompt to the underlying completion model. + /// + /// If the completion model's response is a message, then it is returned as a string. + /// + /// If the completion model's response is a tool call, then the tool is called and + /// the result is returned as a string. + /// + /// If the tool does not exist, or the tool call fails, then an error is returned. fn prompt( &self, prompt: &str, @@ -144,8 +149,13 @@ pub trait Prompt: Send + Sync { /// Trait defining a high-level LLM chat interface (i.e.: prompt and chat history in, response out). pub trait Chat: Send + Sync { /// Send a prompt with optional chat history to the underlying completion model. - /// If the response is a message, then it is returned as a string. If the response - /// is a tool call, then the tool is called and the result is returned as a string. + /// + /// If the completion model's response is a message, then it is returned as a string. + /// + /// If the completion model's response is a tool call, then the tool is called and the result + /// is returned as a string. + /// + /// If the tool does not exist, or the tool call fails, then an error is returned. fn chat( &self, prompt: &str, @@ -207,6 +217,7 @@ pub trait CompletionModel: Clone + Send + Sync { ) -> impl std::future::Future, CompletionError>> + Send; + /// Generates a completion request builder for the given `prompt`. fn completion_request(&self, prompt: &str) -> CompletionRequestBuilder { CompletionRequestBuilder::new(self.clone(), prompt.to_string()) } @@ -231,6 +242,49 @@ pub struct CompletionRequest { } /// Builder struct for constructing a completion request. +/// +/// Example usage: +/// ```rust +/// use rig::{ +/// providers::openai::{Client, self}, +/// completion::CompletionRequestBuilder, +/// }; +/// +/// let openai = Client::new("your-openai-api-key"); +/// let model = openai.completion_model(openai::GPT_4O).build(); +/// +/// // Create the completion request and execute it separately +/// let request = CompletionRequestBuilder::new(model, "Who are you?".to_string()) +/// .preamble("You are Marvin from the Hitchhiker's Guide to the Galaxy.".to_string()) +/// .temperature(0.5) +/// .build(); +/// +/// let response = model.completion(request) +/// .await +/// .expect("Failed to get completion response"); +/// ``` +/// +/// Alternatively, you can execute the completion request directly from the builder: +/// ```rust +/// use rig::{ +/// providers::openai::{Client, self}, +/// completion::CompletionRequestBuilder, +/// }; +/// +/// let openai = Client::new("your-openai-api-key"); +/// let model = openai.completion_model(openai::GPT_4O).build(); +/// +/// // Create the completion request and execute it directly +/// let response = CompletionRequestBuilder::new(model, "Who are you?".to_string()) +/// .preamble("You are Marvin from the Hitchhiker's Guide to the Galaxy.".to_string()) +/// .temperature(0.5) +/// .send() +/// .await +/// .expect("Failed to get completion response"); +/// ``` +/// +/// Note: It is usually unnecessary to create a completion request builder directly. +/// Instead, use the [CompletionModel::completion_request] method. pub struct CompletionRequestBuilder { model: M, prompt: String, diff --git a/rig-core/src/embeddings.rs b/rig-core/src/embeddings.rs index df08f279..aa9772a3 100644 --- a/rig-core/src/embeddings.rs +++ b/rig-core/src/embeddings.rs @@ -28,12 +28,12 @@ //! //! // Create an embeddings builder and add documents //! let embeddings = EmbeddingsBuilder::new(embedding_model) -//! .simple_document("doc1", "This is the first document.") +//! .simple_document("doc1", "This is the first document.") //! .simple_document("doc2", "This is the second document.") //! .build() //! .await //! .expect("Failed to build embeddings."); -//! +//! //! // Use the generated embeddings //! // ... //! ``` @@ -70,6 +70,7 @@ pub enum EmbeddingError { /// Trait for embedding models that can generate embeddings for documents. pub trait EmbeddingModel: Clone + Sync + Send { + /// The maximum number of documents that can be embedded in a single request. const MAX_DOCUMENTS: usize; /// Embed a single document @@ -132,9 +133,12 @@ impl Embedding { /// Struct that holds a document and its embeddings. /// /// The struct is designed to model any kind of documents that can be serialized to JSON -/// (including a simple string). Moreover, it can hold multiple embeddings for the same -/// document, thus allowing a large or non-text document to be "ragged" from various -/// smaller text documents. +/// (including a simple string). +/// +/// Moreover, it can hold multiple embeddings for the same document, thus allowing a +/// large document to be retrieved from a query that matches multiple smaller and +/// distinct text documents. For example, if the document is a textbook, a summary of +/// each chapter could serve as the book's embeddings. #[derive(Clone, Eq, PartialEq, Serialize, Deserialize)] pub struct DocumentEmbeddings { #[serde(rename = "_id")] diff --git a/rig-core/src/extractor.rs b/rig-core/src/extractor.rs index 5eb2de4f..02699496 100644 --- a/rig-core/src/extractor.rs +++ b/rig-core/src/extractor.rs @@ -1,4 +1,7 @@ //! This module provides high-level abstractions for extracting structured data from text using LLMs. +//! +//! Note: The target structure must implement the `serde::Deserialize`, `serde::Serialize`, +//! and `schemars::JsonSchema` traits. Those can be easily derived using the `derive` macro. //! //! # Example //! ``` @@ -8,7 +11,7 @@ //! let openai = openai::Client::new("your-open-ai-api-key"); //! //! // Define the structure of the data you want to extract -//! #[derive(serde::Deserialize)] +//! #[derive(serde::Deserialize, serde::Serialize, schemars::JsonSchema)] //! struct Person { //! name: Option, //! age: Option, @@ -21,7 +24,8 @@ //! //! // Extract structured data from text //! let person = extractor.extract("John Doe is a 30 year old doctor.") -//! .await; +//! .await +//! .expect("Failed to extract data from text"); //! ``` use std::marker::PhantomData; diff --git a/rig-core/src/lib.rs b/rig-core/src/lib.rs index c323fc74..c6e82612 100644 --- a/rig-core/src/lib.rs +++ b/rig-core/src/lib.rs @@ -70,7 +70,7 @@ //! - OpenAI //! - Cohere //! -//! Rig currently has the following integration sub-libraries: +//! Rig currently has the following integration companion crates: //! - `rig-mongodb`: Vector store implementation for MongoDB //! diff --git a/rig-core/src/model.rs b/rig-core/src/model.rs index 1e05930e..b75f2601 100644 --- a/rig-core/src/model.rs +++ b/rig-core/src/model.rs @@ -11,19 +11,41 @@ //! //! # Example //! ```rust -//! use rig::{completion::{Chat, Prompt}, providers::openai}; +//! use rig::{ +//! completion::{Chat, Completion, Prompt}, +//! providers::openai, +//! }; //! -//! let openai_client = openai::Client::from_env(); +//! let openai = openai::Client::from_env(); //! //! // Configure the model -//! let model = client.model("gpt-4o") +//! let model = openai.model("gpt-4o") //! .temperature(0.8) //! .build(); //! //! // Use the model for completions and prompts -//! let completion_req_builder = model.completion("Prompt", chat_history).await; -//! let chat_response = model.chat("Prompt", chat_history).await; -//! let prompt_response = model.prompt("Prompt").await; +//! // Generate a chat completion response from a prompt and chat history +//! let chat_response = agent.chat("Prompt", chat_history) +//! .await +//! .expect("Failed to chat with model"); +//! +//! // Generate a prompt completion response from a simple prompt +//! let chat_response = agent.prompt("Prompt") +//! .await +//! .expect("Failed to prompt the model"); +//! +//! // Generate a completion request builder from a prompt and chat history. The builder +//! // will contain the model's configuration (i.e.: model parameters, etc.), but these +//! // can be overwritten. +//! let completion_req_builder = agent.completion("Prompt", chat_history) +//! .await +//! .expect("Failed to create completion request builder"); +//! +//! let response = completion_req_builder +//! .temperature(0.9) // Overwrite the model's temperature +//! .send() +//! .await +//! .expect("Failed to send completion request"); //! ``` use crate::completion::{ Chat, Completion, CompletionError, CompletionModel, CompletionRequestBuilder, diff --git a/rig-core/src/rag.rs b/rig-core/src/rag.rs index aa86b839..89bb21d1 100644 --- a/rig-core/src/rag.rs +++ b/rig-core/src/rag.rs @@ -18,10 +18,10 @@ //! }; //! //! // Initialize OpenAI client -//! let openai_client = openai::Client::from_env(); +//! let openai = openai::Client::from_env(); //! //! // Initialize OpenAI embedding model -//! let embedding_model = openai_client.embedding_model(openai::TEXT_EMBEDDING_ADA_002); +//! let embedding_model = openai.embedding_model(openai::TEXT_EMBEDDING_ADA_002); //! //! // Create vector store, compute embeddings and load them in the store //! let mut vector_store = InMemoryVectorStore::default(); @@ -41,7 +41,7 @@ //! // Create vector store index //! let index = vector_store.index(embedding_model); //! -//! let rag_agent = openai_client.context_rag_agent(openai::GPT_4O) +//! let rag_agent = openai.context_rag_agent(openai::GPT_4O) //! .preamble(" //! You are a dictionary assistant here to assist the user in understanding the meaning of words. //! You will find additional non-standard word definitions that could be useful below. @@ -212,6 +212,7 @@ impl RagAgent, } @@ -76,10 +77,18 @@ impl VectorStore for MongoDbVectorStore { } impl MongoDbVectorStore { + /// Create a new `MongoDbVectorStore` from a MongoDB collection. pub fn new(collection: mongodb::Collection) -> Self { Self { collection } } + /// Create a new `MongoDbVectorIndex` from an existing `MongoDbVectorStore`. + /// + /// The index (of type "vector") must already exist for the MongoDB collection. + /// See the MongoDB [documentation](https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-type/) for more information on creating indexes. + /// + /// An additional filter can be provided to further restrict the documents that are + /// considered in the search. pub fn index( &self, model: M, @@ -90,6 +99,7 @@ impl MongoDbVectorStore { } } +/// A vector index for a MongoDB collection. pub struct MongoDbVectorIndex { collection: mongodb::Collection, model: M,