From b5ef7e4881871d75c029c78b49048e596033d49f Mon Sep 17 00:00:00 2001 From: LucasDLTG Date: Mon, 11 May 2026 18:28:04 +0200 Subject: [PATCH] feat: generate title --- ...4bf0f09143b2ddf200757b5e8ac53e1e8a993.json | 15 +++++ src/databases/postgres/chat.rs | 38 ++++++++++++- src/routes/v1/chat.rs | 55 ++++++++++++++++++- 3 files changed, 102 insertions(+), 6 deletions(-) create mode 100644 .sqlx/query-ea0755640e64de9eb034de9c1d94bf0f09143b2ddf200757b5e8ac53e1e8a993.json diff --git a/.sqlx/query-ea0755640e64de9eb034de9c1d94bf0f09143b2ddf200757b5e8ac53e1e8a993.json b/.sqlx/query-ea0755640e64de9eb034de9c1d94bf0f09143b2ddf200757b5e8ac53e1e8a993.json new file mode 100644 index 0000000..c4fc210 --- /dev/null +++ b/.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" +} diff --git a/src/databases/postgres/chat.rs b/src/databases/postgres/chat.rs index 70e7c56..2ce598a 100644 --- a/src/databases/postgres/chat.rs +++ b/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, user_id: Uuid, -) -> Result { +) -> Result { 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(()) +} diff --git a/src/routes/v1/chat.rs b/src/routes/v1/chat.rs index 581edc9..b81ad7c 100644 --- a/src/routes/v1/chat.rs +++ b/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 +}