headless_lms_models/
chatbot_conversation_messages_citations.rs

1use crate::prelude::*;
2use utoipa::ToSchema;
3
4#[derive(Clone, PartialEq, Deserialize, Serialize, ToSchema)]
5pub struct ChatbotConversationMessageCitation {
6    pub id: Uuid,
7    pub created_at: DateTime<Utc>,
8    pub updated_at: DateTime<Utc>,
9    pub deleted_at: Option<DateTime<Utc>>,
10    pub conversation_message_id: Uuid,
11    pub conversation_id: Uuid,
12    pub course_material_chapter_number: Option<i32>,
13    pub title: String,
14    pub content: String,
15    pub document_url: String,
16    pub citation_number: i32,
17}
18
19impl Default for ChatbotConversationMessageCitation {
20    fn default() -> Self {
21        Self {
22            id: Uuid::nil(),
23            created_at: Default::default(),
24            updated_at: Default::default(),
25            deleted_at: None,
26            conversation_message_id: Uuid::nil(),
27            conversation_id: Uuid::nil(),
28            course_material_chapter_number: None,
29            title: Default::default(),
30            content: Default::default(),
31            document_url: Default::default(),
32            citation_number: Default::default(),
33        }
34    }
35}
36
37pub async fn insert(
38    conn: &mut PgConnection,
39    input: ChatbotConversationMessageCitation,
40) -> ModelResult<ChatbotConversationMessageCitation> {
41    let res = sqlx::query_as!(
42        ChatbotConversationMessageCitation,
43        r#"
44INSERT INTO chatbot_conversation_messages_citations (
45  conversation_message_id,
46  conversation_id,
47  course_material_chapter_number,
48  title,
49  content,
50  document_url,
51  citation_number)
52VALUES ($1, $2, $3, $4, $5, $6, $7)
53RETURNING *
54        "#,
55        input.conversation_message_id,
56        input.conversation_id,
57        input.course_material_chapter_number,
58        input.title,
59        input.content,
60        input.document_url,
61        input.citation_number
62    )
63    .fetch_one(conn)
64    .await?;
65    Ok(res)
66}
67
68/// Insert a batch of citation from the same conversation
69pub async fn insert_batch(
70    conn: &mut PgConnection,
71    input: Vec<ChatbotConversationMessageCitation>,
72    page_ids: Vec<Option<Uuid>>,
73) -> ModelResult<Vec<ChatbotConversationMessageCitation>> {
74    if input.is_empty() {
75        return Ok(vec![]);
76    }
77    let conv_id = input[0].conversation_id;
78    let cm_ids: Vec<Uuid> = input.iter().map(|x| x.conversation_message_id).collect();
79    let titles: Vec<String> = input.iter().map(|x| x.title.to_owned()).collect();
80    let contents: Vec<String> = input.iter().map(|x| x.content.to_owned()).collect();
81    let document_urls: Vec<String> = input.iter().map(|x| x.document_url.to_owned()).collect();
82    let citation_numbers: Vec<i32> = input.iter().map(|x| x.citation_number).collect();
83
84    let res = sqlx::query_as!(
85        ChatbotConversationMessageCitation,
86        r#"
87INSERT INTO chatbot_conversation_messages_citations (
88    conversation_id,
89    conversation_message_id,
90    title,
91    content,
92    document_url,
93    citation_number,
94    course_material_chapter_number
95  )
96SELECT $1,
97  input.cm_id,
98  input.title,
99  input.content,
100  input.document_url,
101  input.citation_number,
102  c.chapter_number
103FROM (
104    SELECT UNNEST($2::UUID []) cm_id,
105      UNNEST($3::TEXT []) title,
106      UNNEST($4::TEXT []) content,
107      UNNEST($5::TEXT []) document_url,
108      UNNEST($6::INTEGER []) citation_number,
109      UNNEST($7::UUID []) page_id
110  ) AS input
111  JOIN pages p ON p.id = input.page_id
112  LEFT JOIN chapters c ON p.chapter_id = c.id
113WHERE c.deleted_at IS NULL
114  AND p.deleted_at IS NULL
115RETURNING *
116        "#,
117        conv_id,
118        &cm_ids,
119        &titles,
120        &contents,
121        &document_urls,
122        &citation_numbers,
123        &page_ids as _,
124    )
125    .fetch_all(conn)
126    .await?;
127    Ok(res)
128}
129
130pub async fn get_by_message_id(
131    conn: &mut PgConnection,
132    message_id: Uuid,
133) -> ModelResult<Vec<ChatbotConversationMessageCitation>> {
134    let res = sqlx::query_as!(
135        ChatbotConversationMessageCitation,
136        r#"
137SELECT * FROM chatbot_conversation_messages_citations
138WHERE conversation_message_id = $1
139AND deleted_at IS NULL
140        "#,
141        message_id
142    )
143    .fetch_all(conn)
144    .await?;
145    Ok(res)
146}
147
148pub async fn get_by_conversation_id(
149    conn: &mut PgConnection,
150    conversation_id: Uuid,
151) -> ModelResult<Vec<ChatbotConversationMessageCitation>> {
152    let res = sqlx::query_as!(
153        ChatbotConversationMessageCitation,
154        r#"
155SELECT * FROM chatbot_conversation_messages_citations
156WHERE conversation_id = $1
157AND deleted_at IS NULL
158        "#,
159        conversation_id
160    )
161    .fetch_all(conn)
162    .await?;
163    Ok(res)
164}
165
166/// Sets the correct conversation_message_id to citations. The correct id is the
167/// id of the chatbot text message that uses the citations. Update the citations
168/// that currently connected to a conversation_message that contains tool output
169/// and has the same response_id.
170pub async fn update_citation_message_ids(
171    conn: &mut PgConnection,
172    response_id: String,
173    conversation_message_id: Uuid,
174) -> ModelResult<Vec<ChatbotConversationMessageCitation>> {
175    let res = sqlx::query_as!(
176        ChatbotConversationMessageCitation,
177        r#"
178UPDATE chatbot_conversation_messages_citations
179SET conversation_message_id = $1
180WHERE conversation_message_id IN (
181    SELECT id
182    FROM chatbot_conversation_messages
183    WHERE id IN (
184        SELECT conversation_message_id
185        FROM chatbot_conversation_message_tool_outputs
186        WHERE response_id = $2
187          AND deleted_at IS NULL
188      )
189      AND deleted_at IS NULL
190  )
191RETURNING *
192        "#,
193        conversation_message_id,
194        response_id
195    )
196    .fetch_all(conn)
197    .await?;
198    Ok(res)
199}
headless_lms_models/chatbot_conversation_messages_citations.rs

headless_lms_models/
chatbot_conversation_messages_citations.rs
