Rig Tools
Tools are a core concept in Rig that allow agents to perform specific actions or computations. They provide a structured way to extend an agent’s capabilities beyond pure language model interactions.
Overview
Tools in Rig are implemented through two main traits:
Tool
: The base trait for implementing simple toolsToolEmbedding
: An extension trait that allows tools to be stored in vector stores and used with RAG (Retrieval Augmented Generation)
Basic Tool Implementation
A basic tool requires implementing the Tool
trait, which defines:
- A unique name identifier
- Input argument types
- Output types
- Error handling
- Tool definition (description and parameters)
- Execution logic
Here’s a simple example of a tool that adds two numbers:
#[derive(Deserialize)]
struct AddArgs {
x: i32,
y: i32,
}
#[derive(Deserialize, Serialize)]
struct Adder;
impl Tool for Adder {
const NAME: &'static str = "add";
type Error = MathError;
type Args = AddArgs;
type Output = i32;
async fn definition(&self, _prompt: String) -> ToolDefinition {
ToolDefinition {
name: "add".to_string(),
description: "Add x and y together".to_string(),
parameters: json!({
"type": "object",
"properties": {
"x": { "type": "number", "description": "First number" },
"y": { "type": "number", "description": "Second number" }
}
})
}
}
async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
Ok(args.x + args.y)
}
}
RAG-Enabled Tools
Tools can be made RAG-enabled by implementing the ToolEmbedding
trait, which allows them to be:
- Stored in vector stores
- Retrieved based on semantic similarity
- Dynamically added to agent prompts
Reference implementation:
struct Add;
impl Tool for Add {
const NAME: &'static str = "add";
type Error = MathError;
type Args = OperationArgs;
type Output = i32;
async fn definition(&self, _prompt: String) -> ToolDefinition {
serde_json::from_value(json!({
"name": "add",
"description": "Add x and y together",
"parameters": {
"type": "object",
"properties": {
"x": {
"type": "number",
"description": "The first number to add"
},
"y": {
"type": "number",
"description": "The second number to add"
}
}
}
}))
.expect("Tool Definition")
}
async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
let result = args.x + args.y;
Ok(result)
}
}
impl ToolEmbedding for Add {
type InitError = InitError;
type Context = ();
type State = ();
fn init(_state: Self::State, _context: Self::Context) -> Result<Self, Self::InitError> {
Ok(Add)
}
fn embedding_docs(&self) -> Vec<String> {
vec!["Add x and y together".into()]
}
fn context(&self) -> Self::Context {}
Deriving JSON schemas with macros
While implementing your own tools, you may find that using the raw serde_json::json!
macro can be a bit clumsy and error prone.
Fortunately, the schemars
crate has a great way to solve this: with a derive macro (and helpers)! By deriving the schemars::JsonSchema
macro, we can use description helpers to write our JSON schema instead of having to write it all manually.
#[derive(Deserialize, Serialize, schemars::JsonSchema)]
struct OperationArgs {
#[schemars(description = "The first number to add.")]
x: i32,
#[schemars(description = "The first number to add.")]
y: i32,
}
Once you’re ready to implement your tool, now rather than having to write it all out manually you can turn your schema into a serde_json::Value
and slot it into your tool definition:
#[derive(Deserialize, Serialize)]
struct Adder;
impl Tool for Adder {
// .. other trait impl parts here
async fn definition(&self, _prompt: String) -> ToolDefinition {
// this should technically never error out as it's generated from set codegen
let parameters = serde_json::to_value(schema_for!(OperationArgs)).unwrap();
ToolDefinition {
name: "add".to_string(),
description: "Add x and y together".to_string(),
parameters,
}
}
// .. other trait impl parts here
}
Rig currently uses schemars
v0.8.16. If you’re encountering strange errors, you may want to double check the version of your schemars
crate.
Using Tools with Agents
Tools can be added to agents in two ways:
- Static Tools: Always available to the agent
let agent = client
.agent("gpt-4")
.preamble("You are a calculator.")
.tool(Adder)
.tool(Subtract)
.build();
- Dynamic Tools: Retrieved from a vector store based on the query
let agent = client
.agent("gpt-4")
.preamble("You are a calculator.")
.dynamic_tools(2, vector_store_index, toolset)
.build();
Tool Organization
Tools are typically organized in a ToolSet
, which provides:
- Tool registration and management
- Tool lookup by name
- Tool execution routing
- Conversion to embeddings for RAG
Best Practices
- Unique Names: Ensure each tool has a unique name within your application
- Clear Descriptions: Provide clear, detailed descriptions in tool definitions
- Type Safety: Use strong typing for tool arguments and outputs
- Error Handling: Implement proper error types and handling
- RAG Consideration: Consider implementing
ToolEmbedding
if your tool might benefit from semantic retrieval
Integration with LLMs
Tools are automatically integrated with LLM providers through Rig’s agent system. The library handles:
- Converting tool definitions to provider-specific formats
- Parsing LLM outputs into tool calls
- Routing tool calls to appropriate implementations
- Returning tool results to the LLM
For more information on integrating tools with specific LLM providers, see the provider-specific documentation in the providers
module.
API Reference (Tools)