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    let conv_id = input[0].conversation_id;
75    let cm_ids: Vec<Uuid> = input.iter().map(|x| x.conversation_message_id).collect();
76    let titles: Vec<String> = input.iter().map(|x| x.title.to_owned()).collect();
77    let contents: Vec<String> = input.iter().map(|x| x.content.to_owned()).collect();
78    let document_urls: Vec<String> = input.iter().map(|x| x.document_url.to_owned()).collect();
79    let citation_numbers: Vec<i32> = input.iter().map(|x| x.citation_number).collect();
80
81    let res = sqlx::query_as!(
82        ChatbotConversationMessageCitation,
83        r#"
84INSERT INTO chatbot_conversation_messages_citations (
85    conversation_id,
86    conversation_message_id,
87    title,
88    content,
89    document_url,
90    citation_number,
91    course_material_chapter_number
92  )
93SELECT $1,
94  input.cm_id,
95  input.title,
96  input.content,
97  input.document_url,
98  input.citation_number,
99  c.chapter_number
100FROM (
101    SELECT UNNEST($2::UUID []) cm_id,
102      UNNEST($3::TEXT []) title,
103      UNNEST($4::TEXT []) content,
104      UNNEST($5::TEXT []) document_url,
105      UNNEST($6::INTEGER []) citation_number,
106      UNNEST($7::UUID []) page_id
107  ) AS input
108  JOIN pages p ON p.id = input.page_id
109  LEFT JOIN chapters c ON p.chapter_id = c.id
110WHERE c.deleted_at IS NULL
111  AND p.deleted_at IS NULL
112RETURNING *
113        "#,
114        conv_id,
115        &cm_ids,
116        &titles,
117        &contents,
118        &document_urls,
119        &citation_numbers,
120        &page_ids as _,
121    )
122    .fetch_all(conn)
123    .await?;
124    Ok(res)
125}
126
127pub async fn get_by_message_id(
128    conn: &mut PgConnection,
129    message_id: Uuid,
130) -> ModelResult<Vec<ChatbotConversationMessageCitation>> {
131    let res = sqlx::query_as!(
132        ChatbotConversationMessageCitation,
133        r#"
134SELECT * FROM chatbot_conversation_messages_citations
135WHERE conversation_message_id = $1
136AND deleted_at IS NULL
137        "#,
138        message_id
139    )
140    .fetch_all(conn)
141    .await?;
142    Ok(res)
143}
144
145pub async fn get_by_conversation_id(
146    conn: &mut PgConnection,
147    conversation_id: Uuid,
148) -> ModelResult<Vec<ChatbotConversationMessageCitation>> {
149    let res = sqlx::query_as!(
150        ChatbotConversationMessageCitation,
151        r#"
152SELECT * FROM chatbot_conversation_messages_citations
153WHERE conversation_id = $1
154AND deleted_at IS NULL
155        "#,
156        conversation_id
157    )
158    .fetch_all(conn)
159    .await?;
160    Ok(res)
161}
162
163/// Sets the correct conversation_message_id to citations. The correct id is the
164/// id of the chatbot text message that uses the citations. Update the citations
165/// that currently connected to a conversation_message that contains tool output
166/// and has the same response_id.
167pub async fn update_citation_message_ids(
168    conn: &mut PgConnection,
169    response_id: String,
170    conversation_message_id: Uuid,
171) -> ModelResult<Vec<ChatbotConversationMessageCitation>> {
172    let res = sqlx::query_as!(
173        ChatbotConversationMessageCitation,
174        r#"
175UPDATE chatbot_conversation_messages_citations
176SET conversation_message_id = $1
177WHERE conversation_message_id IN (
178    SELECT id
179    FROM chatbot_conversation_messages
180    WHERE id IN (
181        SELECT conversation_message_id
182        FROM chatbot_conversation_message_tool_outputs
183        WHERE response_id = $2
184          AND deleted_at IS NULL
185      )
186      AND deleted_at IS NULL
187  )
188RETURNING *
189        "#,
190        conversation_message_id,
191        response_id
192    )
193    .fetch_all(conn)
194    .await?;
195    Ok(res)
196}
headless_lms_models/chatbot_conversation_messages_citations.rs

headless_lms_models/
chatbot_conversation_messages_citations.rs
