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