Skip to main content

af_agent/
prompt.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
6#[serde(rename_all = "snake_case")]
7pub enum PromptAuthority {
8    Platform,
9    Product,
10    Profile,
11    TrustedCapability,
12    Plugin,
13}
14
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16pub struct PromptSection {
17    pub id: String,
18    pub order: i32,
19    pub authority: PromptAuthority,
20    pub source: String,
21    pub version: String,
22    pub content: String,
23}
24
25#[derive(Debug, Clone, Default)]
26pub struct PromptRegistry {
27    sections: BTreeMap<(i32, String), PromptSection>,
28}
29
30impl PromptRegistry {
31    pub fn insert(&mut self, section: PromptSection) -> Result<(), String> {
32        if section.id.trim().is_empty() || section.content.trim().is_empty() {
33            return Err("prompt section id and content are required".into());
34        }
35        if self
36            .sections
37            .values()
38            .any(|current| current.id == section.id)
39        {
40            return Err(format!("duplicate prompt section '{}'", section.id));
41        }
42        self.sections
43            .insert((section.order, section.id.clone()), section);
44        Ok(())
45    }
46
47    pub fn render(&self) -> String {
48        self.sections
49            .values()
50            .map(|section| section.content.trim())
51            .collect::<Vec<_>>()
52            .join("\n\n")
53    }
54
55    pub fn sections(&self) -> Vec<PromptSection> {
56        self.sections.values().cloned().collect()
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[test]
65    fn renders_stably_and_rejects_duplicate_ids() {
66        let mut prompts = PromptRegistry::default();
67        prompts
68            .insert(PromptSection {
69                id: "profile".into(),
70                order: 20,
71                authority: PromptAuthority::Profile,
72                source: "profile".into(),
73                version: "r1".into(),
74                content: "profile".into(),
75            })
76            .unwrap();
77        prompts
78            .insert(PromptSection {
79                id: "base".into(),
80                order: 10,
81                authority: PromptAuthority::Product,
82                source: "product".into(),
83                version: "1".into(),
84                content: "base".into(),
85            })
86            .unwrap();
87        assert_eq!(prompts.render(), "base\n\nprofile");
88        assert!(prompts
89            .insert(PromptSection {
90                id: "base".into(),
91                order: 30,
92                authority: PromptAuthority::Plugin,
93                source: "plugin".into(),
94                version: "1".into(),
95                content: "duplicate".into(),
96            })
97            .is_err());
98    }
99}