ai_lib_contact/plugins/
base.rs1use 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}