1use 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
33pub const EXTENSION: &str = ".txt.jinja";
35
36pub fn root() -> PathBuf {
38 std::env::current_dir()
39 .unwrap_or_default()
40 .join("app/prompts")
41}
42
43pub 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#[derive(Debug, Clone)]
57pub struct Prompt {
58 name: String,
59 filename: String,
60 config: Arc<Config>,
61}
62
63impl Prompt {
64 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 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 pub fn render(&self, locals: Value) -> Result<String> {
106 render_file(&self.config, &self.name, self.path(), locals_map(locals))
107 }
108}
109
110pub 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
128fn 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(¤t, &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
220fn 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}