headless_lms_chatbot/
citations.rs

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

headless_lms_chatbot/
citations.rs
