headless_lms_chatbot/
citations.rs

1use std::path::PathBuf;
2
3use secrecy::SecretString;
4
5use crate::{llm_utils::build_llm_headers, prelude::*};
6
7use headless_lms_models::chatbot_conversation_messages_citations::{
8    self, ChatbotConversationMessageCitation,
9};
10use headless_lms_utils::strings::truncate_utf8_at_boundary;
11use headless_lms_utils::url_encoding::url_decode;
12use reqwest::Response;
13use serde::{Deserialize, Serialize};
14use tracing::{error, instrument, trace};
15use url::Url;
16
17#[derive(Serialize, Deserialize, Debug, Clone)]
18pub struct CourseMaterialDocument {
19    pub chunk_id: String,
20    pub chunk: String,
21    pub title: String,
22    pub url: String,
23    pub filepath: String,
24}
25
26pub struct DocumentProperties {
27    pub page_id: Uuid,
28}
29
30/// Parse the filepath of a document from the Azure search index and return the page_id of
31/// the document. The page id is the same as the id of the page in our DB.
32pub fn parse_document_filepath(filepath: &str) -> ChatbotResult<DocumentProperties> {
33    let mut page_path = PathBuf::from(filepath);
34    page_path.set_extension("");
35    let page_id_str = page_path.file_name().ok_or(chatbot_err!(
36        ToolUseError,
37        "Failed to parse document filepath"
38    ))?;
39    let page_id = Uuid::parse_str(page_id_str.to_string_lossy().as_ref()).map_err(|_| {
40        chatbot_err!(
41            ToolUseError,
42            format!("Failed to parse document page id: {:?}", page_id_str)
43        )
44    })?;
45
46    Ok(DocumentProperties { page_id })
47}
48
49impl CourseMaterialDocument {
50    /// Converts the document to citation. Returns also the page_id of the cited document
51    /// so we can get the correct chapter_number later.
52    pub fn to_chatbot_conversation_message_citation(
53        &self,
54        conversation_message_id: Uuid,
55        conversation_id: Uuid,
56        citation_number: i32,
57    ) -> ChatbotResult<(ChatbotConversationMessageCitation, Option<Uuid>)> {
58        // Shorten the content if needed
59        let content = if self.chunk.len() < 255 {
60            self.chunk.clone()
61        } else {
62            truncate_utf8_at_boundary(&self.chunk, 255).to_string()
63        };
64
65        // The title and URL come from Azure Blob Storage metadata, which was URL-encoded
66        // (percent-encoded) because Azure Blob Storage metadata values must be ASCII-only.
67        // We decode them back to their original UTF-8 strings before storing in the database.
68        let decoded_title = url_decode(&self.title)?;
69        let decoded_url = url_decode(&self.url)?;
70
71        // Get the page id
72        let page_id = parse_document_filepath(&self.filepath)
73            .ok()
74            .map(|x| x.page_id);
75        Ok((
76            ChatbotConversationMessageCitation {
77                conversation_message_id,
78                conversation_id,
79                title: decoded_title,
80                content,
81                document_url: decoded_url,
82                citation_number,
83                ..Default::default()
84            },
85            page_id,
86        ))
87    }
88}
89
90/// Get documents cited by the chatbot from the search index and save them
91/// as chatbot_conversation_message_citations into the database
92pub async fn chatbot_cited_documents_to_citations(
93    conn: &mut PgConnection,
94    test_chatbot: bool,
95    mut document_urls: Vec<Url>,
96    api_key: &SecretString,
97    conversation_message_id: Uuid,
98    conversation_id: Uuid,
99) -> ChatbotResult<Vec<ChatbotConversationMessageCitation>> {
100    let mut documents: Vec<(CourseMaterialDocument, i32)> = vec![];
101    for (idx, url) in document_urls.iter_mut().enumerate() {
102        let document = get_course_material_document(url, api_key).await?;
103        let citation_number = idx as i32;
104        documents.push((document, citation_number));
105    }
106    let res = save_documents(
107        conn,
108        test_chatbot,
109        documents,
110        conversation_message_id,
111        conversation_id,
112    )
113    .await?;
114
115    Ok(res)
116}
117
118/// Get a document from the search index with a LLM-provided get url
119async fn get_course_material_document(
120    endpoint: &mut Url,
121    api_key: &SecretString,
122) -> ChatbotResult<CourseMaterialDocument> {
123    endpoint.set_query(Some(
124        "api-version=2024-07-01&$select=chunk_id,parent_id,chunk,title,url,filepath,course_id",
125    ));
126    let headers = build_llm_headers(api_key)?;
127
128    let response = REQWEST_CLIENT
129        .get(endpoint.clone())
130        .headers(headers)
131        .send()
132        .await?;
133
134    process_course_material_document_response(response).await
135}
136
137#[instrument(skip(response), fields(status = %response.status()))]
138async fn process_course_material_document_response(
139    response: Response,
140) -> ChatbotResult<CourseMaterialDocument> {
141    if !response.status().is_success() {
142        let status = response.status();
143        let error_text = response.text().await?;
144        error!(
145            status = %status,
146            error = %error_text,
147            "Error fetching document from search index."
148        );
149        return Err(chatbot_err!(
150            FailedAzureResponse,
151            format!(
152                "Error fetching document from search index: Status: {}. Error: {}",
153                status, error_text
154            )
155        ));
156    }
157
158    trace!("Processing successful LLM response");
159    // Parse the response
160    let document: CourseMaterialDocument = response.json().await?;
161
162    Ok(document)
163}
164
165/// Save a course material document into the database as a citation
166async fn save_documents(
167    conn: &mut PgConnection,
168    test_chatbot: bool,
169    documents_with_citation_numbers: Vec<(CourseMaterialDocument, i32)>,
170    conversation_message_id: Uuid,
171    conversation_id: Uuid,
172) -> ChatbotResult<Vec<ChatbotConversationMessageCitation>> {
173    let (citations, page_ids): (Vec<ChatbotConversationMessageCitation>, Vec<Option<Uuid>>) =
174        documents_with_citation_numbers
175            .iter()
176            .map(|(d, citation_number)| {
177                d.to_chatbot_conversation_message_citation(
178                    conversation_message_id,
179                    conversation_id,
180                    citation_number.to_owned(),
181                )
182            })
183            .collect::<ChatbotResult<Vec<(ChatbotConversationMessageCitation, Option<Uuid>)>>>()?
184            .into_iter()
185            .unzip();
186    if test_chatbot {
187        return save_documents_mock(conn, citations).await;
188    };
189    let res =
190        chatbot_conversation_messages_citations::insert_batch(conn, citations, page_ids).await?;
191
192    Ok(res)
193}
194
195async fn save_documents_mock(
196    conn: &mut PgConnection,
197    citations: Vec<ChatbotConversationMessageCitation>,
198) -> ChatbotResult<Vec<ChatbotConversationMessageCitation>> {
199    let mut res = vec![];
200    for input in citations {
201        let a = chatbot_conversation_messages_citations::insert(conn, input).await?;
202        res.push(a)
203    }
204    Ok(res)
205}
headless_lms_chatbot/citations.rs

headless_lms_chatbot/
citations.rs
