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
40impl Default for SystemPrompt {
41    fn default() -> Self {
42        SystemPrompt {
43            role: None,
44            guidelines: Vec::new(),
45            constraints: Vec::new(),
46            template: default_prompt_template(),
47        }
48    }
49}
50
51impl SystemPromptBuilder {
52    pub fn guidelines<T: IntoIterator<Item = S>, S: AsRef<str>>(
53        &mut self,
54        guidelines: T,
55    ) -> &mut Self {
56        self.guidelines = Some(
57            guidelines
58                .into_iter()
59                .map(|s| s.as_ref().to_string())
60                .collect(),
61        );
62        self
63    }
64
65    pub fn constraints<T: IntoIterator<Item = S>, S: AsRef<str>>(
66        &mut self,
67        constraints: T,
68    ) -> &mut Self {
69        self.constraints = Some(
70            constraints
71                .into_iter()
72                .map(|s| s.as_ref().to_string())
73                .collect(),
74        );
75        self
76    }
77}
78
79fn default_prompt_template() -> Prompt {
80    include_str!("system_prompt_template.md").into()
81}
82
83#[allow(clippy::from_over_into)]
84impl Into<Prompt> for SystemPrompt {
85    fn into(self) -> Prompt {
86        let SystemPrompt {
87            role,
88            guidelines,
89            constraints,
90            template,
91        } = self;
92
93        template
94            .with_context_value("role", role)
95            .with_context_value("guidelines", guidelines)
96            .with_context_value("constraints", constraints)
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[tokio::test]
105    async fn test_customization() {
106        let prompt = SystemPrompt::builder()
107            .role("role")
108            .guidelines(["guideline"])
109            .constraints(vec!["constraint".to_string()])
110            .build()
111            .unwrap();
112
113        let prompt: Prompt = prompt.into();
114
115        let rendered = prompt.render().unwrap();
116
117        insta::assert_snapshot!(rendered);
118    }
119}