diff --git a/api/libs/agents/src/tools/categories/file_tools/common.rs b/api/libs/agents/src/tools/categories/file_tools/common.rs index cf3ec7a4f..4e1eb15ad 100644 --- a/api/libs/agents/src/tools/categories/file_tools/common.rs +++ b/api/libs/agents/src/tools/categories/file_tools/common.rs @@ -98,7 +98,7 @@ pub async fn validate_metric_ids(ids: &[Uuid]) -> Result> { } pub const METRIC_YML_SCHEMA: &str = r##" -# METRIC CONFIGURATION - YAML STRUCTURE +# METRIC CONFIGURATION - YML STRUCTURE # ------------------------------------- # Required top-level fields: # @@ -483,7 +483,7 @@ definitions: "##; pub const DASHBOARD_YML_SCHEMA: &str = r##" -# DASHBOARD CONFIGURATION - YAML STRUCTURE +# DASHBOARD CONFIGURATION - YML STRUCTURE # ---------------------------------------- # Required fields: # @@ -590,7 +590,7 @@ pub async fn process_metric_file( let metric_file = MetricFile { id: metric_id, name: metric_yml.title.clone(), - file_name: format!("{}.yml", file_name), + file_name: file_name.clone(), content: serde_json::to_value(metric_yml.clone()) .map_err(|e| format!("Failed to process metric: {}", e))?, created_by: Uuid::new_v4(), diff --git a/api/libs/agents/src/tools/categories/file_tools/create_dashboard_files.rs b/api/libs/agents/src/tools/categories/file_tools/create_dashboard_files.rs index 33038a768..04344dc9d 100644 --- a/api/libs/agents/src/tools/categories/file_tools/create_dashboard_files.rs +++ b/api/libs/agents/src/tools/categories/file_tools/create_dashboard_files.rs @@ -101,7 +101,7 @@ async fn process_dashboard_file( let dashboard_file = DashboardFile { id: dashboard_id, name: dashboard_yml.name.clone(), - file_name: format!("{}.yml", file.name), + file_name: file.name.clone(), content: serde_json::to_value(dashboard_yml.clone()) .map_err(|e| format!("Failed to process dashboard: {}", e))?, filter: None, @@ -260,7 +260,7 @@ impl ToolExecutor for CreateDashboardFilesTool { "properties": { "name": { "type": "string", - "description": "The name of the dashboard file to be created. This should exclude the file extension. (i.e. '.yml')" + "description": "The name of the dashboard file to be created. Do not include the file_extension." }, "yml_content": { "type": "string", diff --git a/api/libs/agents/src/tools/categories/file_tools/create_metric_files.rs b/api/libs/agents/src/tools/categories/file_tools/create_metric_files.rs index c8d79204f..b900374d6 100644 --- a/api/libs/agents/src/tools/categories/file_tools/create_metric_files.rs +++ b/api/libs/agents/src/tools/categories/file_tools/create_metric_files.rs @@ -205,7 +205,7 @@ impl ToolExecutor for CreateMetricFilesTool { "properties": { "name": { "type": "string", - "description": "Name of the file" + "description": "Name of the file. Do not include the file_extension." }, "yml_content": { "type": "string", diff --git a/api/libs/database/src/models.rs b/api/libs/database/src/models.rs index d1982729f..1cf4b10ef 100644 --- a/api/libs/database/src/models.rs +++ b/api/libs/database/src/models.rs @@ -39,8 +39,10 @@ pub struct DashboardFile { #[diesel(table_name = messages)] pub struct Message { pub id: Uuid, - pub request: String, - pub response: Value, + pub request_message: String, + pub response_messages: Value, + pub reasoning: Value, + pub final_reasoning_message: String, pub chat_id: Uuid, pub created_at: DateTime, pub updated_at: DateTime, @@ -89,6 +91,7 @@ pub struct Chat { pub updated_at: DateTime, pub deleted_at: Option>, pub created_by: Uuid, + pub updated_by: Uuid, } #[derive(Queryable, Insertable, Associations, Debug)] diff --git a/api/libs/database/src/schema.rs b/api/libs/database/src/schema.rs index ec55a4e6a..80a738d4e 100644 --- a/api/libs/database/src/schema.rs +++ b/api/libs/database/src/schema.rs @@ -91,6 +91,7 @@ diesel::table! { updated_at -> Timestamptz, deleted_at -> Nullable, created_by -> Uuid, + updated_by -> Uuid, } } @@ -315,8 +316,10 @@ diesel::table! { diesel::table! { messages (id) { id -> Uuid, - request -> Text, - response -> Jsonb, + request_message -> Text, + response_messages -> Jsonb, + reasoning -> Jsonb, + final_reasoning_message -> Text, chat_id -> Uuid, created_at -> Timestamptz, updated_at -> Timestamptz, @@ -590,7 +593,6 @@ diesel::table! { diesel::joinable!(api_keys -> organizations (organization_id)); diesel::joinable!(api_keys -> users (owner_id)); diesel::joinable!(chats -> organizations (organization_id)); -diesel::joinable!(chats -> users (created_by)); diesel::joinable!(collections -> organizations (organization_id)); diesel::joinable!(dashboard_versions -> dashboards (dashboard_id)); diesel::joinable!(dashboards -> organizations (organization_id)); diff --git a/api/libs/handlers/src/chats/context_loaders/chat_context.rs b/api/libs/handlers/src/chats/context_loaders/chat_context.rs index f239b8134..6d44d3fbc 100644 --- a/api/libs/handlers/src/chats/context_loaders/chat_context.rs +++ b/api/libs/handlers/src/chats/context_loaders/chat_context.rs @@ -47,10 +47,10 @@ impl ContextLoader for ChatContextLoader { let mut agent_messages = Vec::new(); for message in messages { // Add user message - agent_messages.push(AgentMessage::user(message.request)); + agent_messages.push(AgentMessage::user(message.request_message)); // Add assistant messages from response - if let Ok(response_messages) = serde_json::from_value::>(message.response) + if let Ok(response_messages) = serde_json::from_value::>(message.response_messages) { agent_messages.extend(response_messages); } diff --git a/api/libs/handlers/src/chats/get_chat_handler.rs b/api/libs/handlers/src/chats/get_chat_handler.rs index 9d922844b..a4e9bf505 100644 --- a/api/libs/handlers/src/chats/get_chat_handler.rs +++ b/api/libs/handlers/src/chats/get_chat_handler.rs @@ -27,8 +27,10 @@ pub struct ThreadWithUser { #[derive(Queryable)] pub struct MessageWithUser { pub id: Uuid, - pub request: String, - pub response: Value, + pub request_message: String, + pub response_messages: Value, + pub reasoning: Value, + pub final_reasoning_message: String, pub created_at: DateTime, pub user_id: Uuid, pub user_name: Option, @@ -83,8 +85,10 @@ pub async fn get_chat_handler(chat_id: &Uuid, user_id: &Uuid) -> Result Result + let response_messages = msg.response_messages .as_array() - .map(|arr| { - arr.iter().fold( - (Vec::new(), Vec::new()), - |(mut reasoning, mut responses), item| { - if let Some(reasoning_obj) = item.get("reasoning") { - reasoning.push(reasoning_obj.clone()); - } else if let Some(response_obj) = item.get("response_message") { - responses.push(response_obj.clone()); - } - (reasoning, responses) - }, - ) - }) + .map(|arr| arr.to_vec()) + .unwrap_or_default(); + + let reasoning = msg.reasoning + .as_array() + .map(|arr| arr.to_vec()) .unwrap_or_default(); ChatMessage { id: msg.id, request_message: crate::messages::types::ChatUserMessage { - request: msg.request, + request: msg.request_message, sender_id: msg.user_id, sender_name: msg.user_name.unwrap_or_else(|| "Unknown".to_string()), sender_avatar, diff --git a/api/libs/handlers/src/chats/post_chat_handler.rs b/api/libs/handlers/src/chats/post_chat_handler.rs index a827495fc..2cf5dd1cd 100644 --- a/api/libs/handlers/src/chats/post_chat_handler.rs +++ b/api/libs/handlers/src/chats/post_chat_handler.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::{collections::HashMap, time::Instant}; use agents::{ tools::file_tools::search_data_catalog::SearchDataCatalogOutput, AgentExt, AgentMessage, @@ -59,6 +59,7 @@ pub async fn post_chat_handler( user: User, tx: Option>>, ) -> Result { + let reasoning_duration = Instant::now(); // Validate context request validate_context_request(request.chat_id, request.metric_id, request.dashboard_id)?; @@ -87,6 +88,7 @@ pub async fn post_chat_handler( created_at: Utc::now(), updated_at: Utc::now(), deleted_at: None, + updated_by: user.id.clone(), }; let mut chat_with_messages = ChatWithMessages { @@ -236,17 +238,22 @@ pub async fn post_chat_handler( } let title = title_handle.await??; + let reasoning_duration = reasoning_duration.elapsed().as_secs(); + + let final_reasoning_message = format!("Reasoned for {} seconds", reasoning_duration); // Create and store message in the database let message = Message { id: message_id, - request: request.prompt, - response: serde_json::to_value(&all_messages)?, - chat_id: chat_id, + request_message: request.prompt, + chat_id, created_by: user.id.clone(), created_at: Utc::now(), updated_at: Utc::now(), deleted_at: None, + response_messages: serde_json::to_value(&all_messages)?, + reasoning: serde_json::to_value(&all_messages)?, + final_reasoning_message, }; // Insert message into database @@ -319,7 +326,7 @@ async fn store_final_message_state( diesel::update(messages::table) .filter(messages::id.eq(message.id)) .set(( - messages::response.eq(&message.response), + messages::response_messages.eq(&message.response_messages), messages::updated_at.eq(message.updated_at), )) .execute(conn) @@ -561,7 +568,7 @@ pub fn transform_message( } => { if let Some(content) = content { let mut containers = Vec::new(); - + // Create the regular content message let chat_messages = match transform_text_message( id.clone().unwrap_or_else(|| Uuid::new_v4().to_string()), @@ -599,11 +606,13 @@ pub fn transform_message( status: Some("completed".to_string()), }); - containers.push(BusterContainer::ReasoningMessage(BusterReasoningMessageContainer { - reasoning: reasoning_message, - chat_id: *chat_id, - message_id: *message_id, - })); + containers.push(BusterContainer::ReasoningMessage( + BusterReasoningMessageContainer { + reasoning: reasoning_message, + chat_id: *chat_id, + message_id: *message_id, + }, + )); } return Ok((containers, ThreadEvent::GeneratingResponseMessage)); diff --git a/api/migrations/2025-01-27-390057_adjust_threads_table/up.sql b/api/migrations/2025-01-27-390057_adjust_threads_table/up.sql index 0affcbab8..7dc443c29 100644 --- a/api/migrations/2025-01-27-390057_adjust_threads_table/up.sql +++ b/api/migrations/2025-01-27-390057_adjust_threads_table/up.sql @@ -10,7 +10,8 @@ CREATE TABLE chats( created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(), updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(), deleted_at TIMESTAMP WITH TIME ZONE, - created_by UUID NOT NULL REFERENCES users(id) + created_by UUID NOT NULL REFERENCES users(id), + updated_by UUID NOT NULL REFERENCES users(id) ); -- Create indexes for common query patterns diff --git a/api/migrations/2025-01-28-174921_adjust_messages_table/up.sql b/api/migrations/2025-01-28-174921_adjust_messages_table/up.sql index 76d148729..790bd8181 100644 --- a/api/migrations/2025-01-28-174921_adjust_messages_table/up.sql +++ b/api/migrations/2025-01-28-174921_adjust_messages_table/up.sql @@ -7,8 +7,10 @@ ALTER TABLE messages RENAME TO messages_deprecated; -- Create new messages table with updated schema CREATE TABLE messages ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), - request TEXT NOT NULL, - response JSONB NOT NULL, + request_message TEXT NOT NULL, + response_messages JSONB NOT NULL, + reasoning JSONB NOT NULL, + final_reasoning_message TEXT NOT NULL, chat_id UUID NOT NULL REFERENCES chats(id), created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(), updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),