feat: generate title

.sqlx/query-ea0755640e64de9eb034de9c1d94bf0f09143b2ddf200757b5e8ac53e1e8a993.json

@@ -0,0 +1,15 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "UPDATE chat.conversation SET title = $1 WHERE id = $2",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Left": [
+        "Text",
+        "Uuid"
+      ]
+    },
+    "nullable": []
+  },
+  "hash": "ea0755640e64de9eb034de9c1d94bf0f09143b2ddf200757b5e8ac53e1e8a993"
+}

src/databases/postgres/chat.rs

@@ -36,11 +36,17 @@ where
     Ok(rec.id)
 }
 
+#[derive(Debug)]
+pub enum ConversationState {
+    Existing(Uuid),
+    Created(Uuid),
+}
+
 pub async fn get_or_create_conversation(
     pool: &PgPool,
     conversation_id: Option<Uuid>,
     user_id: Uuid,
-) -> Result<Uuid, DbError> {
+) -> Result<ConversationState, DbError> {
     tracing::debug!("Testing conversation");
 
     let mut tx = pool.begin().await?;
@@ -56,12 +62,38 @@ pub async fn get_or_create_conversation(
         match conversation_id {
             Some(id) => {
                 validate_conversation(&mut *conn, id).await?;
-                id
+                ConversationState::Existing(id)
             }
-            None => create_conversation(&mut *conn, user_id).await?,
+            None => ConversationState::Created(create_conversation(&mut *conn, user_id).await?),
         }
     };
 
     tx.commit().await?;
     Ok(result)
 }
+
+pub async fn set_conversation_title(
+    pool: &PgPool,
+    conversation_id: Uuid,
+    user_id: Uuid,
+    title: &str,
+) -> Result<(), DbError> {
+    let mut tx = pool.begin().await?;
+    let conn = tx.acquire().await?;
+
+    sqlx::query(r#"SELECT set_config('app.current_user_id', $1::text, true)"#)
+        .bind(user_id.to_string())
+        .execute(&mut *conn)
+        .await?;
+
+    sqlx::query!(
+        r#"UPDATE chat.conversation SET title = $1 WHERE id = $2"#,
+        title,
+        conversation_id
+    )
+    .execute(&mut *conn)
+    .await?;
+
+    tx.commit().await?;
+    Ok(())
+}

src/routes/v1/chat.rs

@@ -1,4 +1,7 @@
-use crate::middlewares::auth::middleware::Auth;
+use crate::{
+    databases::postgres::chat::ConversationState, dto::api::BaseLLMRequest,
+    middlewares::auth::middleware::Auth, providers::ollama::client::OllamaProvider,
+};
 use axum::{
     Json,
     extract::{Extension, State},
@@ -8,8 +11,9 @@ use axum::{
     },
 };
 
-use crate::databases::postgres::chat::get_or_create_conversation;
+use crate::databases::postgres::chat::{get_or_create_conversation, set_conversation_title};
 use crate::dto::api;
+use crate::dto::api::CompletionRequest;
 use crate::providers::ollama::errors::into_http_response;
 use crate::state::app_state::AppState;
 
@@ -128,11 +132,35 @@ pub async fn chat_completions(
             body.conversation_id.is_some()
         );
 
-        let conversation_id =
+        let conversation_state =
             get_or_create_conversation(&state.postgres, body.conversation_id, auth.user_id())
                 .await
                 .map_err(crate::databases::postgres::errors::into_http_response)?;
 
+        let conversation_id = match &conversation_state {
+            ConversationState::Existing(uuid) => *uuid,
+            ConversationState::Created(uuid) => {
+                let conversation_id = *uuid;
+                let title =
+                    generate_conversation_title(&state.ollama, body.messages[0].content.as_str())
+                        .await;
+
+                if let Err(e) =
+                    set_conversation_title(&state.postgres, conversation_id, auth.user_id(), &title)
+                        .await
+                {
+                    tracing::warn!(
+                        conversation_id = %conversation_id,
+                        error = %e,
+                        "Failed to set conversation title, continuing with no title"
+                    );
+                }
+
+                tracing::debug!(conversation_id = %conversation_id, %title, "generated conversation title");
+                conversation_id
+            }
+        };
+
         tracing::debug!("Using conversation_id: {:?}", conversation_id);
     }
 
@@ -156,3 +184,24 @@ pub async fn chat_completions(
         Ok(Json(response).into_response())
     }
 }
+
+async fn generate_conversation_title(ollama: &OllamaProvider, first_message: &str) -> String {
+    let request = CompletionRequest {
+        base: BaseLLMRequest {
+            model: "llama3:latest".to_string(),
+            ..Default::default()
+        },
+        prompt: format!(
+            "Generate a short title (max 6 words) for the following chat conversation: {}",
+            first_message
+        ),
+    };
+
+    ollama
+        .completions(&request)
+        .await
+        .ok()
+        .and_then(|r| r.choices.first().map(|c| c.text.trim().to_string()))
+        .filter(|t| !t.is_empty())
+        .unwrap_or_else(|| "New Conversation".to_string()) // ← default on any error
+}