kovi_plugin_ai/
lib.rs

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

kovi_plugin_ai/
lib.rs
