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 =
62                        anthropic::ClientBuilder::<reqwest::Client>::new(&agent_conf.api_key);
63                    if let Some(api_base_url) = &agent_conf.api_base_url {
64                        client_builder = client_builder.base_url(api_base_url);
65                    }
66                    match client_builder.build() {
67                        Ok(client) => {
68                            let agent = client
69                                .agent(&agent_conf.model_name)
70                                .name(agent_name.as_str())
71                                .preamble(&system_prompt)
72                                .build();
73                            self.agents.push((
74                                agent,
75                                agent_conf.id,
76                                agent_conf.provider.to_string(),
77                                agent_conf.model_name,
78                            ));
79                        }
80                        Err(err) => {
81                            tracing::error!("添加 {} 错误: {}", agent_conf.provider, err);
82                        }
83                    }
84                }
85                ProviderEnum::Cohere => {
86                    let client = cohere::Client::new(&agent_conf.api_key);
87                    let agent = client
88                        .agent(&agent_conf.model_name)
89                        .name(agent_name.as_str())
90                        .preamble(&system_prompt)
91                        .build();
92                    self.agents.push((
93                        agent,
94                        agent_conf.id,
95                        agent_conf.provider.to_string(),
96                        agent_conf.model_name,
97                    ));
98                }
99                ProviderEnum::Gemini => {
100                    let mut client_builder = gemini::Client::builder(&agent_conf.api_key);
101                    if let Some(api_base_url) = &agent_conf.api_base_url {
102                        client_builder = client_builder.base_url(api_base_url);
103                    }
104                    match client_builder.build() {
105                        Ok(client) => {
106                            let agent = client
107                                .agent(&agent_conf.model_name)
108                                .name(agent_name.as_str())
109                                .preamble(&system_prompt)
110                                .build();
111                            self.agents.push((
112                                agent,
113                                agent_conf.id,
114                                agent_conf.provider.to_string(),
115                                agent_conf.model_name,
116                            ));
117                        }
118                        Err(err) => {
119                            tracing::error!("添加 {} 错误: {}", agent_conf.provider, err);
120                        }
121                    }
122                }
123                ProviderEnum::Huggingface => {
124                    let client = huggingface::Client::new(&agent_conf.api_key);
125                    let agent = client
126                        .agent(&agent_conf.model_name)
127                        .name(agent_name.as_str())
128                        .preamble(&system_prompt)
129                        .build();
130                    self.agents.push((
131                        agent,
132                        agent_conf.id,
133                        agent_conf.provider.to_string(),
134                        agent_conf.model_name,
135                    ));
136                }
137                ProviderEnum::Mistral => {
138                    let client = mistral::Client::new(&agent_conf.api_key);
139                    let agent = client
140                        .agent(&agent_conf.model_name)
141                        .name(agent_name.as_str())
142                        .preamble(&system_prompt)
143                        .build();
144                    self.agents.push((
145                        agent,
146                        agent_conf.id,
147                        agent_conf.provider.to_string(),
148                        agent_conf.model_name,
149                    ));
150                }
151                ProviderEnum::OpenAi => {
152                    let mut client_builder = openai::ClientBuilder::new(&agent_conf.api_key);
153                    if let Some(api_base_url) = &agent_conf.api_base_url {
154                        client_builder = client_builder.base_url(api_base_url)
155                    }
156
157                    let client = client_builder.build();
158
159                    let agent =
160                        get_openai_agent(client, &agent_conf.model_name, agent_name, system_prompt);
161                    self.agents.push((
162                        agent,
163                        agent_conf.id,
164                        agent_conf.provider.to_string(),
165                        agent_conf.model_name,
166                    ));
167                }
168                ProviderEnum::OpenRouter => {
169                    let mut client_builder =
170                        openrouter::ClientBuilder::<reqwest::Client>::new(&agent_conf.api_key);
171                    if let Some(api_base_url) = &agent_conf.api_base_url {
172                        client_builder = client_builder.base_url(api_base_url)
173                    }
174                    let client = client_builder.build();
175                    let agent = client
176                        .agent(&agent_conf.model_name)
177                        .name(agent_name.as_str())
178                        .preamble(&system_prompt)
179                        .build();
180                    self.agents.push((
181                        agent,
182                        agent_conf.id,
183                        agent_conf.provider.to_string(),
184                        agent_conf.model_name,
185                    ));
186                }
187                ProviderEnum::Together => {
188                    let client = together::Client::new(&agent_conf.api_key);
189                    let agent = client
190                        .agent(&agent_conf.model_name)
191                        .name(agent_name.as_str())
192                        .preamble(&system_prompt)
193                        .build();
194                    self.agents.push((
195                        agent,
196                        agent_conf.id,
197                        agent_conf.provider.to_string(),
198                        agent_conf.model_name,
199                    ));
200                }
201                ProviderEnum::XAI => {
202                    let client = xai::Client::new(&agent_conf.api_key);
203                    let agent = client
204                        .agent(&agent_conf.model_name)
205                        .name(agent_name.as_str())
206                        .preamble(&system_prompt)
207                        .build();
208                    self.agents.push((
209                        agent,
210                        agent_conf.id,
211                        agent_conf.provider.to_string(),
212                        agent_conf.model_name,
213                    ));
214                }
215                ProviderEnum::Azure => {
216                    tracing::info!("Azure simple_builder暂不支持,参数有点多,可以自行添加........ ")
217                }
218                ProviderEnum::DeepSeek => {
219                    let client = deepseek::Client::new(&agent_conf.api_key);
220                    let agent = client
221                        .agent(&agent_conf.model_name)
222                        .name(agent_name.as_str())
223                        .preamble(&system_prompt)
224                        .build();
225                    self.agents.push((
226                        agent,
227                        agent_conf.id,
228                        agent_conf.provider.to_string(),
229                        agent_conf.model_name,
230                    ));
231                }
232                ProviderEnum::Galadriel => {
233                    let client = galadriel::Client::new(&agent_conf.api_key);
234                    let agent = client
235                        .agent(&agent_conf.model_name)
236                        .name(agent_name.as_str())
237                        .preamble(&system_prompt)
238                        .build();
239                    self.agents.push((
240                        agent,
241                        agent_conf.id,
242                        agent_conf.provider.to_string(),
243                        agent_conf.model_name,
244                    ));
245                }
246                ProviderEnum::Groq => {
247                    let client = groq::Client::new(&agent_conf.api_key);
248                    let agent = client
249                        .agent(&agent_conf.model_name)
250                        .name(agent_name.as_str())
251                        .preamble(&system_prompt)
252                        .build();
253                    self.agents.push((
254                        agent,
255                        agent_conf.id,
256                        agent_conf.provider.to_string(),
257                        agent_conf.model_name,
258                    ));
259                }
260                ProviderEnum::Hyperbolic => {
261                    let client = hyperbolic::Client::new(&agent_conf.api_key);
262                    let agent = client
263                        .agent(&agent_conf.model_name)
264                        .name(agent_name.as_str())
265                        .preamble(&system_prompt)
266                        .build();
267                    self.agents.push((
268                        agent,
269                        agent_conf.id,
270                        agent_conf.provider.to_string(),
271                        agent_conf.model_name,
272                    ));
273                }
274                ProviderEnum::Mira => {
275                    let client = mira::Client::new(&agent_conf.api_key);
276                    let agent = client
277                        .agent(&agent_conf.model_name)
278                        .name(agent_name.as_str())
279                        .preamble(&system_prompt)
280                        .build();
281                    self.agents.push((
282                        agent,
283                        agent_conf.id,
284                        agent_conf.provider.to_string(),
285                        agent_conf.model_name,
286                    ));
287                }
288                ProviderEnum::Mooshot => {
289                    let client = moonshot::Client::new(&agent_conf.api_key);
290                    let agent = client
291                        .agent(&agent_conf.model_name)
292                        .name(agent_name.as_str())
293                        .preamble(&system_prompt)
294                        .build();
295                    self.agents.push((
296                        agent,
297                        agent_conf.id,
298                        agent_conf.provider.to_string(),
299                        agent_conf.model_name,
300                    ));
301                }
302                ProviderEnum::Ollama => {
303                    let mut client_builder = ollama::ClientBuilder::<reqwest::Client>::new();
304                    if let Some(api_base_url) = &agent_conf.api_base_url {
305                        client_builder = client_builder.base_url(api_base_url);
306                    }
307
308                    let client = client_builder.build();
309                    let agent = client
310                        .agent(&agent_conf.model_name)
311                        .name(agent_name.as_str())
312                        .preamble(&system_prompt)
313                        .build();
314                    self.agents.push((
315                        agent,
316                        agent_conf.id,
317                        agent_conf.provider.to_string(),
318                        agent_conf.model_name,
319                    ));
320                }
321                ProviderEnum::Perplexity => {
322                    // let client = perplexity::Client::new(&agent_conf.api_key);
323                    // let agent = client
324                    //     .agent(&agent_conf.model_name)
325                    //     .name(agent_name.as_str())
326                    //     .preamble(&system_prompt)
327                    //     .build();
328                    // self.agents.push((
329                    //     agent,
330                    //     agent_conf.id,
331                    //     agent_conf.provider.to_string(),
332                    //     agent_conf.model_name,
333                    // ));
334                    tracing::info!("Perplexity 暂不支持,没有实现BoxAgent........ ")
335                }
336                ProviderEnum::Bigmodel => {
337                    let client = if let Some(api_base_url) = agent_conf.api_base_url {
338                        bigmodel::Client::from_url(&agent_conf.api_key, &api_base_url)
339                    } else {
340                        bigmodel::Client::new(&agent_conf.api_key)
341                    };
342                    let agent = client
343                        .agent(&agent_conf.model_name)
344                        .name(agent_name.as_str())
345                        .preamble(&system_prompt)
346                        .build();
347                    self.agents.push((
348                        agent,
349                        agent_conf.id,
350                        agent_conf.provider.to_string(),
351                        agent_conf.model_name,
352                    ));
353                }
354            }
355        }
356        self
357    }
358}