rig/agent/
builder.rs

1use std::collections::HashMap;
2
3use crate::{
4    completion::{CompletionModel, Document},
5    tool::{Tool, ToolSet},
6    vector_store::VectorStoreIndexDyn,
7};
8
9#[cfg(feature = "mcp")]
10use crate::tool::McpTool;
11
12use super::Agent;
13
14/// A builder for creating an agent
15///
16/// # Example
17/// ```
18/// use rig::{providers::openai, agent::AgentBuilder};
19///
20/// let openai = openai::Client::from_env();
21///
22/// let gpt4o = openai.completion_model("gpt-4o");
23///
24/// // Configure the agent
25/// let agent = AgentBuilder::new(model)
26///     .preamble("System prompt")
27///     .context("Context document 1")
28///     .context("Context document 2")
29///     .tool(tool1)
30///     .tool(tool2)
31///     .temperature(0.8)
32///     .additional_params(json!({"foo": "bar"}))
33///     .build();
34/// ```
35pub struct AgentBuilder<M: CompletionModel> {
36    /// Completion model (e.g.: OpenAI's gpt-3.5-turbo-1106, Cohere's command-r)
37    model: M,
38    /// System prompt
39    preamble: Option<String>,
40    /// Context documents always available to the agent
41    static_context: Vec<Document>,
42    /// Tools that are always available to the agent (by name)
43    static_tools: Vec<String>,
44    /// Additional parameters to be passed to the model
45    additional_params: Option<serde_json::Value>,
46    /// Maximum number of tokens for the completion
47    max_tokens: Option<u64>,
48    /// List of vector store, with the sample number
49    dynamic_context: Vec<(usize, Box<dyn VectorStoreIndexDyn>)>,
50    /// Dynamic tools
51    dynamic_tools: Vec<(usize, Box<dyn VectorStoreIndexDyn>)>,
52    /// Temperature of the model
53    temperature: Option<f64>,
54    /// Actual tool implementations
55    tools: ToolSet,
56}
57
58impl<M: CompletionModel> AgentBuilder<M> {
59    pub fn new(model: M) -> Self {
60        Self {
61            model,
62            preamble: None,
63            static_context: vec![],
64            static_tools: vec![],
65            temperature: None,
66            max_tokens: None,
67            additional_params: None,
68            dynamic_context: vec![],
69            dynamic_tools: vec![],
70            tools: ToolSet::default(),
71        }
72    }
73
74    /// Set the system prompt
75    pub fn preamble(mut self, preamble: &str) -> Self {
76        self.preamble = Some(preamble.into());
77        self
78    }
79
80    /// Append to the preamble of the agent
81    pub fn append_preamble(mut self, doc: &str) -> Self {
82        self.preamble = Some(format!(
83            "{}\n{}",
84            self.preamble.unwrap_or_else(|| "".into()),
85            doc
86        ));
87        self
88    }
89
90    /// Add a static context document to the agent
91    pub fn context(mut self, doc: &str) -> Self {
92        self.static_context.push(Document {
93            id: format!("static_doc_{}", self.static_context.len()),
94            text: doc.into(),
95            additional_props: HashMap::new(),
96        });
97        self
98    }
99
100    /// Add a static tool to the agent
101    pub fn tool(mut self, tool: impl Tool + 'static) -> Self {
102        let toolname = tool.name();
103        self.tools.add_tool(tool);
104        self.static_tools.push(toolname);
105        self
106    }
107
108    // Add an MCP tool to the agent
109    #[cfg(feature = "mcp")]
110    pub fn mcp_tool<T: mcp_core::transport::Transport>(
111        mut self,
112        tool: mcp_core::types::Tool,
113        client: mcp_core::client::Client<T>,
114    ) -> Self {
115        let toolname = tool.name.clone();
116        self.tools.add_tool(McpTool::from_mcp_server(tool, client));
117        self.static_tools.push(toolname);
118        self
119    }
120
121    /// Add some dynamic context to the agent. On each prompt, `sample` documents from the
122    /// dynamic context will be inserted in the request.
123    pub fn dynamic_context(
124        mut self,
125        sample: usize,
126        dynamic_context: impl VectorStoreIndexDyn + 'static,
127    ) -> Self {
128        self.dynamic_context
129            .push((sample, Box::new(dynamic_context)));
130        self
131    }
132
133    /// Add some dynamic tools to the agent. On each prompt, `sample` tools from the
134    /// dynamic toolset will be inserted in the request.
135    pub fn dynamic_tools(
136        mut self,
137        sample: usize,
138        dynamic_tools: impl VectorStoreIndexDyn + 'static,
139        toolset: ToolSet,
140    ) -> Self {
141        self.dynamic_tools.push((sample, Box::new(dynamic_tools)));
142        self.tools.add_tools(toolset);
143        self
144    }
145
146    /// Set the temperature of the model
147    pub fn temperature(mut self, temperature: f64) -> Self {
148        self.temperature = Some(temperature);
149        self
150    }
151
152    /// Set the maximum number of tokens for the completion
153    pub fn max_tokens(mut self, max_tokens: u64) -> Self {
154        self.max_tokens = Some(max_tokens);
155        self
156    }
157
158    /// Set additional parameters to be passed to the model
159    pub fn additional_params(mut self, params: serde_json::Value) -> Self {
160        self.additional_params = Some(params);
161        self
162    }
163
164    /// Build the agent
165    pub fn build(self) -> Agent<M> {
166        Agent {
167            model: self.model,
168            preamble: self.preamble.unwrap_or_default(),
169            static_context: self.static_context,
170            static_tools: self.static_tools,
171            temperature: self.temperature,
172            max_tokens: self.max_tokens,
173            additional_params: self.additional_params,
174            dynamic_context: self.dynamic_context,
175            dynamic_tools: self.dynamic_tools,
176            tools: self.tools,
177        }
178    }
179}
rig/agent/builder.rs

rig/agent/
builder.rs
