wip: sse implementation chat

src/server/routes/html/chat/mod.rs

@@ -1,18 +1,29 @@
+use std::time::Duration;
+
 use axum::{
-    extract::State,
-    response::{IntoResponse, Redirect},
+    extract::{Path, Query, State},
+    response::{
+        sse::{Event, KeepAlive},
+        Html, IntoResponse, Redirect, Sse,
+    },
     Form,
 };
 use axum_session_auth::AuthSession;
 use axum_session_surreal::SessionSurrealPool;
+use futures::{stream, Stream, StreamExt};
 use surrealdb::{engine::any::Any, Surreal};
+use tokio::time::sleep;
 use tracing::info;
+use uuid::Uuid;
 
 use crate::{
     error::HtmlError,
     page_data,
     server::{routes::html::render_template, AppState},
-    storage::types::user::User,
+    storage::types::{
+        message::{Message, MessageRole},
+        user::User,
+    },
 };
 
 // Update your ChatStartParams struct to properly deserialize the references
@@ -36,27 +47,8 @@ where
 page_data!(ChatData, "chat/base.html", {
     user: User,
     history: Vec<Message>,
-});
-
-#[derive(Deserialize, Debug, Serialize)]
-pub enum MessageRole {
-    User,
-    AI,
-    System,
-}
-
-#[derive(Deserialize, Debug, Serialize)]
-pub struct Message {
     conversation_id: String,
-    role: MessageRole,
-    content: String,
-    references: Option<Vec<String>>,
-}
-
-pub struct Conversation {
-    user_id: String,
-    title: String,
-}
+});
 
 pub async fn show_initialized_chat(
     State(state): State<AppState>,
@@ -72,19 +64,16 @@ pub async fn show_initialized_chat(
 
     info!("{:?}", form);
 
-    let user_message = Message {
-        conversation_id: "test".to_string(),
-        role: MessageRole::User,
-        content: form.user_query,
-        references: None,
-    };
+    let conversation_id = Uuid::new_v4().to_string();
 
-    let ai_message = Message {
-        conversation_id: "test".to_string(),
-        role: MessageRole::AI,
-        content: form.llm_response,
-        references: Some(form.references),
-    };
+    let user_message = Message::new("test".to_string(), MessageRole::User, form.user_query, None);
+
+    let ai_message = Message::new(
+        "test".to_string(),
+        MessageRole::AI,
+        form.llm_response,
+        Some(form.references),
+    );
 
     let messages = vec![user_message, ai_message];
 
@@ -93,6 +82,7 @@ pub async fn show_initialized_chat(
         ChatData {
             history: messages,
             user,
+            conversation_id,
         },
         state.templates.clone(),
     )?;
@@ -111,14 +101,98 @@ pub async fn show_chat_base(
         None => return Ok(Redirect::to("/").into_response()),
     };
 
+    let conversation_id = Uuid::new_v4().to_string();
+
     let output = render_template(
         ChatData::template_name(),
         ChatData {
             history: vec![],
             user,
+            conversation_id,
         },
         state.templates.clone(),
     )?;
 
     Ok(output.into_response())
 }
+
+#[derive(Deserialize)]
+pub struct NewMessageForm {
+    content: String,
+}
+
+pub async fn new_user_message(
+    Path(conversation_id): Path<String>,
+    State(state): State<AppState>,
+    auth: AuthSession<User, String, SessionSurrealPool<Any>, Surreal<Any>>,
+    Form(form): Form<NewMessageForm>,
+) -> Result<impl IntoResponse, HtmlError> {
+    info!("Displaying empty chat start");
+
+    let user = match auth.current_user {
+        Some(user) => user,
+        None => return Ok(Redirect::to("/").into_response()),
+    };
+
+    let query_id = Uuid::new_v4().to_string();
+    let user_message = form.content.clone();
+
+    // Save to database
+    // state
+    //     .db
+    //     .save(conversation_id, query_id.clone(), user_message)
+    //     .await;
+
+    #[derive(Serialize)]
+    struct SSEResponseInitData {
+        user_message: String,
+        query_id: String,
+    }
+
+    let output = render_template(
+        "chat/streaming_response.html",
+        SSEResponseInitData {
+            user_message,
+            query_id,
+        },
+        state.templates.clone(),
+    )?;
+
+    Ok(output.into_response())
+}
+
+#[derive(Deserialize)]
+pub struct QueryParams {
+    query_id: String,
+}
+
+pub async fn get_response_stream(
+    State(_state): State<AppState>,
+    auth: AuthSession<User, String, SessionSurrealPool<Any>, Surreal<Any>>,
+    Query(params): Query<QueryParams>,
+) -> Sse<impl Stream<Item = Result<Event, axum::Error>>> {
+    let stream = stream::iter(vec![
+        Event::default()
+            .event("chat_message")
+            .data("Hello, starting stream!"),
+        Event::default()
+            .event("chat_message")
+            .data("This is message 2"),
+        Event::default().event("chat_message").data("Final message"),
+        Event::default()
+            .event("close_stream")
+            .data("Stream complete"), // Signal to close
+    ])
+    .then(|event| async move {
+        sleep(Duration::from_millis(500)).await; // Delay between messages
+        Ok(event)
+    });
+
+    info!("Streaming started");
+
+    Sse::new(stream).keep_alive(
+        axum::response::sse::KeepAlive::new()
+            .interval(Duration::from_secs(15))
+            .text("keep-alive"),
+    )
+}

src/server/routes/mod.rs

@@ -16,7 +16,7 @@ use axum_session_surreal::SessionSurrealPool;
 use html::{
     account::{delete_account, set_api_key, show_account_page, update_timezone},
     admin_panel::{show_admin_panel, toggle_registration_status},
-    chat::{show_chat_base, show_initialized_chat},
+    chat::{get_response_stream, new_user_message, show_chat_base, show_initialized_chat},
     content::{patch_text_content, show_content_page, show_text_content_edit_form},
     documentation::{
         show_documentation_index, show_get_started, show_mobile_friendly, show_privacy_policy,
@@ -65,6 +65,8 @@ pub fn html_routes(app_state: &AppState) -> Router<AppState> {
         .route("/gdpr/deny", post(deny_gdpr))
         .route("/search", get(search_result_handler))
         .route("/chat", get(show_chat_base).post(show_initialized_chat))
+        .route("/chat/:id", post(new_user_message))
+        .route("/chat/response-stream", get(get_response_stream))
         .route("/signout", get(sign_out_user))
         .route("/signin", get(show_signin_form).post(authenticate_user))
         .route(

src/storage/types/conversation.rs

@@ -0,0 +1,21 @@
+use uuid::Uuid;
+
+use crate::stored_object;
+
+stored_object!(Conversation, "conversation", {
+    user_id: String,
+    title: String
+});
+
+impl Conversation {
+    pub fn new(user_id: String, title: String) -> Self {
+        let now = Utc::now();
+        Self {
+            id: Uuid::new_v4().to_string(),
+            created_at: now,
+            updated_at: now,
+            user_id,
+            title,
+        }
+    }
+}

src/storage/types/message.rs

@@ -0,0 +1,37 @@
+use uuid::Uuid;
+
+use crate::stored_object;
+
+#[derive(Deserialize, Debug, Serialize)]
+pub enum MessageRole {
+    User,
+    AI,
+    System,
+}
+
+stored_object!(Message, "message", {
+    conversation_id: String,
+    role: MessageRole,
+    content: String,
+    references: Option<Vec<String>>
+});
+
+impl Message {
+    pub fn new(
+        conversation_id: String,
+        role: MessageRole,
+        content: String,
+        references: Option<Vec<String>>,
+    ) -> Self {
+        let now = Utc::now();
+        Self {
+            id: Uuid::new_v4().to_string(),
+            created_at: now,
+            updated_at: now,
+            conversation_id,
+            role,
+            content,
+            references,
+        }
+    }
+}

src/storage/types/mod.rs

@@ -1,10 +1,12 @@
 use axum::async_trait;
 use serde::{Deserialize, Serialize};
 pub mod analytics;
+pub mod conversation;
 pub mod file_info;
 pub mod job;
 pub mod knowledge_entity;
 pub mod knowledge_relationship;
+pub mod message;
 pub mod system_settings;
 pub mod text_chunk;
 pub mod text_content;