Skip to main content

ai_lib_rust/plugins/
base.rs

1//! Base plugin types.
2
3use 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}