headless_lms_chatbot/chatbot_tools/custom_tools/
document_lookup.rs

1use indexmap::IndexMap;
2
3use headless_lms_models::chatbot_configurations::ToolCategory;
4use headless_lms_models::{course_page_markdown_content, pages};
5use headless_lms_utils::{
6    document_schema_processor::remove_sensitive_attributes,
7    json_schema_types::{JSONType, JsonItem, Schema, SchemaPropertyType},
8    strings::truncate_utf8_at_boundary,
9};
10
11use crate::{
12    azure_chatbot::azure::tools::{AzureLLMFunctionToolDefinition, LLMToolType},
13    chatbot_tools::{
14        ChatbotTool, ChatbotToolDeclaration, ToolProperties,
15        argument_parsing::deserialize_to_optional_uuid_and_errors_to_none,
16        course_scope::{
17            COURSE_ID_ARGUMENT_DESCRIPTION, material_requirements, resolve_course_scope,
18        },
19        tool_authorization::ToolRequirement,
20    },
21    citations::parse_document_filepath,
22    llm_utils::estimate_tokens,
23    prelude::*,
24    user_context::ChatbotTurnContext,
25};
26
27pub type DocumentLookupTool = ToolProperties<DocumentLookupState>;
28
29pub struct DocumentLookupState {
30    document: Option<String>,
31}
32
33#[derive(Deserialize)]
34pub struct DocumentLookupArguments {
35    /// Required by the tool's schema, but the lookup resolves the document by id or filepath and
36    /// never reads this back; kept only so a call missing it fails to deserialize.
37    #[allow(dead_code)]
38    title: String,
39    filepath: Option<String>,
40    #[serde(deserialize_with = "deserialize_to_optional_uuid_and_errors_to_none")]
41    page_id: Option<Uuid>,
42    format: String,
43    #[serde(deserialize_with = "deserialize_to_optional_uuid_and_errors_to_none")]
44    course_id: Option<Uuid>,
45}
46
47/// Truncates page content until its estimated token count fits the budget we are willing to hand
48/// to the LLM. Scaling the byte length by the token ratio always shrinks the content, so this
49/// terminates in a pass or two.
50fn shorten_page_content(mut content: String) -> String {
51    const MAX_TOKENS: i32 = 25_000;
52    loop {
53        let tokens = estimate_tokens(&content);
54        if tokens <= MAX_TOKENS {
55            return content;
56        }
57        let max_bytes = content.len() * (MAX_TOKENS as usize - 1_000) / tokens as usize;
58        content = truncate_utf8_at_boundary(&content, max_bytes).to_string();
59    }
60}
61
62impl ChatbotToolDeclaration for DocumentLookupTool {
63    const NAME: &'static str = "document_lookup";
64
65    fn offer_requirements(user_context: &ChatbotTurnContext) -> Vec<ToolRequirement> {
66        material_requirements(user_context.course_id)
67    }
68
69    const CATEGORY: ToolCategory = ToolCategory::CourseMaterial;
70
71    fn get_tool_definition() -> AzureLLMFunctionToolDefinition {
72        AzureLLMFunctionToolDefinition {
73            tool_type: LLMToolType::Function,
74            name: Self::NAME.to_string(),
75            description: "Look up the full content of a specific document by the title and filepath or id (page_id). The needed arguments can be found from Azure search results or by using the course_structure tool. Either a filepath or a page_id is required to find the correct document, in addition to the document title. The document can be returned in Markdown or JSON format. The Markdown format is cleaner and preferred, but might have errors: if you suspect it's erroneous, you can request the JSON version.".to_string(),
76            parameters: Schema::strict_object(
77                IndexMap::from([
78                    (
79                        "filepath".to_string(),
80                        SchemaPropertyType::Item(JsonItem {
81                            type_field: JSONType::String,
82                            description: Some("The filepath of the document to look up, as returned from Azure search. Either the filepath or page_id is required.".to_string()),
83                        }),
84                    ),
85                    (
86                        "title".to_string(),
87                        SchemaPropertyType::Item(JsonItem {
88                            type_field: JSONType::String,
89                            description: Some("The title of the document to look up, as returned from Azure search. Optional.".to_string()),
90                        }),
91                    ),
92                    (
93                        "page_id".to_string(),
94                        SchemaPropertyType::Item(JsonItem {
95                            type_field: JSONType::String,
96                            description: Some("The page_id of the document to look up. Either page_id or the filepath is required.".to_string()),
97                        }),
98                    ),
99                    (
100                        "format".to_string(),
101                        SchemaPropertyType::Item(JsonItem {
102                            type_field: JSONType::String,
103                            description: Some("The format of the document. Optional. Valid values are 'json' and 'markdown'. Markdown content is human readable, but might have errors. ".to_string()),
104                        }),
105                    ),
106                    (
107                        "course_id".to_string(),
108                        SchemaPropertyType::Item(JsonItem {
109                            type_field: JSONType::String,
110                            description: Some(COURSE_ID_ARGUMENT_DESCRIPTION.to_string()),
111                        }),
112                    )
113                ]),
114                None,
115            ),
116            strict: true,
117        }
118    }
119}
120
121/// Look up a document (page) from the course the chatbot is on.
122impl ChatbotTool for DocumentLookupTool {
123    type Arguments = DocumentLookupArguments;
124
125    fn call_requirements(
126        arguments: &Self::Arguments,
127        user_context: &ChatbotTurnContext,
128    ) -> Vec<ToolRequirement> {
129        material_requirements(resolve_course_scope(user_context, arguments.course_id).ok())
130    }
131
132    async fn from_db_and_arguments(
133        conn: &mut PgConnection,
134        _app_config: &ApplicationConfiguration,
135        arguments: Self::Arguments,
136        user_context: &ChatbotTurnContext,
137    ) -> ChatbotResult<Self> {
138        let course_id = resolve_course_scope(user_context, arguments.course_id)?;
139
140        let page_id = if let Some(id) = &arguments.page_id {
141            id.to_owned()
142        } else if let Some(f) = &arguments.filepath {
143            let res = parse_document_filepath(f);
144            match res {
145                Ok(d) => d.page_id,
146                Err(e) => Err(chatbot_err!(
147                    InvalidToolArguments,
148                    "Couldn't parse document file path and no valid page id was provided, unable to look up document.",
149                    e
150                ))?,
151            }
152        } else {
153            return Err(chatbot_err!(
154                InvalidToolArguments,
155                format!(
156                    "Unable to call document_lookup tool. No filepath or page id provided. One of them is needed to find the document."
157                )
158            ));
159        };
160        let document = match course_page_markdown_content::get_course_page_content_by_page_id(
161            conn, page_id,
162        )
163        .await
164        {
165            // A page of another course is not the caller's to read, so it reads as not found.
166            Ok(page_content) if page_content.course_id == course_id => {
167                if arguments.format == "json" {
168                    let s =
169                        shorten_page_content(serde_json::to_string(&page_content.json_content)?);
170                    Some(s)
171                } else if let Some(content) = page_content.markdown_content {
172                    let s = shorten_page_content(content);
173                    Some(s)
174                } else {
175                    let base = "Markdown content not found. Page JSON content:\n\n".to_string();
176                    let s =
177                        shorten_page_content(serde_json::to_string(&page_content.json_content)?);
178                    Some(base + &s)
179                }
180            }
181            Ok(_) => None,
182            // No chatbot has ever synced this course's markdown, which covers most courses: fall
183            // back to the page's own blocks, sanitized the way the syncer would before indexing
184            // them, instead of reporting the document not found.
185            Err(e) if e.error_type() == &ModelErrorType::RecordNotFound => {
186                match pages::get_page(conn, page_id).await {
187                    Ok(page) if page.course_id == Some(course_id) && page.deleted_at.is_none() => {
188                        let blocks = remove_sensitive_attributes(page.blocks_cloned()?);
189                        let base = "No converted markdown exists for this course; this is raw block JSON:\n\n".to_string();
190                        let s = shorten_page_content(serde_json::to_string(&blocks)?);
191                        Some(base + &s)
192                    }
193                    _ => None,
194                }
195            }
196            Err(e) => return Err(ChatbotError::from(e)),
197        };
198
199        Ok(DocumentLookupTool {
200            state: DocumentLookupState { document },
201        })
202    }
203
204    fn output(&self) -> String {
205        if let Some(d) = &self.state.document {
206            d.to_string()
207        } else {
208            "Document not found.".to_string()
209        }
210    }
211
212    fn output_description_instructions(&self) -> Option<String> {
213        Some("Do not return the whole document to the user. Use the document as a source of more information for answering the user etc. Cite the course_material_search result the page came from; document_lookup itself produces no citation.".to_string())
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220
221    #[test]
222    fn shorten_page_content_shortens_prose() {
223        let input = "The quick brown fox jumps over the lazy dog. ".repeat(4600);
224        assert!(input.len() > 200_000);
225
226        let shortened = shorten_page_content(input);
227
228        assert!(estimate_tokens(&shortened) <= 25_000);
229    }
230
231    #[test]
232    fn shorten_page_content_shortens_punctuation_heavy_json() {
233        let input = format!(
234            "[{}]",
235            r#"{"id":"1","name":"block","attributes":{"content":"Hei, mitä kuuluu?"}},"#
236                .repeat(2500)
237        );
238        assert!(input.len() > 150_000);
239
240        let shortened = shorten_page_content(input);
241
242        assert!(estimate_tokens(&shortened) <= 25_000);
243    }
244
245    #[test]
246    fn shorten_page_content_leaves_short_content_alone() {
247        let input = "Short enough.".to_string();
248
249        assert_eq!(shorten_page_content(input.clone()), input);
250    }
251}
headless_lms_chatbot/chatbot_tools/custom_tools/document_lookup.rs

headless_lms_chatbot/chatbot_tools/custom_tools/
document_lookup.rs
