Skip to main content

headless_lms_models/
application_task_default_language_models.rs

1use crate::{chatbot_configurations_models::ModelType, prelude::*};
2
3#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, Copy, Type)]
4#[sqlx(type_name = "application_task", rename_all = "kebab-case")]
5pub enum ApplicationTask {
6    ContentCleaning,
7    MessageSuggestion,
8    CmsParagraphSuggestion,
9    SisuDescriptionSummary,
10}
11
12#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
13pub struct ApplicationTaskDefaultLanguageModel {
14    pub id: Uuid,
15    pub created_at: DateTime<Utc>,
16    pub updated_at: DateTime<Utc>,
17    pub deleted_at: Option<DateTime<Utc>>,
18    pub model_id: Uuid,
19    pub task: ApplicationTask,
20    pub context_utilization: f32,
21}
22
23impl Default for ApplicationTaskDefaultLanguageModel {
24    fn default() -> Self {
25        Self {
26            id: Uuid::nil(),
27            created_at: Default::default(),
28            updated_at: Default::default(),
29            deleted_at: None,
30            model_id: Default::default(),
31            task: ApplicationTask::ContentCleaning,
32            context_utilization: 0.75,
33        }
34    }
35}
36
37pub struct TaskLMSpec {
38    pub id: Uuid,
39    pub task: ApplicationTask,
40    pub context_utilization: f32,
41    pub model: String,
42    pub model_type: ModelType,
43    pub context_size: i32,
44}
45
46pub async fn insert(
47    conn: &mut PgConnection,
48    input: ApplicationTaskDefaultLanguageModel,
49) -> ModelResult<ApplicationTaskDefaultLanguageModel> {
50    let res = sqlx::query_as!(
51        ApplicationTaskDefaultLanguageModel,
52        r#"
53INSERT INTO application_task_default_language_models (model_id, task, context_utilization)
54VALUES ($1, $2, $3)
55RETURNING
56    id,
57    created_at,
58    updated_at,
59    deleted_at,
60    model_id,
61    task,
62    context_utilization
63        "#,
64        input.model_id,
65        input.task as ApplicationTask,
66        input.context_utilization
67    )
68    .fetch_one(conn)
69    .await?;
70    Ok(res)
71}
72
73pub async fn delete(conn: &mut PgConnection, id: Uuid) -> ModelResult<()> {
74    sqlx::query!(
75        r#"
76UPDATE application_task_default_language_models
77SET deleted_at = now()
78WHERE id = $1
79AND deleted_at IS NULL
80        "#,
81        id
82    )
83    .execute(conn)
84    .await?;
85    Ok(())
86}
87
88pub async fn get_for_task(
89    conn: &mut PgConnection,
90    task: ApplicationTask,
91) -> ModelResult<TaskLMSpec> {
92    let res = sqlx::query_as!(
93        TaskLMSpec,
94        r#"
95SELECT
96    a.id,
97    a.task,
98    a.context_utilization,
99    model.model,
100    model.model_type as "model_type: ModelType",
101    model.context_size
102FROM application_task_default_language_models AS a
103JOIN chatbot_configurations_models AS model ON model.id = a.model_id
104WHERE a.task = $1
105AND a.deleted_at IS NULL
106AND model.deleted_at IS NULL
107        "#,
108        task as ApplicationTask
109    )
110    .fetch_one(conn)
111    .await?;
112    Ok(res)
113}