ai_lib_rust/plugins/
base.rs1use async_trait::async_trait;
4use std::collections::HashMap;
5use std::sync::Arc;
6use crate::Result;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
9pub enum PluginPriority { Highest = 0, High = 25, Normal = 50, Low = 75, Lowest = 100 }
10impl Default for PluginPriority { fn default() -> Self { PluginPriority::Normal } }
11
12#[derive(Debug, Clone, Default)]
13pub struct PluginContext {
14 pub request: Option<serde_json::Value>,
15 pub response: Option<serde_json::Value>,
16 pub request_id: Option<String>,
17 pub model: Option<String>,
18 pub provider: Option<String>,
19 pub metadata: HashMap<String, serde_json::Value>,
20 pub error: Option<String>,
21 pub skip: bool,
22}
23
24impl PluginContext {
25 pub fn new() -> Self { Self::default() }
26 pub fn with_request(mut self, r: serde_json::Value) -> Self { self.request = Some(r); self }
27 pub fn with_request_id(mut self, id: impl Into<String>) -> Self { self.request_id = Some(id.into()); self }
28 pub fn with_model(mut self, m: impl Into<String>) -> Self { self.model = Some(m.into()); self }
29 pub fn skip(&mut self) { self.skip = true; }
30 pub fn should_skip(&self) -> bool { self.skip }
31 pub fn set_error(&mut self, e: impl Into<String>) { self.error = Some(e.into()); }
32 pub fn has_error(&self) -> bool { self.error.is_some() }
33}
34
35#[async_trait]
36pub trait Plugin: Send + Sync {
37 fn name(&self) -> &str;
38 fn priority(&self) -> PluginPriority { PluginPriority::Normal }
39 async fn on_register(&self) -> Result<()> { Ok(()) }
40 async fn on_unregister(&self) -> Result<()> { Ok(()) }
41 async fn on_before_request(&self, _ctx: &mut PluginContext) -> Result<()> { Ok(()) }
42 async fn on_after_response(&self, _ctx: &mut PluginContext) -> Result<()> { Ok(()) }
43 async fn on_error(&self, _ctx: &mut PluginContext) -> Result<()> { Ok(()) }
44 async fn on_stream_event(&self, _ctx: &mut PluginContext, _event: &serde_json::Value) -> Result<()> { Ok(()) }
45}
46
47pub struct CompositePlugin { name: String, plugins: Vec<Arc<dyn Plugin>> }
48impl CompositePlugin {
49 pub fn new(name: impl Into<String>) -> Self { Self { name: name.into(), plugins: Vec::new() } }
50 pub fn add(mut self, p: Arc<dyn Plugin>) -> Self { self.plugins.push(p); self }
51 pub fn len(&self) -> usize { self.plugins.len() }
52 pub fn is_empty(&self) -> bool { self.plugins.is_empty() }
53}
54
55#[async_trait]
56impl Plugin for CompositePlugin {
57 fn name(&self) -> &str { &self.name }
58 async fn on_register(&self) -> Result<()> { for p in &self.plugins { p.on_register().await?; } Ok(()) }
59 async fn on_unregister(&self) -> Result<()> { for p in &self.plugins { p.on_unregister().await?; } Ok(()) }
60 async fn on_before_request(&self, ctx: &mut PluginContext) -> Result<()> { for p in &self.plugins { if ctx.should_skip() { break; } p.on_before_request(ctx).await?; } Ok(()) }
61 async fn on_after_response(&self, ctx: &mut PluginContext) -> Result<()> { for p in &self.plugins { if ctx.should_skip() { break; } p.on_after_response(ctx).await?; } Ok(()) }
62 async fn on_error(&self, ctx: &mut PluginContext) -> Result<()> { for p in &self.plugins { p.on_error(ctx).await?; } Ok(()) }
63}