chat_prompts/chat/
nvidia.rs

1use crate::{
2    error::{PromptError, Result},
3    BuildChatPrompt,
4};
5use endpoints::chat::{
6    ChatCompletionAssistantMessage, ChatCompletionRequestMessage, ChatCompletionSystemMessage,
7    ChatCompletionToolMessage, ChatCompletionUserMessage, ChatCompletionUserMessageContent,
8    ContentPart, Tool,
9};
10
11/// Generate prompts for the `nemotron-mini-instruct` model.
12#[derive(Debug, Default, Clone)]
13pub struct NemotronChatPrompt;
14impl NemotronChatPrompt {
15    /// Create a system prompt from a chat completion request message.
16    fn create_system_prompt(&self, message: &ChatCompletionSystemMessage) -> String {
17        let content = message.content();
18        match content.is_empty() {
19            true => String::from("<extra_id_0>System\nYou are a helpful, respectful and honest assistant. Always answer as short as possible, while being safe."),
20            false =>format!(
21                "<extra_id_0>System\n{content}"
22            )
23        }
24    }
25
26    /// Create a user prompt from a chat completion request message.
27    fn append_user_message(
28        &self,
29        chat_history: impl AsRef<str>,
30        system_prompt: impl AsRef<str>,
31        message: &ChatCompletionUserMessage,
32    ) -> String {
33        let content = match message.content() {
34            ChatCompletionUserMessageContent::Text(text) => text.to_string(),
35            ChatCompletionUserMessageContent::Parts(parts) => {
36                let mut content = String::new();
37                for part in parts {
38                    if let ContentPart::Text(text_content) = part {
39                        content.push_str(text_content.text());
40                        content.push('\n');
41                    }
42                }
43                content
44            }
45        };
46
47        match chat_history.as_ref().is_empty() {
48            true => format!(
49                "{system_prompt}\n<extra_id_1>User\n{user_message}",
50                system_prompt = system_prompt.as_ref().trim(),
51                user_message = content.trim(),
52            ),
53            false => format!(
54                "{chat_history}\n<extra_id_1>User\n{user_message}",
55                chat_history = chat_history.as_ref().trim(),
56                user_message = content.trim(),
57            ),
58        }
59    }
60
61    /// create an assistant prompt from a chat completion request message.
62    fn append_assistant_message(
63        &self,
64        chat_history: impl AsRef<str>,
65        message: &ChatCompletionAssistantMessage,
66    ) -> Result<String> {
67        let content = match message.content() {
68            Some(content) => content.to_string(),
69            // Note that the content is optional if `tool_calls` is specified.
70            None => match message.tool_calls().is_some() {
71                true => String::new(),
72                false => return Err(PromptError::NoAssistantMessage),
73            },
74        };
75
76        Ok(format!(
77            "{chat_history}<extra_id_1>Assistant\n{assistant_message}",
78            chat_history = chat_history.as_ref().trim(),
79            assistant_message = content.trim(),
80        ))
81    }
82}
83impl BuildChatPrompt for NemotronChatPrompt {
84    fn build(&self, messages: &mut Vec<ChatCompletionRequestMessage>) -> Result<String> {
85        if messages.is_empty() {
86            return Err(crate::error::PromptError::NoMessages);
87        }
88
89        // system prompt
90        let system_prompt = match messages[0] {
91            ChatCompletionRequestMessage::System(ref message) => {
92                self.create_system_prompt(message)
93            }
94            _ => String::from("<extra_id_0>System\nYou are a helpful, respectful and honest assistant. Always answer as short as possible, while being safe."),
95        };
96
97        // append user/assistant messages
98        let mut prompt = String::new();
99        for message in messages {
100            match message {
101                ChatCompletionRequestMessage::User(message) => {
102                    prompt = self.append_user_message(&prompt, &system_prompt, message);
103                }
104                ChatCompletionRequestMessage::Assistant(message) => {
105                    prompt = self.append_assistant_message(&prompt, message)?;
106                }
107                _ => continue,
108            }
109        }
110
111        prompt.push_str("\n<extra_id_1>Assistant\n");
112
113        Ok(prompt)
114    }
115}
116
117/// Generate prompts for the models using ChatML template.
118#[derive(Debug, Default, Clone)]
119pub struct NemotronToolPrompt;
120impl NemotronToolPrompt {
121    /// Create a system prompt from a chat completion request message.
122    fn create_system_prompt(&self, message: &ChatCompletionSystemMessage) -> String {
123        let content = message.content();
124        match content.is_empty() {
125            true => String::from("<|im_start|>system\nAnswer as concisely as possible.<|im_end|>"),
126            false => format!("<|im_start|>system\n{content}<|im_end|>"),
127        }
128    }
129
130    fn create_system_prompt_tool(
131        &self,
132        message: &ChatCompletionSystemMessage,
133        tools: Option<&[Tool]>,
134    ) -> String {
135        let content = message.content();
136        match content.is_empty() {
137            true => match tools {
138                Some(tools) => {
139                    let available_tools = serde_json::to_string(tools).unwrap();
140                    let tools = format!("<tool> {available_tools} </tool>");
141
142                    let begin = r#"<extra_id_0>System\nYou are a helpful, respectful and honest assistant. Always answer as short as possible, while being safe."#;
143
144                    format!("{begin}\n\n{tools}")
145                }
146                None => {
147                    String::from("<extra_id_0>System\nYou are a helpful, respectful and honest assistant. Always answer as short as possible, while being safe.")
148                }
149            },
150            false => match tools {
151                Some(tools) => {
152                    let available_tools = serde_json::to_string(tools).unwrap();
153                    let tools = format!("<tool> {available_tools} </tool>");
154
155                    let begin = format!(
156                        "<extra_id_0>System\n{content}"
157                    );
158
159                    format!("{begin}\n\n{tools}")
160                }
161                None => {
162                    format!(
163                        "<extra_id_0>System\n{content}"
164                    )
165                }
166            },
167        }
168    }
169
170    /// Create a user prompt from a chat completion request message.
171    fn append_user_message(
172        &self,
173        chat_history: impl AsRef<str>,
174        system_prompt: impl AsRef<str>,
175        message: &ChatCompletionUserMessage,
176    ) -> String {
177        let content = match message.content() {
178            ChatCompletionUserMessageContent::Text(text) => text.to_string(),
179            ChatCompletionUserMessageContent::Parts(parts) => {
180                let mut content = String::new();
181                for part in parts {
182                    if let ContentPart::Text(text_content) = part {
183                        content.push_str(text_content.text());
184                        content.push('\n');
185                    }
186                }
187                content
188            }
189        };
190
191        match chat_history.as_ref().is_empty() {
192            true => format!(
193                "{system_prompt}\n\n<extra_id_1>User\n{user_message}",
194                system_prompt = system_prompt.as_ref().trim(),
195                user_message = content.trim(),
196            ),
197            false => format!(
198                "{chat_history}\n<extra_id_1>User\n{user_message}",
199                chat_history = chat_history.as_ref().trim(),
200                user_message = content.trim(),
201            ),
202        }
203    }
204
205    /// create an assistant prompt from a chat completion request message.
206    fn append_assistant_message(
207        &self,
208        chat_history: impl AsRef<str>,
209        message: &ChatCompletionAssistantMessage,
210    ) -> Result<String> {
211        let content = match message.content() {
212            Some(content) => content.to_string(),
213            // Note that the content is optional if `tool_calls` is specified.
214            None => match message.tool_calls().is_some() {
215                true => String::new(),
216                false => return Err(PromptError::NoAssistantMessage),
217            },
218        };
219
220        Ok(format!(
221            "{chat_history}<extra_id_1>Assistant\n{assistant_message}",
222            chat_history = chat_history.as_ref().trim(),
223            assistant_message = content.trim(),
224        ))
225    }
226
227    /// create a tool prompt from a chat completion request message.
228    fn append_tool_message(
229        &self,
230        chat_history: impl AsRef<str>,
231        message: &ChatCompletionToolMessage,
232    ) -> String {
233        format!(
234            "{chat_history}\n<extra_id_1>Tool\n{tool_message}",
235            chat_history = chat_history.as_ref().trim(),
236            tool_message = message.content().trim(),
237        )
238    }
239}
240impl BuildChatPrompt for NemotronToolPrompt {
241    fn build(&self, messages: &mut Vec<ChatCompletionRequestMessage>) -> Result<String> {
242        if messages.is_empty() {
243            return Err(crate::error::PromptError::NoMessages);
244        }
245
246        // system prompt
247        let system_prompt = match messages[0] {
248            ChatCompletionRequestMessage::System(ref message) => self.create_system_prompt(message),
249            _ => String::from("<|im_start|>system\nAnswer as concisely as possible.<|im_end|>"),
250        };
251
252        // append user/assistant messages
253        let mut prompt = String::new();
254        for message in messages {
255            match message {
256                ChatCompletionRequestMessage::User(message) => {
257                    prompt = self.append_user_message(&prompt, &system_prompt, message);
258                }
259                ChatCompletionRequestMessage::Assistant(message) => {
260                    prompt = self.append_assistant_message(&prompt, message)?;
261                }
262                ChatCompletionRequestMessage::Tool(message) => {
263                    prompt = self.append_tool_message(&prompt, message);
264                }
265                _ => continue,
266            }
267        }
268
269        prompt.push_str("\n<|im_start|>assistant");
270
271        Ok(prompt)
272    }
273
274    fn build_with_tools(
275        &self,
276        messages: &mut Vec<ChatCompletionRequestMessage>,
277        tools: Option<&[Tool]>,
278    ) -> Result<String> {
279        if messages.is_empty() {
280            return Err(crate::error::PromptError::NoMessages);
281        }
282
283        // system prompt
284        let system_prompt = match messages[0] {
285            ChatCompletionRequestMessage::System(ref message) => {
286                self.create_system_prompt_tool(message, tools)
287            }
288            _ => match tools {
289                Some(tools) => {
290                    let mut tools_s = String::new();
291                    for tool in tools {
292                        let available_tool = serde_json::to_string(&tool.function).unwrap();
293
294                        let tool = format!("<tool> {available_tool} </tool>\n");
295
296                        tools_s.push_str(&tool);
297                    }
298
299                    let begin = r#"<extra_id_0>System\nYou are a helpful, respectful and honest assistant. Always answer as short as possible, while being safe."#;
300
301                    format!("{}\n{}", begin, tools_s.trim())
302                }
303                None => {
304                    String::from("<|im_start|>system\nAnswer as concisely as possible.<|im_end|>")
305                }
306            },
307        };
308
309        // append user/assistant messages
310        let mut prompt = String::new();
311        for message in messages {
312            match message {
313                ChatCompletionRequestMessage::User(message) => {
314                    prompt = self.append_user_message(&prompt, &system_prompt, message);
315                }
316                ChatCompletionRequestMessage::Assistant(message) => {
317                    prompt = self.append_assistant_message(&prompt, message)?;
318                }
319                ChatCompletionRequestMessage::Tool(message) => {
320                    prompt = self.append_tool_message(&prompt, message);
321                }
322                _ => continue,
323            }
324        }
325
326        prompt.push_str("\n<extra_id_1>Assistant\n");
327
328        Ok(prompt)
329    }
330}
chat_prompts/chat/nvidia.rs

chat_prompts/chat/
nvidia.rs
