Skip to main content

kproc_llm/
prompts.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4use smart_default::SmartDefault as Default;
5
6use crate::Message;
7
8use crate::prelude::*;
9
10/// Format
11#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize)]
12pub enum Format
13{
14  /// Output text
15  #[default]
16  Text,
17  /// Output Json
18  Json,
19}
20
21/// Generation options
22#[derive(Debug, Default, Deserialize, Serialize)]
23pub struct Options
24{
25  /// Requested output format
26  #[default(Format::Text)]
27  pub(crate) format: Format,
28  /// Enable thinking mode (if the model support)
29  #[default(false)]
30  pub(crate) thinking: bool,
31}
32
33/// Prompt for a chat session
34#[derive(Debug, Default, Deserialize, Serialize)]
35pub struct ChatPrompt
36{
37  /// Messages
38  pub(crate) messages: Vec<Message>,
39  /// Generation options
40  #[serde(default)]
41  pub(crate) options: Options,
42  /// Extraxt context passed to the template
43  pub(crate) template_context: HashMap<String, minijinja::Value>,
44}
45
46impl ChatPrompt
47{
48  /// Create a new message, from the given role and content.
49  pub fn message(mut self, role: Role, content: impl Into<String>) -> Self
50  {
51    self.messages.push(Message {
52      role,
53      content: content.into(),
54    });
55    self
56  }
57  /// Create a new message, from the given role and content. Ignore if content is null.
58  pub fn message_opt(mut self, role: Role, content: Option<String>) -> Self
59  {
60    if let Some(content) = content
61    {
62      self.messages.push(Message { role, content });
63    }
64    self
65  }
66  /// Create a new prompt, from the given string.
67  pub fn user(self, content: impl Into<String>) -> Self
68  {
69    self.message(Role::User, content)
70  }
71  /// Set the system hint.
72  pub fn system(self, content: impl Into<String>) -> Self
73  {
74    self.message(Role::System, content)
75  }
76  /// Set the system hint.
77  pub fn system_opt(self, content: Option<String>) -> Self
78  {
79    self.message_opt(Role::System, content)
80  }
81  /// Set the system hint.
82  pub fn assistant(self, content: impl Into<String>) -> Self
83  {
84    self.message(Role::Assistant, content)
85  }
86  /// Set the system hint.
87  pub fn assistant_opt(self, content: Option<String>) -> Self
88  {
89    self.message_opt(Role::Assistant, content)
90  }
91  /// Set the result format
92  pub fn format(mut self, format: impl Into<Format>) -> Self
93  {
94    self.options.format = format.into();
95    self
96  }
97  /// Enable thinking
98  pub fn thiking(mut self, thinking: bool) -> Self
99  {
100    self.options.thinking = thinking;
101    self
102  }
103  /// Set options
104  pub fn options(mut self, options: Options) -> Self
105  {
106    self.options = options;
107    self
108  }
109  /// Add to template context
110  pub fn template_context(mut self, key: String, value: impl Into<minijinja::Value>) -> Self
111  {
112    self.template_context.insert(key, value.into());
113    self
114  }
115}
116
117/// Prompt
118#[derive(Debug)]
119pub struct GenerationPrompt
120{
121  /// Messages
122  pub(crate) user: String,
123  /// Messages
124  pub(crate) system: Option<String>,
125  /// Messages
126  pub(crate) assistant: Option<String>,
127  /// Requested generation options
128  pub(crate) options: Options,
129  /// Extraxt context passed to the template
130  pub(crate) template_context: HashMap<String, minijinja::Value>,
131  /// Optional prompt image (for multi modal models)
132  #[cfg(feature = "image")]
133  pub(crate) image: Option<kproc_values::Image>,
134}
135
136impl GenerationPrompt
137{
138  /// Start creation of generation prompt, with user prompt
139  pub fn prompt(user: impl Into<String>) -> Self
140  {
141    Self {
142      user: user.into(),
143      system: Default::default(),
144      assistant: Default::default(),
145      options: Default::default(),
146      template_context: Default::default(),
147      #[cfg(feature = "image")]
148      image: None,
149    }
150  }
151  /// Set the system hint.
152  pub fn system(mut self, content: impl Into<String>) -> Self
153  {
154    self.system = Some(content.into());
155    self
156  }
157  /// Set the assistant hint.
158  pub fn assistant(mut self, content: impl Into<String>) -> Self
159  {
160    self.assistant = Some(content.into());
161    self
162  }
163  /// Set the result format
164  pub fn format(mut self, format: impl Into<Format>) -> Self
165  {
166    self.options.format = format.into();
167    self
168  }
169  /// Enable thinking
170  pub fn thinking(mut self, thinking: bool) -> Self
171  {
172    self.options.thinking = thinking;
173    self
174  }
175  /// Add to template context
176  pub fn template_context(mut self, key: String, value: impl Into<minijinja::Value>) -> Self
177  {
178    self.template_context.insert(key, value.into());
179    self
180  }
181  /// Set the image, for use in multi-modal models
182  #[cfg(feature = "image")]
183  pub fn image(mut self, image: impl Into<kproc_values::Image>) -> Self
184  {
185    self.image = Some(image.into());
186    self
187  }
188}
189
190impl From<GenerationPrompt> for Vec<Message>
191{
192  fn from(value: GenerationPrompt) -> Self
193  {
194    let mut vec = Self::default();
195    if let Some(system) = value.system
196    {
197      vec.push(Message {
198        role: Role::System,
199        content: system,
200      });
201    }
202    if let Some(assistant) = value.assistant
203    {
204      vec.push(Message {
205        role: Role::Assistant,
206        content: assistant,
207      });
208    }
209    vec.push(Message {
210      role: Role::User,
211      content: value.user,
212    });
213    vec
214  }
215}