Merge branch 'evals' of https://github.com/buster-so/buster into evals

api/libs/handlers/src/chats/get_chat_handler.rs

@@ -1,5 +1,6 @@
 use anyhow::{anyhow, Result};
 use chrono::{DateTime, Utc};
+use database::enums::MessageFeedback;
 use diesel::prelude::Queryable;
 use diesel::{ExpressionMethods, JoinOnDsl, NullableExpressionMethods, QueryDsl};
 use diesel_async::RunQueryDsl;
@@ -41,6 +42,7 @@ pub struct MessageWithUser {
     pub user_id: Uuid,
     pub user_name: Option<String>,
     pub user_attributes: Value,
+    pub feedback: Option<String>,
 }
 
 #[derive(Queryable)]
@@ -122,6 +124,7 @@ pub async fn get_chat_handler(
                     users::id,
                     users::name.nullable(),
                     users::attributes,
+                    messages::feedback.nullable(),
                 ))
                 .load::<MessageWithUser>(&mut conn)
                 .await
@@ -275,6 +278,7 @@ pub async fn get_chat_handler(
                 reasoning,
                 msg.final_reasoning_message,
                 msg.created_at,
+                msg.feedback,
             )
         })
         .collect();

api/libs/handlers/src/chats/post_chat_handler.rs

@@ -298,6 +298,7 @@ pub async fn post_chat_handler(
                 vec![],
                 None,
                 message.created_at,
+                None
             );
 
             chat_with_messages.add_message(chat_message);
@@ -817,6 +818,7 @@ pub async fn post_chat_handler(
         reasoning_messages.clone(),
         Some(formatted_final_reasoning_duration.clone()), // Use formatted reasoning duration
         Utc::now(),
+        None,
     );
 
     chat_with_messages.update_message(final_message);
@@ -2442,6 +2444,7 @@ async fn initialize_chat(
             Vec::new(),
             None,
             Utc::now(),
+            None,
         );
 
         // Add message to existing chat
