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

ok the sub agents are working!

This commit is contained in:
dal 2025-02-21 12:44:00 -07:00
parent ddcb8c7798
commit 9f1d8eff7d
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
14 changed files with 303 additions and 279 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
api/libs/litellm/src/client.rs
View File

@ -66,10 +66,15 @@ impl LiteLLMClient {
.post(&url)
.json(&request)
.send()
.await?
.json::<ChatCompletionResponse>()
.await?;
// Print the raw response text
let response_text = response.text().await?;
println!("DEBUG: Raw response payload: {}", response_text);
// Parse the response text into the expected type
let response: ChatCompletionResponse = serde_json::from_str(&response_text)?;
// Print tool calls if present
if let Some(Message::Assistant {
tool_calls: Some(tool_calls),
@ -128,7 +133,7 @@ impl LiteLLMClient {
match chunk_result {
Ok(chunk) => {
let chunk_str = String::from_utf8_lossy(&chunk);
println!("DEBUG: Received raw stream chunk: {}", chunk_str);
println!("DEBUG: Raw response payload: {}", chunk_str);
buffer.push_str(&chunk_str);
while let Some(pos) = buffer.find("\n\n") {

28
api/src/routes/rest/routes/chats/post_chat.rs
View File

@ -13,6 +13,7 @@ use serde_json::Value;
use uuid::Uuid;
use crate::routes::rest::ApiResponse;
use crate::utils::agent::AgentThread;
use crate::{
database_dep::{
enums::Verification,
@ -88,19 +89,27 @@ async fn process_chat(request: ChatCreateNewChat, user: User) -> Result<ThreadWi
// Initialize agent and process request
let agent = ManagerAgent::new(user.id, chat_id).await?;
let input = ManagerAgentInput {
prompt: request.prompt.clone(),
thread_id: Some(chat_id),
message_id: Some(message_id),
};
let mut thread = AgentThread::new(
Some(chat_id),
user.id,
vec![AgentMessage::user(request.prompt.clone())],
);
let output = agent.process_request(input, user.id).await?;
// Get the receiver and collect all messages
let mut rx = agent.run(&mut thread).await?;
let mut messages = Vec::new();
while let Some(msg_result) = rx.recv().await {
match msg_result {
Ok(msg) => messages.push(msg),
Err(e) => return Err(e.into()),
}
}
// Create and store initial message
let message = Message {
id: message_id,
request: request.prompt,
response: serde_json::to_value(&output.messages)?,
response: serde_json::to_value(&messages)?,
thread_id: chat_id,
created_by: user.id.clone(),
created_at: Utc::now(),
@ -118,7 +127,7 @@ async fn process_chat(request: ChatCreateNewChat, user: User) -> Result<ThreadWi
store_final_message_state(
&mut conn,
&message,
&output.messages,
&messages,
&user_org_id,
&user.id,
)
@ -126,8 +135,7 @@ async fn process_chat(request: ChatCreateNewChat, user: User) -> Result<ThreadWi
// Update thread_with_messages with processed messages
if let Some(thread_message) = thread_with_messages.messages.first_mut() {
let transformed_messages: Vec<BusterContainer> = output
.messages
let transformed_messages: Vec<BusterContainer> = messages
.iter()
.filter_map(|msg| {
transform_message(&chat_id, &message_id, msg.clone())

2
api/src/utils/agent/agent.rs
View File

@ -375,6 +375,7 @@ impl Agent {
if let Some(tool) = self.tools.read().await.get(&tool_call.function.name) {
let params: Value = serde_json::from_str(&tool_call.function.arguments)?;
let result = tool.execute(params).await?;
println!("Tool Call result: {:?}", result);
let result_str = serde_json::to_string(&result)?;
let tool_message = Message::tool(
None,
@ -456,6 +457,7 @@ impl PendingToolCall {
/// A trait that provides convenient access to Agent functionality
/// when the agent is stored behind an Arc
#[async_trait::async_trait]
pub trait AgentExt {
fn get_agent(&self) -> &Arc<Agent>;

111
api/src/utils/agent/agents/dashboard_agent.rs
View File

@ -2,17 +2,15 @@ use std::sync::Arc;
use anyhow::Result;
use std::collections::HashMap;
use tokio::sync::mpsc::Receiver;
use uuid::Uuid;
use crate::utils::{
agent::{Agent, AgentExt, AgentThread},
tools::{
file_tools::{
agents_as_tools::dashboard_agent_tool::DashboardAgentOutput, file_tools::{
CreateDashboardFilesTool, CreateMetricFilesTool, ModifyDashboardFilesTool,
ModifyMetricFilesTool, OpenFilesTool, SearchFilesTool,
},
IntoValueTool, ToolExecutor,
ModifyMetricFilesTool,
}, IntoValueTool, ToolExecutor
},
};
@ -37,8 +35,6 @@ impl DashboardAgent {
let modify_dashboard_tool = ModifyDashboardFilesTool::new(Arc::clone(&agent));
let create_metric_tool = CreateMetricFilesTool::new(Arc::clone(&agent));
let modify_metric_tool = ModifyMetricFilesTool::new(Arc::clone(&agent));
let open_files_tool = OpenFilesTool::new(Arc::clone(&agent));
let search_files_tool = SearchFilesTool::new(Arc::clone(&agent));
// Add tools directly to the Arc<Agent>
agent
@ -65,18 +61,6 @@ impl DashboardAgent {
modify_metric_tool.into_value_tool(),
)
.await;
agent
.add_tool(
open_files_tool.get_name(),
open_files_tool.into_value_tool(),
)
.await;
agent
.add_tool(
search_files_tool.get_name(),
search_files_tool.into_value_tool(),
)
.await;
Ok(Self { agent })
}
@ -90,9 +74,6 @@ impl DashboardAgent {
let modify_dashboard_tool = ModifyDashboardFilesTool::new(Arc::clone(&agent));
let create_metric_tool = CreateMetricFilesTool::new(Arc::clone(&agent));
let modify_metric_tool = ModifyMetricFilesTool::new(Arc::clone(&agent));
let open_files_tool = OpenFilesTool::new(Arc::clone(&agent));
let search_files_tool = SearchFilesTool::new(Arc::clone(&agent));
// Add tools to the agent
agent
@ -119,32 +100,80 @@ impl DashboardAgent {
modify_metric_tool.into_value_tool(),
)
.await;
agent
.add_tool(
open_files_tool.get_name(),
open_files_tool.into_value_tool(),
)
.await;
agent
.add_tool(
search_files_tool.get_name(),
search_files_tool.into_value_tool(),
)
.await;
Ok(Self { agent })
}
pub async fn run(
&self,
thread: &mut AgentThread,
) -> Result<Receiver<Result<AgentMessage, anyhow::Error>>> {
// Process using agent's streaming functionality
fn is_completion_signal(msg: &AgentMessage) -> bool {
matches!(msg, AgentMessage::Assistant { content: Some(content), tool_calls: None, .. }
if content == "AGENT_COMPLETE")
}
pub async fn run(&self, thread: &mut AgentThread) -> Result<DashboardAgentOutput> {
println!("Running dashboard agent");
println!("Setting developer message");
thread.set_developer_message(DASHBOARD_AGENT_PROMPT.to_string());
println!("Dashboard agent thread: {:?}", thread);
println!("Starting stream_process_thread");
let mut rx = self.stream_process_thread(thread).await?;
println!("Got receiver from stream_process_thread");
self.stream_process_thread(thread).await
println!("Starting message processing loop");
// Process messages internally until we determine we're done
while let Some(msg_result) = rx.recv().await {
println!("Received message from channel");
match msg_result {
Ok(msg) => {
println!("Message content: {:?}", msg.get_content());
println!("Message has tool calls: {:?}", msg.get_tool_calls());
println!("Forwarding message to stream sender");
if let Err(e) = self.get_agent().get_stream_sender().await.send(Ok(msg.clone())).await {
println!("Error forwarding message: {:?}", e);
// Continue processing even if we fail to forward
continue;
}
if let Some(content) = msg.get_content() {
println!("Message has content: {}", content);
if content == "AGENT_COMPLETE" {
println!("Found completion signal, breaking loop");
break;
}
}
}
Err(e) => {
println!("Error receiving message: {:?}", e);
println!("Error details: {:?}", e.to_string());
// Log error but continue processing instead of returning error
continue;
}
}
}
println!("Exited message processing loop");
println!("Creating completion signal");
let completion_msg = AgentMessage::assistant(
None,
Some("AGENT_COMPLETE".to_string()),
None,
None,
None,
);
println!("Sending completion signal");
self.get_agent()
.get_stream_sender()
.await
.send(Ok(completion_msg))
.await?;
println!("Sent completion signal, returning output");
Ok(DashboardAgentOutput {
message: "Dashboard processing complete".to_string(),
duration: 0,
files: vec![],
})
}
}

62
api/src/utils/agent/agents/manager_agent.rs
View File

@ -7,6 +7,7 @@ use tracing::{debug, info};
use uuid::Uuid;
use crate::utils::tools::agents_as_tools::{DashboardAgentTool, MetricAgentTool};
use crate::utils::tools::file_tools::{send_assets_to_user, SendAssetsToUserTool};
use crate::utils::{
agent::{Agent, AgentExt, AgentThread},
tools::{
@ -109,6 +110,7 @@ impl ManagerAgent {
let create_or_modify_metrics_tool = MetricAgentTool::new(Arc::clone(&agent));
let create_or_modify_dashboards_tool = DashboardAgentTool::new(Arc::clone(&agent));
let exploratory_agent_tool = ExploratoryAgentTool::new(Arc::clone(&agent));
let send_assets_to_user = SendAssetsToUserTool::new(Arc::clone(&agent));
// Add tools to the agent
agent
@ -141,62 +143,22 @@ impl ManagerAgent {
exploratory_agent_tool.into_value_tool(),
)
.await;
agent
.add_tool(
send_assets_to_user.get_name(),
send_assets_to_user.into_value_tool(),
)
.await;
Ok(Self { agent })
}
pub async fn process_request(
pub async fn run(
&self,
input: ManagerAgentInput,
user_id: Uuid,
) -> Result<ManagerAgentOutput> {
let start_time = std::time::Instant::now();
let thread_id = input.thread_id.unwrap_or_else(Uuid::new_v4);
info!("Starting manager request processing for thread: {}", thread_id);
thread: &mut AgentThread,
) -> Result<Receiver<Result<AgentMessage, anyhow::Error>>> {
thread.set_developer_message(MANAGER_AGENT_PROMPT.to_string());
// Create thread with manager context
let thread = AgentThread::new(
Some(thread_id),
user_id,
vec![
AgentMessage::developer(MANAGER_AGENT_PROMPT.to_string()),
AgentMessage::user(input.prompt),
],
);
// Process using agent's streaming functionality
let mut rx = self.run(&thread).await?;
let mut messages = Vec::new();
while let Some(msg_result) = rx.recv().await {
if let Ok(msg) = msg_result {
messages.push(msg);
}
}
let duration = start_time.elapsed().as_millis() as i64;
let message = format!(
"Completed request processing with {} messages",
messages.len()
);
info!(
duration_ms = duration,
messages_count = messages.len(),
thread_id = %thread_id,
"Completed manager request processing"
);
Ok(ManagerAgentOutput {
message,
duration,
thread_id,
messages,
})
}
pub async fn run(&self, thread: &AgentThread) -> Result<Receiver<Result<AgentMessage, anyhow::Error>>> {
// Process using agent's streaming functionality
self.stream_process_thread(thread).await
}
}

7
api/src/utils/agent/types.rs
View File

@ -34,4 +34,11 @@ impl AgentThread {
self.messages.insert(0, Message::developer(message));
}
}
/// Remove the most recent assistant message from the thread
pub fn remove_last_assistant_message(&mut self) {
if let Some(pos) = self.messages.iter().rposition(|msg| matches!(msg, Message::Assistant { .. })) {
self.messages.remove(pos);
}
}
}

40
api/src/utils/tools/agents_as_tools/dashboard_agent_tool.rs
View File

@ -14,6 +14,11 @@ pub struct DashboardAgentTool {
agent: Arc<Agent>,
}
fn is_completion_signal(msg: &AgentMessage) -> bool {
matches!(msg, AgentMessage::Assistant { content: Some(content), tool_calls: None, .. }
if content == "AGENT_COMPLETE")
}
#[derive(Debug, Serialize, Deserialize)]
pub struct DashboardAgentParams {
pub ticket_description: String,
@ -42,35 +47,34 @@ impl ToolExecutor for DashboardAgentTool {
async fn execute(&self, params: Self::Params) -> Result<Self::Output> {
// Create and initialize the agent
println!("DashboardAgentTool: Creating dashboard agent");
let dashboard_agent = DashboardAgent::from_existing(&self.agent).await?;
println!("DashboardAgentTool: Dashboard agent created");
println!("DashboardAgentTool: Getting current thread");
let mut current_thread = self
.agent
.get_current_thread()
.await
.ok_or_else(|| anyhow::anyhow!("No current thread"))?;
println!("DashboardAgentTool: Current thread retrieved");
// Add the ticket description to the thread
current_thread
.messages
.push(AgentMessage::user(params.ticket_description));
println!("DashboardAgentTool: Removing last assistant message");
current_thread.remove_last_assistant_message();
println!("DashboardAgentTool: Last assistant message removed");
// Run the dashboard agent and get the receiver
let mut rx = dashboard_agent.run(&mut current_thread).await?;
println!("DashboardAgentTool: Starting dashboard agent run");
// Run the dashboard agent and get the output
let output = dashboard_agent.run(&mut current_thread).await?;
println!("DashboardAgentTool: Dashboard agent run completed");
// Process all messages from the receiver
let mut messages = Vec::new();
while let Some(msg_result) = rx.recv().await {
match msg_result {
Ok(msg) => messages.push(msg),
Err(e) => return Err(e.into()),
}
}
// Return the messages as part of the output
println!("DashboardAgentTool: Preparing success response");
// Return success with the output
Ok(serde_json::json!({
"messages": messages,
"status": "completed"
"status": "success",
"message": output.message,
"duration": output.duration,
"files": output.files
}))
}

22
api/src/utils/tools/agents_as_tools/exploratory_agent_tool.rs
View File

@ -43,27 +43,15 @@ impl ToolExecutor for ExploratoryAgentTool {
.await
.ok_or_else(|| anyhow::anyhow!("No current thread"))?;
// Add the ticket description to the thread
current_thread
.messages
.push(AgentMessage::user(params.ticket_description));
current_thread.remove_last_assistant_message();
// Run the exploratory agent and get the receiver
let mut rx = exploratory_agent.run(&mut current_thread).await?;
let _rx = exploratory_agent.run(&mut current_thread).await?;
// Process all messages from the receiver
let mut messages = Vec::new();
while let Some(msg_result) = rx.recv().await {
match msg_result {
Ok(msg) => messages.push(msg),
Err(e) => return Err(e.into()),
}
}
// Return the messages as part of the output
// Return immediately with status
Ok(serde_json::json!({
"messages": messages,
"status": "completed"
"status": "running",
"message": "Exploratory agent started successfully"
}))
}

21
api/src/utils/tools/agents_as_tools/metric_agent_tool.rs
View File

@ -51,26 +51,15 @@ impl ToolExecutor for MetricAgentTool {
ticket_description: params.ticket_description,
};
// current_thread
// .messages
// .push(AgentMessage::user(agent_input.ticket_description));
current_thread.remove_last_assistant_message();
// Run the metric agent and get the receiver
let mut rx = metric_agent.run(&mut current_thread).await?;
let _rx = metric_agent.run(&mut current_thread).await?;
// Process all messages from the receiver
let mut messages = Vec::new();
while let Some(msg_result) = rx.recv().await {
match msg_result {
Ok(msg) => messages.push(msg),
Err(e) => return Err(e.into()),
}
}
// Return the messages as part of the output
// Return immediately with status
Ok(serde_json::json!({
"messages": messages,
"status": "completed"
"status": "running",
"message": "Metric agent started successfully"
}))
}

5
api/src/utils/tools/file_tools/file_types/dashboard_yml.rs
View File

@ -7,15 +7,16 @@ use uuid::Uuid;
pub struct DashboardYml {
pub id: Option<Uuid>,
pub updated_at: Option<DateTime<Utc>>,
#[serde(alias = "title")]
pub name: String,
pub rows: Vec<Row>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Row {
pub items: Vec<RowItem>, // max number of items in a row is 4, min is 1
pub items: Vec<RowItem>, // max number of items in a row is 4, min is 1
#[serde(skip_serializing_if = "Option::is_none")]
pub row_height: Option<u32>, // max is 550, min is 320
pub row_height: Option<u32>, // max is 550, min is 320
#[serde(skip_serializing_if = "Option::is_none")]
pub column_sizes: Option<Vec<u32>>, // max sum of elements is 12 min is 3
}

4
api/src/utils/tools/file_tools/mod.rs
View File

@ -9,14 +9,14 @@ pub mod modify_metric_files;
pub mod open_files;
pub mod search_data_catalog;
pub mod search_files;
pub mod send_files_to_user;
pub mod send_assets_to_user;
pub use create_files::CreateFilesTool;
pub use modify_files::ModifyFilesTool;
pub use open_files::OpenFilesTool;
pub use search_data_catalog::SearchDataCatalogTool;
pub use search_files::SearchFilesTool;
pub use send_files_to_user::SendFilesToUserTool;
pub use send_assets_to_user::SendAssetsToUserTool;
pub use create_dashboard_files::CreateDashboardFilesTool;
pub use create_metric_files::CreateMetricFilesTool;
pub use modify_dashboard_files::ModifyDashboardFilesTool;

120
api/src/utils/tools/file_tools/modify_metric_files.rs
View File

File diff suppressed because one or more lines are too long

89
api/src/utils/tools/file_tools/send_assets_to_user.rs Normal file
View File

@ -0,0 +1,89 @@
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;
use crate::utils::{agent::Agent, tools::ToolExecutor};
use litellm::ToolCall;
#[derive(Debug, Serialize, Deserialize)]
pub struct SendToUserParams {
metric_id: String,
}
#[derive(Debug, Serialize)]
pub struct SendToUserOutput {
success: bool,
}
pub struct SendAssetsToUserTool {
agent: Arc<Agent>,
}
impl SendAssetsToUserTool {
pub fn new(agent: Arc<Agent>) -> Self {
Self { agent }
}
}
#[async_trait]
impl ToolExecutor for SendAssetsToUserTool {
type Output = SendToUserOutput;
type Params = SendToUserParams;
async fn execute(&self, params: Self::Params) -> Result<Self::Output> {
// TODO: Implement actual send to user logic
Ok(SendToUserOutput { success: true })
}
fn get_name(&self) -> String {
"decide_assets_to_return".to_string()
}
fn get_schema(&self) -> Value {
serde_json::json!({
"name": "decide_assets_to_return",
"description": "Use after you have created or modified any assets (metrics or dashboards) to specify exactly which assets to present in the final response. If you have not created or modified any assets, do not call this action.",
"strict": true,
"parameters": {
"type": "object",
"required": [
"assets_to_return",
"ticket_description"
],
"properties": {
"assets_to_return": {
"type": "array",
"description": "List of assets to present in the final response, each with an ID and a name",
"items": {
"type": "object",
"properties": {
"id": {
"type": "string",
"description": "Unique identifier for the asset"
},
"name": {
"type": "string",
"description": "Name of the asset"
}
},
"required": [
"id",
"name"
],
"additionalProperties": false
}
},
"ticket_description": {
"type": "string",
"description": "Description of the ticket related to the assets"
}
},
"additionalProperties": false
}
})
}
}

60
api/src/utils/tools/file_tools/send_files_to_user.rs
View File

@ -1,60 +0,0 @@
use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;
use crate::utils::tools::ToolExecutor;
use litellm::ToolCall;
#[derive(Debug, Serialize, Deserialize)]
pub struct SendToUserParams {
metric_id: String,
}
#[derive(Debug, Serialize)]
pub struct SendToUserOutput {
success: bool,
}
pub struct SendFilesToUserTool;
impl SendFilesToUserTool {
pub fn new() -> Self {
Self
}
}
#[async_trait]
impl ToolExecutor for SendFilesToUserTool {
type Output = SendToUserOutput;
type Params = SendToUserParams;
async fn execute(&self, params: Self::Params) -> Result<Self::Output> {
// TODO: Implement actual send to user logic
Ok(SendToUserOutput { success: true })
}
fn get_name(&self) -> String {
"send_to_user".to_string()
}
fn get_schema(&self) -> Value {
serde_json::json!({
"name": "send_to_user",
"strict": true,
"parameters": {
"type": "object",
"required": ["metric_id"],
"properties": {
"metric_id": {
"type": "string",
"description": "The ID of the metric to send to the user"
}
},
"additionalProperties": false
},
"description": "Sends a metric to the user by its ID."
})
}
}