Skip to main content

ai_lib_contact/plugins/
base.rs

1//! Base plugin types.
2
3use ai_lib_core::Result;
4use async_trait::async_trait;
5use std::collections::HashMap;
6use std::sync::Arc;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
9pub enum PluginPriority {
10    Highest = 0,
11    High = 25,
12    #[default]
13    Normal = 50,
14    Low = 75,
15    Lowest = 100,
16}
17
18#[derive(Debug, Clone, Default)]
19pub struct PluginContext {
20    pub request: Option<serde_json::Value>,
21    pub response: Option<serde_json::Value>,
22    pub request_id: Option<String>,
23    pub model: Option<String>,
24    pub provider: Option<String>,
25    pub metadata: HashMap<String, serde_json::Value>,
26    pub error: Option<String>,
27    pub skip: bool,
28}
29
30impl PluginContext {
31    pub fn new() -> Self {
32        Self::default()
33    }
34    pub fn with_request(mut self, r: serde_json::Value) -> Self {
35        self.request = Some(r);
36        self
37    }
38    pub fn with_request_id(mut self, id: impl Into<String>) -> Self {
39        self.request_id = Some(id.into());
40        self
41    }
42    pub fn with_model(mut self, m: impl Into<String>) -> Self {
43        self.model = Some(m.into());
44        self
45    }
46    pub fn skip(&mut self) {
47        self.skip = true;
48    }
49    pub fn should_skip(&self) -> bool {
50        self.skip
51    }
52    pub fn set_error(&mut self, e: impl Into<String>) {
53        self.error = Some(e.into());
54    }
55    pub fn has_error(&self) -> bool {
56        self.error.is_some()
57    }
58}
59
60#[async_trait]
61pub trait Plugin: Send + Sync {
62    fn name(&self) -> &str;
63    fn priority(&self) -> PluginPriority {
64        PluginPriority::Normal
65    }
66    async fn on_register(&self) -> Result<()> {
67        Ok(())
68    }
69    async fn on_unregister(&self) -> Result<()> {
70        Ok(())
71    }
72    async fn on_before_request(&self, _ctx: &mut PluginContext) -> Result<()> {
73        Ok(())
74    }
75    async fn on_after_response(&self, _ctx: &mut PluginContext) -> Result<()> {
76        Ok(())
77    }
78    async fn on_error(&self, _ctx: &mut PluginContext) -> Result<()> {
79        Ok(())
80    }
81    async fn on_stream_event(
82        &self,
83        _ctx: &mut PluginContext,
84        _event: &serde_json::Value,
85    ) -> Result<()> {
86        Ok(())
87    }
88}
89
90pub struct CompositePlugin {
91    name: String,
92    plugins: Vec<Arc<dyn Plugin>>,
93}
94impl CompositePlugin {
95    pub fn new(name: impl Into<String>) -> Self {
96        Self {
97            name: name.into(),
98            plugins: Vec::new(),
99        }
100    }
101    pub fn with_plugin(mut self, p: Arc<dyn Plugin>) -> Self {
102        self.plugins.push(p);
103        self
104    }
105    pub fn len(&self) -> usize {
106        self.plugins.len()
107    }
108    pub fn is_empty(&self) -> bool {
109        self.plugins.is_empty()
110    }
111}
112
113#[async_trait]
114impl Plugin for CompositePlugin {
115    fn name(&self) -> &str {
116        &self.name
117    }
118    async fn on_register(&self) -> Result<()> {
119        for p in &self.plugins {
120            p.on_register().await?;
121        }
122        Ok(())
123    }
124    async fn on_unregister(&self) -> Result<()> {
125        for p in &self.plugins {
126            p.on_unregister().await?;
127        }
128        Ok(())
129    }
130    async fn on_before_request(&self, ctx: &mut PluginContext) -> Result<()> {
131        for p in &self.plugins {
132            if ctx.should_skip() {
133                break;
134            }
135            p.on_before_request(ctx).await?;
136        }
137        Ok(())
138    }
139    async fn on_after_response(&self, ctx: &mut PluginContext) -> Result<()> {
140        for p in &self.plugins {
141            if ctx.should_skip() {
142                break;
143            }
144            p.on_after_response(ctx).await?;
145        }
146        Ok(())
147    }
148    async fn on_error(&self, ctx: &mut PluginContext) -> Result<()> {
149        for p in &self.plugins {
150            p.on_error(ctx).await?;
151        }
152        Ok(())
153    }
154}