Skip to main content

headless_lms_models/
chatbot_conversation_messages.rs

1use utoipa::ToSchema;
2
3use crate::{
4    chatbot_conversation_message_messages::{self, ChatbotConversationMessageMessage},
5    chatbot_conversation_message_reasoning::{self, ChatbotConversationMessageReasoning},
6    chatbot_conversation_message_tool_calls::{self, ChatbotConversationMessageToolCall},
7    chatbot_conversation_message_tool_outputs::{self, ChatbotConversationMessageToolOutput},
8    prelude::*,
9};
10
11#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
12pub struct ChatbotConversationMessageRow {
13    pub id: Uuid,
14    pub created_at: DateTime<Utc>,
15    pub updated_at: DateTime<Utc>,
16    pub deleted_at: Option<DateTime<Utc>>,
17    pub conversation_id: Uuid,
18    pub order_number: i32,
19}
20
21#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema)]
22#[serde(untagged)]
23pub enum Message {
24    Text(ChatbotConversationMessageMessage),
25    ToolCall(ChatbotConversationMessageToolCall),
26    ToolOutput(ChatbotConversationMessageToolOutput),
27    Reasoning(ChatbotConversationMessageReasoning),
28}
29
30#[derive(Clone, PartialEq, Deserialize, Serialize, Debug, ToSchema)]
31pub struct ChatbotConversationMessage {
32    pub id: Uuid,
33    pub created_at: DateTime<Utc>,
34    pub updated_at: DateTime<Utc>,
35    pub deleted_at: Option<DateTime<Utc>>,
36    pub conversation_id: Uuid,
37    pub order_number: i32,
38    pub message: Message,
39}
40
41impl Default for ChatbotConversationMessage {
42    fn default() -> Self {
43        Self {
44            id: Uuid::nil(),
45            created_at: Default::default(),
46            updated_at: Default::default(),
47            deleted_at: None,
48            conversation_id: Uuid::nil(),
49            order_number: Default::default(),
50            message: Message::Text(ChatbotConversationMessageMessage::default()),
51        }
52    }
53}
54
55impl ChatbotConversationMessage {
56    pub fn from_row(r: ChatbotConversationMessageRow, m: Message) -> Self {
57        ChatbotConversationMessage {
58            id: r.id,
59            created_at: r.created_at,
60            updated_at: r.updated_at,
61            deleted_at: r.deleted_at,
62            conversation_id: r.conversation_id,
63            order_number: r.order_number,
64            message: m,
65        }
66    }
67}
68
69pub async fn insert(
70    conn: &mut PgConnection,
71    input: ChatbotConversationMessage,
72) -> ModelResult<ChatbotConversationMessage> {
73    let mut tx = conn.begin().await?;
74    let msg = sqlx::query_as!(
75        ChatbotConversationMessageRow,
76        r#"
77INSERT INTO chatbot_conversation_messages (conversation_id, order_number)
78VALUES (
79    $1,
80    COALESCE((
81      SELECT order_number
82      FROM chatbot_conversation_messages
83      WHERE conversation_id = $1
84        AND deleted_at IS NULL
85      ORDER BY order_number DESC
86      LIMIT 1
87    ), 0) + 1
88  )
89RETURNING *
90        "#,
91        input.conversation_id,
92    )
93    .fetch_one(&mut *tx)
94    .await?;
95
96    let inner = match input.message {
97        Message::Text(message) => {
98            let res =
99                chatbot_conversation_message_messages::insert(&mut tx, message, msg.id).await?;
100            Message::Text(res)
101        }
102        Message::ToolCall(tool_call) => {
103            let res =
104                chatbot_conversation_message_tool_calls::insert(&mut tx, tool_call, msg.id).await?;
105            Message::ToolCall(res)
106        }
107        Message::ToolOutput(tool_output) => {
108            let res =
109                chatbot_conversation_message_tool_outputs::insert(&mut tx, tool_output, msg.id)
110                    .await?;
111            Message::ToolOutput(res)
112        }
113        Message::Reasoning(reasoning) => {
114            let res =
115                chatbot_conversation_message_reasoning::insert(&mut tx, reasoning, msg.id).await?;
116            Message::Reasoning(res)
117        }
118    };
119
120    let res = ChatbotConversationMessage::from_row(msg, inner);
121    tx.commit().await?;
122    Ok(res)
123}
124
125// todo
126pub async fn insert_for_conversation_user_and_configuration(
127    conn: &mut PgConnection,
128    input: ChatbotConversationMessage,
129    user_id: Uuid,
130    chatbot_configuration_id: Uuid,
131) -> ModelResult<ChatbotConversationMessage> {
132    let mut tx = conn.begin().await?;
133
134    sqlx::query!(
135        r#"
136SELECT id
137FROM chatbot_conversations
138WHERE id = $1
139  AND user_id = $2
140  AND chatbot_configuration_id = $3
141  AND deleted_at IS NULL
142        "#,
143        input.conversation_id,
144        user_id,
145        chatbot_configuration_id
146    )
147    .fetch_one(&mut *tx)
148    .await?;
149
150    let msg = sqlx::query_as!(
151        ChatbotConversationMessageRow,
152        r#"
153INSERT INTO chatbot_conversation_messages (
154    conversation_id,
155    order_number
156)
157VALUES (
158    $1,
159    COALESCE((
160      SELECT order_number
161      FROM chatbot_conversation_messages
162      WHERE conversation_id = $1
163        AND deleted_at IS NULL
164      ORDER BY order_number DESC
165      LIMIT 1
166    ), 0) + 1
167)
168RETURNING *
169        "#,
170        input.conversation_id,
171    )
172    .fetch_one(&mut *tx)
173    .await?;
174
175    let inner = match input.message {
176        Message::Text(message) => {
177            let res =
178                chatbot_conversation_message_messages::insert(&mut tx, message, msg.id).await?;
179            Message::Text(res)
180        }
181        Message::ToolCall(tool_call) => {
182            let res =
183                chatbot_conversation_message_tool_calls::insert(&mut tx, tool_call, msg.id).await?;
184            Message::ToolCall(res)
185        }
186        Message::ToolOutput(tool_output) => {
187            let res =
188                chatbot_conversation_message_tool_outputs::insert(&mut tx, tool_output, msg.id)
189                    .await?;
190            Message::ToolOutput(res)
191        }
192        Message::Reasoning(reasoning) => {
193            let res =
194                chatbot_conversation_message_reasoning::insert(&mut tx, reasoning, msg.id).await?;
195            Message::Reasoning(res)
196        }
197    };
198
199    let res = ChatbotConversationMessage::from_row(msg, inner);
200    tx.commit().await?;
201    Ok(res)
202}
203
204pub async fn get_by_conversation_id(
205    conn: &mut PgConnection,
206    conversation_id: Uuid,
207) -> ModelResult<Vec<ChatbotConversationMessage>> {
208    let mut tx = conn.begin().await?;
209    let mut msgs: Vec<ChatbotConversationMessageRow> = sqlx::query_as!(
210        ChatbotConversationMessageRow,
211        r#"
212SELECT *
213FROM chatbot_conversation_messages
214WHERE conversation_id = $1
215AND deleted_at IS NULL
216        "#,
217        conversation_id
218    )
219    .fetch_all(&mut *tx)
220    .await?;
221    // Should have the same order as in the conversation.
222    msgs.sort_by_key(|a| a.order_number);
223    let mut res = vec![];
224    for m in msgs {
225        let msg = message_row_to_message(&mut tx, m).await?;
226        res.push(msg);
227    }
228    tx.commit().await?;
229    Ok(res)
230}
231
232pub async fn delete(conn: &mut PgConnection, id: Uuid) -> ModelResult<ChatbotConversationMessage> {
233    let mut tx = conn.begin().await?;
234
235    let row = sqlx::query_as!(
236        ChatbotConversationMessageRow,
237        r#"
238UPDATE chatbot_conversation_messages
239SET deleted_at = NOW()
240WHERE id = $1
241  AND deleted_at IS NULL
242RETURNING *
243        "#,
244        id
245    )
246    .fetch_one(&mut *tx)
247    .await?;
248
249    // delete the child
250    let child = delete_message_fields(&mut tx, row.id).await?;
251
252    let res = ChatbotConversationMessage::from_row(row, child);
253    tx.commit().await?;
254    Ok(res)
255}
256
257/// Sometimes during chatbot conversation streaming, the stream ends unexpectedly while a
258/// tool call has been made but not answered. This happens also with provider tools that
259/// we can't control. In this case, the conversation is left in a state which is invalid,
260/// so we need to answer the un-answered tool calls to inform of the failure and continue
261/// the conversation.
262pub async fn answer_hanging_tool_call_messages_for_conversation(
263    conn: &mut PgConnection,
264    conversation_id: Uuid,
265) -> ModelResult<Vec<ChatbotConversationMessage>> {
266    let mut res = vec![];
267
268    let hanging_children =
269        chatbot_conversation_message_tool_calls::get_hanging_tool_calls_for_conversation(
270            conn,
271            conversation_id,
272        )
273        .await?;
274
275    if !hanging_children.is_empty() {
276        for tool_call in hanging_children.into_iter() {
277            let tool_output = ChatbotConversationMessage {
278                conversation_id,
279                message: Message::ToolOutput(ChatbotConversationMessageToolOutput {
280                    output: "Unexpected error encountered, tool call aborted.".to_string(),
281                    tool_call_id: tool_call.tool_call_id,
282                    tool_kind: tool_call.tool_kind,
283                    response_id: tool_call.response_id,
284                    ..Default::default()
285                }),
286                ..Default::default()
287            };
288            let inserted = insert(conn, tool_output).await?;
289            res.push(inserted);
290        }
291    }
292
293    Ok(res)
294}
295
296pub async fn update(
297    conn: &mut PgConnection,
298    id: Uuid,
299    text: &str,
300    message_is_complete: bool,
301    used_tokens: i32,
302) -> ModelResult<ChatbotConversationMessage> {
303    let mut tx = conn.begin().await?;
304
305    let row = sqlx::query_as!(
306        ChatbotConversationMessageRow,
307        r#"
308UPDATE chatbot_conversation_messages
309SET updated_at = NOW()
310WHERE id = $1
311  AND deleted_at IS NULL
312RETURNING *
313        "#,
314        id
315    )
316    .fetch_one(&mut *tx)
317    .await?;
318
319    // update the parent
320    let child = chatbot_conversation_message_messages::update(
321        &mut tx,
322        row.id,
323        text,
324        message_is_complete,
325        used_tokens,
326    )
327    .await?;
328
329    let res = ChatbotConversationMessage::from_row(row, Message::Text(child));
330    tx.commit().await?;
331
332    Ok(res)
333}
334
335pub async fn message_row_to_message(
336    conn: &mut PgConnection,
337    row: ChatbotConversationMessageRow,
338) -> ModelResult<ChatbotConversationMessage> {
339    let inner_message = get_message_fields(conn, row.id).await?;
340    let res = ChatbotConversationMessage::from_row(row, inner_message);
341    Ok(res)
342}
343
344pub async fn get_message_fields(conn: &mut PgConnection, message_id: Uuid) -> ModelResult<Message> {
345    if let Some(message) =
346        chatbot_conversation_message_messages::get_by_message_id(conn, message_id).await?
347    {
348        Ok(Message::Text(message))
349    } else if let Some(tool_call) =
350        chatbot_conversation_message_tool_calls::get_by_message_id(conn, message_id).await?
351    {
352        Ok(Message::ToolCall(tool_call))
353    } else if let Some(tool_output) =
354        chatbot_conversation_message_tool_outputs::get_by_message_id(conn, message_id).await?
355    {
356        Ok(Message::ToolOutput(tool_output))
357    } else if let Some(reasoning) =
358        chatbot_conversation_message_reasoning::get_by_message_id(conn, message_id).await?
359    {
360        Ok(Message::Reasoning(reasoning))
361    } else {
362        Err(ModelError::new(
363            ModelErrorType::RecordNotFound,
364            "No inner message found for this ChatbotConversationMessage",
365            None,
366        ))
367    }
368}
369
370pub async fn delete_message_fields(
371    conn: &mut PgConnection,
372    message_id: Uuid,
373) -> ModelResult<Message> {
374    if let Some(message) =
375        chatbot_conversation_message_messages::get_by_message_id(conn, message_id).await?
376    {
377        let res = chatbot_conversation_message_messages::delete(conn, message.id).await?;
378        Ok(Message::Text(res))
379    } else if let Some(tool_call) =
380        chatbot_conversation_message_tool_calls::get_by_message_id(conn, message_id).await?
381    {
382        let res = chatbot_conversation_message_tool_calls::delete(conn, tool_call.id).await?;
383        Ok(Message::ToolCall(res))
384    } else if let Some(tool_output) =
385        chatbot_conversation_message_tool_outputs::get_by_message_id(conn, message_id).await?
386    {
387        let res = chatbot_conversation_message_tool_outputs::delete(conn, tool_output.id).await?;
388        Ok(Message::ToolOutput(res))
389    } else if let Some(reasoning) =
390        chatbot_conversation_message_reasoning::get_by_message_id(conn, message_id).await?
391    {
392        let res = chatbot_conversation_message_reasoning::delete(conn, reasoning.id).await?;
393        Ok(Message::Reasoning(res))
394    } else {
395        Err(ModelError::new(
396            ModelErrorType::RecordNotFound,
397            "No inner message found for this ChatbotConversationMessage",
398            None,
399        ))
400    }
401}