Skip to main content

headless_lms_chatbot/
message_suggestion.rs

1use std::collections::HashMap;
2
3use crate::{
4    azure_chatbot::{
5        ArrayItem, ArrayProperty, InputItem, JSONType, LLMRequest, LLMRequestParams,
6        LLMRequestResponseFormatParam, NonThinkingParams, RequestTextOptions, Schema,
7        ThinkingParams,
8    },
9    chatbot_error::chatbot_err,
10    content_cleaner::calculate_safe_token_limit,
11    llm_utils::{
12        APIInputMessage, MessageContent, estimate_tokens, make_blocking_llm_request,
13        model_is_thinking, parse_text_completion,
14    },
15    prelude::{ChatbotError, ChatbotErrorType, ChatbotResult},
16};
17use headless_lms_base::config::ApplicationConfiguration;
18use headless_lms_base::error::backend_error::BackendError;
19use headless_lms_models::{
20    application_task_default_language_models::TaskLMSpec,
21    chatbot_conversation_message_messages::MessageRole,
22    chatbot_conversation_messages::{ChatbotConversationMessage, Message},
23};
24use rand::seq::IndexedRandom;
25
26/// Shape of the structured LLM output response, defined by the JSONSchema in
27/// generate_suggested_messages
28#[derive(serde::Deserialize)]
29struct ChatbotNextMessageSuggestionResponse {
30    suggestions: Vec<String>,
31}
32
33/// System prompt instructions for generating suggested next messages
34const SYSTEM_PROMPT: &str = r#"You are given a conversation between a helpful teaching assistant chatbot and a student. Your task is to analyze the conversation and suggest what messages the user could send to the teaching assistant chatbot next to best support the user's learning.
35
36When generating suggestions:
37- Base them strictly on the content and tone of the conversation so far
38- Think of questions that will help deepen the user's understanding, aid in learning, clear misconceptions and motivate the user to think and reason about the subject at hand
39- Avoid introducing unrelated topics or information not motivated by the conversation
40- Maintain a supportive, respectful, and clear tone
41- Keep the suggested messages short and concise
42- The teaching assistant's messages are marked to have been said by 'assistant'.
43
44Steps:
451. Analyze the student's current level of understanding, confusion, and engagement.
462. Identify any misconceptions, gaps, or opportunities for deeper learning.
473. Propose next messages that:
48   - Advance understanding
49   - Encourage active thinking
50   - Are appropriate in difficulty and tone
51   - Are appropriate in the context of the conversation so far
52   - Are relevant to the topic of the conversation, context and course
53
54Constraints:
55- Do not continue the conversation yourself.
56- Role-play a student who wants to learn.
57- Only output the suggested messages, nothing else.
58- Suggest exactly 3 alternate next user messages.
59- Be brief, concise and clear. Use as few words and sentences as possible.
60- Generate messages that fit some of the categories 'clarification', 'elaboration', 'placing in context', 'how to learn this', and 'practical use'
61
62"#;
63
64/// User prompt instructions for generating suggested next messages
65pub const USER_PROMPT: &str = r#"Suggest exactly three messages that the user could send next."#;
66
67/// Calls an LLM and generates suggested messages for a chatbot conversation
68pub async fn generate_suggested_messages(
69    app_config: &ApplicationConfiguration,
70    task_lm: TaskLMSpec,
71    conversation_messages: &[ChatbotConversationMessage],
72    initial_suggested_messages: Option<Vec<String>>,
73    course_name: &str,
74    course_desc: Option<String>,
75) -> ChatbotResult<Vec<String>> {
76    let prompt = SYSTEM_PROMPT.to_owned()
77        + &format!("The course is: {}\n\n", course_name)
78        // if there are initial suggested messages, then include <=5 of them as examples
79        + &(if let Some(ism) = initial_suggested_messages {
80            let mut rng = rand::rng();
81            let examples = ism
82                .sample(&mut rng, 5)
83                .cloned()
84                .collect::<Vec<String>>()
85                .join(" ");
86            format!("Example suggested messages: {}\n\n", examples)} else {"".to_string()})
87        + &(if let Some(c_d) = course_desc {format!("Description for course: {}\n\n", c_d)} else {"".to_string()})
88        + "The conversation so far:\n";
89
90    let used_tokens = estimate_tokens(&prompt) + estimate_tokens(USER_PROMPT);
91    let token_budget =
92        calculate_safe_token_limit(task_lm.context_size, task_lm.context_utilization);
93
94    let conversation =
95        &create_conversation_from_msgs(conversation_messages, used_tokens, token_budget)?;
96
97    let system_prompt = APIInputMessage {
98        message_type: InputItem::Message {
99            role: MessageRole::System,
100            content: MessageContent::Text(prompt + conversation),
101        },
102    };
103
104    let user_prompt = APIInputMessage {
105        message_type: InputItem::Message {
106            role: MessageRole::User,
107            content: MessageContent::Text(USER_PROMPT.to_string()),
108        },
109    };
110
111    let (params, max_output_tokens) = if model_is_thinking(task_lm.model_type) {
112        (
113            LLMRequestParams::GPTThinking(ThinkingParams { reasoning: None }),
114            Some(7000),
115        )
116    } else {
117        (
118            LLMRequestParams::GPTNonThinking(NonThinkingParams {
119                temperature: None,
120                top_p: None,
121                frequency_penalty: None,
122                presence_penalty: None,
123            }),
124            Some(4000),
125        )
126    };
127
128    let chat_request = LLMRequest {
129        input: vec![system_prompt, user_prompt],
130        model: task_lm.model.to_owned(),
131        max_output_tokens,
132        tools: vec![],
133        tool_choice: None,
134        params,
135        text: Some(RequestTextOptions {
136            verbosity: None,
137            format: Some(LLMRequestResponseFormatParam {
138                format_type: JSONType::JsonSchema,
139                name: "ChatbotNextMessageSuggestionResponse".to_string(),
140                schema: Schema {
141                    type_field: JSONType::Object,
142                    properties: HashMap::from([(
143                        "suggestions".to_string(),
144                        ArrayProperty {
145                            type_field: JSONType::Array,
146                            items: ArrayItem {
147                                type_field: JSONType::String,
148                            },
149                        },
150                    )]),
151                    required: Vec::from(["suggestions".to_string()]),
152                    additional_properties: false,
153                },
154                strict: true,
155            }),
156        }),
157    };
158
159    let completion = make_blocking_llm_request(chat_request, app_config).await?;
160
161    let completion_content: &String = &parse_text_completion(completion)?;
162    let suggestions: ChatbotNextMessageSuggestionResponse =
163        serde_json::from_str(completion_content).map_err(|_| {
164            chatbot_err!(
165                ChatbotMessageSuggestError,
166                "The message suggestion LLM returned an incorrectly formatted response."
167                    .to_string()
168            )
169        })?;
170
171    Ok(suggestions.suggestions)
172}
173
174/// Take ChatbotConversationMessages from a list until no more fit into the token budget.
175/// Transcribe the conversation's messages' content into a string.
176fn create_conversation_from_msgs(
177    conversation_messages: &[ChatbotConversationMessage],
178    mut used_tokens: i32,
179    token_budget: i32,
180) -> ChatbotResult<String> {
181    // todo empty messages?
182    let mut sorted_conversation_messages = conversation_messages.to_vec();
183    sorted_conversation_messages.sort_by_key(|el| el.order_number);
184    let conversation: Vec<ChatbotConversationMessage> = sorted_conversation_messages
185        .iter()
186        .filter_map(|el| match &el.message {
187            Message::Text(_) => Some(el.to_owned()),
188            _ => None,
189        })
190        .collect();
191
192    let conv_len = conversation.len();
193    // calculate how many messages to include in the conversation context
194    let cutoff = conversation
195        .iter()
196        .enumerate()
197        // we want to take messages starting from the newest (=last)
198        .rev()
199        .map_while(|(idx, el)| {
200            match el.message.to_owned()  {
201                Message::Text(msg) => {
202                // add this message's tokens and extra 5 tokens for newline and
203                // tag to used_tokens
204                used_tokens += msg.used_tokens + 5;}
205                Message::ToolOutput(o) => {
206                // add the tokens needed for the tool info to used_tokens
207                let s = format!("Output:\n{}\n\n", o.output); // todo: which tool call was it
208                used_tokens += estimate_tokens(&s);
209                }
210                _ => {
211                    // if there is no message or tool output, skip this element.
212                    // putting in conv_len won't affect the truncation later as we take min.
213                    return Some(conv_len);
214                }
215            }
216
217            if used_tokens > token_budget {
218                return None;
219            }
220            // include this element's index as a potential cutoff point
221            Some(idx)
222        })
223        // select the minimum index i.e. oldest message to include
224        .min()
225        .ok_or(chatbot_err!(
226            ChatbotMessageSuggestError,
227            "Failed to create context for message suggestion LLM, there were no conversation messages or none of them fit into the context."
228        ))?;
229    // cut off messages older than order_n from the conversation to keep context short
230    let conversation = conversation[cutoff..]
231        .iter()
232        .map(create_msg_string)
233        .collect::<Vec<String>>()
234        .join("");
235    if conversation.trim().is_empty() {
236        // this happens only if the conversation contains only ChatbotConversationMessages
237        // that have no message property, which should never happen in practice
238        return Err(chatbot_err!(
239            ChatbotMessageSuggestError,
240            "Failed to create context for message suggestion LLM, there were no conversation messages or no content in any messages. There should be some messages before messages are suggested."
241        ));
242    };
243
244    Ok(conversation)
245}
246
247fn create_msg_string(m: &ChatbotConversationMessage) -> String {
248    match m.message.to_owned() {
249        Message::Text(text_message) => {
250            match text_message.message_role {
251                MessageRole::Assistant => {
252                    format!("Assistant:\n{}\n\n", text_message.text)
253                }
254                MessageRole::User => {
255                    format!("Student:\n{}\n\n", text_message.text)
256                }
257                _ => "".to_string(), //todo error?
258            }
259        }
260        Message::ToolCall(tool_call) => {
261            format!(
262                "Tool call: Name: {} Arguments: {}\n\n",
263                tool_call.tool_name, tool_call.tool_arguments
264            )
265        }
266        Message::ToolOutput(tool_output) => {
267            format!("Tool output: {}\n\n", tool_output.output) // todo: get the tool name
268        }
269        Message::Reasoning(..) => "".to_string(),
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use headless_lms_models::chatbot_conversation_message_messages::ChatbotConversationMessageMessage;
276
277    use super::*;
278
279    fn get_msgs() -> Vec<ChatbotConversationMessage> {
280        vec![ChatbotConversationMessage {
281            order_number: 0,
282            message: Message::Text(ChatbotConversationMessageMessage {
283                message_role: MessageRole::System,
284                text: "You are a helpful assistant.".to_string(),
285                used_tokens: 6,
286                ..Default::default()
287            }),
288            ..Default::default()
289        },
290        ChatbotConversationMessage {
291            order_number: 0,
292            message: Message::Text(ChatbotConversationMessageMessage {
293            message_role: MessageRole::Assistant,
294            text: "Hello! What can I help you with?".to_string(),
295            used_tokens: 9,
296        ..Default::default()}),
297            ..Default::default()
298        },
299        ChatbotConversationMessage {
300            order_number: 1,
301                        message: Message::Text(ChatbotConversationMessageMessage {
302            message_role: MessageRole::User,
303
304            text: "Hi, I’m stuck on solving quadratic equations. I don’t really get when to factor and when to use the quadratic formula.".to_string(),
305            used_tokens: 26,
306            ..Default::default()}),
307            ..Default::default()
308        },
309        ChatbotConversationMessage {
310            order_number: 2,
311                        message: Message::Text(ChatbotConversationMessageMessage {
312
313            message_role: MessageRole::Assistant,
314            text: "No worries — that’s a super common question 😊
315Let’s start simple: do you remember the standard form of a quadratic equation?".to_string(),
316            used_tokens: 27,
317            ..Default::default()}),
318            ..Default::default()
319        },
320        ChatbotConversationMessage {
321            order_number: 3,
322                        message: Message::Text(ChatbotConversationMessageMessage {
323
324            message_role: MessageRole::User,
325            text: "Is it something like ax² + bx + c = 0?".to_string(),
326            used_tokens: 14,
327            ..Default::default()}),
328            ..Default::default()
329        },
330        ChatbotConversationMessage {
331            order_number: 4,
332                        message: Message::Text(ChatbotConversationMessageMessage {
333
334            message_role: MessageRole::Assistant,
335            text: "Exactly! Great start.
336Now, factoring works best when the numbers are “nice” — meaning you can easily find two numbers that multiply to ac and add to b.
337
338The quadratic formula works every time, but it can be more work.
339
340Let’s try this one:
341x² + 5x + 6 = 0
342
343Would you like to try factoring it?".to_string(),
344            used_tokens: 79,
345            ..Default::default()}),
346            ..Default::default()
347        },
348        ChatbotConversationMessage {
349            order_number: 5,
350                        message: Message::Text(ChatbotConversationMessageMessage {
351
352            message_role: MessageRole::User,
353            text: "Okay… I need two numbers that multiply to 6 and add to 5. So 2 and 3?".to_string(),
354            used_tokens: 24,
355            ..Default::default()}),
356            ..Default::default()
357        },
358        ChatbotConversationMessage {
359            order_number: 6,
360                        message: Message::Text(ChatbotConversationMessageMessage {
361
362            message_role: MessageRole::Assistant,
363            text: "Perfect! 🎯
364
365So we rewrite it as:
366(x + 2)(x + 3) = 0
367
368Now what do we do?".to_string(),
369            used_tokens: 56,
370            ..Default::default()}),
371            ..Default::default()
372        },
373        ChatbotConversationMessage {
374            order_number: 7,
375                        message: Message::Text(ChatbotConversationMessageMessage {
376
377            message_role: MessageRole::User,
378            text: "Set each part equal to zero?
379x + 2 = 0 and x + 3 = 0?".to_string(),
380            used_tokens: 23,
381            ..Default::default()}),
382            ..Default::default()
383        },
384        ChatbotConversationMessage {
385            order_number: 8,
386                        message: Message::Text(ChatbotConversationMessageMessage {
387
388            message_role: MessageRole::Assistant,
389            text: "Yes!
390
391So your solutions are:
392x = -2
393x = -3
394
395See how clean that was? That’s when factoring is ideal.
396
397Now let’s try one that’s harder:
3982x² + 3x + 1 = 0
399
400Do you want to try factoring again, or use the quadratic formula?".to_string(),
401            used_tokens: 67,
402            ..Default::default()}),
403            ..Default::default()
404        },
405        ChatbotConversationMessage {
406            order_number: 9,
407                        message: Message::Text(ChatbotConversationMessageMessage {
408
409            message_role: MessageRole::User,
410            text: "Factoring looks harder with the 2 in front… maybe quadratic formula?".to_string(),
411            used_tokens: 15,
412            ..Default::default()}),
413            ..Default::default()
414        },
415        ChatbotConversationMessage {
416            order_number: 10,
417                        message: Message::Text(ChatbotConversationMessageMessage {
418
419            message_role: MessageRole::Assistant,
420            text: "Good instinct. When a ≠ 1, the quadratic formula is often safer.
421
422The formula is:
423
424x = (-b ± √(b² - 4ac)) / 2a
425
426For this equation:
427a = 2
428b = 3
429c = 1
430
431First step: calculate the discriminant (b² - 4ac). Want to try that?".to_string(),
432            used_tokens: 80,
433            ..Default::default()}),
434            ..Default::default()
435        },
436        ChatbotConversationMessage {
437            order_number: 11,
438                        message: Message::Text(ChatbotConversationMessageMessage {
439
440            message_role: MessageRole::User,
441            text: "Okay
4423² - 4(2)(1)
443= 9 - 8
444= 1".to_string(),
445            used_tokens: 22,
446            ..Default::default()}),
447            ..Default::default()
448
449        },
450        ChatbotConversationMessage {
451            order_number: 12,
452                        message: Message::Text(ChatbotConversationMessageMessage {
453
454            message_role: MessageRole::Assistant,
455            text: "Excellent!
456
457Now plug it into the formula:
458
459x = (-3 ± √1) / 4
460
461Since √1 = 1, what are the two answers?".to_string(),
462            used_tokens: 34,
463            ..Default::default()}),
464            ..Default::default()
465        },
466        ChatbotConversationMessage {
467            order_number: 13,
468                        message: Message::Text(ChatbotConversationMessageMessage {
469
470            message_role: MessageRole::User,
471            text: "(-3 + 1)/4 = -2/4 = -1/2
472(-3 - 1)/4 = -4/4 = -1".to_string(),
473            used_tokens: 33,
474            ..Default::default()}),
475            ..Default::default()
476        },
477                ChatbotConversationMessage {
478            order_number: 14,
479                        message: Message::Text(ChatbotConversationMessageMessage {
480
481            message_role: MessageRole::Assistant,
482            text: "YES 👏
483
484So your solutions are:
485x = -1/2
486x = -1
487
488You handled that really well.
489
490Quick check:
491When would you choose factoring over the quadratic formula?".to_string(),
492            used_tokens: 40,
493            ..Default::default()}),
494            ..Default::default()
495        },
496        ChatbotConversationMessage {
497            order_number: 15,
498                        message: Message::Text(ChatbotConversationMessageMessage {
499
500            message_role: MessageRole::User,
501            text: "If the numbers are simple and easy to factor. Otherwise use the formula since it always works.".to_string(),
502            used_tokens: 19,
503            ..Default::default()}),
504            ..Default::default()
505        },
506                ChatbotConversationMessage {
507            order_number: 16,
508                        message: Message::Text(ChatbotConversationMessageMessage {
509
510            message_role: MessageRole::Assistant,
511            text: "Exactly right. You’ve got it!
512
513Want to try a challenge problem next time with completing the square?".to_string(),
514            used_tokens: 21,
515            ..Default::default()}),
516            ..Default::default()
517        },
518        ]
519    }
520
521    #[test]
522    fn create_context_it_fits() {
523        let used_tokens = 2000;
524        let token_budget = 3000;
525        let conv_msgs = get_msgs();
526
527        let convo = create_conversation_from_msgs(&conv_msgs, used_tokens, token_budget)
528            .expect("Creating conversation string failed!");
529
530        let expected_string = r#"Assistant:
531Hello! What can I help you with?
532
533Student:
534Hi, I’m stuck on solving quadratic equations. I don’t really get when to factor and when to use the quadratic formula.
535
536Assistant:
537No worries — that’s a super common question 😊
538Let’s start simple: do you remember the standard form of a quadratic equation?
539
540Student:
541Is it something like ax² + bx + c = 0?
542
543Assistant:
544Exactly! Great start.
545Now, factoring works best when the numbers are “nice” — meaning you can easily find two numbers that multiply to ac and add to b.
546
547The quadratic formula works every time, but it can be more work.
548
549Let’s try this one:
550x² + 5x + 6 = 0
551
552Would you like to try factoring it?
553
554Student:
555Okay… I need two numbers that multiply to 6 and add to 5. So 2 and 3?
556
557Assistant:
558Perfect! 🎯
559
560So we rewrite it as:
561(x + 2)(x + 3) = 0
562
563Now what do we do?
564
565Student:
566Set each part equal to zero?
567x + 2 = 0 and x + 3 = 0?
568
569Assistant:
570Yes!
571
572So your solutions are:
573x = -2
574x = -3
575
576See how clean that was? That’s when factoring is ideal.
577
578Now let’s try one that’s harder:
5792x² + 3x + 1 = 0
580
581Do you want to try factoring again, or use the quadratic formula?
582
583Student:
584Factoring looks harder with the 2 in front… maybe quadratic formula?
585
586Assistant:
587Good instinct. When a ≠ 1, the quadratic formula is often safer.
588
589The formula is:
590
591x = (-b ± √(b² - 4ac)) / 2a
592
593For this equation:
594a = 2
595b = 3
596c = 1
597
598First step: calculate the discriminant (b² - 4ac). Want to try that?
599
600Student:
601Okay
6023² - 4(2)(1)
603= 9 - 8
604= 1
605
606Assistant:
607Excellent!
608
609Now plug it into the formula:
610
611x = (-3 ± √1) / 4
612
613Since √1 = 1, what are the two answers?
614
615Student:
616(-3 + 1)/4 = -2/4 = -1/2
617(-3 - 1)/4 = -4/4 = -1
618
619Assistant:
620YES 👏
621
622So your solutions are:
623x = -1/2
624x = -1
625
626You handled that really well.
627
628Quick check:
629When would you choose factoring over the quadratic formula?
630
631Student:
632If the numbers are simple and easy to factor. Otherwise use the formula since it always works.
633
634Assistant:
635Exactly right. You’ve got it!
636
637Want to try a challenge problem next time with completing the square?
638
639"#.to_string();
640
641        assert_eq!(convo, expected_string);
642    }
643
644    #[test]
645    fn create_context_it_wont_fit() {
646        let used_tokens = 2000;
647        let token_budget = 2300;
648        let conv_msgs = get_msgs();
649
650        let convo = create_conversation_from_msgs(&conv_msgs, used_tokens, token_budget)
651            .expect("Creating conversation string failed!");
652
653        let expected_string = r#"Assistant:
654Good instinct. When a ≠ 1, the quadratic formula is often safer.
655
656The formula is:
657
658x = (-b ± √(b² - 4ac)) / 2a
659
660For this equation:
661a = 2
662b = 3
663c = 1
664
665First step: calculate the discriminant (b² - 4ac). Want to try that?
666
667Student:
668Okay
6693² - 4(2)(1)
670= 9 - 8
671= 1
672
673Assistant:
674Excellent!
675
676Now plug it into the formula:
677
678x = (-3 ± √1) / 4
679
680Since √1 = 1, what are the two answers?
681
682Student:
683(-3 + 1)/4 = -2/4 = -1/2
684(-3 - 1)/4 = -4/4 = -1
685
686Assistant:
687YES 👏
688
689So your solutions are:
690x = -1/2
691x = -1
692
693You handled that really well.
694
695Quick check:
696When would you choose factoring over the quadratic formula?
697
698Student:
699If the numbers are simple and easy to factor. Otherwise use the formula since it always works.
700
701Assistant:
702Exactly right. You’ve got it!
703
704Want to try a challenge problem next time with completing the square?
705
706"#
707        .to_string();
708        assert_eq!(convo, expected_string);
709    }
710}