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<Prompt> for SystemPrompt {
67    fn from(prompt: Prompt) -> Self {
68        SystemPrompt {
69            role: None,
70            guidelines: Vec::new(),
71            constraints: Vec::new(),
72            template: prompt,
73        }
74    }
75}
76
77impl Default for SystemPrompt {
78    fn default() -> Self {
79        SystemPrompt {
80            role: None,
81            guidelines: Vec::new(),
82            constraints: Vec::new(),
83            template: default_prompt_template(),
84        }
85    }
86}
87
88impl SystemPromptBuilder {
89    pub fn add_guideline(&mut self, guideline: &str) -> &mut Self {
90        self.guidelines
91            .get_or_insert_with(Vec::new)
92            .push(guideline.to_string());
93        self
94    }
95
96    pub fn add_constraint(&mut self, constraint: &str) -> &mut Self {
97        self.constraints
98            .get_or_insert_with(Vec::new)
99            .push(constraint.to_string());
100        self
101    }
102
103    pub fn guidelines<T: IntoIterator<Item = S>, S: AsRef<str>>(
104        &mut self,
105        guidelines: T,
106    ) -> &mut Self {
107        self.guidelines = Some(
108            guidelines
109                .into_iter()
110                .map(|s| s.as_ref().to_string())
111                .collect(),
112        );
113        self
114    }
115
116    pub fn constraints<T: IntoIterator<Item = S>, S: AsRef<str>>(
117        &mut self,
118        constraints: T,
119    ) -> &mut Self {
120        self.constraints = Some(
121            constraints
122                .into_iter()
123                .map(|s| s.as_ref().to_string())
124                .collect(),
125        );
126        self
127    }
128}
129
130fn default_prompt_template() -> Prompt {
131    include_str!("system_prompt_template.md").into()
132}
133
134#[allow(clippy::from_over_into)]
135impl Into<Prompt> for SystemPrompt {
136    fn into(self) -> Prompt {
137        let SystemPrompt {
138            role,
139            guidelines,
140            constraints,
141            template,
142        } = self;
143
144        template
145            .with_context_value("role", role)
146            .with_context_value("guidelines", guidelines)
147            .with_context_value("constraints", constraints)
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[tokio::test]
156    async fn test_customization() {
157        let prompt = SystemPrompt::builder()
158            .role("role")
159            .guidelines(["guideline"])
160            .constraints(vec!["constraint".to_string()])
161            .build()
162            .unwrap();
163
164        let prompt: Prompt = prompt.into();
165
166        let rendered = prompt.render().unwrap();
167
168        insta::assert_snapshot!(rendered);
169    }
170}