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

some updates to better align with types

This commit is contained in:
dal 2025-03-04 10:47:16 -07:00
parent ab0538c737
commit e5f25a6356
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
10 changed files with 61 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
api/libs/agents/src/tools/categories/file_tools/common.rs
View File

@ -98,7 +98,7 @@ pub async fn validate_metric_ids(ids: &[Uuid]) -> Result<Vec<Uuid>> {
}
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(),

4
api/libs/agents/src/tools/categories/file_tools/create_dashboard_files.rs
View File

@ -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",

2
api/libs/agents/src/tools/categories/file_tools/create_metric_files.rs
View File

@ -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",

7
api/libs/database/src/models.rs
View File

@ -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<Utc>,
pub updated_at: DateTime<Utc>,
@ -89,6 +91,7 @@ pub struct Chat {
pub updated_at: DateTime<Utc>,
pub deleted_at: Option<DateTime<Utc>>,
pub created_by: Uuid,
pub updated_by: Uuid,
}
#[derive(Queryable, Insertable, Associations, Debug)]

8
api/libs/database/src/schema.rs
View File

@ -91,6 +91,7 @@ diesel::table! {
updated_at -> Timestamptz,
deleted_at -> Nullable<Timestamptz>,
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));

4
api/libs/handlers/src/chats/context_loaders/chat_context.rs
View File

@ -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::<Vec<AgentMessage>>(message.response)
if let Ok(response_messages) = serde_json::from_value::<Vec<AgentMessage>>(message.response_messages)
{
agent_messages.extend(response_messages);
}

38
api/libs/handlers/src/chats/get_chat_handler.rs
View File

@ -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<Utc>,
pub user_id: Uuid,
pub user_name: Option<String>,
@ -83,8 +85,10 @@ pub async fn get_chat_handler(chat_id: &Uuid, user_id: &Uuid) -> Result<ChatWith
.order_by(messages::created_at.desc())
.select((
messages::id,
messages::request,
messages::response,
messages::request_message,
messages::response_messages,
messages::reasoning,
messages::final_reasoning_message,
messages::created_at,
users::id,
users::name.nullable(),
@ -124,29 +128,21 @@ pub async fn get_chat_handler(chat_id: &Uuid, user_id: &Uuid) -> Result<ChatWith
.and_then(|v| v.as_str())
.map(String::from);
// Parse the response JSON to extract reasoning and response messages
let (reasoning, response_messages) = msg
.response
// Convert response_messages and reasoning to Vec<Value>
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,

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

@ -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<mpsc::Sender<Result<(BusterContainer, ThreadEvent)>>>,
) -> Result<ChatWithMessages> {
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));

3
api/migrations/2025-01-27-390057_adjust_threads_table/up.sql
View File

@ -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

6
api/migrations/2025-01-28-174921_adjust_messages_table/up.sql
View File

@ -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(),