90 lines · 2.4 KB · Rust
RawDownload
use anyhow::Result;use gpui::SharedString;use handlebars::Handlebars;use rust_embed::RustEmbed;use serde::Serialize;use std::sync::Arc; #[derive(RustEmbed)]#[folder = "src/templates"]#[include = "*.hbs"]struct Assets; pub struct Templates(Handlebars<'static>); impl Templates {    pub fn new() -> Arc<Self> {        let mut handlebars = Handlebars::new();        handlebars.set_strict_mode(true);        handlebars.register_helper("contains", Box::new(contains));        handlebars.register_embed_templates::<Assets>().unwrap();        Arc::new(Self(handlebars))    }} pub trait Template: Sized {    const TEMPLATE_NAME: &'static str;     fn render(&self, templates: &Templates) -> Result<String>    where        Self: Serialize + Sized,    {        Ok(templates.0.render(Self::TEMPLATE_NAME, self)?)    }} #[derive(Serialize)]pub struct SystemPromptTemplate<'a> {    #[serde(flatten)]    pub project: &'a prompt_store::ProjectContext,    pub available_tools: Vec<SharedString>,    pub model_name: Option<String>,} impl Template for SystemPromptTemplate<'_> {    const TEMPLATE_NAME: &'static str = "system_prompt.hbs";} /// Handlebars helper for checking if an item is in a listfn contains(    h: &handlebars::Helper,    _: &handlebars::Handlebars,    _: &handlebars::Context,    _: &mut handlebars::RenderContext,    out: &mut dyn handlebars::Output,) -> handlebars::HelperResult {    let list = h        .param(0)        .and_then(|v| v.value().as_array())        .ok_or_else(|| {            handlebars::RenderError::new("contains: missing or invalid list parameter")        })?;    let query = h.param(1).map(|v| v.value()).ok_or_else(|| {        handlebars::RenderError::new("contains: missing or invalid query parameter")    })?;     if list.contains(query) {        out.write("true")?;    }     Ok(())} #[cfg(test)]mod tests {    use super::*;     #[test]    fn test_system_prompt_template() {        let project = prompt_store::ProjectContext::default();        let template = SystemPromptTemplate {            project: &project,            available_tools: vec!["echo".into()],            model_name: Some("test-model".to_string()),        };        let templates = Templates::new();        let rendered = template.render(&templates).unwrap();        assert!(rendered.contains("## Fixing Diagnostics"));        assert!(rendered.contains("test-model"));    }}