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, JsonItem, LLMRequest, LLMRequestParams,
6        LLMRequestResponseFormatParam, NonThinkingParams, RequestTextOptions, Schema,
7        SchemaPropertyType, 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        parallel_tool_calls: None,
135        params,
136        text: Some(RequestTextOptions {
137            verbosity: None,
138            format: Some(LLMRequestResponseFormatParam {
139                format_type: JSONType::JsonSchema,
140                name: "ChatbotNextMessageSuggestionResponse".to_string(),
141                schema: Schema {
142                    type_field: JSONType::Object,
143                    properties: HashMap::from([(
144                        "suggestions".to_string(),
145                        SchemaPropertyType::ArrayProperty(ArrayProperty {
146                            type_field: JSONType::Array,
147                            items: ArrayItem::JsonItem(JsonItem {
148                                type_field: JSONType::String,
149                            }),
150                        }),
151                    )]),
152                    required: vec!["suggestions".to_string()],
153                    additional_properties: false,
154                },
155                strict: true,
156            }),
157        }),
158    };
159
160    let completion = make_blocking_llm_request(chat_request, app_config).await?;
161
162    let completion_content: &String = &parse_text_completion(completion)?;
163    let suggestions: ChatbotNextMessageSuggestionResponse =
164        serde_json::from_str(completion_content).map_err(|_| {
165            chatbot_err!(
166                ChatbotMessageSuggestError,
167                "The message suggestion LLM returned an incorrectly formatted response."
168                    .to_string()
169            )
170        })?;
171
172    Ok(suggestions.suggestions)
173}
174
175/// Take ChatbotConversationMessages from a list until no more fit into the token budget.
176/// Transcribe the conversation's messages' content into a string.
177fn create_conversation_from_msgs(
178    conversation_messages: &[ChatbotConversationMessage],
179    mut used_tokens: i32,
180    token_budget: i32,
181) -> ChatbotResult<String> {
182    // todo empty messages?
183    let mut sorted_conversation_messages = conversation_messages.to_vec();
184    sorted_conversation_messages.sort_by_key(|el| el.order_number);
185    let conversation: Vec<ChatbotConversationMessage> = sorted_conversation_messages
186        .iter()
187        .filter_map(|el| match &el.message {
188            Message::Text(_) => Some(el.to_owned()),
189            _ => None,
190        })
191        .collect();
192
193    let conv_len = conversation.len();
194    // calculate how many messages to include in the conversation context
195    let cutoff = conversation
196        .iter()
197        .enumerate()
198        // we want to take messages starting from the newest (=last)
199        .rev()
200        .map_while(|(idx, el)| {
201            match el.message.to_owned()  {
202                Message::Text(msg) => {
203                // add this message's tokens and extra 5 tokens for newline and
204                // tag to used_tokens
205                used_tokens += msg.used_tokens + 5;}
206                Message::ToolOutput(o) => {
207                // add the tokens needed for the tool info to used_tokens
208                let s = format!("Output:\n{}\n\n", o.output); // todo: which tool call was it
209                used_tokens += estimate_tokens(&s);
210                }
211                _ => {
212                    // if there is no message or tool output, skip this element.
213                    // putting in conv_len won't affect the truncation later as we take min.
214                    return Some(conv_len);
215                }
216            }
217
218            if used_tokens > token_budget {
219                return None;
220            }
221            // include this element's index as a potential cutoff point
222            Some(idx)
223        })
224        // select the minimum index i.e. oldest message to include
225        .min()
226        .ok_or(chatbot_err!(
227            ChatbotMessageSuggestError,
228            "Failed to create context for message suggestion LLM, there were no conversation messages or none of them fit into the context."
229        ))?;
230    // cut off messages older than order_n from the conversation to keep context short
231    let conversation = conversation[cutoff..]
232        .iter()
233        .map(create_msg_string)
234        .collect::<Vec<String>>()
235        .join("");
236    if conversation.trim().is_empty() {
237        // this happens only if the conversation contains only ChatbotConversationMessages
238        // that have no message property, which should never happen in practice
239        return Err(chatbot_err!(
240            ChatbotMessageSuggestError,
241            "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."
242        ));
243    };
244
245    Ok(conversation)
246}
247
248fn create_msg_string(m: &ChatbotConversationMessage) -> String {
249    match m.message.to_owned() {
250        Message::Text(text_message) => {
251            match text_message.message_role {
252                MessageRole::Assistant => {
253                    format!("Assistant:\n{}\n\n", text_message.text)
254                }
255                MessageRole::User => {
256                    format!("Student:\n{}\n\n", text_message.text)
257                }
258                _ => "".to_string(), //todo error?
259            }
260        }
261        Message::ToolCall(tool_call) => {
262            format!(
263                "Tool call: Name: {} Arguments: {}\n\n",
264                tool_call.tool_name, tool_call.tool_arguments
265            )
266        }
267        Message::ToolOutput(tool_output) => {
268            format!("Tool output: {}\n\n", tool_output.output) // todo: get the tool name
269        }
270        Message::Reasoning(..) => "".to_string(),
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use headless_lms_models::chatbot_conversation_message_messages::ChatbotConversationMessageMessage;
277
278    use super::*;
279
280    fn get_msgs() -> Vec<ChatbotConversationMessage> {
281        vec![ChatbotConversationMessage {
282            order_number: 0,
283            message: Message::Text(ChatbotConversationMessageMessage {
284                message_role: MessageRole::System,
285                text: "You are a helpful assistant.".to_string(),
286                used_tokens: 6,
287                ..Default::default()
288            }),
289            ..Default::default()
290        },
291        ChatbotConversationMessage {
292            order_number: 0,
293            message: Message::Text(ChatbotConversationMessageMessage {
294            message_role: MessageRole::Assistant,
295            text: "Hello! What can I help you with?".to_string(),
296            used_tokens: 9,
297        ..Default::default()}),
298            ..Default::default()
299        },
300        ChatbotConversationMessage {
301            order_number: 1,
302                        message: Message::Text(ChatbotConversationMessageMessage {
303            message_role: MessageRole::User,
304
305            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(),
306            used_tokens: 26,
307            ..Default::default()}),
308            ..Default::default()
309        },
310        ChatbotConversationMessage {
311            order_number: 2,
312                        message: Message::Text(ChatbotConversationMessageMessage {
313
314            message_role: MessageRole::Assistant,
315            text: "No worries — that’s a super common question 😊
316Let’s start simple: do you remember the standard form of a quadratic equation?".to_string(),
317            used_tokens: 27,
318            ..Default::default()}),
319            ..Default::default()
320        },
321        ChatbotConversationMessage {
322            order_number: 3,
323                        message: Message::Text(ChatbotConversationMessageMessage {
324
325            message_role: MessageRole::User,
326            text: "Is it something like ax² + bx + c = 0?".to_string(),
327            used_tokens: 14,
328            ..Default::default()}),
329            ..Default::default()
330        },
331        ChatbotConversationMessage {
332            order_number: 4,
333                        message: Message::Text(ChatbotConversationMessageMessage {
334
335            message_role: MessageRole::Assistant,
336            text: "Exactly! Great start.
337Now, factoring works best when the numbers are “nice” — meaning you can easily find two numbers that multiply to ac and add to b.
338
339The quadratic formula works every time, but it can be more work.
340
341Let’s try this one:
342x² + 5x + 6 = 0
343
344Would you like to try factoring it?".to_string(),
345            used_tokens: 79,
346            ..Default::default()}),
347            ..Default::default()
348        },
349        ChatbotConversationMessage {
350            order_number: 5,
351                        message: Message::Text(ChatbotConversationMessageMessage {
352
353            message_role: MessageRole::User,
354            text: "Okay… I need two numbers that multiply to 6 and add to 5. So 2 and 3?".to_string(),
355            used_tokens: 24,
356            ..Default::default()}),
357            ..Default::default()
358        },
359        ChatbotConversationMessage {
360            order_number: 6,
361                        message: Message::Text(ChatbotConversationMessageMessage {
362
363            message_role: MessageRole::Assistant,
364            text: "Perfect! 🎯
365
366So we rewrite it as:
367(x + 2)(x + 3) = 0
368
369Now what do we do?".to_string(),
370            used_tokens: 56,
371            ..Default::default()}),
372            ..Default::default()
373        },
374        ChatbotConversationMessage {
375            order_number: 7,
376                        message: Message::Text(ChatbotConversationMessageMessage {
377
378            message_role: MessageRole::User,
379            text: "Set each part equal to zero?
380x + 2 = 0 and x + 3 = 0?".to_string(),
381            used_tokens: 23,
382            ..Default::default()}),
383            ..Default::default()
384        },
385        ChatbotConversationMessage {
386            order_number: 8,
387                        message: Message::Text(ChatbotConversationMessageMessage {
388
389            message_role: MessageRole::Assistant,
390            text: "Yes!
391
392So your solutions are:
393x = -2
394x = -3
395
396See how clean that was? That’s when factoring is ideal.
397
398Now let’s try one that’s harder:
3992x² + 3x + 1 = 0
400
401Do you want to try factoring again, or use the quadratic formula?".to_string(),
402            used_tokens: 67,
403            ..Default::default()}),
404            ..Default::default()
405        },
406        ChatbotConversationMessage {
407            order_number: 9,
408                        message: Message::Text(ChatbotConversationMessageMessage {
409
410            message_role: MessageRole::User,
411            text: "Factoring looks harder with the 2 in front… maybe quadratic formula?".to_string(),
412            used_tokens: 15,
413            ..Default::default()}),
414            ..Default::default()
415        },
416        ChatbotConversationMessage {
417            order_number: 10,
418                        message: Message::Text(ChatbotConversationMessageMessage {
419
420            message_role: MessageRole::Assistant,
421            text: "Good instinct. When a ≠ 1, the quadratic formula is often safer.
422
423The formula is:
424
425x = (-b ± √(b² - 4ac)) / 2a
426
427For this equation:
428a = 2
429b = 3
430c = 1
431
432First step: calculate the discriminant (b² - 4ac). Want to try that?".to_string(),
433            used_tokens: 80,
434            ..Default::default()}),
435            ..Default::default()
436        },
437        ChatbotConversationMessage {
438            order_number: 11,
439                        message: Message::Text(ChatbotConversationMessageMessage {
440
441            message_role: MessageRole::User,
442            text: "Okay
4433² - 4(2)(1)
444= 9 - 8
445= 1".to_string(),
446            used_tokens: 22,
447            ..Default::default()}),
448            ..Default::default()
449
450        },
451        ChatbotConversationMessage {
452            order_number: 12,
453                        message: Message::Text(ChatbotConversationMessageMessage {
454
455            message_role: MessageRole::Assistant,
456            text: "Excellent!
457
458Now plug it into the formula:
459
460x = (-3 ± √1) / 4
461
462Since √1 = 1, what are the two answers?".to_string(),
463            used_tokens: 34,
464            ..Default::default()}),
465            ..Default::default()
466        },
467        ChatbotConversationMessage {
468            order_number: 13,
469                        message: Message::Text(ChatbotConversationMessageMessage {
470
471            message_role: MessageRole::User,
472            text: "(-3 + 1)/4 = -2/4 = -1/2
473(-3 - 1)/4 = -4/4 = -1".to_string(),
474            used_tokens: 33,
475            ..Default::default()}),
476            ..Default::default()
477        },
478                ChatbotConversationMessage {
479            order_number: 14,
480                        message: Message::Text(ChatbotConversationMessageMessage {
481
482            message_role: MessageRole::Assistant,
483            text: "YES 👏
484
485So your solutions are:
486x = -1/2
487x = -1
488
489You handled that really well.
490
491Quick check:
492When would you choose factoring over the quadratic formula?".to_string(),
493            used_tokens: 40,
494            ..Default::default()}),
495            ..Default::default()
496        },
497        ChatbotConversationMessage {
498            order_number: 15,
499                        message: Message::Text(ChatbotConversationMessageMessage {
500
501            message_role: MessageRole::User,
502            text: "If the numbers are simple and easy to factor. Otherwise use the formula since it always works.".to_string(),
503            used_tokens: 19,
504            ..Default::default()}),
505            ..Default::default()
506        },
507                ChatbotConversationMessage {
508            order_number: 16,
509                        message: Message::Text(ChatbotConversationMessageMessage {
510
511            message_role: MessageRole::Assistant,
512            text: "Exactly right. You’ve got it!
513
514Want to try a challenge problem next time with completing the square?".to_string(),
515            used_tokens: 21,
516            ..Default::default()}),
517            ..Default::default()
518        },
519        ]
520    }
521
522    #[test]
523    fn create_context_it_fits() {
524        let used_tokens = 2000;
525        let token_budget = 3000;
526        let conv_msgs = get_msgs();
527
528        let convo = create_conversation_from_msgs(&conv_msgs, used_tokens, token_budget)
529            .expect("Creating conversation string failed!");
530
531        let expected_string = r#"Assistant:
532Hello! What can I help you with?
533
534Student:
535Hi, I’m stuck on solving quadratic equations. I don’t really get when to factor and when to use the quadratic formula.
536
537Assistant:
538No worries — that’s a super common question 😊
539Let’s start simple: do you remember the standard form of a quadratic equation?
540
541Student:
542Is it something like ax² + bx + c = 0?
543
544Assistant:
545Exactly! Great start.
546Now, factoring works best when the numbers are “nice” — meaning you can easily find two numbers that multiply to ac and add to b.
547
548The quadratic formula works every time, but it can be more work.
549
550Let’s try this one:
551x² + 5x + 6 = 0
552
553Would you like to try factoring it?
554
555Student:
556Okay… I need two numbers that multiply to 6 and add to 5. So 2 and 3?
557
558Assistant:
559Perfect! 🎯
560
561So we rewrite it as:
562(x + 2)(x + 3) = 0
563
564Now what do we do?
565
566Student:
567Set each part equal to zero?
568x + 2 = 0 and x + 3 = 0?
569
570Assistant:
571Yes!
572
573So your solutions are:
574x = -2
575x = -3
576
577See how clean that was? That’s when factoring is ideal.
578
579Now let’s try one that’s harder:
5802x² + 3x + 1 = 0
581
582Do you want to try factoring again, or use the quadratic formula?
583
584Student:
585Factoring looks harder with the 2 in front… maybe quadratic formula?
586
587Assistant:
588Good instinct. When a ≠ 1, the quadratic formula is often safer.
589
590The formula is:
591
592x = (-b ± √(b² - 4ac)) / 2a
593
594For this equation:
595a = 2
596b = 3
597c = 1
598
599First step: calculate the discriminant (b² - 4ac). Want to try that?
600
601Student:
602Okay
6033² - 4(2)(1)
604= 9 - 8
605= 1
606
607Assistant:
608Excellent!
609
610Now plug it into the formula:
611
612x = (-3 ± √1) / 4
613
614Since √1 = 1, what are the two answers?
615
616Student:
617(-3 + 1)/4 = -2/4 = -1/2
618(-3 - 1)/4 = -4/4 = -1
619
620Assistant:
621YES 👏
622
623So your solutions are:
624x = -1/2
625x = -1
626
627You handled that really well.
628
629Quick check:
630When would you choose factoring over the quadratic formula?
631
632Student:
633If the numbers are simple and easy to factor. Otherwise use the formula since it always works.
634
635Assistant:
636Exactly right. You’ve got it!
637
638Want to try a challenge problem next time with completing the square?
639
640"#.to_string();
641
642        assert_eq!(convo, expected_string);
643    }
644
645    #[test]
646    fn create_context_it_wont_fit() {
647        let used_tokens = 2000;
648        let token_budget = 2300;
649        let conv_msgs = get_msgs();
650
651        let convo = create_conversation_from_msgs(&conv_msgs, used_tokens, token_budget)
652            .expect("Creating conversation string failed!");
653
654        let expected_string = r#"Assistant:
655Good instinct. When a ≠ 1, the quadratic formula is often safer.
656
657The formula is:
658
659x = (-b ± √(b² - 4ac)) / 2a
660
661For this equation:
662a = 2
663b = 3
664c = 1
665
666First step: calculate the discriminant (b² - 4ac). Want to try that?
667
668Student:
669Okay
6703² - 4(2)(1)
671= 9 - 8
672= 1
673
674Assistant:
675Excellent!
676
677Now plug it into the formula:
678
679x = (-3 ± √1) / 4
680
681Since √1 = 1, what are the two answers?
682
683Student:
684(-3 + 1)/4 = -2/4 = -1/2
685(-3 - 1)/4 = -4/4 = -1
686
687Assistant:
688YES 👏
689
690So your solutions are:
691x = -1/2
692x = -1
693
694You handled that really well.
695
696Quick check:
697When would you choose factoring over the quadratic formula?
698
699Student:
700If the numbers are simple and easy to factor. Otherwise use the formula since it always works.
701
702Assistant:
703Exactly right. You’ve got it!
704
705Want to try a challenge problem next time with completing the square?
706
707"#
708        .to_string();
709        assert_eq!(convo, expected_string);
710    }
711}