systemprompt_templates/
core_provider.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3
4use serde::Deserialize;
5use systemprompt_template_provider::{TemplateDefinition, TemplateProvider, TemplateSource};
6use tokio::fs;
7use tracing::{debug, warn};
8
9#[derive(Debug, Deserialize, Default)]
10struct TemplateManifest {
11    #[serde(default)]
12    templates: HashMap<String, TemplateConfig>,
13}
14
15#[derive(Debug, Deserialize)]
16struct TemplateConfig {
17    #[serde(default)]
18    content_types: Vec<String>,
19}
20
21#[derive(Debug)]
22pub struct CoreTemplateProvider {
23    template_dir: PathBuf,
24    templates: Vec<TemplateDefinition>,
25    priority: u32,
26}
27
28impl CoreTemplateProvider {
29    pub const DEFAULT_PRIORITY: u32 = 1000;
30    pub const EXTENSION_PRIORITY: u32 = 500;
31
32    #[must_use]
33    pub fn new(template_dir: impl Into<PathBuf>) -> Self {
34        Self {
35            template_dir: template_dir.into(),
36            templates: Vec::new(),
37            priority: Self::DEFAULT_PRIORITY,
38        }
39    }
40
41    #[must_use]
42    pub fn with_priority(template_dir: impl Into<PathBuf>, priority: u32) -> Self {
43        Self {
44            template_dir: template_dir.into(),
45            templates: Vec::new(),
46            priority,
47        }
48    }
49
50    pub async fn discover(&mut self) -> anyhow::Result<()> {
51        self.templates = discover_templates(&self.template_dir, self.priority).await?;
52        Ok(())
53    }
54
55    pub async fn discover_from(template_dir: impl Into<PathBuf>) -> anyhow::Result<Self> {
56        let mut provider = Self::new(template_dir);
57        provider.discover().await?;
58        Ok(provider)
59    }
60
61    pub async fn discover_with_priority(
62        template_dir: impl Into<PathBuf>,
63        priority: u32,
64    ) -> anyhow::Result<Self> {
65        let mut provider = Self::with_priority(template_dir, priority);
66        provider.discover().await?;
67        Ok(provider)
68    }
69}
70
71impl TemplateProvider for CoreTemplateProvider {
72    fn provider_id(&self) -> &'static str {
73        "core"
74    }
75
76    fn priority(&self) -> u32 {
77        self.priority
78    }
79
80    fn templates(&self) -> Vec<TemplateDefinition> {
81        self.templates.clone()
82    }
83}
84
85async fn load_manifest(dir: &Path) -> TemplateManifest {
86    let manifest_path = dir.join("templates.yaml");
87
88    let Ok(content) = fs::read_to_string(&manifest_path).await else {
89        return TemplateManifest::default();
90    };
91
92    match serde_yaml::from_str(&content) {
93        Ok(manifest) => {
94            debug!(path = %manifest_path.display(), "Loaded template manifest");
95            manifest
96        },
97        Err(e) => {
98            warn!(
99                path = %manifest_path.display(),
100                error = %e,
101                "Failed to parse template manifest, using defaults"
102            );
103            TemplateManifest::default()
104        },
105    }
106}
107
108async fn discover_templates(dir: &Path, priority: u32) -> anyhow::Result<Vec<TemplateDefinition>> {
109    let mut templates = Vec::new();
110
111    if !dir.exists() {
112        return Ok(templates);
113    }
114
115    let manifest = load_manifest(dir).await;
116    let mut entries = fs::read_dir(dir).await?;
117
118    while let Some(entry) = entries.next_entry().await? {
119        let path = entry.path();
120
121        if path.extension().is_some_and(|ext| ext == "html") {
122            let Some(file_stem) = path.file_stem() else {
123                continue;
124            };
125            let template_name = file_stem.to_string_lossy().to_string();
126
127            debug!(
128                template = %template_name,
129                path = %path.display(),
130                priority = priority,
131                "Discovered template"
132            );
133
134            let content_types = manifest.templates.get(&template_name).map_or_else(
135                || infer_content_types(&template_name),
136                |config| config.content_types.clone(),
137            );
138
139            let filename = path.file_name().map_or_else(|| path.clone(), PathBuf::from);
140
141            templates.push(TemplateDefinition {
142                name: template_name,
143                source: TemplateSource::File(filename),
144                priority,
145                content_types,
146            });
147        }
148    }
149
150    Ok(templates)
151}
152
153fn infer_content_types(name: &str) -> Vec<String> {
154    match name {
155        _ if name.ends_with("-post") => {
156            let content_type = name.trim_end_matches("-post");
157            vec![content_type.into()]
158        },
159        _ if name.ends_with("-list") => {
160            vec![name.into()]
161        },
162        _ => vec![],
163    }
164}
systemprompt_templates/core_provider.rs

systemprompt_templates/
core_provider.rs
