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#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize)]
12pub enum Format
13{
14 #[default]
16 Text,
17 Json,
19}
20
21#[derive(Debug, Default, Deserialize, Serialize)]
23pub struct Options
24{
25 #[default(Format::Text)]
27 pub(crate) format: Format,
28 #[default(false)]
30 pub(crate) thinking: bool,
31}
32
33#[derive(Debug, Default, Deserialize, Serialize)]
35pub struct ChatPrompt
36{
37 pub(crate) messages: Vec<Message>,
39 #[serde(default)]
41 pub(crate) options: Options,
42 pub(crate) template_context: HashMap<String, minijinja::Value>,
44}
45
46impl ChatPrompt
47{
48 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 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 pub fn user(self, content: impl Into<String>) -> Self
68 {
69 self.message(Role::User, content)
70 }
71 pub fn system(self, content: impl Into<String>) -> Self
73 {
74 self.message(Role::System, content)
75 }
76 pub fn system_opt(self, content: Option<String>) -> Self
78 {
79 self.message_opt(Role::System, content)
80 }
81 pub fn assistant(self, content: impl Into<String>) -> Self
83 {
84 self.message(Role::Assistant, content)
85 }
86 pub fn assistant_opt(self, content: Option<String>) -> Self
88 {
89 self.message_opt(Role::Assistant, content)
90 }
91 pub fn format(mut self, format: impl Into<Format>) -> Self
93 {
94 self.options.format = format.into();
95 self
96 }
97 pub fn thiking(mut self, thinking: bool) -> Self
99 {
100 self.options.thinking = thinking;
101 self
102 }
103 pub fn options(mut self, options: Options) -> Self
105 {
106 self.options = options;
107 self
108 }
109 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#[derive(Debug)]
119pub struct GenerationPrompt
120{
121 pub(crate) user: String,
123 pub(crate) system: Option<String>,
125 pub(crate) assistant: Option<String>,
127 pub(crate) options: Options,
129 pub(crate) template_context: HashMap<String, minijinja::Value>,
131 #[cfg(feature = "image")]
133 pub(crate) image: Option<kproc_values::Image>,
134}
135
136impl GenerationPrompt
137{
138 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 pub fn system(mut self, content: impl Into<String>) -> Self
153 {
154 self.system = Some(content.into());
155 self
156 }
157 pub fn assistant(mut self, content: impl Into<String>) -> Self
159 {
160 self.assistant = Some(content.into());
161 self
162 }
163 pub fn format(mut self, format: impl Into<Format>) -> Self
165 {
166 self.options.format = format.into();
167 self
168 }
169 pub fn thinking(mut self, thinking: bool) -> Self
171 {
172 self.options.thinking = thinking;
173 self
174 }
175 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 #[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}