Files
keep/src/meta_plugin/read_rate.rs
Andrew Phillips 79fdf05d84 refactor: replace meta_name with MetaPluginType from strum
Co-authored-by: aider (openai/andrew/openrouter/deepseek/deepseek-chat-v3.1) <aider@aider.chat>
2025-08-27 21:04:52 -03:00

143 lines
4.2 KiB
Rust

use std::time::Instant;
use crate::meta_plugin::MetaPlugin;
#[derive(Debug, Clone, Default)]
pub struct ReadRateMetaPlugin {
start_time: Option<Instant>,
bytes_read: u64,
is_finalized: bool,
meta_name: String,
outputs: std::collections::HashMap<String, serde_yaml::Value>,
options: std::collections::HashMap<String, serde_yaml::Value>,
}
impl ReadRateMetaPlugin {
pub fn new(
_options: Option<std::collections::HashMap<String, serde_yaml::Value>>,
outputs: Option<std::collections::HashMap<String, serde_yaml::Value>>,
) -> ReadRateMetaPlugin {
// Start with default options
let mut final_options = std::collections::HashMap::new();
if let Some(opts) = _options {
for (key, value) in opts {
final_options.insert(key, value);
}
}
// Start with default outputs
let mut final_outputs = std::collections::HashMap::new();
let default_outputs = Self::default().default_outputs();
for output_name in default_outputs {
final_outputs.insert(output_name.clone(), serde_yaml::Value::String(output_name));
}
if let Some(outs) = outputs {
for (key, value) in outs {
final_outputs.insert(key, value);
}
}
ReadRateMetaPlugin {
start_time: None,
bytes_read: 0,
is_finalized: false,
meta_name: "read_rate".to_string(),
outputs: final_outputs,
options: final_options,
}
}
}
impl MetaPlugin for ReadRateMetaPlugin {
fn is_finalized(&self) -> bool {
self.is_finalized
}
fn set_finalized(&mut self, finalized: bool) {
self.is_finalized = finalized;
}
fn finalize(&mut self) -> crate::meta_plugin::MetaPluginResponse {
// If already finalized, don't process again
if self.is_finalized {
return crate::meta_plugin::MetaPluginResponse {
metadata: Vec::new(),
is_finalized: true,
};
}
let mut metadata = Vec::new();
if let Some(start_time) = self.start_time {
let duration = start_time.elapsed();
let rate = if duration.as_secs_f64() > 0.0 {
format!("{:.2} KB/s", (self.bytes_read as f64 / 1024.0) / duration.as_secs_f64())
} else {
"N/A".to_string()
};
// Use process_metadata_outputs to handle output mapping
if let Some(meta_data) = crate::meta_plugin::process_metadata_outputs(
"read_rate",
serde_yaml::Value::String(rate),
&self.outputs
) {
metadata.push(meta_data);
}
}
// Mark as finalized
self.is_finalized = true;
crate::meta_plugin::MetaPluginResponse {
metadata,
is_finalized: true,
}
}
fn update(&mut self, data: &[u8]) -> crate::meta_plugin::MetaPluginResponse {
// If already finalized, don't process more data
if self.is_finalized {
return crate::meta_plugin::MetaPluginResponse {
metadata: Vec::new(),
is_finalized: true,
};
}
if self.start_time.is_none() {
self.start_time = Some(Instant::now());
}
self.bytes_read += data.len() as u64;
crate::meta_plugin::MetaPluginResponse {
metadata: Vec::new(),
is_finalized: false,
}
}
fn meta_type(&self) -> MetaPluginType {
MetaPluginType::ReadRate
}
fn outputs(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.outputs
}
fn outputs_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.outputs
}
fn default_outputs(&self) -> Vec<String> {
vec!["read_rate".to_string()]
}
fn options(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.options
}
fn options_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.options
}
}