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 = §ion.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;