feat: log user message

src/routes/v1/chat.rs

@@ -7,11 +7,17 @@ use axum::{
     extract::{Extension, State},
     response::{
         IntoResponse, Response,
-        sse::{KeepAlive, Sse},
+        sse::{Event, KeepAlive, Sse},
     },
 };
+use sqlx::PgPool;
+use tokio_stream::StreamExt;
+use uuid::Uuid;
 
-use crate::databases::postgres::chat::{get_or_create_conversation, set_conversation_title};
+use crate::databases::postgres::chat::{
+    get_or_create_conversation, set_conversation_title, update_message_tokens,
+};
+use crate::databases::postgres::{chat, errors};
 use crate::dto::api;
 use crate::dto::api::CompletionRequest;
 use crate::providers::ollama::errors::into_http_response;
@@ -126,6 +132,7 @@ pub async fn chat_completions(
 ) -> Result<Response, (axum::http::StatusCode, String)> {
     tracing::debug!("Received /chat/completion with body {:?}", body);
 
+    // Get current conversation or create it
     let conversation_id = if matches!(&auth, Auth::Jwt(_)) {
         tracing::debug!(
             "Is conversation_id existing: {:?}",
@@ -173,16 +180,157 @@ pub async fn chat_completions(
             .chat_completions_stream(&body)
             .await
             .map_err(into_http_response)?;
-        Ok(Sse::new(stream)
-            .keep_alive(KeepAlive::default())
-            .into_response())
+
+        if let Some(conversation_id) = conversation_id {
+            let user_msg_id = log_user_message(
+                &state.postgres,
+                auth.user_id(),
+                conversation_id,
+                body.parent_id,
+                body.messages
+                    .last()
+                    .map(|m| m.content.as_str())
+                    .unwrap_or(""),
+                None, // not known at this stage
+            )
+            .await
+            .map_err(errors::into_http_response)?;
+
+            let start_event = serde_json::json!({
+                "type": "start",
+                "model": body.base.model,
+                "conversation_id": conversation_id,
+                "created": chrono::Utc::now().timestamp() as u64,
+            });
+
+            let start_stream = futures::stream::once(async move {
+                Ok::<Event, crate::providers::ollama::errors::OllamaError>(
+                    Event::default()
+                        .event("metadata")
+                        .data(start_event.to_string()),
+                )
+            });
+
+            let mut accumulated = String::new();
+
+            let wrapped_stream = stream.map(move |item| match item {
+                Ok(chunk) => {
+                    let is_done = chunk.choices[0].finish_reason == Some(api::FinishReason::Stop);
+
+                    if let Some(content) = chunk.choices[0].delta.content.as_ref() {
+                        accumulated.push_str(content);
+                    }
+
+                    if is_done {
+                        let content = accumulated.clone();
+                        let pool = state.postgres.clone();
+                        let user_id = auth.user_id();
+                        let prompt_tokens = chunk.usage.as_ref().map(|u| u.prompt_tokens as i32);
+                        let completion_tokens =
+                            chunk.usage.as_ref().map(|u| u.completion_tokens as i32);
+
+                        let end_event = serde_json::json!({
+                            "type": "end",
+                            "usage": {
+                                "prompt_tokens": prompt_tokens,
+                                "completion_tokens": completion_tokens,
+                                "total_tokens": chunk.usage.as_ref().map(|u| u.total_tokens),
+                            }
+                        });
+
+                        tokio::spawn(async move {
+                            if let Some(prompt_tokens) = prompt_tokens {
+                                let _ = update_message_tokens(
+                                    &pool,
+                                    user_id,
+                                    user_msg_id,
+                                    prompt_tokens,
+                                )
+                                .await;
+                            }
+
+                            let _ = log_assistant_message(
+                                &pool,
+                                user_id,
+                                conversation_id,
+                                user_msg_id,
+                                &content,
+                                completion_tokens,
+                            )
+                            .await;
+                        });
+
+                        return Ok(Event::default()
+                            .event("metadata")
+                            .data(end_event.to_string()));
+                    }
+
+                    let data = serde_json::to_string(&chunk).unwrap_or_default();
+                    Ok(Event::default().data(data))
+                }
+                Err(e) => Err(e),
+            });
+
+            let full_stream = start_stream.chain(wrapped_stream);
+
+            Ok(Sse::new(full_stream)
+                .keep_alive(KeepAlive::default())
+                .into_response())
+        } else {
+            let plain_stream = stream.map(
+                |item| -> Result<Event, crate::providers::ollama::errors::OllamaError> {
+                    match item {
+                        Ok(chunk) => {
+                            let data = serde_json::to_string(&chunk).unwrap_or_default();
+                            Ok(Event::default().data(data))
+                        }
+                        Err(e) => Err(e),
+                    }
+                },
+            );
+
+            Ok(Sse::new(plain_stream)
+                .keep_alive(KeepAlive::default())
+                .into_response())
+        }
     } else {
         let mut response = state
             .ollama
             .chat_completions(&body)
             .await
             .map_err(into_http_response)?;
+
         response.conversation_id = conversation_id;
+
+        if let Some(conversation_id) = conversation_id {
+            // log user message, get its id as parent for assistant
+            let user_msg_id = log_user_message(
+                &state.postgres,
+                auth.user_id(),
+                conversation_id,
+                body.parent_id,
+                body.messages
+                    .last()
+                    .map(|m| m.content.as_str())
+                    .unwrap_or(""),
+                response.usage.map(|u| u.prompt_tokens as i32),
+            )
+            .await
+            .map_err(errors::into_http_response)?;
+
+            // log assistant response
+            log_assistant_message(
+                &state.postgres,
+                auth.user_id(),
+                conversation_id,
+                user_msg_id,
+                &response.choices[0].message.content,
+                response.usage.map(|u| u.completion_tokens as i32),
+            )
+            .await
+            .map_err(errors::into_http_response)?;
+        }
+
         Ok(Json(response).into_response())
     }
 }
@@ -207,3 +355,43 @@ async fn generate_conversation_title(ollama: &OllamaProvider, first_message: &st
         .filter(|t| !t.is_empty())
         .unwrap_or_else(|| "New Conversation".to_string()) // ← default on any error
 }
+
+async fn log_user_message(
+    pool: &PgPool,
+    user_id: Uuid,
+    conversation_id: Uuid,
+    parent_id: Option<Uuid>,
+    content: &str,
+    tokens: Option<i32>,
+) -> Result<Uuid, errors::DbError> {
+    chat::insert_message(
+        pool,
+        user_id,
+        conversation_id,
+        parent_id,
+        chat::MessageRole::User,
+        content,
+        tokens,
+    )
+    .await
+}
+
+async fn log_assistant_message(
+    pool: &PgPool,
+    user_id: Uuid,
+    conversation_id: Uuid,
+    parent_id: Uuid,
+    content: &str,
+    tokens: Option<i32>,
+) -> Result<Uuid, errors::DbError> {
+    chat::insert_message(
+        pool,
+        user_id,
+        conversation_id,
+        Some(parent_id),
+        chat::MessageRole::Assistant,
+        content,
+        tokens,
+    )
+    .await
+}