Skip to main content

headless_lms_models/
chatbot_conversation_message_tool_calls.rs

1use serde_json::Value;
2use utoipa::ToSchema;
3
4use crate::prelude::*;
5
6#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, Copy, Type, ToSchema)]
7#[sqlx(type_name = "tool_kind", rename_all = "kebab-case")]
8#[serde(rename_all = "snake_case")]
9pub enum ToolKind {
10    Function,
11    AzureAiSearch,
12}
13
14#[derive(Clone, PartialEq, Deserialize, Serialize, Debug, ToSchema)]
15pub struct ChatbotConversationMessageToolCall {
16    pub id: Uuid,
17    pub created_at: DateTime<Utc>,
18    pub updated_at: DateTime<Utc>,
19    pub deleted_at: Option<DateTime<Utc>>,
20    pub chatbot_conversation_message_id: Uuid,
21    pub tool_name: String,
22    /// Is type = string in the OpenApi schema, because there is no known shape
23    /// for this JSON value/object so we treat it as 'string' instead of
24    /// 'unknown' in the API autogenerated types. Since it's valid JSON, we save
25    /// it as JSON in our DB.
26    #[schema(value_type = String)]
27    pub tool_arguments: Value,
28    pub tool_call_id: String,
29    pub tool_kind: ToolKind,
30    pub response_id: String,
31}
32
33impl Default for ChatbotConversationMessageToolCall {
34    fn default() -> Self {
35        Self {
36            id: Uuid::nil(),
37            created_at: Default::default(),
38            updated_at: Default::default(),
39            deleted_at: None,
40            chatbot_conversation_message_id: Uuid::nil(),
41            tool_name: Default::default(),
42            tool_arguments: Default::default(),
43            tool_call_id: Default::default(),
44            tool_kind: ToolKind::Function,
45            response_id: Default::default(),
46        }
47    }
48}
49
50pub async fn insert(
51    conn: &mut PgConnection,
52    input: ChatbotConversationMessageToolCall,
53    msg_id: Uuid,
54) -> ModelResult<ChatbotConversationMessageToolCall> {
55    let res = sqlx::query_as!(
56        ChatbotConversationMessageToolCall,
57        r#"
58INSERT INTO chatbot_conversation_message_tool_calls (
59    chatbot_conversation_message_id,
60    tool_name,
61    tool_arguments,
62    tool_call_id,
63    tool_kind,
64    response_id
65  )
66VALUES ($1, $2, $3, $4, $5, $6)
67RETURNING *
68        "#,
69        msg_id,
70        input.tool_name,
71        input.tool_arguments,
72        input.tool_call_id,
73        input.tool_kind as ToolKind,
74        input.response_id
75    )
76    .fetch_one(conn)
77    .await?;
78    Ok(res)
79}
80
81pub async fn get_by_id(
82    conn: &mut PgConnection,
83    id: Uuid,
84) -> ModelResult<ChatbotConversationMessageToolCall> {
85    let res = sqlx::query_as!(
86        ChatbotConversationMessageToolCall,
87        r#"
88SELECT *
89FROM chatbot_conversation_message_tool_calls
90WHERE id = $1
91  AND deleted_at IS NULL
92        "#,
93        id
94    )
95    .fetch_one(conn)
96    .await?;
97    Ok(res)
98}
99
100pub async fn get_by_message_id(
101    conn: &mut PgConnection,
102    msg_id: Uuid,
103) -> ModelResult<Option<ChatbotConversationMessageToolCall>> {
104    let res = sqlx::query_as!(
105        ChatbotConversationMessageToolCall,
106        r#"
107SELECT *
108FROM chatbot_conversation_message_tool_calls
109WHERE chatbot_conversation_message_id = $1
110  AND deleted_at IS NULL
111        "#,
112        msg_id
113    )
114    .fetch_optional(conn)
115    .await?;
116    Ok(res)
117}
118
119pub async fn delete(
120    conn: &mut PgConnection,
121    id: Uuid,
122) -> ModelResult<ChatbotConversationMessageToolCall> {
123    let res = sqlx::query_as!(
124        ChatbotConversationMessageToolCall,
125        r#"
126UPDATE chatbot_conversation_message_tool_calls
127SET deleted_at = NOW()
128WHERE id = $1
129  AND deleted_at IS NULL
130RETURNING *
131        "#,
132        id
133    )
134    .fetch_one(conn)
135    .await?;
136    Ok(res)
137}
138
139/// Sometimes during chatbot conversation streaming, the stream ends unexpectedly while
140/// a tool call has been made but not answered. This happens with provider tools that we
141/// can't control. In this case, the conversation is left in a state which is invalid,
142/// so we need to delete the un-answered tool call(s).
143pub async fn get_hanging_tool_calls_for_conversation(
144    conn: &mut PgConnection,
145    conversation_id: Uuid,
146) -> ModelResult<Vec<ChatbotConversationMessageToolCall>> {
147    let res = sqlx::query_as!(
148        ChatbotConversationMessageToolCall,
149        r#"
150SELECT *
151FROM chatbot_conversation_message_tool_calls AS ccmtc
152WHERE ccmtc.chatbot_conversation_message_id IN (
153    SELECT id
154    FROM chatbot_conversation_messages
155    WHERE conversation_id = $1
156      AND deleted_at IS NULL
157  )
158  AND ccmtc.deleted_at IS NULL
159  AND NOT EXISTS (
160    SELECT id
161    FROM chatbot_conversation_message_tool_outputs
162    WHERE tool_call_id = ccmtc.tool_call_id
163      AND deleted_at IS NULL
164  )
165        "#,
166        conversation_id
167    )
168    .fetch_all(conn)
169    .await?;
170    Ok(res)
171}