rig_extra/
simple_rand_builder.rs

1use crate::extra_providers::bigmodel;
2use crate::get_openai_agent::get_openai_agent;
3use crate::rand_agent::RandAgentBuilder;
4use rig::client::completion::CompletionClientDyn;
5use rig::providers::*;
6use serde::{Deserialize, Serialize};
7use strum_macros::Display;
8
9#[derive(Debug, Display, Deserialize, Serialize)]
10#[serde(rename_all = "lowercase")]
11pub enum ProviderEnum {
12    Anthropic,
13    Cohere,
14    Gemini,
15    Huggingface,
16    Mistral,
17    OpenAi,
18    OpenRouter,
19    Together,
20    XAI,
21    Azure,
22    DeepSeek,
23    Galadriel,
24    Groq,
25    Hyperbolic,
26    Mira,
27    Mooshot,
28    Ollama,
29    Perplexity,
30    // embedding模型
31    // Voyageai,
32    Bigmodel,
33}
34
35#[derive(Debug, Deserialize)]
36pub struct AgentConfig {
37    pub id: i32,
38    pub provider: ProviderEnum,
39    pub model_name: String,
40    pub api_key: String,
41    pub api_base_url: Option<String>,
42    pub system_prompt: Option<String>,
43    pub agent_name: Option<String>,
44}
45
46impl RandAgentBuilder {
47    /// 简单构建器
48    pub fn simple_builder(
49        mut self,
50        agent_configs: Vec<AgentConfig>,
51        global_system_prompt: String,
52    ) -> Self {
53        for agent_conf in agent_configs {
54            let agent_name = agent_conf.agent_name.unwrap_or("rand agent".to_string());
55            let system_prompt = agent_conf
56                .system_prompt
57                .unwrap_or(global_system_prompt.clone());
58
59            match agent_conf.provider {
60                ProviderEnum::Anthropic => {
61                    let mut client_builder = anthropic::ClientBuilder::new(&agent_conf.api_key);
62                    if let Some(api_base_url) = &agent_conf.api_base_url {
63                        client_builder = client_builder.base_url(api_base_url);
64                    }
65                    match client_builder.build() {
66                        Ok(client) => {
67                            let agent = client
68                                .agent(&agent_conf.model_name)
69                                .name(agent_name.as_str())
70                                .preamble(&system_prompt)
71                                .build();
72                            self.agents.push((
73                                agent,
74                                agent_conf.id,
75                                agent_conf.provider.to_string(),
76                                agent_conf.model_name,
77                            ));
78                        }
79                        Err(err) => {
80                            tracing::error!("添加 {} 错误: {}", agent_conf.provider, err);
81                        }
82                    }
83                }
84                ProviderEnum::Cohere => {
85                    let client = cohere::Client::new(&agent_conf.api_key);
86                    let agent = client
87                        .agent(&agent_conf.model_name)
88                        .name(agent_name.as_str())
89                        .preamble(&system_prompt)
90                        .build();
91                    self.agents.push((
92                        agent,
93                        agent_conf.id,
94                        agent_conf.provider.to_string(),
95                        agent_conf.model_name,
96                    ));
97                }
98                ProviderEnum::Gemini => {
99                    let mut client_builder = gemini::Client::builder(&agent_conf.api_key);
100                    if let Some(api_base_url) = &agent_conf.api_base_url {
101                        client_builder = client_builder.base_url(api_base_url);
102                    }
103                    match client_builder.build() {
104                        Ok(client) => {
105                            let agent = client
106                                .agent(&agent_conf.model_name)
107                                .name(agent_name.as_str())
108                                .preamble(&system_prompt)
109                                .build();
110                            self.agents.push((
111                                agent,
112                                agent_conf.id,
113                                agent_conf.provider.to_string(),
114                                agent_conf.model_name,
115                            ));
116                        }
117                        Err(err) => {
118                            tracing::error!("添加 {} 错误: {}", agent_conf.provider, err);
119                        }
120                    }
121                }
122                ProviderEnum::Huggingface => {
123                    let client = huggingface::Client::new(&agent_conf.api_key);
124                    let agent = client
125                        .agent(&agent_conf.model_name)
126                        .name(agent_name.as_str())
127                        .preamble(&system_prompt)
128                        .build();
129                    self.agents.push((
130                        agent,
131                        agent_conf.id,
132                        agent_conf.provider.to_string(),
133                        agent_conf.model_name,
134                    ));
135                }
136                ProviderEnum::Mistral => {
137                    let client = mistral::Client::new(&agent_conf.api_key);
138                    let agent = client
139                        .agent(&agent_conf.model_name)
140                        .name(agent_name.as_str())
141                        .preamble(&system_prompt)
142                        .build();
143                    self.agents.push((
144                        agent,
145                        agent_conf.id,
146                        agent_conf.provider.to_string(),
147                        agent_conf.model_name,
148                    ));
149                }
150                ProviderEnum::OpenAi => {
151                    let mut client_builder = openai::ClientBuilder::new(&agent_conf.api_key);
152                    if let Some(api_base_url) = &agent_conf.api_base_url {
153                        client_builder = client_builder.base_url(api_base_url)
154                    }
155
156                    let client = client_builder.build();
157
158                    let agent =
159                        get_openai_agent(client, &agent_conf.model_name, agent_name, system_prompt);
160                    self.agents.push((
161                        agent,
162                        agent_conf.id,
163                        agent_conf.provider.to_string(),
164                        agent_conf.model_name,
165                    ));
166                }
167                ProviderEnum::OpenRouter => {
168                    let mut client_builder = openrouter::ClientBuilder::new(&agent_conf.api_key);
169                    if let Some(api_base_url) = &agent_conf.api_base_url {
170                        client_builder = client_builder.base_url(api_base_url)
171                    }
172                    let client = client_builder.build();
173                    let agent = client
174                        .agent(&agent_conf.model_name)
175                        .name(agent_name.as_str())
176                        .preamble(&system_prompt)
177                        .build();
178                    self.agents.push((
179                        agent,
180                        agent_conf.id,
181                        agent_conf.provider.to_string(),
182                        agent_conf.model_name,
183                    ));
184                }
185                ProviderEnum::Together => {
186                    let client = together::Client::new(&agent_conf.api_key);
187                    let agent = client
188                        .agent(&agent_conf.model_name)
189                        .name(agent_name.as_str())
190                        .preamble(&system_prompt)
191                        .build();
192                    self.agents.push((
193                        agent,
194                        agent_conf.id,
195                        agent_conf.provider.to_string(),
196                        agent_conf.model_name,
197                    ));
198                }
199                ProviderEnum::XAI => {
200                    let client = xai::Client::new(&agent_conf.api_key);
201                    let agent = client
202                        .agent(&agent_conf.model_name)
203                        .name(agent_name.as_str())
204                        .preamble(&system_prompt)
205                        .build();
206                    self.agents.push((
207                        agent,
208                        agent_conf.id,
209                        agent_conf.provider.to_string(),
210                        agent_conf.model_name,
211                    ));
212                }
213                ProviderEnum::Azure => {
214                    tracing::info!("Azure simple_builder暂不支持,参数有点多,可以自行添加........ ")
215                }
216                ProviderEnum::DeepSeek => {
217                    let client = deepseek::Client::new(&agent_conf.api_key);
218                    let agent = client
219                        .agent(&agent_conf.model_name)
220                        .name(agent_name.as_str())
221                        .preamble(&system_prompt)
222                        .build();
223                    self.agents.push((
224                        agent,
225                        agent_conf.id,
226                        agent_conf.provider.to_string(),
227                        agent_conf.model_name,
228                    ));
229                }
230                ProviderEnum::Galadriel => {
231                    let client = galadriel::Client::new(&agent_conf.api_key);
232                    let agent = client
233                        .agent(&agent_conf.model_name)
234                        .name(agent_name.as_str())
235                        .preamble(&system_prompt)
236                        .build();
237                    self.agents.push((
238                        agent,
239                        agent_conf.id,
240                        agent_conf.provider.to_string(),
241                        agent_conf.model_name,
242                    ));
243                }
244                ProviderEnum::Groq => {
245                    let client = groq::Client::new(&agent_conf.api_key);
246                    let agent = client
247                        .agent(&agent_conf.model_name)
248                        .name(agent_name.as_str())
249                        .preamble(&system_prompt)
250                        .build();
251                    self.agents.push((
252                        agent,
253                        agent_conf.id,
254                        agent_conf.provider.to_string(),
255                        agent_conf.model_name,
256                    ));
257                }
258                ProviderEnum::Hyperbolic => {
259                    let client = hyperbolic::Client::new(&agent_conf.api_key);
260                    let agent = client
261                        .agent(&agent_conf.model_name)
262                        .name(agent_name.as_str())
263                        .preamble(&system_prompt)
264                        .build();
265                    self.agents.push((
266                        agent,
267                        agent_conf.id,
268                        agent_conf.provider.to_string(),
269                        agent_conf.model_name,
270                    ));
271                }
272                ProviderEnum::Mira => {
273                    let client = mira::Client::new(&agent_conf.api_key);
274                    let agent = client
275                        .agent(&agent_conf.model_name)
276                        .name(agent_name.as_str())
277                        .preamble(&system_prompt)
278                        .build();
279                    self.agents.push((
280                        agent,
281                        agent_conf.id,
282                        agent_conf.provider.to_string(),
283                        agent_conf.model_name,
284                    ));
285                }
286                ProviderEnum::Mooshot => {
287                    let client = moonshot::Client::new(&agent_conf.api_key);
288                    let agent = client
289                        .agent(&agent_conf.model_name)
290                        .name(agent_name.as_str())
291                        .preamble(&system_prompt)
292                        .build();
293                    self.agents.push((
294                        agent,
295                        agent_conf.id,
296                        agent_conf.provider.to_string(),
297                        agent_conf.model_name,
298                    ));
299                }
300                ProviderEnum::Ollama => {
301                    let mut client_builder = ollama::ClientBuilder::new();
302                    if let Some(api_base_url) = &agent_conf.api_base_url {
303                        client_builder = client_builder.base_url(api_base_url);
304                    }
305
306                    let client = client_builder.build();
307                    let agent = client
308                        .agent(&agent_conf.model_name)
309                        .name(agent_name.as_str())
310                        .preamble(&system_prompt)
311                        .build();
312                    self.agents.push((
313                        agent,
314                        agent_conf.id,
315                        agent_conf.provider.to_string(),
316                        agent_conf.model_name,
317                    ));
318                }
319                ProviderEnum::Perplexity => {
320                    // let client = perplexity::Client::new(&agent_conf.api_key);
321                    // let agent = client
322                    //     .agent(&agent_conf.model_name)
323                    //     .name(agent_name.as_str())
324                    //     .preamble(&system_prompt)
325                    //     .build();
326                    // self.agents.push((
327                    //     agent,
328                    //     agent_conf.id,
329                    //     agent_conf.provider.to_string(),
330                    //     agent_conf.model_name,
331                    // ));
332                    tracing::info!("Perplexity 暂不支持,没有实现BoxAgent........ ")
333                }
334                ProviderEnum::Bigmodel => {
335                    let client = if let Some(api_base_url) = agent_conf.api_base_url {
336                        bigmodel::Client::from_url(&agent_conf.api_key, &api_base_url)
337                    } else {
338                        bigmodel::Client::new(&agent_conf.api_key)
339                    };
340                    let agent = client
341                        .agent(&agent_conf.model_name)
342                        .name(agent_name.as_str())
343                        .preamble(&system_prompt)
344                        .build();
345                    self.agents.push((
346                        agent,
347                        agent_conf.id,
348                        agent_conf.provider.to_string(),
349                        agent_conf.model_name,
350                    ));
351                }
352            }
353        }
354        self
355    }
356}