refactor: implemented state machines for retrieval pipeline, improved tracing

html-router/src/routes/account/handlers.rs

@@ -24,7 +24,10 @@ pub async fn show_account_page(
     RequireUser(user): RequireUser,
     State(state): State<HtmlState>,
 ) -> Result<impl IntoResponse, HtmlError> {
-    let timezones = TZ_VARIANTS.iter().map(std::string::ToString::to_string).collect();
+    let timezones = TZ_VARIANTS
+        .iter()
+        .map(std::string::ToString::to_string)
+        .collect();
     let conversation_archive = User::get_user_conversations(&user.id, &state.db).await?;
 
     Ok(TemplateResponse::new_template(
@@ -102,7 +105,10 @@ pub async fn update_timezone(
         ..user.clone()
     };
 
-    let timezones = TZ_VARIANTS.iter().map(std::string::ToString::to_string).collect();
+    let timezones = TZ_VARIANTS
+        .iter()
+        .map(std::string::ToString::to_string)
+        .collect();
 
     // Render the API key section block
     Ok(TemplateResponse::new_partial(

html-router/src/routes/auth/signin.rs

@@ -27,11 +27,15 @@ pub async fn show_signin_form(
     if auth.is_authenticated() {
         return Ok(TemplateResponse::redirect("/"));
     }
-    if boosted { Ok(TemplateResponse::new_partial(
-        "auth/signin_base.html",
-        "body",
-        (),
-    )) } else { Ok(TemplateResponse::new_template("auth/signin_base.html", ())) }
+    if boosted {
+        Ok(TemplateResponse::new_partial(
+            "auth/signin_base.html",
+            "body",
+            (),
+        ))
+    } else {
+        Ok(TemplateResponse::new_template("auth/signin_base.html", ()))
+    }
 }
 
 pub async fn authenticate_user(

html-router/src/routes/auth/signup.rs

@@ -29,11 +29,15 @@ pub async fn show_signup_form(
         return Ok(TemplateResponse::redirect("/"));
     }
 
-    if boosted { Ok(TemplateResponse::new_partial(
-        "auth/signup_form.html",
-        "body",
-        (),
-    )) } else { Ok(TemplateResponse::new_template("auth/signup_form.html", ())) }
+    if boosted {
+        Ok(TemplateResponse::new_partial(
+            "auth/signup_form.html",
+            "body",
+            (),
+        ))
+    } else {
+        Ok(TemplateResponse::new_template("auth/signup_form.html", ()))
+    }
 }
 
 pub async fn process_signup_and_show_verification(

html-router/src/routes/chat/message_response_stream.rs

@@ -9,11 +9,8 @@ use axum::{
     },
 };
 use composite_retrieval::{
-    answer_retrieval::{
-        create_chat_request, create_user_message_with_history, format_entities_json,
-        LLMResponseFormat,
-    },
-    retrieve_entities,
+    answer_retrieval::{create_chat_request, create_user_message_with_history, LLMResponseFormat},
+    retrieve_entities, retrieved_entities_to_json,
 };
 use futures::{
     stream::{self, once},
@@ -136,7 +133,7 @@ pub async fn get_response_stream(
     };
 
     // 3. Create the OpenAI request
-    let entities_json = format_entities_json(&entities);
+    let entities_json = retrieved_entities_to_json(&entities);
     let formatted_user_message =
         create_user_message_with_history(&entities_json, &history, &user_message.content);
     let settings = match SystemSettings::get_current(&state.db).await {
@@ -260,7 +257,11 @@ pub async fn get_response_stream(
         .chain(stream::once(async move {
             if let Some(message) = rx_final.recv().await {
                 // Don't send any event if references is empty
-                if message.references.as_ref().is_some_and(std::vec::Vec::is_empty) {
+                if message
+                    .references
+                    .as_ref()
+                    .is_some_and(std::vec::Vec::is_empty)
+                {
                     return Ok(Event::default().event("empty")); // This event won't be sent
                 }