headless_lms_chatbot/azure_chatbot/
search_grounding.rs1use crate::chatbot_tools::ChatbotToolDeclaration;
6use crate::chatbot_tools::custom_tools::course_structure::CourseStructureTool;
7use crate::chatbot_tools::custom_tools::document_lookup::DocumentLookupTool;
8use crate::chatbot_tools::tool_category::EnabledToolCategories;
9use headless_lms_models::chatbot_configurations::ToolCategory;
10
11use super::azure::tools::AZURE_AI_SEARCH_TOOL_NAME;
12
13struct ToolMention {
15 category: ToolCategory,
16 sentence: fn() -> String,
17}
18
19const TOOL_MENTIONS: &[ToolMention] = &[
20 ToolMention {
21 category: ToolCategory::CourseMaterial,
22 sentence: || {
23 format!(
24 " If you need more information about a specific document or a topic covered in it, use the {} tool to retrieve the full document.",
25 DocumentLookupTool::NAME
26 )
27 },
28 },
29 ToolMention {
30 category: ToolCategory::CourseInfo,
31 sentence: || {
32 format!(
33 " If you need more information about the course, like what pages and chapters are in it, use the {} tool.",
34 CourseStructureTool::NAME
35 )
36 },
37 },
38];
39
40pub fn build_search_grounding_instruction(
44 enabled_tool_categories: &EnabledToolCategories,
45) -> String {
46 let mut instruction = format!(
47 "\n\nSearch the course material with the {AZURE_AI_SEARCH_TOOL_NAME} tool before answering, and ground your answer in the results with citations. Put only what you want to find in the query; the search is already limited to this course, so don't include the course name. Searching more than once is fine when it helps — to cover distinct sub-questions or angles, to refine when the first results don't answer, or when a follow-up or new instruction needs material you don't already have. When one search already answers, stop there."
48 );
49 for mention in TOOL_MENTIONS {
50 if enabled_tool_categories.contains(mention.category) {
51 instruction.push_str(&(mention.sentence)());
52 }
53 }
54 instruction.push_str(
55 " Skip searching only for messages that don't need course material, like greetings or thanks.",
56 );
57 instruction
58}
59
60#[cfg(test)]
61mod tests {
62 use super::*;
63
64 #[test]
65 fn mentions_no_tools_when_no_categories_enabled() {
66 let instruction = build_search_grounding_instruction(&EnabledToolCategories::only(&[]));
67 assert!(instruction.contains("Search the course material"));
68 assert!(instruction.contains("Skip searching only for messages"));
69 assert!(!instruction.contains(DocumentLookupTool::NAME));
70 assert!(!instruction.contains(CourseStructureTool::NAME));
71 }
72
73 #[test]
74 fn mentions_document_lookup_when_course_material_enabled() {
75 let instruction = build_search_grounding_instruction(&EnabledToolCategories::only(&[
76 ToolCategory::CourseMaterial,
77 ]));
78 assert!(instruction.contains(DocumentLookupTool::NAME));
79 assert!(!instruction.contains(CourseStructureTool::NAME));
80 }
81
82 #[test]
83 fn mentions_course_structure_when_course_info_enabled() {
84 let instruction = build_search_grounding_instruction(&EnabledToolCategories::only(&[
85 ToolCategory::CourseInfo,
86 ]));
87 assert!(!instruction.contains(DocumentLookupTool::NAME));
88 assert!(instruction.contains(CourseStructureTool::NAME));
89 }
90
91 #[test]
92 fn mentions_both_tools_in_original_order_when_both_enabled() {
93 let instruction = build_search_grounding_instruction(&EnabledToolCategories::only(&[
94 ToolCategory::CourseMaterial,
95 ToolCategory::CourseInfo,
96 ]));
97 let document_lookup_index = instruction
98 .find(DocumentLookupTool::NAME)
99 .expect("document_lookup mention missing");
100 let course_structure_index = instruction
101 .find(CourseStructureTool::NAME)
102 .expect("course_structure mention missing");
103 assert!(document_lookup_index < course_structure_index);
104 }
105}