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