1use crate::{Error, Result};
2use pforge_config::{ParamType, PromptDef};
3use rustc_hash::FxHashMap;
4use serde_json::Value;
5
6pub struct PromptManager {
8 prompts: FxHashMap<String, PromptEntry>,
9}
10
11struct PromptEntry {
12 description: String,
13 template: String,
14 arguments: FxHashMap<String, ParamType>,
15}
16
17impl PromptManager {
18 pub fn new() -> Self {
19 Self {
20 prompts: FxHashMap::default(),
21 }
22 }
23
24 pub fn register(&mut self, def: PromptDef) -> Result<()> {
26 if self.prompts.contains_key(&def.name) {
27 return Err(Error::Handler(format!(
28 "Prompt '{}' already registered",
29 def.name
30 )));
31 }
32
33 self.prompts.insert(
34 def.name.clone(),
35 PromptEntry {
36 description: def.description,
37 template: def.template,
38 arguments: def.arguments,
39 },
40 );
41
42 Ok(())
43 }
44
45 pub fn render(&self, name: &str, args: FxHashMap<String, Value>) -> Result<String> {
47 let entry = self
48 .prompts
49 .get(name)
50 .ok_or_else(|| Error::Handler(format!("Prompt '{}' not found", name)))?;
51
52 self.validate_arguments(entry, &args)?;
54
55 self.interpolate(&entry.template, &args)
57 }
58
59 pub fn get_prompt(&self, name: &str) -> Option<PromptMetadata> {
61 self.prompts.get(name).map(|entry| PromptMetadata {
62 description: entry.description.clone(),
63 arguments: entry.arguments.clone(),
64 })
65 }
66
67 pub fn list_prompts(&self) -> Vec<String> {
69 self.prompts.keys().cloned().collect()
70 }
71
72 fn validate_arguments(
74 &self,
75 entry: &PromptEntry,
76 args: &FxHashMap<String, Value>,
77 ) -> Result<()> {
78 for (arg_name, param_type) in &entry.arguments {
80 let is_required = match param_type {
81 ParamType::Complex { required, .. } => *required,
82 _ => false,
83 };
84
85 if is_required && !args.contains_key(arg_name) {
86 return Err(Error::Handler(format!(
87 "Required argument '{}' not provided",
88 arg_name
89 )));
90 }
91 }
92
93 Ok(())
95 }
96
97 fn interpolate(&self, template: &str, args: &FxHashMap<String, Value>) -> Result<String> {
100 let mut result = template.to_string();
101
102 for (key, value) in args {
103 let placeholder = format!("{{{{{}}}}}", key);
104 let replacement = match value {
105 Value::String(s) => s.clone(),
106 Value::Number(n) => n.to_string(),
107 Value::Bool(b) => b.to_string(),
108 Value::Null => String::new(),
109 _ => serde_json::to_string(value)
110 .map_err(|e| Error::Handler(format!("Failed to serialize value: {}", e)))?,
111 };
112
113 result = result.replace(&placeholder, &replacement);
114 }
115
116 if result.contains("{{") && result.contains("}}") {
118 let unresolved: Vec<&str> = result
120 .split("{{")
121 .skip(1)
122 .filter_map(|s| s.split("}}").next())
123 .collect();
124
125 if !unresolved.is_empty() {
126 return Err(Error::Handler(format!(
127 "Unresolved template variables: {}",
128 unresolved.join(", ")
129 )));
130 }
131 }
132
133 Ok(result)
134 }
135}
136
137impl Default for PromptManager {
138 fn default() -> Self {
139 Self::new()
140 }
141}
142
143#[derive(Debug, Clone)]
145pub struct PromptMetadata {
146 pub description: String,
147 pub arguments: FxHashMap<String, ParamType>,
148}
149
150#[cfg(test)]
151mod tests {
152 use super::*;
153 use pforge_config::SimpleType;
154 use serde_json::json;
155
156 #[test]
157 fn test_prompt_registration() {
158 let mut manager = PromptManager::new();
159
160 let def = PromptDef {
161 name: "greeting".to_string(),
162 description: "A simple greeting prompt".to_string(),
163 template: "Hello, {{name}}!".to_string(),
164 arguments: FxHashMap::default(),
165 };
166
167 manager.register(def).unwrap();
168 assert_eq!(manager.list_prompts(), vec!["greeting"]);
169 }
170
171 #[test]
172 fn test_duplicate_prompt_registration() {
173 let mut manager = PromptManager::new();
174
175 let def = PromptDef {
176 name: "test".to_string(),
177 description: "Test".to_string(),
178 template: "{{x}}".to_string(),
179 arguments: FxHashMap::default(),
180 };
181
182 manager.register(def.clone()).unwrap();
183 let result = manager.register(def);
184 assert!(result.is_err());
185 assert!(result
186 .unwrap_err()
187 .to_string()
188 .contains("already registered"));
189 }
190
191 #[test]
192 fn test_simple_interpolation() {
193 let mut manager = PromptManager::new();
194
195 let def = PromptDef {
196 name: "greeting".to_string(),
197 description: "Greeting".to_string(),
198 template: "Hello, {{name}}! You are {{age}} years old.".to_string(),
199 arguments: FxHashMap::default(),
200 };
201
202 manager.register(def).unwrap();
203
204 let mut args = FxHashMap::default();
205 args.insert("name".to_string(), json!("Alice"));
206 args.insert("age".to_string(), json!(30));
207
208 let result = manager.render("greeting", args).unwrap();
209 assert_eq!(result, "Hello, Alice! You are 30 years old.");
210 }
211
212 #[test]
213 fn test_required_argument_validation() {
214 let mut manager = PromptManager::new();
215
216 let mut arguments = FxHashMap::default();
217 arguments.insert(
218 "name".to_string(),
219 ParamType::Complex {
220 ty: SimpleType::String,
221 required: true,
222 default: None,
223 description: None,
224 validation: None,
225 },
226 );
227
228 let def = PromptDef {
229 name: "greeting".to_string(),
230 description: "Greeting".to_string(),
231 template: "Hello, {{name}}!".to_string(),
232 arguments,
233 };
234
235 manager.register(def).unwrap();
236
237 let args = FxHashMap::default();
238 let result = manager.render("greeting", args);
239 assert!(result.is_err());
240 assert!(result
241 .unwrap_err()
242 .to_string()
243 .contains("Required argument"));
244 }
245
246 #[test]
247 fn test_unresolved_placeholder() {
248 let mut manager = PromptManager::new();
249
250 let def = PromptDef {
251 name: "test".to_string(),
252 description: "Test".to_string(),
253 template: "Hello, {{name}}! Welcome to {{location}}.".to_string(),
254 arguments: FxHashMap::default(),
255 };
256
257 manager.register(def).unwrap();
258
259 let mut args = FxHashMap::default();
260 args.insert("name".to_string(), json!("Alice"));
261 let result = manager.render("test", args);
264 assert!(result.is_err());
265 assert!(result
266 .unwrap_err()
267 .to_string()
268 .contains("Unresolved template variables"));
269 }
270
271 #[test]
272 fn test_get_prompt_metadata() {
273 let mut manager = PromptManager::new();
274
275 let mut arguments = FxHashMap::default();
276 arguments.insert(
277 "name".to_string(),
278 ParamType::Complex {
279 ty: SimpleType::String,
280 required: true,
281 default: None,
282 description: Some("User name".to_string()),
283 validation: None,
284 },
285 );
286
287 let def = PromptDef {
288 name: "greeting".to_string(),
289 description: "A greeting prompt".to_string(),
290 template: "Hello, {{name}}!".to_string(),
291 arguments,
292 };
293
294 manager.register(def).unwrap();
295
296 let metadata = manager.get_prompt("greeting").unwrap();
297 assert_eq!(metadata.description, "A greeting prompt");
298 assert!(metadata.arguments.contains_key("name"));
299 }
300
301 #[test]
302 fn test_complex_value_interpolation() {
303 let mut manager = PromptManager::new();
304
305 let def = PromptDef {
306 name: "test".to_string(),
307 description: "Test".to_string(),
308 template: "String: {{str}}, Number: {{num}}, Bool: {{bool}}".to_string(),
309 arguments: FxHashMap::default(),
310 };
311
312 manager.register(def).unwrap();
313
314 let mut args = FxHashMap::default();
315 args.insert("str".to_string(), json!("hello"));
316 args.insert("num".to_string(), json!(42));
317 args.insert("bool".to_string(), json!(true));
318
319 let result = manager.render("test", args).unwrap();
320 assert_eq!(result, "String: hello, Number: 42, Bool: true");
321 }
322
323 #[test]
324 fn test_required_argument_provided_succeeds() {
325 let mut manager = PromptManager::new();
328
329 let mut arguments = FxHashMap::default();
330 arguments.insert(
331 "name".to_string(),
332 ParamType::Complex {
333 ty: SimpleType::String,
334 required: true,
335 default: None,
336 description: None,
337 validation: None,
338 },
339 );
340
341 let def = PromptDef {
342 name: "greeting".to_string(),
343 description: "Greeting".to_string(),
344 template: "Hello, {{name}}!".to_string(),
345 arguments,
346 };
347
348 manager.register(def).unwrap();
349
350 let mut args = FxHashMap::default();
351 args.insert("name".to_string(), json!("Alice"));
352
353 let result = manager.render("greeting", args).unwrap();
355 assert_eq!(result, "Hello, Alice!");
356 }
357
358 #[test]
359 fn test_null_value_interpolation() {
360 let mut manager = PromptManager::new();
362
363 let def = PromptDef {
364 name: "test".to_string(),
365 description: "Test".to_string(),
366 template: "Value is: {{val}}.".to_string(),
367 arguments: FxHashMap::default(),
368 };
369
370 manager.register(def).unwrap();
371
372 let mut args = FxHashMap::default();
373 args.insert("val".to_string(), Value::Null);
374
375 let result = manager.render("test", args).unwrap();
376 assert_eq!(result, "Value is: .");
378 }
379
380 #[test]
381 fn test_array_value_interpolation() {
382 let mut manager = PromptManager::new();
383
384 let def = PromptDef {
385 name: "test".to_string(),
386 description: "Test".to_string(),
387 template: "Items: {{items}}".to_string(),
388 arguments: FxHashMap::default(),
389 };
390
391 manager.register(def).unwrap();
392
393 let mut args = FxHashMap::default();
394 args.insert("items".to_string(), json!(["a", "b", "c"]));
395
396 let result = manager.render("test", args).unwrap();
397 assert_eq!(result, "Items: [\"a\",\"b\",\"c\"]");
398 }
399}