@@ -2481,6 +2484,7 @@ async fn initialize_chat(
             Vec::new(),
             None,
             Utc::now(),
+            None,
         );
 
         let mut chat_with_messages = ChatWithMessages::new(

api/libs/handlers/src/messages/helpers/update_message_handler.rs

@@ -4,11 +4,8 @@ use database::{pool::get_pg_pool, schema::messages};
 use diesel::prelude::*;
 use diesel_async::RunQueryDsl;
 use middleware::AuthenticatedUser;
-use std::str::FromStr;
 use uuid::Uuid;
 
-use crate::messages::types::MessageFeedback;
-
 /// Update a message with new properties
 ///
 /// # Arguments
@@ -47,15 +44,11 @@ pub async fn update_message_handler(
 
     // Add feedback if provided
     if let Some(fb_str) = feedback {
-        // Validate feedback value
-        let feedback = MessageFeedback::from_str(&fb_str)
-            .map_err(|e| anyhow!(e))?;
-        
         // Update the feedback column directly
         update_statement
             .set((
                 messages::updated_at.eq(Utc::now()),
-                messages::feedback.eq(feedback.to_string())
+                messages::feedback.eq(fb_str)
             ))
             .execute(&mut conn)
             .await?;

api/libs/handlers/src/messages/types.rs

@@ -6,7 +6,6 @@ use serde_json::Value;
 use uuid::Uuid;
 
 pub mod message_feedback;
-pub use message_feedback::MessageFeedback;
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct ChatMessage {
@@ -20,6 +19,7 @@ pub struct ChatMessage {
     pub reasoning_messages: HashMap<String, Value>,
     pub created_at: chrono::DateTime<chrono::Utc>,
     pub final_reasoning_message: Option<String>,
+    pub feedback: Option<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
@@ -51,6 +51,7 @@ impl ChatMessage {
             reasoning_messages: HashMap::new(),
             created_at: Utc::now(),
             final_reasoning_message: None,
+            feedback: None,
         }
     }
 
@@ -61,6 +62,7 @@ impl ChatMessage {
         reasoning_messages: Vec<Value>,
         final_reasoning_message: Option<String>,
         created_at: chrono::DateTime<chrono::Utc>,
+        feedback: Option<String>,
     ) -> Self {
         let response_message_ids: Vec<String> = response_messages
             .iter()
@@ -97,6 +99,7 @@ impl ChatMessage {
             reasoning_messages: reasoning_messages_map,
             created_at,
             final_reasoning_message,
+            feedback,
         }
     }
 }

api/libs/handlers/src/users/invite_user_handler.rs

@@ -2,11 +2,11 @@ use anyhow::{Context, Result};
 use chrono::Utc;
 use database::{
     self,
-    enums::{UserOrganizationRole, UserOrganizationStatus, SharingSetting},
+    enums::{SharingSetting, UserOrganizationRole, UserOrganizationStatus},
     models::{User, UserToOrganization},
+    pool::get_pg_pool,
     schema::{users, users_to_organizations},
 };
-use diesel::prelude::*;
 use diesel_async::{AsyncPgConnection, RunQueryDsl};
 use middleware::AuthenticatedUser;
 use serde_json::json;
@@ -17,7 +17,6 @@ use uuid::Uuid;
 pub async fn invite_user_handler(
     inviting_user: &AuthenticatedUser,
     emails: Vec<String>,
-    conn: &mut AsyncPgConnection, // Accept the connection directly
 ) -> Result<()> {
     let organization_id = inviting_user
         .organizations
@@ -42,7 +41,7 @@ pub async fn invite_user_handler(
 
         // 2. Create User struct instance using the generated ID and attributes
         let user_to_insert = User {
-            id: new_user_id, // Use the generated ID
+            id: new_user_id,      // Use the generated ID
             email: email.clone(), // Use the original email variable again or the cloned one
             name: None,
             config: json!({}),
@@ -52,23 +51,35 @@ pub async fn invite_user_handler(
             avatar_url: None,
         };
 
+        let mut conn = match get_pg_pool().get().await {
+            Ok(mut conn) => conn,
+            Err(e) => {
+                return Err(e.into());
+            }
+        };
+
         // 3. Insert user
-        diesel::insert_into(users::table)
+        match diesel::insert_into(users::table)
             .values(&user_to_insert)
-            .execute(conn)
+            .execute(&mut conn)
             .await
-            .context("Failed to insert new user")?;
+        {
+            Ok(_) => (),
+            Err(e) => {
+                return Err(e.into());
+            }
+        };
 
         // 4. Create UserToOrganization struct instance
         let user_org_to_insert = UserToOrganization {
             user_id: new_user_id, // Use the generated ID
             organization_id,
-            role: assigned_role, // Use the role variable
+            role: assigned_role,                   // Use the role variable
             sharing_setting: SharingSetting::None, // Default setting
-            edit_sql: false, // Default permission
-            upload_csv: false, // Default permission
-            export_assets: false, // Default permission
-            email_slack_enabled: false, // Default setting
+            edit_sql: false,                       // Default permission
+            upload_csv: false,                     // Default permission
+            export_assets: false,                  // Default permission
+            email_slack_enabled: false,            // Default setting
             created_at: now,
             updated_at: now,
             deleted_at: None,
@@ -79,11 +90,16 @@ pub async fn invite_user_handler(
         };
 
         // 5. Insert user organization mapping
-        diesel::insert_into(users_to_organizations::table)
+        match diesel::insert_into(users_to_organizations::table)
             .values(&user_org_to_insert)
-            .execute(conn)
+            .execute(&mut conn)
             .await
-            .context("Failed to map user to organization")?;
+        {
+            Ok(_) => (),
+            Err(e) => {
+                return Err(e.into());
+            }
+        };
     }
 
     Ok(())

api/libs/handlers/src/users/mod.rs

@@ -1 +1,3 @@
-pub mod invite_user_handler; 
+pub mod invite_user_handler; 
+
+pub use invite_user_handler::*;

api/server/src/routes/rest/routes/users/invite_users.rs

@@ -1,11 +1,12 @@
 use anyhow::Result;
 use axum::{Extension, Json};
+use handlers::users::invite_user_handler;
 
 use crate::routes::rest::ApiResponse;
 use axum::http::StatusCode;
-use database::enums::UserOrganizationRole;
 use middleware::AuthenticatedUser;
 use serde::Deserialize;
+use tracing::error;
 
 #[derive(Deserialize)]
 pub struct InviteUsersRequest {
@@ -16,5 +17,16 @@ pub async fn invite_users(
     Extension(user): Extension<AuthenticatedUser>,
     Json(body): Json<InviteUsersRequest>,
 ) -> Result<ApiResponse<()>, (StatusCode, &'static str)> {
-    Ok(ApiResponse::NoContent)
+    let result = invite_user_handler(&user, body.emails).await;
+
+    match result {
+        Ok(_) => Ok(ApiResponse::NoContent),
+        Err(e) => {
+            error!("Failed to invite users: {}", e);
+            Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Failed to process invitation request",
+            ))
+        }
+    }
 }