kovi_plugin_pet_cat/
lib.rs

1mod config;
2
3use kovi::{
4    Message, PluginBuilder as plugin, RuntimeBot,
5    bot::runtimebot::kovi_api::SetAccessControlList,
6    event::GroupMsgEvent,
7    log::{error, info},
8    serde_json::{Value, json},
9    tokio::task::JoinSet,
10};
11use reqwest::Client;
12use std::sync::Arc;
13
14use crate::config::Condition;
15
16const PLUGIN_NAME: &str = "kovi-plugin-pet-cat";
17
18#[kovi::plugin]
19async fn main() {
20    let bot = plugin::get_runtime_bot();
21    let client = Arc::new(reqwest::ClientBuilder::new().build().unwrap());
22
23    let config = config::init(&bot).await.unwrap();
24
25    if let Some(groups) = &config.allow_groups {
26        bot.set_plugin_access_control(PLUGIN_NAME, true).unwrap();
27        bot.set_plugin_access_control_list(
28            PLUGIN_NAME,
29            true,
30            SetAccessControlList::Adds(groups.clone()),
31        )
32        .unwrap();
33    } else {
34        bot.set_plugin_access_control(PLUGIN_NAME, false).unwrap();
35    }
36
37    plugin::on_group_msg({
38        let bot = bot.clone();
39        let client = client.clone();
40        move |msg| on_group_msg(msg, bot.clone(), client.clone())
41    });
42
43    info!("[pet-cat] Ready to pet cats!");
44}
45
46async fn on_group_msg(event: Arc<GroupMsgEvent>, bot: Arc<RuntimeBot>, client: Arc<Client>) {
47    let imgs = event.message.get("image");
48
49    for img in imgs {
50        let map = img.data.as_object();
51        if img.data.as_object().is_none() {
52            info!("[pet-cat] No data provided by image segment. (Strange!)");
53            continue;
54        }
55
56        let url = map.unwrap().get("url");
57        if url.is_none() {
58            info!("[pet-cat] No url provided by image segment. (Strange!)");
59            continue;
60        }
61
62        let mut url = url.unwrap().as_str().unwrap().to_string();
63        if url.starts_with("https") {
64            url = url.replace("https", "http");
65        }
66        if predict_cat(&url, &client).await {
67            info!("[pet-cat] Cat detected, sending pet cat meme...");
68            send_pet_cat(event.group_id, &bot).await;
69        } else {
70            info!("[pet-cat] No cat detected.")
71        }
72    }
73}
74
75async fn predict_cat(url: &str, client: &Arc<Client>) -> bool {
76    info!("[pet-cat] Predicting cat for image: {url}");
77
78    let config = config::CONFIG.get().unwrap();
79    let mut set = JoinSet::new();
80
81    for c in &config.conditions {
82        set.spawn(predict_cond(url.to_string(), client.clone(), &c));
83    }
84
85    let mut flag = true;
86    while let Some(res) = set.join_next().await {
87        match res {
88            Ok(res) => flag &= res,
89            Err(e) => error!("[pet-cat] Error when querying llm: {e}"),
90        }
91    }
92
93    flag
94}
95
96async fn predict_cond(url: String, client: Arc<Client>, cond: &Condition) -> bool {
97    let config = config::CONFIG.get().unwrap();
98    let req = match client
99        .post(&config.api_url)
100        .bearer_auth(&config.api_key)
101        .json(&json!({
102            "model": config.model,
103            "messages": [
104                {
105                    "role": "system",
106                    "content": [
107                        {
108                            "type": "text",
109                            "text": "你是一个专业的图片分辨专家，可以精确地依据用户的指示，分辨图片中是否包含某一特定物体。**你只能回答 是 或 否，不要做出多余的回答或进行解释**。"
110                        }
111                    ]
112                },
113                {
114                    "role": "user",
115                    "content": [
116                        {
117                            "type": "image_url",
118                            "image_url": {
119                                "url": url
120                            }
121                        },
122                        {
123                            "type": "text",
124                            "text": cond.prompt
125                        }
126                    ]
127                }
128            ],
129            "stream": false,
130        }))
131        .build(){
132            Ok(req) => req,
133            Err(e) => {
134                error!("[pet-cat] Failed to build request: {e}");
135                return false;
136            }
137        };
138
139    let resp = match client.execute(req).await {
140        Ok(resp) => resp,
141        Err(e) => {
142            error!("[pet-cat] Failed to get response: {e}");
143            return false;
144        }
145    };
146
147    let resp: Value = match resp.json().await {
148        Ok(resp) => resp,
149        Err(e) => {
150            error!("[pet-cat] Failed to parse response: {e}");
151            return false;
152        }
153    };
154
155    let resp = resp.as_object().unwrap();
156
157    let Some(result) = resp.get("choices") else {
158        info!("[pet-cat] Invalid response: {resp:?}");
159        return false;
160    };
161
162    let Some(result) = result.as_array().unwrap().first() else {
163        info!("[pet-cat] No choice provided: {resp:?}");
164        return false;
165    };
166
167    let result = result["message"]["content"].as_str();
168
169    let mut flag = false;
170
171    if let Some(s) = result {
172        flag = s.trim() == cond.prediction;
173    }
174
175    info!("[pet-cat] Predicting \"{}\": {flag}", cond.name);
176    flag
177}
178
179async fn send_pet_cat(group: i64, bot: &Arc<RuntimeBot>) {
180    let config = config::CONFIG.get().unwrap();
181    bot.send_group_msg(
182        group,
183        Message::from_value(json!([
184            {
185                "type":"image",
186                "data": {
187                    "file": config.pet_cat_img,
188                }
189            }
190        ]))
191        .unwrap(),
192    );
193}
kovi_plugin_pet_cat/lib.rs

kovi_plugin_pet_cat/
lib.rs
