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

ok types are matched with streaam

This commit is contained in:
dal 2025-03-03 10:42:29 -07:00
parent 71ab5891a9
commit 9aa1cbbab7
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
2 changed files with 65 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -212,8 +212,9 @@ pub async fn post_chat_handler(
} }
BusterContainer::ReasoningMessage(reasoning) => { BusterContainer::ReasoningMessage(reasoning) => {
let reasoning_value = match &reasoning.reasoning { let reasoning_value = match &reasoning.reasoning {
ReasoningMessage::Thought(thought) => serde_json::to_value(thought).ok(), ReasoningMessage::Pill(thought) => serde_json::to_value(thought).ok(),
ReasoningMessage::File(file) => serde_json::to_value(file).ok(), ReasoningMessage::File(file) => serde_json::to_value(file).ok(),
ReasoningMessage::Text(text) => serde_json::to_value(text).ok(),
}; };
if let Some(value) = reasoning_value { if let Some(value) = reasoning_value {
@ -362,8 +363,8 @@ async fn store_final_message_state(
#[serde(untagged)] #[serde(untagged)]
pub enum BusterChatContainer { pub enum BusterChatContainer {
ChatMessage(BusterChatMessage), ChatMessage(BusterChatMessage),
Thought(BusterThought), Thought(BusterReasoningPill),
File(BusterFileMessage), File(BusterReasoningFile),
} }
#[derive(Debug, Serialize, Clone)] #[derive(Debug, Serialize, Clone)]
@ -376,8 +377,9 @@ pub struct BusterChatMessageContainer {
#[derive(Debug, Serialize, Clone)] #[derive(Debug, Serialize, Clone)]
#[serde(untagged)] #[serde(untagged)]
pub enum ReasoningMessage { pub enum ReasoningMessage {
Thought(BusterThought), Pill(BusterReasoningPill),
File(BusterFileMessage), File(BusterReasoningFile),
Text(BusterReasoningText),
} }
#[derive(Debug, Serialize, Clone)] #[derive(Debug, Serialize, Clone)]
@ -398,23 +400,33 @@ pub struct BusterChatMessage {
} }
#[derive(Debug, Serialize, Clone)] #[derive(Debug, Serialize, Clone)]
pub struct BusterThought { pub struct BusterReasoningPill {
pub id: String, pub id: String,
#[serde(rename = "type")] #[serde(rename = "type")]
pub thought_type: String, pub thought_type: String,
#[serde(rename = "title")] pub title: String,
pub thought_title: String, pub secondary_title: String,
#[serde(rename = "secondary_title")] pub pill_containers: Option<Vec<BusterThoughtPillContainer>>,
pub thought_secondary_title: String,
#[serde(rename = "pill_containers")]
pub thoughts: Option<Vec<BusterThoughtPillContainer>>,
pub status: String, pub status: String,
} }
#[derive(Debug, Serialize, Clone)]
pub struct BusterReasoningText {
pub id: String,
#[serde(rename = "type")]
pub reasoning_type: String,
pub title: String,
pub secondary_title: Option<String>,
pub message: Option<String>,
pub message_chunk: Option<String>,
pub status: Option<String>,
}
#[derive(Debug, Serialize, Clone)] #[derive(Debug, Serialize, Clone)]
pub struct BusterThoughtPillContainer { pub struct BusterThoughtPillContainer {
pub title: String, pub title: String,
pub thought_pills: Vec<BusterThoughtPill>, pub pills: Vec<BusterThoughtPill>,
} }
#[derive(Debug, Serialize, Clone)] #[derive(Debug, Serialize, Clone)]
@ -426,7 +438,7 @@ pub struct BusterThoughtPill {
} }
#[derive(Debug, Serialize, Clone)] #[derive(Debug, Serialize, Clone)]
pub struct BusterFileMessage { pub struct BusterReasoningFile {
pub id: String, pub id: String,
#[serde(rename = "type")] #[serde(rename = "type")]
pub message_type: String, pub message_type: String,
@ -819,23 +831,23 @@ fn tool_data_catalog_search(
}; };
let buster_thought = if result_count > 0 { let buster_thought = if result_count > 0 {
BusterChatContainer::Thought(BusterThought { BusterChatContainer::Thought(BusterReasoningPill {
id: id, // Use required ID directly id: id, // Use required ID directly
thought_type: "thought".to_string(), thought_type: "thought".to_string(),
thought_title: format!("Found {} results", result_count), title: format!("Found {} results", result_count),
thought_secondary_title: format!("{} seconds", duration), secondary_title: format!("{} seconds", duration),
thoughts: Some(thought_pill_containters), pill_containers: Some(thought_pill_containters),
status: "completed".to_string(), status: "completed".to_string(),
}) })
} else { } else {
BusterChatContainer::Thought(BusterThought { BusterChatContainer::Thought(BusterReasoningPill {
id: id, // Use required ID directly id: id, // Use required ID directly
thought_type: "thought".to_string(), thought_type: "thought".to_string(),
thought_title: "No data catalog items found".to_string(), title: "No data catalog items found".to_string(),
thought_secondary_title: format!("{} seconds", duration), secondary_title: format!("{} seconds", duration),
thoughts: Some(vec![BusterThoughtPillContainer { pill_containers: Some(vec![BusterThoughtPillContainer {
title: "No results found".to_string(), title: "No results found".to_string(),
thought_pills: vec![BusterThoughtPill { pills: vec![BusterThoughtPill {
id: "".to_string(), id: "".to_string(),
text: query_params, text: query_params,
thought_file_type: "empty".to_string(), thought_file_type: "empty".to_string(),
@ -854,7 +866,7 @@ fn proccess_data_catalog_search_results(
if results.results.is_empty() { if results.results.is_empty() {
return Ok(vec![BusterThoughtPillContainer { return Ok(vec![BusterThoughtPillContainer {
title: "No results found".to_string(), title: "No results found".to_string(),
thought_pills: vec![], pills: vec![],
}]); }]);
} }
@ -880,7 +892,7 @@ fn proccess_data_catalog_search_results(
"{count} {} found", "{count} {} found",
title.chars().next().unwrap().to_uppercase().to_string() + &title[1..] title.chars().next().unwrap().to_uppercase().to_string() + &title[1..]
), ),
thought_pills, pills: thought_pills,
} }
}) })
.collect(); .collect();
@ -946,7 +958,7 @@ fn transform_to_reasoning_container(
.into_iter() .into_iter()
.map(|container| match container { .map(|container| match container {
BusterChatContainer::Thought(thought) => BusterReasoningMessageContainer { BusterChatContainer::Thought(thought) => BusterReasoningMessageContainer {
reasoning: ReasoningMessage::Thought(thought), reasoning: ReasoningMessage::Pill(thought),
chat_id, chat_id,
message_id, message_id,
}, },
@ -1057,7 +1069,7 @@ fn transform_assistant_tool_message(
.into_iter() .into_iter()
.map(|message| BusterReasoningMessageContainer { .map(|message| BusterReasoningMessageContainer {
reasoning: match message { reasoning: match message {
BusterChatContainer::Thought(thought) => ReasoningMessage::Thought(thought), BusterChatContainer::Thought(thought) => ReasoningMessage::Pill(thought),
BusterChatContainer::File(file) => ReasoningMessage::File(file), BusterChatContainer::File(file) => ReasoningMessage::File(file),
_ => unreachable!("Assistant tool messages should only return Thought or File"), _ => unreachable!("Assistant tool messages should only return Thought or File"),
}, },
@ -1076,12 +1088,12 @@ fn assistant_data_catalog_search(
match progress { match progress {
MessageProgress::InProgress => { MessageProgress::InProgress => {
if initial { if initial {
Ok(vec![BusterChatContainer::Thought(BusterThought { Ok(vec![BusterChatContainer::Thought(BusterReasoningPill {
id, id,
thought_type: "thought".to_string(), thought_type: "thought".to_string(),
thought_title: "Searching your data catalog...".to_string(), title: "Searching your data catalog...".to_string(),
thought_secondary_title: "".to_string(), secondary_title: "".to_string(),
thoughts: None, pill_containers: None,
status: "loading".to_string(), status: "loading".to_string(),
})]) })])
} else { } else {
@ -1107,12 +1119,12 @@ fn assistant_create_metrics(
if let Some(progress) = progress { if let Some(progress) = progress {
match progress { match progress {
MessageProgress::InProgress => { MessageProgress::InProgress => {
Ok(vec![BusterChatContainer::Thought(BusterThought { Ok(vec![BusterChatContainer::Thought(BusterReasoningPill {
id, id,
thought_type: "thought".to_string(), thought_type: "thought".to_string(),
thought_title: "Creating metrics...".to_string(), title: "Creating metrics...".to_string(),
thought_secondary_title: "".to_string(), secondary_title: "".to_string(),
thoughts: None, pill_containers: None,
status: "loading".to_string(), status: "loading".to_string(),
})]) })])
} }
@ -1191,12 +1203,12 @@ fn assistant_create_dashboards(
if let Some(progress) = progress { if let Some(progress) = progress {
match progress { match progress {
MessageProgress::InProgress => { MessageProgress::InProgress => {
Ok(vec![BusterChatContainer::Thought(BusterThought { Ok(vec![BusterChatContainer::Thought(BusterReasoningPill {
id, id,
thought_type: "thought".to_string(), thought_type: "thought".to_string(),
thought_title: "Creating dashboards...".to_string(), title: "Creating dashboards...".to_string(),
thought_secondary_title: "".to_string(), secondary_title: "".to_string(),
thoughts: None, pill_containers: None,
status: "loading".to_string(), status: "loading".to_string(),
})]) })])
} }
@ -1247,12 +1259,12 @@ fn assistant_modify_dashboards(
if let Some(progress) = progress { if let Some(progress) = progress {
match progress { match progress {
MessageProgress::InProgress => { MessageProgress::InProgress => {
Ok(vec![BusterChatContainer::Thought(BusterThought { Ok(vec![BusterChatContainer::Thought(BusterReasoningPill {
id, id,
thought_type: "thought".to_string(), thought_type: "thought".to_string(),
thought_title: "Updating dashboards...".to_string(), title: "Updating dashboards...".to_string(),
thought_secondary_title: "".to_string(), secondary_title: "".to_string(),
thoughts: None, pill_containers: None,
status: "loading".to_string(), status: "loading".to_string(),
})]) })])
} }
@ -1303,12 +1315,12 @@ fn assistant_create_plan(
if let Some(progress) = progress { if let Some(progress) = progress {
match progress { match progress {
MessageProgress::InProgress => { MessageProgress::InProgress => {
Ok(vec![BusterChatContainer::Thought(BusterThought { Ok(vec![BusterChatContainer::Thought(BusterReasoningPill {
id, id,
thought_type: "thought".to_string(), thought_type: "thought".to_string(),
thought_title: "Creating plan...".to_string(), title: "Creating plan...".to_string(),
thought_secondary_title: "".to_string(), secondary_title: "".to_string(),
thoughts: None, pill_containers: None,
status: "loading".to_string(), status: "loading".to_string(),
})]) })])
} }
@ -1359,12 +1371,12 @@ fn assistant_modify_metrics(
if let Some(progress) = progress { if let Some(progress) = progress {
match progress { match progress {
MessageProgress::InProgress => { MessageProgress::InProgress => {
Ok(vec![BusterChatContainer::Thought(BusterThought { Ok(vec![BusterChatContainer::Thought(BusterReasoningPill {
id, id,
thought_type: "thought".to_string(), thought_type: "thought".to_string(),
thought_title: "Updating metrics...".to_string(), title: "Updating metrics...".to_string(),
thought_secondary_title: "".to_string(), secondary_title: "".to_string(),
thoughts: None, pill_containers: None,
status: "loading".to_string(), status: "loading".to_string(),
})]) })])
} }

4
api/libs/handlers/src/chats/streaming_parser.rs
View File

@ -2,7 +2,7 @@ use anyhow::Result;
use serde_json::Value; use serde_json::Value;
use uuid::Uuid; use uuid::Uuid;
use super::post_chat_handler::{BusterChatContainer, BusterFileLine, BusterFileMessage}; use super::post_chat_handler::{BusterChatContainer, BusterFileLine, BusterReasoningFile};
pub struct StreamingParser { pub struct StreamingParser {
buffer: String, buffer: String,
@ -170,7 +170,7 @@ impl StreamingParser {
}); });
} }
return Ok(Some(BusterChatContainer::File(BusterFileMessage { return Ok(Some(BusterChatContainer::File(BusterReasoningFile {
id, id,
message_type: "file".to_string(), message_type: "file".to_string(),
file_type: file_type.to_string(), file_type: file_type.to_string(),