templates.rs

 1use anyhow::Result;
 2use gpui::SharedString;
 3use handlebars::Handlebars;
 4use rust_embed::RustEmbed;
 5use serde::Serialize;
 6use std::sync::Arc;
 7
 8#[derive(RustEmbed)]
 9#[folder = "src/templates"]
10#[include = "*.hbs"]
11struct Assets;
12
13pub struct Templates(Handlebars<'static>);
14
15impl Templates {
16    pub fn new() -> Arc<Self> {
17        let mut handlebars = Handlebars::new();
18        handlebars.set_strict_mode(true);
19        handlebars.register_helper("contains", Box::new(contains));
20        handlebars.register_embed_templates::<Assets>().unwrap();
21        Arc::new(Self(handlebars))
22    }
23}
24
25pub trait Template: Sized {
26    const TEMPLATE_NAME: &'static str;
27
28    fn render(&self, templates: &Templates) -> Result<String>
29    where
30        Self: Serialize + Sized,
31    {
32        Ok(templates.0.render(Self::TEMPLATE_NAME, self)?)
33    }
34}
35
36#[derive(Serialize)]
37pub struct SystemPromptTemplate<'a> {
38    #[serde(flatten)]
39    pub project: &'a prompt_store::ProjectContext,
40    pub available_tools: Vec<SharedString>,
41    pub model_name: Option<String>,
42}
43
44impl Template for SystemPromptTemplate<'_> {
45    const TEMPLATE_NAME: &'static str = "system_prompt.hbs";
46}
47
48/// Handlebars helper for checking if an item is in a list
49fn contains(
50    h: &handlebars::Helper,
51    _: &handlebars::Handlebars,
52    _: &handlebars::Context,
53    _: &mut handlebars::RenderContext,
54    out: &mut dyn handlebars::Output,
55) -> handlebars::HelperResult {
56    let list = h
57        .param(0)
58        .and_then(|v| v.value().as_array())
59        .ok_or_else(|| {
60            handlebars::RenderError::new("contains: missing or invalid list parameter")
61        })?;
62    let query = h.param(1).map(|v| v.value()).ok_or_else(|| {
63        handlebars::RenderError::new("contains: missing or invalid query parameter")
64    })?;
65
66    if list.contains(query) {
67        out.write("true")?;
68    }
69
70    Ok(())
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn test_system_prompt_template() {
79        let project = prompt_store::ProjectContext::default();
80        let template = SystemPromptTemplate {
81            project: &project,
82            available_tools: vec!["echo".into()],
83            model_name: Some("test-model".to_string()),
84        };
85        let templates = Templates::new();
86        let rendered = template.render(&templates).unwrap();
87        assert!(rendered.contains("## Fixing Diagnostics"));
88        assert!(rendered.contains("test-model"));
89    }
90}