Skip to main content

headless_lms_chatbot/azure_chatbot/
search_grounding.rs

1//! Builds the system-prompt addendum appended when Azure Search is offered, grounding answers in
2//! retrieved course material and pointing the model at whichever other course-material tools this
3//! chatbot's configuration actually offers.
4
5use 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
13/// One sentence pointing the model at another tool, included only when its category is enabled.
14struct 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
40/// Only called when [`AZURE_AI_SEARCH_TOOL_NAME`] itself is in the request's tool list. Each row
41/// in [`TOOL_MENTIONS`] is appended only when its category is enabled, so the instruction never
42/// points the model at a tool this configuration doesn't actually offer.
43pub 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}