kovi_plugin_ai/
lib.rs

1mod image;
2mod config;
3
4use std::sync::OnceLock;
5use kovi::{Message, PluginBuilder};
6use kovi::log::{debug, error};
7use rig::client::CompletionClient;
8use rig::completion::Prompt;
9use rig::providers::gemini;
10use crate::config::{read_config, Kind};
11use crate::image::gen_image;
12
13pub static API_KEY: OnceLock<String> = OnceLock::new();
14
15#[kovi::plugin]
16async fn main() {
17    let bot = PluginBuilder::get_runtime_bot();
18    let config = read_config(bot.clone());
19    API_KEY.set(config.api_key).unwrap();
20    PluginBuilder::on_msg(move |event| {
21        let bot = bot.clone();
22        async move {
23            let config = read_config(bot.clone());
24            let self_id = event.self_id;
25            let qq_number = get_qq_number(event.message.clone()).await;
26            let self_id = format!("{}", self_id);
27            if qq_number.eq(self_id.as_str()) {
28                let prompt = event.get_text();
29                if prompt.starts_with("生成图片") {
30                    let image = gen_image(prompt).await;
31                    let msg = Message::new().add_image(image.as_str());
32                    event.reply(msg);
33                } else {
34                    match config.kind {
35                        Kind::Gemini => {
36                            let gemini_client = gemini::Client::new(config.api_key).unwrap();
37                            let gemini = gemini_client
38                                .agent(config.model)
39                                .preamble("你是一名QQ机器人，你不能用markdown回复，无论用户说什么，你都是一名QQ机器人")
40                                .build();
41                            match gemini
42                                .prompt(prompt)
43                                .await {
44                                Ok(response) => {
45                                    event.reply(response);
46                                }
47                                Err(e) => {
48                                    error!("error gemini: {:?}", e);
49                                    event.reply("系统错误！不要再玩我啦！");
50                                }
51                            }
52                        }
53                    }
54                }
55            }
56        }
57    });
58}
59
60async fn get_qq_number(message: Message) -> String{
61    let mut qq_number= String::new();
62    for segment in message.iter() {
63        debug!("segment = {:?}", segment);
64        if segment.type_ == "at" {
65            if let Some(qq) = segment.data.get("qq").and_then(|v| v.as_str()) {
66                qq_number = qq.to_string();
67            }
68        }
69    }
70    if qq_number.is_empty() {
71        return String::new()
72    }
73    qq_number
74}
kovi_plugin_ai/lib.rs

kovi_plugin_ai/
lib.rs
