llm_kernel/llm/
template.rs

1//! Prompt templates with variable substitution and few-shot examples.
2//!
3//! A [`PromptTemplate`] bundles a `{{variable}}`-substituted body with optional
4//! few-shot example strings rendered before the body. Substitution reuses
5//! [`crate::llm::prompt::render_prompt`], so missing variables are left as-is.
6
7use serde::{Deserialize, Serialize};
8
9/// A prompt template with optional few-shot examples.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct PromptTemplate {
12    /// Template body containing `{{variable}}` placeholders.
13    pub template: String,
14    /// Few-shot example strings rendered before the template body.
15    #[serde(default)]
16    pub few_shot: Vec<String>,
17}
18
19impl PromptTemplate {
20    /// Create a template from a body string.
21    pub fn new(template: impl Into<String>) -> Self {
22        Self {
23            template: template.into(),
24            few_shot: Vec::new(),
25        }
26    }
27
28    /// Set the few-shot examples, returning the updated template.
29    pub fn with_few_shot(mut self, examples: Vec<String>) -> Self {
30        self.few_shot = examples;
31        self
32    }
33
34    /// Render the few-shot examples then the substituted template body.
35    ///
36    /// Each few-shot example is emitted on its own line before the template
37    /// body. Variable substitution reuses [`crate::llm::prompt::render_prompt`],
38    /// so any `{{variable}}` without a matching entry is left as-is.
39    pub fn render(&self, vars: &[(&str, &str)]) -> String {
40        let mut out = String::new();
41        for ex in &self.few_shot {
42            out.push_str(ex);
43            out.push('\n');
44        }
45        out.push_str(&crate::llm::prompt::render_prompt(&self.template, vars));
46        out
47    }
48
49    /// Extract the `{{variable}}` names referenced by the template, in first-seen order.
50    ///
51    /// Scans the template body for tokens between the literal strings `"{{"`
52    /// and `"}}"`, returning the unique names in order of first appearance.
53    /// Few-shot examples are not scanned.
54    pub fn variables(&self) -> Vec<String> {
55        let mut seen = Vec::new();
56        let body = self.template.as_str();
57        let mut rest = body;
58        while let Some(start) = rest.find("{{") {
59            rest = &rest[start + "{{".len()..];
60            let Some(end) = rest.find("}}") else { break };
61            let name = &rest[..end];
62            if !seen.iter().any(|s: &String| s == name) {
63                seen.push(name.to_string());
64            }
65            rest = &rest[end + "}}".len()..];
66        }
67        seen
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn render_substitutes_variables() {
77        let t = PromptTemplate::new("Hello {{name}}, your topic is {{topic}}.");
78        let out = t.render(&[("name", "Alice"), ("topic", "Rust")]);
79        assert_eq!(out, "Hello Alice, your topic is Rust.");
80    }
81
82    #[test]
83    fn few_shot_renders_before_body() {
84        let t = PromptTemplate::new("Classify: {{text}}").with_few_shot(vec![
85            "Q: apples\nA: fruit".to_string(),
86            "Q: rover\nA: dog".to_string(),
87        ]);
88        let out = t.render(&[("text", "carrot")]);
89        let body_pos = out.find("Classify:").expect("body present");
90        let ex1_pos = out.find("Q: apples").expect("example 1 present");
91        let ex2_pos = out.find("Q: rover").expect("example 2 present");
92        assert!(ex1_pos < body_pos, "first example must precede body");
93        assert!(ex2_pos < body_pos, "second example must precede body");
94        assert!(out.ends_with("Classify: carrot"));
95    }
96
97    #[test]
98    fn serde_roundtrip_equal() {
99        let t = PromptTemplate::new("Summarize: {{input}}")
100            .with_few_shot(vec!["Example one".to_string(), "Example two".to_string()]);
101        let json = serde_json::to_string(&t).expect("serialize");
102        let back: PromptTemplate = serde_json::from_str(&json).expect("deserialize");
103        assert_eq!(back.template, t.template);
104        assert_eq!(back.few_shot, t.few_shot);
105        assert!(!back.few_shot.is_empty());
106    }
107
108    #[test]
109    fn missing_variable_left_as_is() {
110        let t = PromptTemplate::new("Hello {{name}}, age {{age}}.");
111        let out = t.render(&[("name", "Bob")]);
112        assert_eq!(out, "Hello Bob, age {{age}}.");
113    }
114
115    #[test]
116    fn variables_in_first_seen_order() {
117        let t = PromptTemplate::new("{{b}} {{a}} {{b}} {{c}}");
118        assert_eq!(t.variables(), vec!["b", "a", "c"]);
119    }
120}
llm_kernel/llm/template.rs

llm_kernel/llm/
template.rs
