sys/sysmod/openai/basicfuncs/
system.rs

1//! システム情報取得。
2
3use crate::sysmod::health::{
4    ThrottleFlags, get_cpu_cores, get_cpu_info, get_current_freq, get_freq_conf,
5    get_throttle_status,
6};
7use crate::sysmod::openai::ParameterType;
8use crate::sysmod::openai::function::{
9    BasicContext, FuncArgs, Function, FunctionTable, ParameterElement, Parameters,
10    get_arg_bool_opt, get_arg_str,
11};
12use anyhow::{Result, anyhow, bail};
13use chrono::{DateTime, Local, Utc};
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16use std::sync::Arc;
17use std::sync::atomic::Ordering;
18
19/// このモジュールの関数をすべて登録する。
20pub fn register_all<T: 'static>(func_table: &mut FunctionTable<T>) {
21    register_debug_mode(func_table);
22    register_get_model(func_table);
23    register_get_rate_limit(func_table);
24    register_get_version(func_table);
25    register_get_cpu_status(func_table);
26    register_get_current_datetime(func_table);
27}
28
29/// デバッグモード取得/設定。
30async fn debug_mode(bctx: Arc<BasicContext>, args: &FuncArgs) -> Result<String> {
31    let enabled = get_arg_bool_opt(args, "enabled")?;
32
33    #[derive(Serialize)]
34    struct FuncResult {
35        current: bool,
36        #[serde(skip_serializing_if = "Option::is_none")]
37        previous: Option<bool>,
38    }
39
40    let result = if let Some(enabled) = enabled {
41        let old = bctx.debug_mode.swap(enabled, Ordering::SeqCst);
42
43        FuncResult {
44            current: enabled,
45            previous: Some(old),
46        }
47    } else {
48        let current = bctx.debug_mode.load(Ordering::SeqCst);
49
50        FuncResult {
51            current,
52            previous: None,
53        }
54    };
55
56    Ok(serde_json::to_string(&result).unwrap())
57}
58
59fn register_debug_mode<T: 'static>(func_table: &mut FunctionTable<T>) {
60    let mut properties = HashMap::new();
61    properties.insert(
62        "enabled".to_string(),
63        ParameterElement {
64            type_: vec![ParameterType::Boolean, ParameterType::Null],
65            description: Some(
66                "New value. If not specified, just get the current value.".to_string(),
67            ),
68            ..Default::default()
69        },
70    );
71
72    func_table.register_function(
73        Function {
74            name: "debug_mode".to_string(),
75            description: Some("Get/Set debug mode of function calls".to_string()),
76            parameters: Parameters {
77                properties,
78                required: vec!["enabled".to_string()],
79                ..Default::default()
80            },
81            ..Default::default()
82        },
83        |bctx, _ctx, args| Box::pin(debug_mode(bctx, args)),
84    );
85}
86
87/// モデル情報取得。
88async fn get_model(bctx: Arc<BasicContext>, _args: &FuncArgs) -> Result<String> {
89    let model = bctx.ctrl.sysmods().openai.lock().await.model_info().await?;
90
91    Ok(serde_json::to_string(&model).unwrap())
92}
93
94fn register_get_model<T: 'static>(func_table: &mut FunctionTable<T>) {
95    func_table.register_function(
96        Function {
97            name: "get_model".to_string(),
98            description: Some("Get GPT model info of the assistant".to_string()),
99            parameters: Parameters {
100                properties: Default::default(),
101                required: Default::default(),
102                ..Default::default()
103            },
104            ..Default::default()
105        },
106        |bctx, _ctx, args| Box::pin(get_model(bctx, args)),
107    );
108}
109
110/// レートリミット情報取得。
111async fn get_rate_limit(bctx: Arc<BasicContext>, _args: &FuncArgs) -> Result<String> {
112    let exp = bctx
113        .ctrl
114        .sysmods()
115        .openai
116        .lock()
117        .await
118        .get_expected_rate_limit();
119
120    exp.ok_or_else(|| anyhow!("No data")).map(|exp| {
121        format!(
122            "Remaining\nRequests: {} / {}\nTokens: {} / {}",
123            exp.remaining_requests, exp.limit_requests, exp.remaining_tokens, exp.limit_tokens,
124        )
125    })
126}
127
128fn register_get_rate_limit<T: 'static>(func_table: &mut FunctionTable<T>) {
129    func_table.register_function(
130        Function {
131            name: "get_rate_limit".to_string(),
132            description: Some("Get rate limit info of GPT usage".to_string()),
133            parameters: Parameters {
134                properties: Default::default(),
135                required: Default::default(),
136                ..Default::default()
137            },
138            ..Default::default()
139        },
140        |bctx, _ctx, args| Box::pin(get_rate_limit(bctx, args)),
141    );
142}
143
144/// バージョン情報取得。
145async fn get_version(_args: &FuncArgs) -> Result<String> {
146    Ok(verinfo::version_info().to_string())
147}
148
149fn register_get_version<T: 'static>(func_table: &mut FunctionTable<T>) {
150    func_table.register_function(
151        Function {
152            name: "get_version".to_string(),
153            description: Some("Get version of the assistant program".to_string()),
154            parameters: Parameters {
155                properties: Default::default(),
156                required: Default::default(),
157                ..Default::default()
158            },
159            ..Default::default()
160        },
161        |_, _, args| Box::pin(get_version(args)),
162    );
163}
164
165#[derive(Serialize, Deserialize)]
166struct CpuStatus {
167    number_of_cores: u32,
168    cpu_usage_percent: f64,
169    #[serde(skip_serializing_if = "Option::is_none")]
170    current_frequency: Option<String>,
171    #[serde(skip_serializing_if = "Option::is_none")]
172    original_frequency: Option<String>,
173    #[serde(skip_serializing_if = "Option::is_none")]
174    temperature_celsius: Option<f64>,
175    #[serde(skip_serializing_if = "Option::is_none")]
176    throttle_status: Option<Vec<String>>,
177}
178
179/// CPU 使用率情報取得。
180async fn get_cpu_status(_args: &FuncArgs) -> Result<String> {
181    let cpu_info = get_cpu_info().await?;
182
183    let number_of_cores = get_cpu_cores().await?;
184    let current_frequency = get_current_freq()
185        .await?
186        .map(|hz| format!("{} MHz", hz / 1_000_000));
187    let original_frequency = get_freq_conf()
188        .await?
189        .map(|hz| format!("{} MHz", hz / 1_000_000));
190    let throttle_status = get_throttle_status().await?.map(|st| {
191        let mut v = vec![];
192        if st.contains(ThrottleFlags::UNDER_VOLTAGE) {
193            v.push("Under Voltage".to_string());
194        }
195        if st.contains(ThrottleFlags::SOFT_TEMP_LIMIT) {
196            v.push("Soft Throttled".to_string());
197        }
198        if st.contains(ThrottleFlags::THROTTLED) {
199            v.push("Hard Throttled".to_string());
200        }
201        v
202    });
203
204    let obj = CpuStatus {
205        number_of_cores,
206        cpu_usage_percent: cpu_info.cpu_percent_total,
207        current_frequency,
208        original_frequency,
209        temperature_celsius: cpu_info.temp,
210        throttle_status,
211    };
212
213    Ok(serde_json::to_string(&obj)?)
214}
215
216fn register_get_cpu_status<T: 'static>(func_table: &mut FunctionTable<T>) {
217    func_table.register_function(
218        Function {
219            name: "get_cpu_status".to_string(),
220            description: Some("Get current status of assistant's CPU".to_string()),
221            parameters: Parameters {
222                properties: Default::default(),
223                required: Default::default(),
224                ..Default::default()
225            },
226            ..Default::default()
227        },
228        |_, _, args| Box::pin(get_cpu_status(args)),
229    );
230}
231
232/// 現在の日時を取得する。
233async fn get_current_datetime(args: &FuncArgs) -> Result<String> {
234    let tz = get_arg_str(args, "tz")?;
235    match tz {
236        "JST" => {
237            let dt: DateTime<Local> = Local::now();
238            Ok(dt.to_string())
239        }
240        "UTC" => {
241            let dt: DateTime<Utc> = Utc::now();
242            Ok(dt.to_string())
243        }
244        _ => {
245            bail!("Parameter tz must be JST or UTC")
246        }
247    }
248}
249
250fn register_get_current_datetime<T: 'static>(func_table: &mut FunctionTable<T>) {
251    let mut properties = HashMap::new();
252    properties.insert(
253        "tz".to_string(),
254        ParameterElement {
255            type_: vec![ParameterType::String],
256            description: Some("Time zone".to_string()),
257            enum_: Some(vec!["JST".to_string(), "UTC".to_string()]),
258        },
259    );
260
261    func_table.register_function(
262        Function {
263            name: "get_current_datetime".to_string(),
264            description: Some("Get the current date and time".to_string()),
265            parameters: Parameters {
266                properties,
267                required: vec!["tz".to_string()],
268                ..Default::default()
269            },
270            ..Default::default()
271        },
272        |_, _, args| Box::pin(get_current_datetime(args)),
273    );
274}