AI models

crates/common/src/scripts/functions/misc.rs

@@ -94,7 +94,7 @@ pub fn fn_hash<'x>(_: &'x Context<'x>, v: Vec<Variable>) -> Variable {
     })
 }
 
-pub fn fn_is_var_names<'x>(ctx: &'x Context<'x>, _: Vec<Variable>) -> Variable {
+pub fn fn_get_var_names<'x>(ctx: &'x Context<'x>, _: Vec<Variable>) -> Variable {
     Variable::Array(
         ctx.global_variable_names()
             .map(|v| Variable::from(v.to_uppercase()))

crates/common/src/scripts/functions/mod.rs

@@ -20,7 +20,7 @@ use self::{
     array::*, email::*, header::*, html::*, image::*, misc::*, text::*, unicode::*, url::*,
 };
 
-pub fn register_functions() -> FunctionMap {
+pub fn register_functions_trusted() -> FunctionMap {
     FunctionMap::new()
         .with_function("trim", fn_trim)
         .with_function("trim_start", fn_trim_start)
@@ -80,6 +80,7 @@ pub fn register_functions() -> FunctionMap {
         .with_function_args("rsplit", fn_rsplit, 2)
         .with_function_args("split_once", fn_split_once, 2)
         .with_function_args("rsplit_once", fn_rsplit_once, 2)
+        .with_function_args("split_n", fn_split_n, 3)
         .with_function_args("strip_prefix", fn_strip_prefix, 2)
         .with_function_args("strip_suffix", fn_strip_suffix, 2)
         .with_function_args("is_intersect", fn_is_intersect, 2)
@@ -87,11 +88,58 @@ pub fn register_functions() -> FunctionMap {
         .with_function_no_args("is_encoding_problem", fn_is_encoding_problem)
         .with_function_no_args("is_attachment", fn_is_attachment)
         .with_function_no_args("is_body", fn_is_body)
-        .with_function_no_args("var_names", fn_is_var_names)
+        .with_function_no_args("var_names", fn_get_var_names)
         .with_function_no_args("attachment_name", fn_attachment_name)
         .with_function_no_args("mime_part_len", fn_mime_part_len)
 }
 
+pub fn register_functions_untrusted() -> FunctionMap {
+    FunctionMap::new()
+        .with_function("trim", fn_trim)
+        .with_function("trim_start", fn_trim_start)
+        .with_function("trim_end", fn_trim_end)
+        .with_function("len", fn_len)
+        .with_function("count", fn_count)
+        .with_function("is_empty", fn_is_empty)
+        .with_function("is_number", fn_is_number)
+        .with_function("is_ascii", fn_is_ascii)
+        .with_function("to_lowercase", fn_to_lowercase)
+        .with_function("to_uppercase", fn_to_uppercase)
+        .with_function("is_email", fn_is_email)
+        .with_function("thread_name", fn_thread_name)
+        .with_function("html_to_text", fn_html_to_text)
+        .with_function("is_uppercase", fn_is_uppercase)
+        .with_function("is_lowercase", fn_is_lowercase)
+        .with_function("has_digits", fn_has_digits)
+        .with_function("count_spaces", fn_count_spaces)
+        .with_function("count_uppercase", fn_count_uppercase)
+        .with_function("count_lowercase", fn_count_lowercase)
+        .with_function("count_chars", fn_count_chars)
+        .with_function("dedup", fn_dedup)
+        .with_function("lines", fn_lines)
+        .with_function("is_ip_addr", fn_is_ip_addr)
+        .with_function("is_ipv4_addr", fn_is_ipv4_addr)
+        .with_function("is_ipv6_addr", fn_is_ipv6_addr)
+        .with_function("winnow", fn_winnow)
+        .with_function_args("sort", fn_sort, 2)
+        .with_function_args("email_part", fn_email_part, 2)
+        .with_function_args("eq_ignore_case", fn_eq_ignore_case, 2)
+        .with_function_args("contains", fn_contains, 2)
+        .with_function_args("contains_ignore_case", fn_contains_ignore_case, 2)
+        .with_function_args("starts_with", fn_starts_with, 2)
+        .with_function_args("ends_with", fn_ends_with, 2)
+        .with_function_args("uri_part", fn_uri_part, 2)
+        .with_function_args("substring", fn_substring, 3)
+        .with_function_args("split", fn_split, 2)
+        .with_function_args("rsplit", fn_rsplit, 2)
+        .with_function_args("split_once", fn_split_once, 2)
+        .with_function_args("rsplit_once", fn_rsplit_once, 2)
+        .with_function_args("split_n", fn_split_n, 3)
+        .with_function_args("strip_prefix", fn_strip_prefix, 2)
+        .with_function_args("strip_suffix", fn_strip_suffix, 2)
+        .with_function_args("is_intersect", fn_is_intersect, 2)
+}
+
 pub trait ApplyString<'x> {
     fn transform(&self, f: impl Fn(&'_ str) -> Variable) -> Variable;
 }

crates/common/src/scripts/functions/text.rs

@@ -191,6 +191,25 @@ pub fn fn_rsplit<'x>(_: &'x Context<'x>, v: Vec<Variable>) -> Variable {
         .into()
 }
 
+pub fn fn_split_n<'x>(_: &'x Context<'x>, v: Vec<Variable>) -> Variable {
+    let value = v[0].to_string();
+    let arg = v[1].to_string();
+    let num = v[2].to_integer() as usize;
+    let mut result = Vec::new();
+
+    let mut s = value.as_ref();
+    for _ in 0..num {
+        if let Some((a, b)) = s.split_once(arg.as_ref()) {
+            result.push(Variable::from(a.to_string()));
+            s = b;
+        } else {
+            break;
+        }
+    }
+    result.push(Variable::from(s.to_string()));
+    result.into()
+}
+
 pub fn fn_split_once<'x>(_: &'x Context<'x>, v: Vec<Variable>) -> Variable {
     v[0].to_string()
         .split_once(v[1].to_string().as_ref())

crates/common/src/scripts/plugins/llm_prompt.rs

@@ -0,0 +1,80 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd <hello@stalw.art>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use std::time::Instant;
+
+use directory::Permission;
+use sieve::{runtime::Variable, FunctionMap};
+use trc::{AiEvent, SecurityEvent};
+
+use super::PluginContext;
+
+pub fn register(plugin_id: u32, fnc_map: &mut FunctionMap) {
+    fnc_map.set_external_function("llm_prompt", plugin_id, 2);
+}
+
+pub async fn exec(ctx: PluginContext<'_>) -> trc::Result<Variable> {
+    // SPDX-SnippetBegin
+    // SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd <hello@stalw.art>
+    // SPDX-License-Identifier: LicenseRef-SEL
+
+    #[cfg(feature = "enterprise")]
+    if let (Variable::String(name), Variable::String(prompt)) =
+        (&ctx.arguments[0], &ctx.arguments[1])
+    {
+        #[cfg(feature = "test_mode")]
+        if name.as_ref() == "echo-test" {
+            return Ok(prompt.to_string().into());
+        }
+
+        if let Some(ai_api) = ctx.server.core.enterprise.as_ref().and_then(|e| {
+            if ctx.access_token.map_or(true, |token| {
+                if token.has_permission(Permission::AiModelInteract) {
+                    true
+                } else {
+                    trc::event!(
+                        Security(SecurityEvent::Unauthorized),
+                        AccountId = token.primary_id(),
+                        Details = Permission::AiModelInteract.name(),
+                        SpanId = ctx.session_id,
+                    );
+                    false
+                }
+            }) {
+                if e.ai_apis.len() == 1 && name.is_empty() {
+                    e.ai_apis.values().next()
+                } else {
+                    e.ai_apis.get(name.as_ref())
+                }
+            } else {
+                None
+            }
+        }) {
+            let time = Instant::now();
+            match ai_api.send_request(prompt.as_ref(), None).await {
+                Ok(response) => {
+                    trc::event!(
+                        Ai(AiEvent::LlmResponse),
+                        Id = ai_api.id.clone(),
+                        Value = prompt.to_string(),
+                        Details = response.clone(),
+                        Elapsed = time.elapsed(),
+                        SpanId = ctx.session_id,
+                    );
+
+                    return Ok(response.into());
+                }
+                Err(err) => {
+                    trc::error!(err.span_id(ctx.session_id));
+                }
+            }
+        }
+    }
+
+    // SPDX-SnippetEnd
+
+    Ok(false.into())
+}

crates/common/src/scripts/plugins/mod.rs

@@ -9,6 +9,7 @@ pub mod dns;
 pub mod exec;
 pub mod headers;
 pub mod http;
+pub mod llm_prompt;
 pub mod lookup;
 pub mod pyzor;
 pub mod query;
@@ -17,7 +18,7 @@ pub mod text;
 use mail_parser::Message;
 use sieve::{runtime::Variable, FunctionMap, Input};
 
-use crate::{Core, Server};
+use crate::{auth::AccessToken, Core, Server};
 
 use super::ScriptModification;
 
@@ -25,13 +26,14 @@ type RegisterPluginFnc = fn(u32, &mut FunctionMap) -> ();
 
 pub struct PluginContext<'x> {
     pub session_id: u64,
+    pub access_token: Option<&'x AccessToken>,
     pub server: &'x Server,
     pub message: &'x Message<'x>,
     pub modifications: &'x mut Vec<ScriptModification>,
     pub arguments: Vec<Variable>,
 }
 
-const PLUGINS_REGISTER: [RegisterPluginFnc; 18] = [
+const PLUGINS_REGISTER: [RegisterPluginFnc; 19] = [
     query::register,
     exec::register,
     lookup::register,
@@ -50,14 +52,16 @@ const PLUGINS_REGISTER: [RegisterPluginFnc; 18] = [
     headers::register,
     text::register_tokenize,
     text::register_domain_part,
+    llm_prompt::register,
 ];
 
 pub trait RegisterSievePlugins {
-    fn register_plugins(self) -> Self;
+    fn register_plugins_trusted(self) -> Self;
+    fn register_plugins_untrusted(self) -> Self;
 }
 
 impl RegisterSievePlugins for FunctionMap {
-    fn register_plugins(mut self) -> Self {
+    fn register_plugins_trusted(mut self) -> Self {
         #[cfg(feature = "test_mode")]
         {
             self.set_external_function("print", PLUGINS_REGISTER.len() as u32, 1)
@@ -68,6 +72,11 @@ impl RegisterSievePlugins for FunctionMap {
         }
         self
     }
+
+    fn register_plugins_untrusted(mut self) -> Self {
+        llm_prompt::register(18, &mut self);
+        self
+    }
 }
 
 impl Core {
@@ -97,6 +106,7 @@ impl Core {
             15 => headers::exec(ctx),
             16 => text::exec_tokenize(ctx),
             17 => text::exec_domain_part(ctx),
+            18 => llm_prompt::exec(ctx).await,
             _ => unreachable!(),
         };