Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

This commit is contained in:
Nate Kelley 2025-04-09 14:00:50 -06:00
parent c66b1878f3 5ffaf3fc74
commit bcbe5d6f97
No known key found for this signature in database
GPG Key ID: FD90372AB8D98B4F
1 changed files with 460 additions and 88 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

446
api/libs/handlers/src/chats/post_chat_handler.rs
View File

@ -1,5 +1,6 @@
use dashmap::DashMap;
use middleware::AuthenticatedUser;
use std::collections::HashSet;
use std::{collections::HashMap, time::Instant};
use agents::{
@ -48,6 +49,15 @@ use crate::messages::types::{ChatMessage, ChatUserMessage};
use super::types::ChatWithMessages;
use tokio::sync::mpsc;
use database::types::dashboard_yml::DashboardYml;
use database::pool::PgPool;
use diesel::OptionalExtension;
use diesel_async::AsyncPgConnection; // Import PgPool
// Add imports for version history types
use database::types::version_history::{VersionContent, VersionHistory};
// Define the helper struct at the module level
#[derive(Debug, Clone)]
struct CompletedFileInfo {
@ -55,6 +65,7 @@ struct CompletedFileInfo {
file_type: String, // "metric" or "dashboard"
file_name: String,
version_number: i32,
content: String, // Added to store file content for parsing
}
// Define ThreadEvent
@ -224,13 +235,15 @@ pub async fn post_chat_handler(
for response in response_arr {
if response.get("type").map_or(false, |t| t == "file") {
// Extract version_number from response, default to 1 if not found
let asset_version_number = response.get("version_number")
let asset_version_number = response
.get("version_number")
.and_then(|v| v.as_i64())
.map(|v| v as i32)
.unwrap_or(1);
// Ensure the response id matches the asset_id
let response_id = response.get("id")
let response_id = response
.get("id")
.and_then(|id| id.as_str())
.and_then(|id_str| Uuid::parse_str(id_str).ok())
.unwrap_or(asset_id_value);
@ -244,7 +257,8 @@ pub async fn post_chat_handler(
asset_version_number,
asset_type_value,
)
.await {
.await
{
tracing::warn!("Failed to create message file association: {}", e);
}
}
@ -305,7 +319,9 @@ pub async fn post_chat_handler(
tracing::info!(
"Updated chat {} with most_recent_file_id: {}, most_recent_file_type: {}",
chat_id, asset_id_value, file_type
chat_id,
asset_id_value,
file_type
);
}
@ -378,6 +394,42 @@ pub async fn post_chat_handler(
let mut sent_initial_files = false; // Flag to track if initial files have been sent
let mut early_sent_file_messages: Vec<Value> = Vec::new(); // Store file messages sent early
// --- START: Load History and Find Context Dashboard ---
let mut context_dashboard_id: Option<Uuid> = None;
if let Some(existing_chat_id) = request.chat_id {
// Fetch the most recent message for this chat to find the last dashboard shown
let pool = get_pg_pool();
let mut conn = pool.get().await?;
let last_message_result = messages::table
.filter(messages::chat_id.eq(existing_chat_id))
.order(messages::created_at.desc())
.first::<Message>(&mut conn)
.await
.optional()?; // Use optional() to handle chats with no previous messages gracefully
if let Some(last_message) = last_message_result {
if let Ok(last_response_values) =
serde_json::from_value::<Vec<Value>>(last_message.response_messages)
{
for value in last_response_values {
if let Ok(response_msg) = serde_json::from_value::<BusterChatMessage>(value) {
if let BusterChatMessage::File { id, file_type, .. } = response_msg {
if file_type == "dashboard" {
if let Ok(uuid) = Uuid::parse_str(&id) {
context_dashboard_id = Some(uuid);
tracing::debug!("Found context dashboard ID: {}", uuid);
break; // Found the most recent dashboard
}
}
}
}
}
}
}
}
// --- END: Load History and Find Context Dashboard ---
// Process all messages from the agent
while let Ok(message_result) = rx.recv().await {
match message_result {
@ -451,36 +503,74 @@ pub async fn post_chat_handler(
// Store transformed containers BEFORE potential early file sending
// This ensures the files are based on the most up-to-date reasoning
let transformed_results = transform_message(&chat_id, &message_id, msg.clone(), tx.as_ref(), &chunk_tracker).await;
let transformed_results = transform_message(
&chat_id,
&message_id,
msg.clone(),
tx.as_ref(),
&chunk_tracker,
)
.await;
match transformed_results {
Ok(containers) => {
// Store all transformed containers first
all_transformed_containers.extend(containers.iter().map(|(c, _)| c.clone()));
all_transformed_containers
.extend(containers.iter().map(|(c, _)| c.clone()));
// --- START: Early File Sending Logic ---
// Check if this is the first text chunk and we haven't sent files yet
if !sent_initial_files {
// Look for an incoming text chunk within the *current* message `msg`
if let AgentMessage::Assistant { content: Some(_), progress: MessageProgress::InProgress, .. } = &msg {
if let AgentMessage::Assistant {
content: Some(_),
progress: MessageProgress::InProgress,
..
} = &msg
{
if let Some(tx_channel) = &tx {
// Set flag immediately to prevent re-entry
sent_initial_files = true;
// Perform filtering based on containers received SO FAR
let current_completed_files = collect_completed_files(&all_transformed_containers);
let filtered_files = apply_file_filtering_rules(&current_completed_files);
early_sent_file_messages = generate_file_response_values(&filtered_files);
let current_completed_files =
collect_completed_files(&all_transformed_containers);
// Pass context_dashboard_id and pool, await the async function
match apply_file_filtering_rules(
&current_completed_files,
context_dashboard_id,
&get_pg_pool(),
)
.await
{
Ok(filtered_files_info) => {
early_sent_file_messages =
generate_file_response_values(&filtered_files_info);
// Send the filtered file messages FIRST
for file_value in &early_sent_file_messages {
if let Ok(buster_chat_message) = serde_json::from_value::<BusterChatMessage>(file_value.clone()) {
let file_container = BusterContainer::ChatMessage(BusterChatMessageContainer {
response_message: buster_chat_message,
if let Ok(buster_chat_message) =
serde_json::from_value::<BusterChatMessage>(
file_value.clone(),
)
{
let file_container =
BusterContainer::ChatMessage(
BusterChatMessageContainer {
response_message:
buster_chat_message,
chat_id,
message_id,
});
if tx_channel.send(Ok((file_container, ThreadEvent::GeneratingResponseMessage))).await.is_err() {
},
);
if tx_channel
.send(Ok((
file_container,
ThreadEvent::GeneratingResponseMessage,
)))
.await
.is_err()
{
tracing::warn!("Client disconnected while sending early file messages");
// Setting the flag ensures we don't retry, but allows loop to continue processing other messages if needed
// Potentially break here if sending is critical: break;
@ -490,6 +580,17 @@ pub async fn post_chat_handler(
}
}
}
Err(e) => {
tracing::error!(
"Error applying file filtering rules early: {}",
e
);
// Optionally send an error over tx_channel or handle otherwise
// For now, proceed without sending early files if filtering fails
early_sent_file_messages = vec![]; // Ensure list is empty
}
}
}
}
}
// --- END: Early File Sending Logic ---
@ -497,8 +598,14 @@ pub async fn post_chat_handler(
// Now send the transformed containers for the current message
if let Some(tx_channel) = &tx {
for (container, thread_event) in containers {
if tx_channel.send(Ok((container, thread_event))).await.is_err() {
tracing::warn!("Client disconnected, but continuing to process messages");
if tx_channel
.send(Ok((container, thread_event)))
.await
.is_err()
{
tracing::warn!(
"Client disconnected, but continuing to process messages"
);
// Don't break immediately, allow storing final state
}
}
@ -750,24 +857,37 @@ async fn process_completed_files(
for msg in messages {
// Use a temporary tracker instance if needed, or reuse the main one
if let Ok(containers) = transform_message(
&message.chat_id, &message.id, msg.clone(), None, chunk_tracker
).await {
&message.chat_id,
&message.id,
msg.clone(),
None,
chunk_tracker,
)
.await
{
transformed_messages_for_db.extend(containers.into_iter().map(|(c, _)| c));
}
}
let mut processed_file_ids = std::collections::HashSet::new();
for container in transformed_messages_for_db { // Use the re-transformed messages
for container in transformed_messages_for_db {
// Use the re-transformed messages
if let BusterContainer::ReasoningMessage(msg) = container {
match &msg.reasoning {
BusterReasoningMessage::File(file) if file.message_type == "files" && file.status == "completed" => {
BusterReasoningMessage::File(file)
if file.message_type == "files" && file.status == "completed" =>
{
for (file_id_key, file_content) in &file.files {
if file_content.status == "completed" { // Ensure inner file is also complete
if file_content.status == "completed" {
// Ensure inner file is also complete
let file_uuid = match Uuid::parse_str(file_id_key) {
Ok(uuid) => uuid,
Err(_) => {
tracing::warn!("Invalid UUID format for file ID in reasoning: {}", file_id_key);
tracing::warn!(
"Invalid UUID format for file ID in reasoning: {}",
file_id_key
);
continue; // Skip this file
}
};
@ -793,12 +913,16 @@ async fn process_completed_files(
if let Err(e) = diesel::insert_into(messages_to_files::table)
.values(&message_to_file)
.execute(conn)
.await {
tracing::error!("Failed to insert message_to_file link for file {}: {}", file_uuid, e);
.await
{
tracing::error!(
"Failed to insert message_to_file link for file {}: {}",
file_uuid,
e
);
continue; // Skip chat update if DB link fails
}
// Determine file type for chat update
let file_type_for_chat = match file_content.file_type.as_str() {
"dashboard" => Some("dashboard".to_string()),
@ -815,7 +939,8 @@ async fn process_completed_files(
chats::updated_at.eq(Utc::now()),
))
.execute(conn)
.await {
.await
{
tracing::error!("Failed to update chat {} with most recent file info for {}: {}", message.chat_id, file_uuid, e);
}
}
@ -2157,11 +2282,15 @@ fn collect_completed_files(containers: &[BusterContainer]) -> Vec<CompletedFileI
if file_reasoning.message_type == "files" && file_reasoning.status == "completed" {
for (_file_id_key, file_detail) in &file_reasoning.files {
if file_detail.status == "completed" {
// Extract content, default to empty string if None
let content = file_detail.file.text.clone().unwrap_or_default();
completed_files.push(CompletedFileInfo {
id: file_detail.id.clone(),
file_type: file_detail.file_type.clone(),
file_name: file_detail.file_name.clone(),
version_number: file_detail.version_number,
content, // Populate the content field
});
}
}
@ -2172,17 +2301,260 @@ fn collect_completed_files(containers: &[BusterContainer]) -> Vec<CompletedFileI
completed_files
}
// Helper function to encapsulate filtering rules
fn apply_file_filtering_rules(completed_files: &[CompletedFileInfo]) -> Vec<CompletedFileInfo> {
let contains_metrics = completed_files.iter().any(|f| f.file_type == "metric");
let contains_dashboards = completed_files.iter().any(|f| f.file_type == "dashboard");
if contains_dashboards {
completed_files.iter().filter(|f| f.file_type == "dashboard").cloned().collect()
} else if contains_metrics {
completed_files.iter().filter(|f| f.file_type == "metric").cloned().collect()
// --- START: New Helper Function ---
// Fetches dashboard details from the database
async fn fetch_dashboard_details(
id: Uuid,
conn: &mut AsyncPgConnection,
) -> Result<Option<CompletedFileInfo>> {
match dashboard_files::table
.filter(dashboard_files::id.eq(id))
// Select id, name, and version_history
.select((dashboard_files::id, dashboard_files::name, dashboard_files::version_history))
// Adjust expected tuple type
.first::<(Uuid, String, VersionHistory)>(conn)
.await
.optional()? // Handle case where dashboard might not be found
{
Some((db_id, name, version_history)) => {
if let Some(latest_version) = version_history.get_latest_version() {
// Extract dashboard_yml content and serialize it back to string
if let VersionContent::DashboardYml(dashboard_yml) = &latest_version.content {
match serde_json::to_string(dashboard_yml) { // Serialize to JSON string
Ok(content_string) => {
Ok(Some(CompletedFileInfo {
id: db_id.to_string(),
file_type: "dashboard".to_string(),
file_name: name,
version_number: latest_version.version_number,
content: content_string, // Use serialized string
}))
}
Err(e) => {
tracing::error!("Failed to serialize DashboardYml content for {}: {}", db_id, e);
Ok(None) // Treat serialization error as unable to fetch details
}
}
} else {
vec![]
// Content was not DashboardYml, unexpected
tracing::warn!("Latest version content for dashboard {} is not DashboardYml type.", db_id);
Ok(None)
}
} else {
// Version history exists but has no versions
tracing::warn!("Version history for dashboard {} has no versions.", db_id);
Ok(None)
}
}
None => Ok(None), // Dashboard not found in the table
}
}
// --- END: New Helper Function ---
// Helper function to encapsulate filtering rules - NOW ASYNC
async fn apply_file_filtering_rules(
completed_files_this_turn: &[CompletedFileInfo],
context_dashboard_id: Option<Uuid>,
pool: &PgPool, // Pass pool
) -> Result<Vec<CompletedFileInfo>> {
// Return Result
let metrics_this_turn: Vec<_> = completed_files_this_turn
.iter()
.filter(|f| f.file_type == "metric")
.cloned()
.collect();
let dashboards_this_turn: Vec<_> = completed_files_this_turn
.iter()
.filter(|f| f.file_type == "dashboard")
.cloned()
.collect();
match context_dashboard_id {
// --- Context Exists ---
Some(ctx_id) => {
// Fetch context dashboard details once upfront
let mut conn = pool.get().await?;
let context_dashboard_info_opt = fetch_dashboard_details(ctx_id, &mut conn).await?;
// If context dashboard couldn't be fetched (e.g., deleted), treat as no context
if context_dashboard_info_opt.is_none() {
tracing::warn!(
"Context dashboard ID {} not found, falling back to current turn logic.",
ctx_id
);
return process_current_turn_files(&metrics_this_turn, &dashboards_this_turn);
}
let context_dashboard_info = context_dashboard_info_opt.unwrap(); // Safe unwrap due to check above
// Case 2 Check: Only context metrics modified, no new dashboards?
if dashboards_this_turn.is_empty() && !metrics_this_turn.is_empty() {
// Parse context dashboard to see if *all* modified metrics belong to it
let mut all_metrics_belong_to_context = true;
let context_metric_ids = match DashboardYml::new(
context_dashboard_info.content.clone(),
) {
Ok(yml) => yml
.rows
.iter()
.flat_map(|r| r.items.iter().map(|i| i.id))
.collect::<HashSet<Uuid>>(),
Err(e) => {
tracing::warn!("Failed to parse context dashboard {} for Case 2 check: {}. Assuming metrics might not belong.", ctx_id, e);
all_metrics_belong_to_context = false; // Cannot confirm, assume they don't all belong
HashSet::new()
}
};
if all_metrics_belong_to_context {
// Only proceed if context dashboard parsed
for metric in &metrics_this_turn {
if let Ok(metric_uuid) = Uuid::parse_str(&metric.id) {
if !context_metric_ids.contains(&metric_uuid) {
all_metrics_belong_to_context = false;
break;
}
}
}
}
// If all modified metrics seem to belong to the context dashboard, return only it.
if all_metrics_belong_to_context {
tracing::debug!(
"Context dashboard {} exists, only metrics belonging to it were modified. Returning context dashboard.",
ctx_id
);
return Ok(vec![context_dashboard_info]);
}
// If not all metrics belong, fall through to Case 3 logic below
}
// Case 1 & 3: New dashboard created OR complex state (metrics modified don't all belong to context).
tracing::debug!(
"New dashboard created or metrics modified don't all belong to context dashboard {}. Checking for modified context metrics.",
ctx_id
);
// Check if any metric modified *this turn* belongs to the context dashboard.
let mut modified_context_metric_this_turn = false;
let context_metric_ids = match DashboardYml::new(context_dashboard_info.content.clone())
{
Ok(yml) => yml
.rows
.iter()
.flat_map(|r| r.items.iter().map(|i| i.id))
.collect::<HashSet<Uuid>>(),
Err(e) => {
tracing::warn!("Failed to parse context dashboard {} for Case 3 check: {}. Assuming no context metrics modified.", ctx_id, e);
HashSet::new() // Assume no overlap if parsing fails
}
};
for metric in &metrics_this_turn {
if let Ok(metric_uuid) = Uuid::parse_str(&metric.id) {
if context_metric_ids.contains(&metric_uuid) {
modified_context_metric_this_turn = true;
break;
}
}
}
// If a context metric was modified AND other assets exist this turn...
if modified_context_metric_this_turn {
tracing::debug!("Context metric modified alongside other assets. Combining context dashboard with current turn processing.");
// Process current turn files
let new_filtered_assets =
process_current_turn_files(&metrics_this_turn, &dashboards_this_turn)?;
// Return context dashboard first, then the processed new assets
Ok(vec![context_dashboard_info]
.into_iter()
.chain(new_filtered_assets.into_iter())
.collect())
} else {
// No context metric modified, or context parsing failed. Process current turn only.
tracing::debug!("No context metric modified (or context parse failed). Processing current turn files only.");
process_current_turn_files(&metrics_this_turn, &dashboards_this_turn)
}
}
// --- No Context ---
None => {
// Case 1 (No context): Process current turn's files normally.
tracing::debug!("No context dashboard ID found. Processing current turn files only.");
process_current_turn_files(&metrics_this_turn, &dashboards_this_turn)
}
}
}
// Helper for the previous filtering logic (refactored, kept synchronous as it doesn't do IO)
fn process_current_turn_files(
metrics: &[CompletedFileInfo],
dashboards: &[CompletedFileInfo],
) -> Result<Vec<CompletedFileInfo>> {
// Return Result for consistency
let contains_metrics = !metrics.is_empty();
let contains_dashboards = !dashboards.is_empty();
if contains_metrics && contains_dashboards {
// Parse dashboards, find referenced metrics, filter unreferenced, combine
let mut metric_uuids = HashSet::new();
for metric in metrics {
if let Ok(uuid) = Uuid::parse_str(&metric.id) {
metric_uuids.insert(uuid);
}
}
let mut referenced_metric_uuids = HashSet::new();
for dashboard_info in dashboards {
match DashboardYml::new(dashboard_info.content.clone()) {
Ok(dashboard_yml) => {
for row in dashboard_yml.rows {
for item in row.items {
referenced_metric_uuids.insert(item.id);
}
}
}
Err(e) => {
tracing::warn!(
"Failed to parse dashboard YML content for ID {} during current turn processing: {}. Skipping for metric reference check.",
dashboard_info.id,
e
);
}
}
}
let unreferenced_metric_uuids: HashSet<_> = metric_uuids
.difference(&referenced_metric_uuids)
.copied()
.collect();
if unreferenced_metric_uuids.is_empty() {
// All metrics referenced, return only dashboards
Ok(dashboards.to_vec())
} else {
let unreferenced_metrics: Vec<_> = metrics
.iter()
.filter(|m| {
Uuid::parse_str(&m.id)
.map_or(false, |uuid| unreferenced_metric_uuids.contains(&uuid))
})
.cloned()
.collect();
// Return unreferenced metrics first, then dashboards
let mut combined = unreferenced_metrics;
combined.extend(dashboards.iter().cloned());
Ok(combined)
}
} else if contains_dashboards {
// Only dashboards
Ok(dashboards.to_vec())
} else if contains_metrics {
// Only metrics
Ok(metrics.to_vec())
} else {
// Neither
Ok(vec![])
}
}