Skip to main content

trustformers_wasm/plugins/
example_plugin.rs

1// Example plugin implementation for TrustformeRS WASM
2// Demonstrates how to create community plugins
3
4use crate::plugin_framework::{
5    ExecutionMetrics, Plugin, PluginConfig, PluginContext, PluginError, PluginErrorCode,
6    PluginMetadata, PluginPermission, PluginResult, PluginType,
7};
8use std::collections::HashMap;
9
10/// Example text processor plugin
11pub struct TextProcessorPlugin {
12    enabled: bool,
13    settings: HashMap<String, String>,
14}
15
16impl Default for TextProcessorPlugin {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl TextProcessorPlugin {
23    pub fn new() -> Self {
24        Self {
25            enabled: false,
26            settings: HashMap::new(),
27        }
28    }
29}
30
31impl Plugin for TextProcessorPlugin {
32    fn metadata(&self) -> PluginMetadata {
33        PluginMetadata {
34            name: "Text Processor".to_string(),
35            version: "1.0.0".to_string(),
36            author: "TrustformeRS Community".to_string(),
37            description: "Example plugin for text preprocessing".to_string(),
38            plugin_type: PluginType::Preprocessor,
39            dependencies: vec![],
40            permissions: vec![
41                PluginPermission::ReadModelData,
42                PluginPermission::DebugAccess,
43            ],
44        }
45    }
46
47    fn initialize(&mut self, config: PluginConfig) -> Result<(), PluginError> {
48        self.settings = config.settings;
49        self.enabled = true;
50
51        // Validate required settings
52        if !self.settings.contains_key("processing_mode") {
53            return Err(PluginError {
54                code: PluginErrorCode::InvalidConfiguration,
55                message: "Missing 'processing_mode' setting".to_string(),
56                details: Some("Valid modes: 'lowercase', 'uppercase', 'trim'".to_string()),
57            });
58        }
59
60        Ok(())
61    }
62
63    fn execute(&self, context: &PluginContext) -> Result<PluginResult, PluginError> {
64        if !self.enabled {
65            return Err(PluginError {
66                code: PluginErrorCode::ExecutionFailed,
67                message: "Plugin not initialized".to_string(),
68                details: None,
69            });
70        }
71
72        let start_time = std::time::Instant::now();
73        let mut result_data = HashMap::new();
74
75        // Get input text from context
76        let input_text = context.request_data.get("input_text").cloned().unwrap_or_default();
77
78        // Process text based on settings
79        let default_mode = "trim".to_string();
80        let processing_mode = self.settings.get("processing_mode").unwrap_or(&default_mode);
81
82        let processed_text = match processing_mode.as_str() {
83            "lowercase" => input_text.to_lowercase(),
84            "uppercase" => input_text.to_uppercase(),
85            "trim" => input_text.trim().to_string(),
86            _ => input_text.clone(), // No processing
87        };
88
89        result_data.insert("processed_text".to_string(), processed_text.clone());
90        result_data.insert("original_length".to_string(), input_text.len().to_string());
91        result_data.insert(
92            "processed_length".to_string(),
93            processed_text.len().to_string(),
94        );
95        result_data.insert("processing_mode".to_string(), processing_mode.clone());
96
97        let execution_time = start_time.elapsed().as_millis() as f64;
98
99        Ok(PluginResult {
100            success: true,
101            data: result_data,
102            metrics: ExecutionMetrics {
103                execution_time_ms: execution_time,
104                memory_used_mb: 0.1, // Minimal memory usage
105                cpu_usage_percent: 5.0,
106                gpu_memory_used_mb: None,
107            },
108            messages: vec![format!(
109                "Processed {} characters using {} mode",
110                processed_text.len(),
111                processing_mode
112            )],
113        })
114    }
115
116    fn cleanup(&mut self) {
117        self.enabled = false;
118        self.settings.clear();
119    }
120}
121
122/// Example model optimizer plugin
123pub struct ModelOptimizerPlugin {
124    enabled: bool,
125    optimization_level: u8,
126}
127
128impl Default for ModelOptimizerPlugin {
129    fn default() -> Self {
130        Self::new()
131    }
132}
133
134impl ModelOptimizerPlugin {
135    pub fn new() -> Self {
136        Self {
137            enabled: false,
138            optimization_level: 1,
139        }
140    }
141}
142
143impl Plugin for ModelOptimizerPlugin {
144    fn metadata(&self) -> PluginMetadata {
145        PluginMetadata {
146            name: "Model Optimizer".to_string(),
147            version: "1.2.0".to_string(),
148            author: "TrustformeRS Team".to_string(),
149            description: "Example plugin for model optimization".to_string(),
150            plugin_type: PluginType::Optimizer,
151            dependencies: vec![],
152            permissions: vec![
153                PluginPermission::ReadModelData,
154                PluginPermission::WriteModelData,
155                PluginPermission::GpuAccess,
156                PluginPermission::ProfilingAccess,
157            ],
158        }
159    }
160
161    fn initialize(&mut self, config: PluginConfig) -> Result<(), PluginError> {
162        // Parse optimization level from config
163        if let Some(level_str) = config.settings.get("optimization_level") {
164            match level_str.parse::<u8>() {
165                Ok(level) if level <= 3 => self.optimization_level = level,
166                _ => {
167                    return Err(PluginError {
168                        code: PluginErrorCode::InvalidConfiguration,
169                        message: "Invalid optimization level (must be 0-3)".to_string(),
170                        details: None,
171                    })
172                },
173            }
174        }
175
176        self.enabled = true;
177        Ok(())
178    }
179
180    fn execute(&self, context: &PluginContext) -> Result<PluginResult, PluginError> {
181        if !self.enabled {
182            return Err(PluginError {
183                code: PluginErrorCode::ExecutionFailed,
184                message: "Plugin not initialized".to_string(),
185                details: None,
186            });
187        }
188
189        let start_time = std::time::Instant::now();
190        let mut result_data = HashMap::new();
191
192        // Simulate model optimization based on level
193        let optimization_applied = match self.optimization_level {
194            0 => "None",
195            1 => "Basic quantization",
196            2 => "Advanced quantization + pruning",
197            3 => "Full optimization suite",
198            _ => "Unknown",
199        };
200
201        // Simulate optimization metrics
202        let size_reduction = match self.optimization_level {
203            0 => 0.0,
204            1 => 15.0,
205            2 => 35.0,
206            3 => 55.0,
207            _ => 0.0,
208        };
209
210        let speed_improvement = match self.optimization_level {
211            0 => 0.0,
212            1 => 10.0,
213            2 => 25.0,
214            3 => 40.0,
215            _ => 0.0,
216        };
217
218        result_data.insert(
219            "optimization_applied".to_string(),
220            optimization_applied.to_string(),
221        );
222        result_data.insert(
223            "size_reduction_percent".to_string(),
224            size_reduction.to_string(),
225        );
226        result_data.insert(
227            "speed_improvement_percent".to_string(),
228            speed_improvement.to_string(),
229        );
230        result_data.insert(
231            "optimization_level".to_string(),
232            self.optimization_level.to_string(),
233        );
234
235        // Add model metadata if available
236        if let Some(ref metadata) = context.model_metadata {
237            result_data.insert(
238                "original_model_size_mb".to_string(),
239                metadata.size_mb.to_string(),
240            );
241            result_data.insert("model_type".to_string(), metadata.model_type.clone());
242        }
243
244        let execution_time = start_time.elapsed().as_millis() as f64;
245
246        Ok(PluginResult {
247            success: true,
248            data: result_data,
249            metrics: ExecutionMetrics {
250                execution_time_ms: execution_time,
251                memory_used_mb: 5.0,
252                cpu_usage_percent: 15.0,
253                gpu_memory_used_mb: Some(2.0),
254            },
255            messages: vec![
256                format!("Applied {optimization_applied} optimization"),
257                format!(
258                    "Achieved {:.1}% size reduction and {:.1}% speed improvement",
259                    size_reduction, speed_improvement
260                ),
261            ],
262        })
263    }
264
265    fn cleanup(&mut self) {
266        self.enabled = false;
267        self.optimization_level = 1;
268    }
269}
270
271/// Example visualization plugin
272pub struct VisualizationPlugin {
273    enabled: bool,
274    chart_type: String,
275}
276
277impl Default for VisualizationPlugin {
278    fn default() -> Self {
279        Self::new()
280    }
281}
282
283impl VisualizationPlugin {
284    pub fn new() -> Self {
285        Self {
286            enabled: false,
287            chart_type: "bar".to_string(),
288        }
289    }
290}
291
292impl Plugin for VisualizationPlugin {
293    fn metadata(&self) -> PluginMetadata {
294        PluginMetadata {
295            name: "Data Visualizer".to_string(),
296            version: "0.9.0".to_string(),
297            author: "Community Developer".to_string(),
298            description: "Create charts and graphs from model data".to_string(),
299            plugin_type: PluginType::Visualizer,
300            dependencies: vec![],
301            permissions: vec![PluginPermission::ReadModelData, PluginPermission::UiAccess],
302        }
303    }
304
305    fn initialize(&mut self, config: PluginConfig) -> Result<(), PluginError> {
306        if let Some(chart_type) = config.settings.get("chart_type") {
307            let valid_types = ["bar", "line", "pie", "scatter"];
308            if valid_types.contains(&chart_type.as_str()) {
309                self.chart_type = chart_type.clone();
310            } else {
311                return Err(PluginError {
312                    code: PluginErrorCode::InvalidConfiguration,
313                    message: format!("Invalid chart type: {chart_type}"),
314                    details: Some("Valid types: bar, line, pie, scatter".to_string()),
315                });
316            }
317        }
318
319        self.enabled = true;
320        Ok(())
321    }
322
323    fn execute(&self, _context: &PluginContext) -> Result<PluginResult, PluginError> {
324        if !self.enabled {
325            return Err(PluginError {
326                code: PluginErrorCode::ExecutionFailed,
327                message: "Plugin not initialized".to_string(),
328                details: None,
329            });
330        }
331
332        let start_time = std::time::Instant::now();
333        let mut result_data = HashMap::new();
334
335        // Generate mock chart data based on chart type
336        let chart_data = match self.chart_type.as_str() {
337            "bar" => {
338                r#"{"type":"bar","data":{"labels":["Q1","Q2","Q3","Q4"],"values":[10,20,15,25]}}"#
339            },
340            "line" => r#"{"type":"line","data":{"x":[1,2,3,4,5],"y":[2,4,3,5,4]}}"#,
341            "pie" => r#"{"type":"pie","data":{"labels":["A","B","C"],"values":[30,40,30]}}"#,
342            "scatter" => r#"{"type":"scatter","data":{"points":[[1,2],[2,4],[3,3],[4,5]]}}"#,
343            _ => r#"{"type":"unknown","data":{}}"#,
344        };
345
346        result_data.insert("chart_type".to_string(), self.chart_type.clone());
347        result_data.insert("chart_data".to_string(), chart_data.to_string());
348        result_data.insert("width".to_string(), "800".to_string());
349        result_data.insert("height".to_string(), "600".to_string());
350
351        let execution_time = start_time.elapsed().as_millis() as f64;
352
353        Ok(PluginResult {
354            success: true,
355            data: result_data,
356            metrics: ExecutionMetrics {
357                execution_time_ms: execution_time,
358                memory_used_mb: 1.0,
359                cpu_usage_percent: 8.0,
360                gpu_memory_used_mb: None,
361            },
362            messages: vec![
363                format!("Generated {chart_type} chart", chart_type = self.chart_type),
364                "Chart data ready for rendering".to_string(),
365            ],
366        })
367    }
368
369    fn cleanup(&mut self) {
370        self.enabled = false;
371        self.chart_type = "bar".to_string();
372    }
373}