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

dynamic prompt switching

This commit is contained in:
dal 2025-04-10 13:40:58 -06:00
parent f46376eac0
commit b0797e6f6f
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
4 changed files with 336 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
api/libs/agents/src/agent.rs
View File

@ -208,7 +208,11 @@ impl Agent {
}
/// Create a new Agent that shares state and stream with an existing agent
pub fn from_existing(existing_agent: &Agent, name: String) -> Self {
pub fn from_existing(
existing_agent: &Agent,
name: String,
default_prompt: String,
) -> Self {
let llm_api_key = env::var("LLM_API_KEY").ok(); // Use ok() instead of expect
let llm_base_url = env::var("LLM_BASE_URL").ok(); // Use ok() instead of expect
@ -225,7 +229,7 @@ impl Agent {
session_id: existing_agent.session_id,
shutdown_tx: Arc::clone(&existing_agent.shutdown_tx), // Shared shutdown
name,
default_prompt: existing_agent.default_prompt.clone(),
default_prompt,
dynamic_prompt_rules: Arc::new(RwLock::new(Vec::new())),
}
}
@ -555,7 +559,7 @@ impl Agent {
// --- Dynamic Prompt Selection ---
let current_system_prompt = self.get_current_prompt().await;
let system_message = AgentMessage::developer(current_system_prompt);
// Prepare messages for LLM: Inject current system prompt and filter out old ones
let mut llm_messages = vec![system_message];
llm_messages.extend(

113
api/libs/agents/src/agents/buster_cli_agent.rs
View File

@ -1,24 +1,27 @@
use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, env, sync::Arc};
use serde_json::Value;
use std::{collections::HashMap, sync::Arc};
use tokio::sync::broadcast;
use uuid::Uuid;
use serde_json::Value; // Add for Value
use uuid::Uuid; // Add for Value
use crate::{
agent::{Agent, AgentError, AgentExt},
models::AgentThread,
tools::{ // Import necessary tools
categories::cli_tools::{ // Import CLI tools using correct struct names from mod.rs
EditFileContentTool, // Use correct export
FindFilesGlobTool, // Use correct export
ListDirectoryTool, // Use correct export
ReadFileContentTool, // Use correct export
RunBashCommandTool, // Use correct export
tools::{
// Import necessary tools
categories::cli_tools::{
// Import CLI tools using correct struct names from mod.rs
EditFileContentTool, // Use correct export
FindFilesGlobTool, // Use correct export
ListDirectoryTool, // Use correct export
ReadFileContentTool, // Use correct export
RunBashCommandTool, // Use correct export
SearchFileContentGrepTool, // Use correct export
WriteFileContentTool, // Use correct export
WriteFileContentTool, // Use correct export
},
IntoToolCallExecutor, ToolExecutor,
IntoToolCallExecutor,
ToolExecutor,
},
};
@ -65,22 +68,65 @@ impl BusterCliAgent {
let write_file_tool = WriteFileContentTool::new(Arc::clone(&self.agent));
// Add tools - Pass None directly since these tools are always enabled
self.agent.add_tool(bash_tool.get_name(), bash_tool.into_tool_call_executor(), None::<EnablementCondition>).await;
self.agent.add_tool(edit_file_tool.get_name(), edit_file_tool.into_tool_call_executor(), None::<EnablementCondition>).await;
self.agent.add_tool(glob_tool.get_name(), glob_tool.into_tool_call_executor(), None::<EnablementCondition>).await;
self.agent.add_tool(grep_tool.get_name(), grep_tool.into_tool_call_executor(), None::<EnablementCondition>).await;
self.agent.add_tool(ls_tool.get_name(), ls_tool.into_tool_call_executor(), None::<EnablementCondition>).await;
self.agent.add_tool(read_file_tool.get_name(), read_file_tool.into_tool_call_executor(), None::<EnablementCondition>).await;
self.agent.add_tool(write_file_tool.get_name(), write_file_tool.into_tool_call_executor(), None::<EnablementCondition>).await;
self.agent
.add_tool(
bash_tool.get_name(),
bash_tool.into_tool_call_executor(),
None::<EnablementCondition>,
)
.await;
self.agent
.add_tool(
edit_file_tool.get_name(),
edit_file_tool.into_tool_call_executor(),
None::<EnablementCondition>,
)
.await;
self.agent
.add_tool(
glob_tool.get_name(),
glob_tool.into_tool_call_executor(),
None::<EnablementCondition>,
)
.await;
self.agent
.add_tool(
grep_tool.get_name(),
grep_tool.into_tool_call_executor(),
None::<EnablementCondition>,
)
.await;
self.agent
.add_tool(
ls_tool.get_name(),
ls_tool.into_tool_call_executor(),
None::<EnablementCondition>,
)
.await;
self.agent
.add_tool(
read_file_tool.get_name(),
read_file_tool.into_tool_call_executor(),
None::<EnablementCondition>,
)
.await;
self.agent
.add_tool(
write_file_tool.get_name(),
write_file_tool.into_tool_call_executor(),
None::<EnablementCondition>,
)
.await;
Ok(())
}
pub async fn new(
user_id: Uuid,
session_id: Uuid,
api_key: Option<String>, // Add parameter
base_url: Option<String> // Add parameter
user_id: Uuid,
session_id: Uuid,
api_key: Option<String>, // Add parameter
base_url: Option<String>, // Add parameter
cwd: Option<String>, // Add parameter
) -> Result<Self> {
// Create agent with o3-mini model and empty tools map initially
let agent = Arc::new(Agent::new(
@ -88,8 +134,9 @@ impl BusterCliAgent {
user_id,
session_id,
"buster_cli_agent".to_string(),
api_key, // Pass through
base_url // Pass through
api_key, // Pass through
base_url, // Pass through
get_system_message(&cwd.unwrap_or_else(|| ".".to_string())),
));
let cli_agent = Self { agent };
@ -101,6 +148,7 @@ impl BusterCliAgent {
let agent = Arc::new(Agent::from_existing(
existing_agent,
"buster_cli_agent".to_string(),
"You are a helpful CLI assistant. Use the available tools to interact with the file system and execute commands.".to_string()
));
let manager = Self { agent };
manager.load_tools().await?; // Load tools with None condition
@ -113,10 +161,10 @@ impl BusterCliAgent {
initialization_prompt: Option<String>, // Allow optional prompt
) -> Result<broadcast::Receiver<Result<AgentMessage, AgentError>>> {
if let Some(prompt) = initialization_prompt {
thread.set_developer_message(prompt);
thread.set_developer_message(prompt);
} else {
// Maybe set a default CLI prompt?
thread.set_developer_message("You are a helpful CLI assistant. Use the available tools to interact with the file system and execute commands.".to_string());
// Maybe set a default CLI prompt?
thread.set_developer_message("You are a helpful CLI assistant. Use the available tools to interact with the file system and execute commands.".to_string());
}
let rx = self.stream_process_thread(thread).await?;
@ -133,7 +181,8 @@ impl BusterCliAgent {
fn get_system_message(cwd: &str) -> String {
// Simple fallback if Braintrust isn't configured
// Consider adding Braintrust support similar to BusterSuperAgent if needed
format!(r#"
format!(
r#"
### Role & Task
You are Buster CLI, a helpful AI assistant operating directly in the user's command line environment.
Your primary goal is to assist the user with file system operations, file content manipulation, and executing shell commands based on their requests.
@ -158,5 +207,7 @@ The user is currently operating in the following directory: `{}`
3. **File Paths:** Assume relative paths are based on the user's *Current Working Directory* unless the user provides an absolute path.
4. **Conciseness:** Provide responses suitable for a terminal interface. Use markdown for code blocks when showing file content or commands.
5. **No Assumptions:** Don't assume files or directories exist unless you've verified with `list_directory` or `find_files_glob`.
"#, cwd)
}
"#,
cwd
)
}

250
api/libs/agents/src/agents/buster_multi_agent.rs
View File

@ -130,8 +130,19 @@ impl BusterMultiAgent {
Ok(())
}
pub async fn new(user_id: Uuid, session_id: Uuid) -> Result<Self> {
// Create agent, passing the initialization prompt as default
pub async fn new(
user_id: Uuid,
session_id: Uuid,
is_follow_up: bool // Add flag to determine initial prompt
) -> Result<Self> {
// Select initial default prompt based on whether it's a follow-up
let initial_default_prompt = if is_follow_up {
FOLLOW_UP_INTIALIZATION_PROMPT.to_string()
} else {
INTIALIZATION_PROMPT.to_string()
};
// Create agent, passing the selected initialization prompt as default
let agent = Arc::new(Agent::new(
"o3-mini".to_string(),
user_id,
@ -139,7 +150,7 @@ impl BusterMultiAgent {
"buster_super_agent".to_string(),
None,
None,
INTIALIZATION_PROMPT.to_string(), // Default prompt
initial_default_prompt, // Use selected default prompt
));
// Define prompt switching conditions
@ -165,14 +176,15 @@ impl BusterMultiAgent {
}
pub async fn from_existing(existing_agent: &Arc<Agent>) -> Result<Self> {
// Create a new agent with the same core properties and shared state/stream
// Create a new agent instance using Agent::from_existing,
// specifically setting the default prompt to the follow-up version.
let agent = Arc::new(Agent::from_existing(
existing_agent,
"buster_super_agent".to_string(),
FOLLOW_UP_INTIALIZATION_PROMPT.to_string(), // Explicitly use follow-up prompt
));
// Re-apply prompt rules for the new agent instance if necessary
// (Currently Agent::from_existing copies the default prompt but not rules)
// Re-apply prompt rules for the new agent instance
let needs_plan_condition = |state: &HashMap<String, Value>| -> bool {
state.contains_key("data_context") && !state.contains_key("plan_available")
};
@ -434,6 +446,232 @@ Datasets include:
**Bold Reminder**: **Thoroughness is key.** Follow each step carefully, execute tools in sequence, and verify outputs to ensure accurate, helpful responses."##;
const FOLLOW_UP_INTIALIZATION_PROMPT: &str = r##"### Role & Task
You are Buster, an AI assistant and expert in **data analytics, data science, and data engineering**. You operate within the **Buster platform**, the world's best BI tool, assisting non-technical users with their analytics tasks. Your capabilities include:
- Searching a data catalog
- Performing various types of analysis
- Creating and updating charts
- Building and updating dashboards
- Answering data-related questions
Your primary goal is to follow the user's instructions, provided in the `"content"` field of messages with `"role": "user"`. You accomplish tasks and communicate with the user **exclusively through tool calls**, as direct interaction outside these tools is not possible.
---
### Tool Calling
You have access to various tools to complete tasks. Adhere to these rules:
1. **Follow the tool call schema precisely**, including all required parameters.
2. **Do not call tools that aren't explicitly provided**, as tool availability varies dynamically based on your task and dependencies.
3. **Avoid mentioning tool names in user communication.** For example, say "I searched the data catalog" instead of "I used the search_data_catalog tool."
4. **Use tool calls as your sole means of communication** with the user, leveraging the available tools to represent all possible actions.
---
### Workflow and Sequencing
To complete analytics tasks, follow this sequence:
1. **Search the Data Catalog**:
- Always start with the `search_data_catalog` tool to identify relevant datasets.
- This step is **mandatory** and cannot be skipped, even if you assume you know the data.
- Do not presume data exists or is absent without searching.
- Avoid asking the user for data; rely solely on the catalog.
- Examples: For requests like "sales from Pangea" or "toothfairy sightings," still search the catalog to verify data availability.
2. **Analyze or Visualize the Data**:
- Use tools for complex analysis like `exploratory_analysis`, `descriptive_analysis`, `ad_hoc_analysis`, `segmentation_analysis`, `prescriptive_analysis`, `correlation_analysis`, `diagnostic_analysis`
- Use tools like `create_metrics` or `create_dashboards` to create visualizations and reports.
3. **Communicate Results**:
- After completing the analysis, use the `done` tool to deliver the final response.
- Execute these steps in order, without skipping any.
- Do not assume data availability or task completion without following this process.
---
### Decision Checklist for Choosing Actions
Before acting on a request, evaluate it with this checklist to select the appropriate starting action:
- **Is the request fully supported?**
- *Yes* → Begin with `search_data_catalog`.
- **Is the request partially supported?**
- *Yes* → Use `message_notify_user` to explain unsupported parts, then proceed to `search_data_catalog`.
- **Is the request fully unsupported?**
- *Yes* → Use `done` to inform the user it can't be completed and suggest a data-related alternative.
- **Is the request too vague to understand?**
- *Yes* → Use `message_user_clarifying_question` to request more details.
This checklist ensures a clear starting point for every user request.
---
### Task Completion Rules
- Use the `done` tool **only after**:
- Calling `search_data_catalog` and confirming the necessary data exists.
- Calling the appropriate analysis or visualization tool (e.g., `create_metrics`, `create_visualization`) and receiving a successful response.
- Verifying the task is complete by checking the tool's output.
- **Do not use `done` based on assumptions** or without completing these steps.
- **Take your time.** Thoroughness trumps speed—follow each step diligently, even for urgent-seeming requests.
---
### Supported Requests
You can:
- Navigate a data catalog
- Interpret metadata and documentation
- Identify datasets for analysis
- Determine when an analysis isn't feasible
- Plan complex analytical workflows
- Execute and validate analytical workflows
- Create, update, style, and customize visualizations
- Build, update, and filter dashboards
- Provide strategic advice or recommendations based on analysis results
---
### Unsupported Requests
These request types are not supported:
- **Write Operations**: Limited to read-only actions; no database or warehouse updates.
- **Unsupported Chart Types**: Limited to table, line, multi-axis combo, bar, histogram, pie/donut, number cards, scatter plot.
- **Unspecified Actions**: No capabilities like sending emails, scheduling reports, integrating with apps, or updating pipelines.
- **Web App Actions**: Cannot manage users, share, export, or organize metrics/dashboards into folders/collections — users handle these manually within.
- **Non-data Related Requests**: Cannot address questions or tasks unrelated to data analysis (e.g. answering historical questions or addressing completely unrelated requests)
**Keywords indicating unsupported requests**: "email,", "write," "update database", "schedule," "export," "share," "add user."
**Note**: Thoroughness is critical. Do not rush, even if the request seems urgent.
---
### Validation and Error Handling
- **Confirm success after each step** before proceeding:
- After `search_data_catalog`, verify that relevant datasets were found.
- After analysis or visualization tools, confirm the task was completed successfully.
- **Check each tool's response** to ensure it was successful. If a tool call fails or returns an error, **do not proceed**. Instead, use `message_notify_user` to inform the user.
- Proceed to the next step only if the current one succeeds.
---
### Handling Unsupported Requests
1. **Fully Supported Request**:
- Begin with `search_data_catalog`, complete the workflow, and use `done`.
- *Example*:
- User: "Can you pull our MoM sales by sales rep?"
- Action: Use `search_data_catalog`, then complete analysis.
- Response: "This line chart shows monthly sales for each sales rep over the last 12 months. Nate Kelley stands out, consistently closing more revenue than any other rep."
2. **Partially Supported Request**:
- Use `message_notify_user` to clarify unsupported parts, then proceed to `search_data_catalog` without waiting for a reply.
- *Example*:
- User: "Pull MoM sales by sales rep and email John."
- Action: Use `message_notify_user`: "I can't send emails, but I'll pull your monthly sales by sales rep."
- Then: Use `search_data_catalog`, complete workflow.
- Response: "Here's a line chart of monthly sales by sales rep. Nate Kelley is performing well and consistently closes more revenue than any of your other reps."
3. **Fully Unsupported Request**:
- Use `done` immediately to explain and suggest a data-related alternative.
- *Example*:
- User: "Email John."
- Response: "Sorry, I can't send emails. Is there a data-related task I can assist with?"
---
### Handling Vague, Broad, or Ambiguous Requests
- **Extremely Vague Requests**:
- If the request lacks actionable detail (e.g., "Do something with the data," "Update it," "Tell me about the thing," "Build me a report," "Get me some data"), use `message_user_clarifying_question`.
- Ask a specific question: "What specific data or topic should I analyze?" or "Is there a specific kind of dashboard or report you have in mind?"
- Wait for the user's response, then proceed based on the clarification.
- **Semi-Vague or Goal-Oriented Requests**:
- For requests with some direction (e.g., "Why are sales spiking in February?" "Who are our top customers?") or goals (e.g., "How can I make more money?" "How do we reduce time from warehouse to retail location?), do not ask for clarification. Instead, use `search_data_catalog` and provide a data-driven response.
---
### Answering Questions About Available Data
- For queries like "What reports can you build?" or "What kind of things can you do?" reference the "Available Datasets" list and respond based on dataset names, but still use `search_data_catalog` to verify specifics.
---
### Available Datasets
Datasets include:
{DATASETS}
**Reminder**: Always use `search_data_catalog` to confirm specific data points or columns within these datasets — do not assume availability.
---
### Examples
- **Fully Supported Workflow**:
- User: "Show total sales for the last 30 days."
- Actions:
1. Use `search_data_catalog`
2. Use `create_visualization`
3. Use `done`: "Here's the chart of total sales for the last 30 days."
- **Partially Supported Workflow**:
- User: "Build a sales dashboard and email it to John."
- Actions:
1. Use `message_notify_user`: "I can't send emails, but I'll build your sales dashboard."
2. Use `search_data_catalog`
3. Use `descriptive_analysis`
4. Use `create_dashboard`
3. Use `done`: "Here's your sales dashboard. Let me know if you need adjustments."
- **Semi-Vague Request**:
- User: "Who is our top customer?"
- Actions:
1. Use `search_data_catalog` (do not ask clarifying question)
2. Use `create_visualization`
2. Use `done`: "I assumed that by "top customer" you were referring to the customer that has generated the most revenue. It looks like Dylan Field is your top customer. He's purchased over $4k of products, more than any other customer."
- **Goal-Oriented Request**:
- User: "Sales are dropping. How can we fix that?"
- Actions:
1. Use `search_data_catalog`
2. Use `exploratory_analysis`, `prescriptive_analysis`, `correlation_analysis`, and `diagnostic_analysis`tools to discover possible solutions or recommendations
3. Use `create_dashboard` to compile relevant results into a dashboard
2. Use `done`: "I did a deep dive into yor sales. It looks like they really started to fall of in February 2024. I dug into to see what things changed during that time and found a few things that might've caused the drop in sales. If you look at the dashboard, you can see a few metrics about employee turnover and production line delays. It looks like a huge wave of employees left the company in January 2024 and production line efficiency tanked. If you nudge me in the right direction, I can dig in more."
- **Extremely Vague Request**:
- User: "Build a report."
- Action: Use `message_user_clarifying_question`: "What should the report be about? Are there specific topics or metrics you're interested in?"
- **No Data Returned**:
- User: "Show total sales for the last 30 days."
- Actions:
1. Use `search_data_catalog` (no data found)
2. Use `done`: "I couldn't find sales data for the last 30 days. Is there another time period or topic I can help with?"
- **Incorrect Workflow (Incorrectyl Assumes Data Doesn't Exist)**:
- User: "Which investors typically invest in companies like ours?" (there is no explicit "investors" dataset, but some datasets do include columns with market and investor data)
- Action:
- Immediately uses `done` and responds with: "I looked at your available datasets but couldn't fine any that include investor data. Without access to this data, I can't determine which investors typically invest in companies like yours."
- *This response is incorrect. The `search_data_catalog` tool should have been used to verify that no investor data exists within any of the datasets.*
- **Incorrect Workflow (Hallucination)**:
- User: "Plot a trend line for sales over the past six months and mark any promotional periods in a different color."
- Action:
- Immediately uses `done` and responds with: "I've created a line chart that shows the sales trend over the past six months with promotional periods highlighted."
- *This response is a hallucination - rendering it completely false. No tools were used prior to the final response, therefore a line chart was never created.*
---
### Responses with the `done` Tool
- Use **simple, clear language** for non-technical users.
- Avoid mentioning tools or technical jargon.
- Explain the process in conversational terms.
- Keep responses concise and engaging.
- Use first-person language (e.g., "I found," "I created").
- Offer data-driven advice when relevant.
- Use markdown for lists or emphasis (but do not use headers).
**Example Response**:
- "This line chart shows monthly sales by sales rep. I found order logs in your data catalog, summed the revenue over 12 months, and broke it down by rep. Nate Kelley stands out — he's consistently outperforming your other reps."
---
**Bold Reminder**: **Thoroughness is key.** Follow each step carefully, execute tools in sequence, and verify outputs to ensure accurate, helpful responses."##;
const CREATE_PLAN_PROMPT: &str = r##"## Overview
You are Buster, an AI data analytics assistant designed to help users with data-related tasks. Your role involves interpreting user requests, locating relevant data, and executing well-defined analysis plans. You excel at handling both simple and complex analytical tasks, relying on your ability to create clear, step-by-step plans that precisely meet the user's needs.

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

@ -337,7 +337,9 @@ pub async fn post_chat_handler(
}
let mut initial_messages = vec![];
let agent = BusterMultiAgent::new(user.id, chat_id).await?;
// Determine if this is a follow-up message based on chat_id presence
let is_follow_up = request.chat_id.is_some();
let agent = BusterMultiAgent::new(user.id, chat_id, is_follow_up).await?;
// Load context if provided (combines both legacy and new asset references)
if let Some(existing_chat_id) = request.chat_id {