rpc_agent/
builder.rs

1use std::{net::SocketAddr, sync::Arc};
2
3use rig::tool::Tool;
4use schemars::JsonSchema;
5
6use crate::{
7    AgentServer, Error, Providers,
8    tools::{NoTool, ToolWrapper},
9};
10
11pub struct AgentServerBuilder<'a> {
12    port: u16,
13    provider: Providers,
14    system_message: &'a str,
15    model: &'a str,
16    api_key: Option<&'a str>,
17    temperature: Option<f64>,
18    max_tokens: Option<u64>,
19}
20
21impl<'a> AgentServerBuilder<'a> {
22    pub fn new(port: u16, provider: Providers, system_message: &'a str, model: &'a str) -> Self {
23        Self {
24            port,
25            provider,
26            system_message,
27            model,
28            api_key: None,
29            temperature: None,
30            max_tokens: None,
31        }
32    }
33
34    #[inline]
35    pub fn api_key(mut self, api_key: &'a str) -> Self {
36        self.api_key = Some(api_key);
37        self
38    }
39
40    #[inline]
41    pub fn temperature(mut self, temperature: f64) -> Self {
42        self.temperature = Some(temperature);
43        self
44    }
45
46    #[inline]
47    pub fn max_tokens(mut self, max_tokens: u64) -> Self {
48        self.max_tokens = Some(max_tokens);
49        self
50    }
51
52    pub fn build(self) -> Result<AgentServer, Error> {
53        let providers = Providers::init::<NoTool>(
54            self.provider,
55            self.model,
56            self.api_key,
57            self.system_message.to_string(),
58            self.temperature,
59            self.max_tokens,
60            None,
61        )?;
62
63        Ok(AgentServer {
64            socket_addr: SocketAddr::from(([0, 0, 0, 0], self.port)),
65            providers: Arc::new(providers),
66        })
67    }
68
69    pub fn build_with_schema<J: JsonSchema>(self) -> Result<AgentServer, Error> {
70        let providers = Providers::init_with_schema::<J, NoTool>(
71            self.provider,
72            self.model,
73            self.api_key,
74            self.system_message.to_string(),
75            self.temperature,
76            self.max_tokens,
77            None,
78        )?;
79
80        Ok(AgentServer {
81            socket_addr: SocketAddr::from(([0, 0, 0, 0], self.port)),
82            providers: Arc::new(providers),
83        })
84    }
85
86    pub fn build_with_tool<T: Tool + 'static>(
87        self,
88        tool: ToolWrapper<T>,
89    ) -> Result<AgentServer, Error> {
90        let providers = Providers::init::<T>(
91            self.provider,
92            self.model,
93            self.api_key,
94            self.system_message.to_string(),
95            self.temperature,
96            self.max_tokens,
97            Some(tool),
98        )?;
99
100        Ok(AgentServer {
101            socket_addr: SocketAddr::from(([0, 0, 0, 0], self.port)),
102            providers: Arc::new(providers),
103        })
104    }
105}
rpc_agent/builder.rs

rpc_agent/
builder.rs
