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