agent-io 0.3.2

A Rust SDK for building AI agents with multi-provider LLM support
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

//! Agent builder

use std::sync::Arc;

use crate::Result;
use crate::llm::BaseChatModel;
use crate::memory::MemoryManager;
use crate::tools::Tool;

use super::config::{AgentConfig, build_ephemeral_config};
use super::service::Agent;

/// Agent builder
#[derive(Default)]
pub struct AgentBuilder {
    llm: Option<Arc<dyn BaseChatModel>>,
    tools: Vec<Arc<dyn Tool>>,
    config: Option<AgentConfig>,
    memory: Option<Arc<RwLock<MemoryManager>>>,
}

use tokio::sync::RwLock;

impl AgentBuilder {
    pub fn with_llm(mut self, llm: Arc<dyn BaseChatModel>) -> Self {
        self.llm = Some(llm);
        self
    }

    pub fn tool(mut self, tool: Arc<dyn Tool>) -> Self {
        self.tools.push(tool);
        self
    }

    pub fn tools(mut self, tools: Vec<Arc<dyn Tool>>) -> Self {
        self.tools = tools;
        self
    }

    pub fn config(mut self, config: AgentConfig) -> Self {
        self.config = Some(config);
        self
    }

    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
        let mut config = self.config.unwrap_or_default();
        config.system_prompt = Some(prompt.into());
        self.config = Some(config);
        self
    }

    pub fn max_iterations(mut self, max: usize) -> Self {
        let mut config = self.config.unwrap_or_default();
        config.max_iterations = max;
        self.config = Some(config);
        self
    }

    /// Add memory manager to the agent
    pub fn with_memory(mut self, memory: Arc<RwLock<MemoryManager>>) -> Self {
        self.memory = Some(memory);
        self
    }

    pub fn build(self) -> Result<Agent> {
        let llm = self
            .llm
            .ok_or_else(|| crate::Error::Config("LLM is required".into()))?;

        let ephemeral_config = build_ephemeral_config(&self.tools);

        Ok(Agent::new_with_config(
            llm,
            self.tools,
            self.config.unwrap_or_default(),
            ephemeral_config,
            self.memory,
        ))
    }
}