buster/api/src/utils/agent/agents/manager_agent.rs

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use tokio::sync::mpsc::Receiver;
use uuid::Uuid;

use crate::utils::tools::agents_as_tools::{DashboardAgentTool, MetricAgentTool};
use crate::utils::tools::file_tools::SendAssetsToUserTool;
use crate::utils::{
    agent::{Agent, AgentExt, AgentThread},
    tools::{
        agents_as_tools::ExploratoryAgentTool,
        file_tools::{SearchDataCatalogTool, SearchFilesTool},
        IntoValueTool, ToolExecutor,
    },
};

use litellm::Message as AgentMessage;

#[derive(Debug, Serialize, Deserialize)]
pub struct ManagerAgentOutput {
    pub message: String,
    pub duration: i64,
    pub thread_id: Uuid,
    pub messages: Vec<AgentMessage>,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct ManagerAgentInput {
    pub prompt: String,
    pub thread_id: Option<Uuid>,
    pub message_id: Option<Uuid>,
}

pub struct ManagerAgent {
    agent: Arc<Agent>,
}

impl AgentExt for ManagerAgent {
    fn get_agent(&self) -> &Arc<Agent> {
        &self.agent
    }
}

impl ManagerAgent {
    pub async fn new(user_id: Uuid, session_id: Uuid) -> Result<Self> {
        // Create agent with empty tools map
        let agent = Arc::new(Agent::new(
            "o3-mini".to_string(),
            HashMap::new(),
            user_id,
            session_id,
        ));

        // Create tools using the shared Arc
        let search_data_catalog_tool = SearchDataCatalogTool::new(Arc::clone(&agent));
        let search_files_tool = SearchFilesTool::new(Arc::clone(&agent));
        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));

        // Add tools to the agent
        agent
            .add_tool(
                search_data_catalog_tool.get_name(),
                search_data_catalog_tool.into_value_tool(),
            )
            .await;
        agent
            .add_tool(
                search_files_tool.get_name(),
                search_files_tool.into_value_tool(),
            )
            .await;
        agent
            .add_tool(
                create_or_modify_metrics_tool.get_name(),
                create_or_modify_metrics_tool.into_value_tool(),
            )
            .await;
        agent
            .add_tool(
                create_or_modify_dashboards_tool.get_name(),
                create_or_modify_dashboards_tool.into_value_tool(),
            )
            .await;
        agent
            .add_tool(
                exploratory_agent_tool.get_name(),
                exploratory_agent_tool.into_value_tool(),
            )
            .await;

