Skip to main content

rig_core/completion/
prompt.rs

1//! A system prompt as data: a role and named sections, rendered into one
2//! system message only when the request is sent, so a layer in between can
3//! read, change or move a section by its tag.
4//!
5//! ```
6//! use rig_core::completion::{CompletionRequest, Section, SystemPrompt};
7//!
8//! let mut prompt = SystemPrompt {
9//!     role: "You are a coding agent.".into(),
10//!     sections: vec![Section::new("environment", "Today is Monday.")],
11//! };
12//! assert_eq!(
13//!     prompt.text(),
14//!     "You are a coding agent.\n\n<environment>\nToday is Monday.\n</environment>"
15//! );
16//! let request = CompletionRequest::new("Hello").with_system(&prompt);
17//! assert_eq!(request.system_instructions(), Some(prompt.text().as_str()));
18//!
19//! let moved = prompt.take("environment");
20//! assert_eq!(moved.map(|section| section.text), Some("Today is Monday.".into()));
21//! assert_eq!(prompt.text(), "You are a coding agent.");
22//! ```
23
24use serde::{Deserialize, Serialize};
25
26use super::request::CompletionRequest;
27
28/// A system prompt in parts: the role, then each non-empty section between
29/// `<tag>` and `</tag>`, sent as one system message.
30#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
31pub struct SystemPrompt {
32    /// The text the prompt opens with: who the model is and how it works.
33    pub role: String,
34    /// The sections after the role, in the order they are sent.
35    pub sections: Vec<Section>,
36}
37
38/// A named part of a [`SystemPrompt`].
39#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
40pub struct Section {
41    /// The tag around the text, such as `project_instructions`.
42    pub tag: String,
43    /// The text. Blank text leaves the section out of the rendering.
44    pub text: String,
45}
46
47impl Section {
48    /// A section of `text` between `<tag>` and `</tag>`.
49    pub fn new(tag: impl Into<String>, text: impl Into<String>) -> Self {
50        Self {
51            tag: tag.into(),
52            text: text.into(),
53        }
54    }
55}
56
57impl SystemPrompt {
58    /// The prompt as one text: the role, then each non-empty section tagged,
59    /// a blank line apart, each trimmed. The one rendering; a user who wants
60    /// another renders `role` and `sections` itself.
61    pub fn text(&self) -> String {
62        let mut prompt = self.role.trim().to_owned();
63        for section in &self.sections {
64            let text = section.text.trim();
65            if text.is_empty() {
66                continue;
67            }
68            if !prompt.is_empty() {
69                prompt.push_str("\n\n");
70            }
71            let tag = &section.tag;
72            prompt.push_str(&format!("<{tag}>\n{text}\n</{tag}>"));
73        }
74        prompt
75    }
76
77    /// The first section tagged `tag`, if there is one.
78    pub fn section(&self, tag: &str) -> Option<&Section> {
79        self.sections.iter().find(|section| section.tag == tag)
80    }
81
82    /// Takes the first section tagged `tag` out, if there is one.
83    pub fn take(&mut self, tag: &str) -> Option<Section> {
84        let at = self
85            .sections
86            .iter()
87            .position(|section| section.tag == tag)?;
88        Some(self.sections.remove(at))
89    }
90}
91
92impl CompletionRequest {
93    /// This request with `prompt`, rendered by [`SystemPrompt::text`], as its
94    /// leading system message, ahead of any system message already there.
95    #[must_use]
96    pub fn with_system(self, prompt: &SystemPrompt) -> Self {
97        self.preamble(prompt.text())
98    }
99}
100
101#[cfg(test)]
102mod tests;