Skip to main content

rust_llm/
prompt.rs

1//! Port of `lib/ruby_llm/prompt.rb` and `lib/ruby_llm/prompt/context.rb`: prompt templates under
2//! `app/prompts`, rendered with locals.
3//!
4//! RubyLLM renders ERB (`name.txt.erb`). ERB evaluates Ruby, so RustLLM uses Jinja through
5//! [minijinja](https://docs.rs/minijinja) instead, with files named `name.txt.jinja`. The
6//! behavior is RubyLLM's: the same lookup roots, partials named `_name`, bare partial names
7//! resolved next to the current template, path names from the roots, no locals leaking into
8//! partials, and an error for an unknown variable. Rewrite a template like this:
9//!
10//! | ERB (`.txt.erb`) | Jinja (`.txt.jinja`) |
11//! |---|---|
12//! | `<%= name %>` | `{{ name }}` |
13//! | `<% if admin %>...<% end %>` | `{% if admin %}...{% endif %}` |
14//! | `<% items.each do \|i\| %>...<% end %>` | `{% for i in items %}...{% endfor %}` |
15//! | `<%= render "tone", name: name %>` | `{{ render("tone", name=name) }}` |
16//! | `<%= render partial: "shared/safety", locals: { name: name } %>` | `{{ render(partial="shared/safety", locals={"name": name}) }}` |
17//! | `<%= local_assigns[:name] \|\| "friend" %>` | `{{ local_assigns.name \| default("friend") }}` |
18//! | `<%= local_assigns["x-y"] %>` | `{{ local_assigns["x-y"] }}` |
19//!
20//! Output is not HTML-escaped, like ERB's. A name that is not a valid identifier (`x-y`, `Name`)
21//! reaches the template only through `local_assigns`, like Ruby's local-variable rule.
22
23use std::path::{Path, PathBuf};
24use std::sync::Arc;
25
26use minijinja::value::{Kwargs, Rest};
27use minijinja::{AutoEscape, Environment, ErrorKind, UndefinedBehavior, Value as Jinja};
28use serde_json::{Map, Value};
29
30use crate::config::Config;
31use crate::error::{Error, Result};
32
33/// The template extension (Ruby's `.txt.erb`).
34pub const EXTENSION: &str = ".txt.jinja";
35
36/// `Prompt.root`: `app/prompts` under the working directory.
37pub fn root() -> PathBuf {
38    std::env::current_dir()
39        .unwrap_or_default()
40        .join("app/prompts")
41}
42
43/// `Prompt.roots`: the application root first, then `config.prompt_roots` in order. An
44/// application overrides an engine's prompt by shipping a file at the same relative path.
45pub fn roots(config: &Config) -> Vec<PathBuf> {
46    let app = config
47        .get("prompt_root")
48        .map(PathBuf::from)
49        .unwrap_or_else(root);
50    std::iter::once(app)
51        .chain(config.prompt_roots.iter().cloned())
52        .collect()
53}
54
55/// `RubyLLM::Prompt`: one named template.
56#[derive(Debug, Clone)]
57pub struct Prompt {
58    name: String,
59    filename: String,
60    config: Arc<Config>,
61}
62
63impl Prompt {
64    /// `Prompt.new(name)`: `name` without the extension, like `"support/instructions"`.
65    pub fn new(name: impl Into<String>) -> Prompt {
66        Prompt::with_config(crate::config(), name)
67    }
68
69    pub fn with_config(config: Arc<Config>, name: impl Into<String>) -> Prompt {
70        let name = name.into();
71        let filename = if name.ends_with(EXTENSION) {
72            name.clone()
73        } else {
74            format!("{name}{EXTENSION}")
75        };
76        Prompt {
77            name,
78            filename,
79            config,
80        }
81    }
82
83    pub fn name(&self) -> &str {
84        &self.name
85    }
86
87    /// `Prompt#path`: the first root that has the file, else where the application root would.
88    pub fn path(&self) -> PathBuf {
89        let candidates: Vec<PathBuf> = roots(&self.config)
90            .iter()
91            .map(|r| r.join(&self.filename))
92            .collect();
93        candidates
94            .iter()
95            .find(|c| c.exists())
96            .cloned()
97            .unwrap_or_else(|| candidates[0].clone())
98    }
99
100    pub fn exists(&self) -> bool {
101        self.path().exists()
102    }
103
104    /// `Prompt#render(**locals)`: `locals` is a JSON object (or `Null` for none).
105    pub fn render(&self, locals: Value) -> Result<String> {
106        render_file(&self.config, &self.name, self.path(), locals_map(locals))
107    }
108}
109
110/// `RubyLLM.render_prompt(name, **locals)`: renders `app/prompts/<name>.txt.jinja` with `locals`
111/// (a JSON object). Fails with `Error::PromptNotFound` when no root has the file.
112///
113/// ```ignore
114/// let instructions = rust_llm::render_prompt("support/instructions", json!({ "product_name": "BillingHub" }))?;
115/// chat.with_instructions(instructions);
116/// ```
117pub fn render_prompt(name: &str, locals: Value) -> Result<String> {
118    Prompt::new(name).render(locals)
119}
120
121fn locals_map(locals: Value) -> Map<String, Value> {
122    match locals {
123        Value::Object(m) => m,
124        _ => Map::new(),
125    }
126}
127
128/// Ruby's `LOCAL_VARIABLE_NAME`: lowercase or `_` first, then word characters.
129fn is_local_name(name: &str) -> bool {
130    let mut chars = name.chars();
131    chars
132        .next()
133        .is_some_and(|c| c == '_' || c.is_lowercase() || !c.is_ascii())
134        && name.chars().all(|c| c == '_' || c.is_alphanumeric())
135}
136
137fn render_file(
138    config: &Arc<Config>,
139    name: &str,
140    path: PathBuf,
141    locals: Map<String, Value>,
142) -> Result<String> {
143    let source = std::fs::read_to_string(&path).map_err(|e| match e.kind() {
144        std::io::ErrorKind::NotFound => {
145            Error::PromptNotFound(format!("Prompt file not found: {}", path.display()))
146        }
147        _ => Error::Io(e),
148    })?;
149    let mut env = Environment::new();
150    env.set_undefined_behavior(UndefinedBehavior::Strict);
151    env.set_keep_trailing_newline(true);
152    env.set_auto_escape_callback(|_| AutoEscape::None);
153    let (config_for_render, current) = (config.clone(), name.to_string());
154    env.add_function(
155        "render",
156        move |args: Rest<Jinja>, kwargs: Kwargs| -> std::result::Result<String, minijinja::Error> {
157            let (partial, locals) = match args.first() {
158                Some(first) => {
159                    let mut locals = Map::new();
160                    for key in kwargs.args() {
161                        let v: Jinja = kwargs.get(key)?;
162                        locals.insert(key.to_string(), to_json(&v)?);
163                    }
164                    (first.to_string(), locals)
165                }
166                None => {
167                    let partial: String = kwargs.get("partial")?;
168                    let locals: Option<Jinja> = kwargs.get("locals")?;
169                    (
170                        partial,
171                        locals
172                            .map(|v| to_json(&v))
173                            .transpose()?
174                            .map(locals_map)
175                            .unwrap_or_default(),
176                    )
177                }
178            };
179            let name = partial_name(&current, &partial);
180            let path = Prompt::with_config(config_for_render.clone(), &name).path();
181            render_file(&config_for_render, &name, path, locals)
182                .map_err(|e| minijinja::Error::new(ErrorKind::InvalidOperation, e.to_string()))
183        },
184    );
185    let template = env
186        .template_from_str(&source)
187        .map_err(|e| Error::Prompt(describe(&e)))?;
188    let mut context: Map<String, Value> = locals
189        .iter()
190        .filter(|(k, _)| is_local_name(k))
191        .map(|(k, v)| (k.clone(), v.clone()))
192        .collect();
193    context.insert("local_assigns".into(), Value::Object(locals));
194    template
195        .render(Jinja::from_serialize(&context))
196        .map_err(|e| match e.kind() {
197            ErrorKind::InvalidOperation
198                if e.detail()
199                    .is_some_and(|d| d.starts_with("Prompt file not found")) =>
200            {
201                Error::PromptNotFound(e.detail().unwrap_or_default().to_string())
202            }
203            _ => Error::Prompt(describe(&e)),
204        })
205}
206
207fn to_json(v: &Jinja) -> std::result::Result<Value, minijinja::Error> {
208    serde_json::to_value(v)
209        .map_err(|e| minijinja::Error::new(ErrorKind::InvalidOperation, e.to_string()))
210}
211
212fn describe(e: &minijinja::Error) -> String {
213    let mut message = e.to_string();
214    if e.kind() == ErrorKind::UndefinedError {
215        message = format!("undefined local variable in prompt: {message}");
216    }
217    message
218}
219
220/// `Context#partial`: `_name` next to the current prompt for a bare name, from the roots for a
221/// path.
222fn partial_name(current: &str, name: &str) -> String {
223    let partial = match name.rsplit_once('/') {
224        Some((dir, file)) => format!("{dir}/_{file}"),
225        None => format!("_{name}"),
226    };
227    let directory = Path::new(current)
228        .parent()
229        .map(|p| p.to_string_lossy().into_owned())
230        .unwrap_or_default();
231    if name.contains('/') || directory.is_empty() {
232        partial
233    } else {
234        format!("{directory}/{partial}")
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    #[test]
243    fn partials_resolve_next_to_the_prompt_or_from_the_root() {
244        assert_eq!(
245            partial_name("work_assistant/instructions", "tone"),
246            "work_assistant/_tone"
247        );
248        assert_eq!(
249            partial_name("work_assistant/instructions", "shared/safety"),
250            "shared/_safety"
251        );
252        assert_eq!(partial_name("instructions", "tone"), "_tone");
253    }
254
255    #[test]
256    fn only_ruby_local_names_become_variables() {
257        assert!(is_local_name("name"));
258        assert!(is_local_name("_x1"));
259        assert!(!is_local_name("x-y"));
260        assert!(!is_local_name("Name"));
261        assert!(!is_local_name("1x"));
262    }
263}