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