rpc_agent/
builder.rs

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

rpc_agent/
builder.rs
