swiftide_core/
template.rs

1use anyhow::{Context as _, Result};
2use tokio::sync::RwLock;
3
4use lazy_static::lazy_static;
5use tera::Tera;
6use uuid::Uuid;
7
8use crate::prompt::Prompt;
9
10lazy_static! {
11    /// Tera repository for templates
12    static ref TEMPLATE_REPOSITORY: RwLock<Tera> = {
13        let prefix = env!("CARGO_MANIFEST_DIR");
14        let path = format!("{prefix}/src/transformers/prompts/**/*.prompt.md");
15
16        match Tera::new(&path)
17        {
18            Ok(t) => RwLock::new(t),
19            Err(e) => {
20                tracing::error!("Parsing error(s): {e}");
21                ::std::process::exit(1);
22            }
23        }
24    };
25}
26/// A `Template` defines a template for a prompt
27#[derive(Clone, Debug)]
28pub enum Template {
29    CompiledTemplate(String),
30    String(String),
31    Static(&'static str),
32}
33
34impl Template {
35    /// Creates a reference to a template already stored in the repository
36    pub fn from_compiled_template_name(name: impl Into<String>) -> Template {
37        Template::CompiledTemplate(name.into())
38    }
39
40    pub fn from_string(template: impl Into<String>) -> Template {
41        Template::String(template.into())
42    }
43
44    /// Extends the prompt repository with a custom [`tera::Tera`] instance.
45    ///
46    /// If you have your own prompt templates or want to add other functionality, you can extend
47    /// the repository with your own [`tera::Tera`] instance.
48    ///
49    /// WARN: Do not use this inside a pipeline or any form of load, as it will lock the repository
50    ///
51    /// # Errors
52    ///
53    /// Errors if the repository could not be extended
54    pub async fn extend(tera: &Tera) -> Result<()> {
55        TEMPLATE_REPOSITORY
56            .write()
57            .await
58            .extend(tera)
59            .context("Could not extend prompt repository with custom Tera instance")
60    }
61
62    /// Compiles a template from a string and returns a `Template` with a reference to the
63    /// string.
64    ///
65    /// WARN: Do not use this inside a pipeline or any form of load, as it will lock the repository
66    ///
67    /// # Errors
68    ///
69    /// Errors if the template fails to compile
70    pub async fn try_compiled_from_str(
71        template: impl AsRef<str> + Send + 'static,
72    ) -> Result<Template> {
73        let id = Uuid::new_v4().to_string();
74        let mut lock = TEMPLATE_REPOSITORY.write().await;
75        lock.add_raw_template(&id, template.as_ref())
76            .context("Failed to add raw template")?;
77
78        Ok(Template::CompiledTemplate(id))
79    }
80
81    /// Renders a template with an optional `tera::Context`
82    ///
83    /// # Errors
84    ///
85    /// - Template cannot be found
86    /// - One-off template has errors
87    /// - Context is missing that is required by the template
88    pub async fn render(&self, context: &tera::Context) -> Result<String> {
89        use Template::{CompiledTemplate, Static, String};
90
91        let template = match self {
92            CompiledTemplate(id) => {
93                let lock = TEMPLATE_REPOSITORY.read().await;
94                tracing::debug!(
95                    id,
96                    available = ?lock.get_template_names().collect::<Vec<_>>(),
97                    "Rendering template ..."
98                );
99                let result = lock.render(id, context);
100
101                if result.is_err() {
102                    tracing::error!(
103                        error = result.as_ref().unwrap_err().to_string(),
104                        available = ?lock.get_template_names().collect::<Vec<_>>(),
105                        "Error rendering template {id}"
106                    );
107                }
108                result.with_context(|| format!("Failed to render template '{id}'"))?
109            }
110            String(template) => Tera::one_off(template, context, false)
111                .context("Failed to render one-off template")?,
112            Static(template) => Tera::one_off(template, context, false)
113                .context("Failed to render one-off template")?,
114        };
115        Ok(template)
116    }
117
118    /// Builds a Prompt from a template with an empty context
119    pub fn to_prompt(&self) -> Prompt {
120        self.into()
121    }
122}
123
124impl From<&'static str> for Template {
125    fn from(template: &'static str) -> Self {
126        Template::Static(template)
127    }
128}
129
130impl From<String> for Template {
131    fn from(template: String) -> Self {
132        Template::String(template)
133    }
134}