swiftide_agents/
system_prompt.rs

1//! The system prompt is the initial role and constraint defining message the LLM will receive for
2//! completion.
3//!
4//! The builder provides an accessible way to build a system prompt.
5//!
6//! The agent will convert the system prompt into a prompt, adding it to the messages list the
7//! first time it is called.
8//!
9//! For customization, either the builder can be used to profit from defaults, or an override can
10//! be provided on the agent level.
11
12use derive_builder::Builder;
13use swiftide_core::prompt::Prompt;
14
15#[derive(Clone, Debug, Builder)]
16#[builder(setter(into, strip_option))]
17pub struct SystemPrompt {
18    /// The role the agent is expected to fulfil.
19    #[builder(default)]
20    role: Option<String>,
21
22    /// Additional guidelines for the agent to follow
23    #[builder(default, setter(custom))]
24    guidelines: Vec<String>,
25    /// Additional constraints
26    #[builder(default, setter(custom))]
27    constraints: Vec<String>,
28
29    /// The template to use for the system prompt
30    #[builder(default = default_prompt_template())]
31    template: Prompt,
32}
33
34impl SystemPrompt {
35    pub fn builder() -> SystemPromptBuilder {
36        SystemPromptBuilder::default()
37    }
38
39    pub fn to_prompt(&self) -> Prompt {
40        self.clone().into()
41    }
42}
43
44impl From<String> for SystemPrompt {
45    fn from(text: String) -> Self {
46        SystemPrompt {
47            role: None,
48            guidelines: Vec::new(),
49            constraints: Vec::new(),
50            template: text.into(),
51        }
52    }
53}
54
55impl From<&'static str> for SystemPrompt {
56    fn from(text: &'static str) -> Self {
57        SystemPrompt {
58            role: None,
59            guidelines: Vec::new(),
60            constraints: Vec::new(),
61            template: text.into(),
62        }
63    }
64}
65
66impl From<SystemPrompt> for SystemPromptBuilder {
67    fn from(val: SystemPrompt) -> Self {
68        SystemPromptBuilder {
69            role: Some(val.role),
70            guidelines: Some(val.guidelines),
71            constraints: Some(val.constraints),
72            template: Some(val.template),
73        }
74    }
75}
76
77impl From<Prompt> for SystemPrompt {
78    fn from(prompt: Prompt) -> Self {
79        SystemPrompt {
80            role: None,
81            guidelines: Vec::new(),
82            constraints: Vec::new(),
83            template: prompt,
84        }
85    }
86}
87
88impl Default for SystemPrompt {
89    fn default() -> Self {
90        SystemPrompt {
91            role: None,
92            guidelines: Vec::new(),
93            constraints: Vec::new(),
94            template: default_prompt_template(),
95        }
96    }
97}
98
99impl SystemPromptBuilder {
100    pub fn add_guideline(&mut self, guideline: &str) -> &mut Self {
101        self.guidelines
102            .get_or_insert_with(Vec::new)
103            .push(guideline.to_string());
104        self
105    }
106
107    pub fn add_constraint(&mut self, constraint: &str) -> &mut Self {
108        self.constraints
109            .get_or_insert_with(Vec::new)
110            .push(constraint.to_string());
111        self
112    }
113
114    pub fn guidelines<T: IntoIterator<Item = S>, S: AsRef<str>>(
115        &mut self,
116        guidelines: T,
117    ) -> &mut Self {
118        self.guidelines = Some(
119            guidelines
120                .into_iter()
121                .map(|s| s.as_ref().to_string())
122                .collect(),
123        );
124        self
125    }
126
127    pub fn constraints<T: IntoIterator<Item = S>, S: AsRef<str>>(
128        &mut self,
129        constraints: T,
130    ) -> &mut Self {
131        self.constraints = Some(
132            constraints
133                .into_iter()
134                .map(|s| s.as_ref().to_string())
135                .collect(),
136        );
137        self
138    }
139}
140
141fn default_prompt_template() -> Prompt {
142    include_str!("system_prompt_template.md").into()
143}
144
145#[allow(clippy::from_over_into)]
146impl Into<Prompt> for SystemPrompt {
147    fn into(self) -> Prompt {
148        let SystemPrompt {
149            role,
150            guidelines,
151            constraints,
152            template,
153        } = self;
154
155        template
156            .with_context_value("role", role)
157            .with_context_value("guidelines", guidelines)
158            .with_context_value("constraints", constraints)
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    #[tokio::test]
167    async fn test_customization() {
168        let prompt = SystemPrompt::builder()
169            .role("role")
170            .guidelines(["guideline"])
171            .constraints(vec!["constraint".to_string()])
172            .build()
173            .unwrap();
174
175        let prompt: Prompt = prompt.into();
176
177        let rendered = prompt.render().unwrap();
178
179        insta::assert_snapshot!(rendered);
180    }
181
182    #[tokio::test]
183    async fn test_system_prompt_to_builder() {
184        let sp = SystemPrompt {
185            role: Some("Assistant".to_string()),
186            guidelines: vec!["Be concise".to_string()],
187            constraints: vec!["No personal opinions".to_string()],
188            template: "Hello, {{role}}! Guidelines: {{guidelines}}, Constraints: {{constraints}}"
189                .into(),
190        };
191
192        let builder = SystemPromptBuilder::from(sp.clone());
193
194        assert_eq!(builder.role, Some(Some("Assistant".to_string())));
195        assert_eq!(builder.guidelines, Some(vec!["Be concise".to_string()]));
196        assert_eq!(
197            builder.constraints,
198            Some(vec!["No personal opinions".to_string()])
199        );
200        // For template, compare the rendered string
201        assert_eq!(
202            builder.template.as_ref().unwrap().render().unwrap(),
203            sp.template.render().unwrap()
204        );
205    }
206}