swiftide_core/
template.rs

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