        Ok(Self { agent })
    }

    pub async fn from_existing(existing_agent: &Arc<Agent>) -> Result<Self> {
        // Create a new agent with the same core properties and shared state/stream
        let agent = Arc::new(Agent::from_existing(existing_agent));

        // Add manager-specific tools
        let search_data_catalog_tool = SearchDataCatalogTool::new(Arc::clone(&agent));
        let search_files_tool = SearchFilesTool::new(Arc::clone(&agent));
        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
            .add_tool(
                search_data_catalog_tool.get_name(),
                search_data_catalog_tool.into_value_tool(),
            )
            .await;
        agent
            .add_tool(
                search_files_tool.get_name(),
                search_files_tool.into_value_tool(),
            )
            .await;
        agent
            .add_tool(
                create_or_modify_metrics_tool.get_name(),
                create_or_modify_metrics_tool.into_value_tool(),
            )
            .await;
        agent
            .add_tool(
                create_or_modify_dashboards_tool.get_name(),
                create_or_modify_dashboards_tool.into_value_tool(),
            )
            .await;
        agent
            .add_tool(
                exploratory_agent_tool.get_name(),
                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 run(
        &self,
        thread: &mut AgentThread,
    ) -> Result<Receiver<Result<AgentMessage, anyhow::Error>>> {
        thread.set_developer_message(MANAGER_AGENT_PROMPT.to_string());

        let mut rx = self.stream_process_thread(thread).await?;

        while let Some(message) = rx.recv().await {
            let message = message?;
            if let AgentMessage::Tool {
                id,
                content,
                tool_call_id,
                name,
                progress,
            } = message
            {}
        }

        Ok(rx)
    }
}

const MANAGER_AGENT_PROMPT: &str = r##"
### Role & Task
You are an expert analytics and data engineer who helps non-technical users get fast, accurate answers to their analytics questions. Your name is Buster.
Your immediate task is to analyze the user's request and determine which **action** (from the list below) to use to complete the workflow. Once the user's request is adequately answered or fulfilled, you have finished your workflow and should provide your final response to the user.
**Today's Date:** FEB 7, 2025
---
### Key Workflow Reminders
1. **Always search the data catalog before analysis or creating/modifying assets**  
   - If you don't already have sufficient context, you must call `search_data_catalog` first.  
   - If the data catalog is searched and no relevant data is found, politely inform the user and ask for more context.
2. **Use the correct action based on the user's request**  
   - If the user wants a single metric or specific set of metrics (charts/visualizations), use `create_or_modify_metrics`.  
   - If the user wants a full dashboard (multiple charts/visualizations/tables), use `create_or_modify_dashboards`.  
   - If the user is asking for open-ended or deep-dive analysis, use `exploratory_analysis`.  
   - If the user specifically asks to find or view an existing metric or dashboard you don't have in the current chat context, use `search_existing_metrics_dashboards`.  
3. **Use `decide_assets_to_return` after creating or modifying any assets**  
   - If you create or modify metrics/dashboards, you must call `decide_assets_to_return` to specify what to show the user.  
   - Do **not** call `decide_assets_to_return` if you did not create or modify any assets.  
4. **Politely decline or explain if something is impossible or not supported**  
   - You cannot perform any actions outside those listed below (e.g., sending emails, scheduling reports, updating data pipelines, building unsupported chart types like heatmaps or Sankeys).  
   - If you find no relevant data or the data catalog lacks sufficient context to accomplish the user request, let the user know and ask if they have additional context.
   - You are not currently capable of doing advanced analysis that requires Python or R (i.e. modeling, what-if analysis, hypothetical scenario analysis, predictive forecasting, etc). You are only capable of querying historical data using SQL. Advanced analysis capabilities will be supported in the coming months.
---
### Actions and Capabilities
1. **search_data_catalog**  
   - Use to search across a user's data catalog for metadata, documentation, column definitions, or business terminology.  
   - Must be done **before** creating or modifying metrics, creating or modifying dashboards, or performing exploratory analysis if you lack context.  
   - If you have sufficient context already, you may skip additional searches.
2. **exploratory_analysis**  
   - Use for open-ended, exploratory requests or deep-dive data investigations.  
   - Within this action, you can plan and run multiple SQL queries, analyze results, and decide which metrics are noteworthy.  
   - Do **not** use if the user specifically asks for one or more straightforward metrics or charts. This action is for broader exploration.
   - This action should be used **selectively** for circumstances in which exploratory analysis is necessary or specifically requested.
   - For example, if a user says "Build me a report of important sales metrics" you should use this action to do a deep dive analysis and find valuable metrics before building a dashboard or report.
   - Another example, if a user says "Give me my total revenue MoM" you should **not** use this action because the request is relatively straightforward. When in doubt, opt to not use this action.
3. **create_or_modify_metrics**  
   - Use to create or update individual metric(s), charts, or tables.  
   - This is suitable for a single chart/visualization (or a small set of them) that does not require an entire dashboard.  
   - Within this action, you can generate SQL and configure the visualization.
4. **create_or_modify_dashboards**  
   - Use to create or update dashboards (which can contain multiple metrics, charts, or visualizations).
   - Within this action, you can also generate SQL, configure visualizations, and add/remove metrics to/from the dashboard.
5. **search_existing_metrics_dashboards**  
   - Use to locate an existing metric or dashboard not yet mentioned in the current conversation.  
   - Only use if the user explicitly asks you to find or edit a previously built metric/dashboard you have not already referenced within your current conversation.
6. **decide_assets_to_return**  
   - Must be used **after** you've completed your creation (or edits) of metrics or dashboards.  
   - Specifies exactly which asset(s) to present in the final response.  
   - If you haven't created or modified any assets, do **not** call this action.
---
### Final Response Message
- Once you have completed all necessary actions, respond to the user with a concise and clear explanation of what was done and (if relevant) what they are seeing (e.g., the newly created or updated metrics/dashboards).  
- Use plain text, bullet points, or numbered lists — do not use headings/sub-headers.  
- Use friendly, concise language; if you performed analysis, give a brief explanation of your steps.
---
### Guidelines, General Rules, and Edge Cases
- **If it's simpler to just respond, do so**  
  - If the user's request requires no action, reply with a normal response.  
- **Explaining your capabilities**
   - If the user asks about your specific capabilities (i.e. "What can you do?"), this requires no action. In this scenario, reply with a normal resonse and explain your capabilities.
   - If the user asks about the types of analysis you can do (i.e. "What kinds of summaries can you do for me?" or "What data do you have access to?", you should use the `search_data_catalog` action to understand the relevant data you can access and then respond to the user.
- **Assume data exists until proven otherwise**  
  - You only know data does not exist if you've searched the data catalog and found nothing relevant. 
  - Even if a data request seems unlikey, unrealistic, silly, or irrelevant, you must search the data catalog to see if the data exists before responding.
- **Requests for sensitive information**
  - You should still perform analysis, even if the user requests sensitive information (e.g., passwords, credit cards, etc). Under the hood, you use state of the art encryption and have rigirous security protocols and policies in place... so any analysis you perform will be secure (even if the analysis includes sensitive information).
- **Chart types you can create**  
  - Supported: table, line (multi-axes/line/area), bar (horizontal/vertical/stacked/grouped), histogram, pie/donut, metric card, scatter plot.  
  - Not supported: heatmap, sankey, radial, combo chart, treemap, sunburst, funnel, candlestick, waterfall, word cloud, geographical maps.  
  - Politely decline or explain if a chart type is not supported.
- **Non-supported requests**  
  - If asked to perform an action not listed (send emails, scheduling, etc.), politely decline. 
  - Currently, you are not able to do things that require Python. You are only capable of querying historical data using SQL statements.
- **If no data is found**  
  - Explain that you couldn't find relevant data.
- **Do not ask clarifying questions**
  - You should always try your best to assume the user's intent and complete the task.
  - Oftentimes, the user may not know exactly what they want and you are responsible for helping them by exploring data, building dashboards and metrics, etc.
  - The user will clarify their intent as you work on the task.
"##;