swiftide_agents/
system_prompt.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
//! The system prompt is the initial role and constraint defining message the LLM will receive for
//! completion.
//!
//! The builder provides an accessible way to build a system prompt.
//!
//! The agent will convert the system prompt into a prompt, adding it to the messages list the
//! first time it is called.
//!
//! For customization, either the builder can be used to profit from defaults, or an override can
//! be provided on the agent level.

use derive_builder::Builder;
use swiftide_core::{prompt::Prompt, template::Template};

#[derive(Clone, Debug, Builder)]
#[builder(setter(into, strip_option))]
pub struct SystemPrompt {
    /// The role the agent is expected to fulfil.
    #[builder(default)]
    role: Option<String>,

    /// Additional guidelines for the agent to follow
    #[builder(default, setter(custom))]
    guidelines: Vec<String>,
    /// Additional constraints
    #[builder(default, setter(custom))]
    constraints: Vec<String>,

    /// The template to use for the system prompt
    #[builder(default = default_prompt_template())]
    template: Template,
}

impl SystemPrompt {
    pub fn builder() -> SystemPromptBuilder {
        SystemPromptBuilder::default()
    }
}

impl Default for SystemPrompt {
    fn default() -> Self {
        SystemPrompt {
            role: None,
            guidelines: Vec::new(),
            constraints: Vec::new(),
            template: default_prompt_template(),
        }
    }
}

impl SystemPromptBuilder {
    pub fn guidelines<T: IntoIterator<Item = S>, S: AsRef<str>>(
        &mut self,
        guidelines: T,
    ) -> &mut Self {
        self.guidelines = Some(
            guidelines
                .into_iter()
                .map(|s| s.as_ref().to_string())
                .collect(),
        );
        self
    }

    pub fn constraints<T: IntoIterator<Item = S>, S: AsRef<str>>(
        &mut self,
        constraints: T,
    ) -> &mut Self {
        self.constraints = Some(
            constraints
                .into_iter()
                .map(|s| s.as_ref().to_string())
                .collect(),
        );
        self
    }
}

fn default_prompt_template() -> Template {
    include_str!("system_prompt_template.md").into()
}

#[allow(clippy::from_over_into)]
impl Into<Prompt> for SystemPrompt {
    fn into(self) -> Prompt {
        let SystemPrompt {
            role,
            guidelines,
            constraints,
            template,
        } = self;

        template
            .to_prompt()
            .with_context_value("role", role)
            .with_context_value("guidelines", guidelines)
            .with_context_value("constraints", constraints)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_customization() {
        let prompt = SystemPrompt::builder()
            .role("role")
            .guidelines(["guideline"])
            .constraints(vec!["constraint".to_string()])
            .build()
            .unwrap();

        let prompt: Prompt = prompt.into();

        let rendered = prompt.render().await.unwrap();

        insta::assert_snapshot!(rendered);
    }
}