headless_lms_chatbot/chatbot_tools/custom_tools/
document_lookup.rs

1use std::collections::HashMap;
2use std::str::FromStr;
3
4use headless_lms_utils::strings::truncate_utf8_at_boundary;
5use serde::{Deserialize, Deserializer};
6use sqlx::PgConnection;
7use uuid::Uuid;
8
9use crate::{
10    azure_chatbot::ChatbotUserContext,
11    chatbot_tools::{
12        AzureLLMFunctionToolDefinition, ChatbotTool, LLMToolParamProperties, LLMToolParamType,
13        LLMToolParams, LLMToolType, ToolProperties,
14    },
15    citations::parse_document_filepath,
16    llm_utils::estimate_tokens,
17    prelude::{BackendError, ChatbotError, ChatbotErrorType, ChatbotResult, chatbot_err},
18};
19
20pub type DocumentLookupTool = ToolProperties<DocumentLookupState, DocumentLookupArguments>;
21
22pub struct DocumentLookupState {
23    document: Option<String>,
24}
25
26#[derive(serde::Serialize, serde::Deserialize)]
27pub struct DocumentLookupArguments {
28    title: String,
29    filepath: Option<String>,
30    #[serde(deserialize_with = "deserialize_to_optional_uuid_and_errors_to_none")]
31    page_id: Option<Uuid>,
32    format: String,
33}
34
35/// Deserializes an optional string field and parses it into an Uuid, doing this
36/// optionally without failure. If an error occurs, a None is returned. This is
37/// desired for the page_id field in DocumentLookupArguments, which is generated
38/// by an LLM and can be None or a non-Uuid string in some cases. If any errors
39/// should occur, they are emitted in ChatbotTools's from_db_and_arguments.
40fn deserialize_to_optional_uuid_and_errors_to_none<'de, D>(
41    deserializer: D,
42) -> Result<Option<Uuid>, D::Error>
43where
44    D: Deserializer<'de>,
45{
46    let res = String::deserialize(deserializer)
47        .ok()
48        .and_then(|s| Uuid::from_str(&s).ok());
49    Ok(res)
50}
51
52fn shorten_page_content(content: String) -> String {
53    let page_tokens = estimate_tokens(&content);
54    if page_tokens <= 25000 {
55        return content.to_string();
56    }
57    // a token is ~4 chars and a char is ~4 bytes
58    let max_bytes = ((page_tokens - 1000) * 4 * 4) as usize;
59
60    let shortened = truncate_utf8_at_boundary(&content, max_bytes);
61    shorten_page_content(shortened.to_string())
62}
63
64/// Look up a document (page) from the course the chatbot is on.
65impl ChatbotTool for DocumentLookupTool {
66    type State = DocumentLookupState;
67    type Arguments = DocumentLookupArguments;
68
69    fn parse_arguments(args_string: String) -> ChatbotResult<Self::Arguments> {
70        serde_json::from_str::<Self::Arguments>(&args_string).map_err(|e| {
71            chatbot_err!(
72                InvalidToolArguments,
73                format!("Couldn't parse tool arguments. Arguments: {args_string}"),
74                e
75            )
76        })
77    }
78
79    async fn from_db_and_arguments(
80        conn: &mut PgConnection,
81        mut arguments: Self::Arguments,
82        user_context: &ChatbotUserContext,
83    ) -> ChatbotResult<Self> {
84        let Some(course_id) = user_context.course_id else {
85            return Err(chatbot_err!(
86                ToolUseError,
87                "Course id is missing.".to_string()
88            ));
89        };
90
91        let page_id = if let Some(id) = &arguments.page_id {
92            id.to_owned()
93        } else if let Some(f) = &arguments.filepath {
94            let res = parse_document_filepath(f);
95            match res {
96                Ok(d) => d.page_id,
97                Err(e) => Err(chatbot_err!(
98                    InvalidToolArguments,
99                    "Couldn't parse document file path and no valid page id was provided, unable to look up document.",
100                    e
101                ))?,
102            }
103        } else {
104            return Err(chatbot_err!(
105                InvalidToolArguments,
106                format!(
107                    "Unable to call document_lookup tool. No filepath or page id provided. One of them is needed to find the document."
108                )
109            ));
110        };
111        let page_content =
112            headless_lms_models::course_page_markdown_content::get_course_page_content_by_page_id(
113                conn, page_id,
114            )
115            .await?;
116
117        let document =
118            // Check if the titles match and the page is part of the same course as
119            // the one the user is on.
120            if page_content.course_id == course_id {
121                arguments.title = page_content.title;
122                if arguments.format == "json" {
123                    let s = shorten_page_content(serde_json::to_string(&page_content.json_content)?);
124                    Some(s)
125                } else {
126                // use markdown content if there is any. else use json as string
127                if let Some(content) = page_content.markdown_content {
128                    let s = shorten_page_content(content);
129                    Some(s)
130                } else {
131                    let base = "Markdown content not found. Page JSON content:\n\n".to_string();
132                    let s = shorten_page_content(serde_json::to_string(&page_content.json_content)?);
133                    Some(base + &s)
134                }
135                }
136
137            } else {
138                None
139            };
140
141        Ok(DocumentLookupTool {
142            state: DocumentLookupState { document },
143            arguments,
144        })
145    }
146
147    fn output(&self) -> String {
148        if let Some(d) = &self.state.document {
149            d.to_string()
150        } else {
151            "Document not found.".to_string()
152        }
153    }
154
155    fn output_description_instructions(&self) -> Option<String> {
156        Some("Do not return the whole document to the user. Use the document as a source of more information for answering the user etc. If you need to cite the content of this document, cite the Azure search result of the document.".to_string())
157    }
158
159    fn get_arguments(&self) -> &Self::Arguments {
160        &self.arguments
161    }
162
163    fn get_tool_definition() -> AzureLLMFunctionToolDefinition {
164        AzureLLMFunctionToolDefinition {
165            tool_type: LLMToolType::Function,
166            name: "document_lookup".to_string(),
167            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(),
168            parameters: LLMToolParams {
169                tool_type: LLMToolParamType::Object,
170                properties: HashMap::from([
171                    (
172                        "filepath".to_string(),
173                        LLMToolParamProperties {
174                            param_type: "string".to_string(),
175                            description: "The filepath of the document to look up, as returned from Azure search. Either the filepath or page_id is required.".to_string(),
176                        },
177                    ),
178                    (
179                        "title".to_string(),
180                        LLMToolParamProperties {
181                            param_type: "string".to_string(),
182                            description: "The title of the document to look up, as returned from Azure search. Optional.".to_string(),
183                        },
184                    ),
185                    (
186                        "page_id".to_string(),
187                        LLMToolParamProperties {
188                            param_type: "string".to_string(),
189                            description: "The page_id of the document to look up. Either page_id or the filepath is required.".to_string(),
190                        },
191                    ),
192                    (
193                        "format".to_string(),
194                        LLMToolParamProperties {
195                            param_type: "string".to_string(),
196                            description: "The format of the document. Optional. Valid values are 'json' and 'markdown'. Markdown content is human readable, but might have errors. ".to_string(),
197                        },
198                    )
199                ]),
200                required: vec!["title".to_string(), "page_id".to_string(), "filepath".to_string(), "format".to_string()],
201                additional_properties: false,
202            },
203            strict: true,
204        }
205    }
206}
headless_lms_chatbot/chatbot_tools/custom_tools/document_lookup.rs

headless_lms_chatbot/chatbot_tools/custom_tools/
document_lookup.rs
