pforge_runtime/
prompt.rs

1use crate::{Error, Result};
2use pforge_config::{ParamType, PromptDef};
3use rustc_hash::FxHashMap;
4use serde_json::Value;
5
6/// Prompt manager handles prompt rendering with template interpolation
7pub 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    /// Register a prompt definition
25    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    /// Render a prompt with given arguments
46    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        // Validate arguments
53        self.validate_arguments(entry, &args)?;
54
55        // Perform template interpolation
56        self.interpolate(&entry.template, &args)
57    }
58
59    /// Get prompt metadata
60    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    /// List all registered prompts
68    pub fn list_prompts(&self) -> Vec<String> {
69        self.prompts.keys().cloned().collect()
70    }
71
72    /// Validate arguments against schema
73    fn validate_arguments(
74        &self,
75        entry: &PromptEntry,
76        args: &FxHashMap<String, Value>,
77    ) -> Result<()> {
78        // Check required arguments
79        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        // Type validation could be added here
94        Ok(())
95    }
96
97    /// Interpolate template with argument values
98    /// Supports {{variable}} syntax
99    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        // Check for unresolved placeholders
117        if result.contains("{{") && result.contains("}}") {
118            // Extract unresolved variable names for better error message
119            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/// Prompt metadata for discovery
144#[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        // Missing 'location' argument
262
263        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        // This test catches the && to || mutation in validate_arguments
326        // When is_required=true AND arg IS provided, should succeed
327        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        // This should succeed - required arg is provided
354        let result = manager.render("greeting", args).unwrap();
355        assert_eq!(result, "Hello, Alice!");
356    }
357
358    #[test]
359    fn test_null_value_interpolation() {
360        // This test catches the deletion of Value::Null match arm
361        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        // Null should be rendered as empty string
377        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}