gitdataai/lib/ai/agent/subagent.rs

use rig::agent::AgentBuilder;
use rig::client::CompletionClient;
use rig::completion::Prompt;
use tracing::{debug, warn};

use super::config::AgentConfig;
use super::helpers::with_retry;
use super::persistence::{
    ActiveAgentRun, AgentRealtime, AgentRuntime, AgentStreamEvent,
    estimate_output_tokens,
};
use super::request::{AgentExpert, AgentExpertOutput};
use crate::client::AiClient;
use crate::error::{AiError, AiResult};

pub async fn run_experts(
    client: &AiClient,
    config: &AgentConfig,
    experts: &[AgentExpert],
    realtime: Option<&AgentRealtime>,
    run: &ActiveAgentRun,
) -> AiResult<Vec<AgentExpertOutput>> {
    let mut outputs = Vec::with_capacity(experts.len());
    let mut failed_count = 0;

    for expert in experts {
        match run_single(client, config, expert, realtime, run).await {
            Ok(output) => {
                debug!(subagent_id = %output.id, role = %output.role, "subagent completed");
                outputs.push(output);
            }
            Err(error) => {
                warn!(subagent_id = %expert.id, role = %expert.role, error = %error, "subagent failed");
                let _ = publish_subagent_failed(
                    realtime,
                    run,
                    expert,
                    &error.to_string(),
                )
                .await;
                failed_count += 1;
            }
        }
    }

    debug!(
        total = experts.len(),
        ok = outputs.len(),
        failed = failed_count,
        "experts done"
    );
    Ok(outputs)
}

async fn run_single(
    client: &AiClient,
    config: &AgentConfig,
    expert: &AgentExpert,
    realtime: Option<&AgentRealtime>,
    run: &ActiveAgentRun,
) -> AiResult<AgentExpertOutput> {
    publish_subagent_started(realtime, run, config, expert).await?;

    let rig_client = client.llm_client().clone();
    let model_name = config.model.clone();
    let temperature = expert.temperature.or(config.temperature);
    let max_completion_tokens = expert
        .max_completion_tokens
        .or(config.max_completion_tokens);
    let retry_attempts = config.retry_max_attempts;
    let retry_delay_ms = config.retry_base_delay_ms;

    let prompt = expert.system_prompt.clone().unwrap_or_else(|| {
        format!(
            "You are a specialist subagent. Role: {}. Produce a concise expert answer for the parent chat agent.",
            expert.role
        )
    });

    let task = build_expert_task(expert);

    let (output, input_tokens_usage, output_tokens_usage) =
        with_retry(retry_attempts, retry_delay_ms, || {
            let rig_client = rig_client.clone();
            let model_name = model_name.clone();
            let prompt = prompt.clone();
            let task = task.clone();
            async move {
                let model = rig_client.completion_model(&model_name);
                let mut builder = AgentBuilder::new(model).preamble(&prompt);
                if let Some(temp) = temperature {
                    builder = builder.temperature(temp);
                }
                if let Some(mt) = max_completion_tokens {
                    builder = builder.max_tokens(mt);
                }
                let agent = builder.build();

                let response =
                    agent.prompt(&task).extended_details().await.map_err(
                        |e: rig::completion::PromptError| {
                            AiError::Api(e.to_string())
                        },
                    )?;

                Ok((
                    response.output,
                    response.usage.input_tokens,
                    response.usage.output_tokens,
                ))
            }
        })
        .await?;

    let input_tokens = input_tokens_usage as i64;
    let output_tokens = if output_tokens_usage > 0 {
        output_tokens_usage as i64
    } else {
        estimate_output_tokens(&output)
    };

    let result = AgentExpertOutput {
        id: expert.id.clone(),
        role: expert.role.clone(),
        task: expert.task.clone(),
        output,
        input_tokens,
        output_tokens,
    };

    publish_subagent_completed(realtime, run, config, &result).await?;
    Ok(result)
}

fn build_expert_task(expert: &AgentExpert) -> String {
    let mut task = String::new();

    if !expert.context.is_empty() {
        task.push_str("Retrieved context for this specialist task:\n");
        for (index, chunk) in expert.context.iter().enumerate() {
            task.push_str(&format!(
                "\n[{}] id={} source={}\n{}\n",
                index + 1,
                chunk.id,
                chunk.source.as_deref().unwrap_or("unknown"),
                chunk.content
            ));
        }
        task.push('\n');
    }

    task.push_str(&expert.task);
    task
}

async fn publish_subagent_started(
    realtime: Option<&AgentRealtime>,
    run: &ActiveAgentRun,
    config: &AgentConfig,
    expert: &AgentExpert,
) -> AiResult<()> {
    AgentRuntime::default()
        .publish(
            realtime,
            &AgentStreamEvent::SubagentStarted {
                conversation_id: run.conversation_id,
                message_id: run.message_id,
                subagent_id: expert.id.clone(),
                role: expert.role.clone(),
                task: expert.task.clone(),
                model: config.model.clone(),
            },
        )
        .await
}

async fn publish_subagent_completed(
    realtime: Option<&AgentRealtime>,
    run: &ActiveAgentRun,
    config: &AgentConfig,
    output: &AgentExpertOutput,
) -> AiResult<()> {
    AgentRuntime::default()
        .publish(
            realtime,
            &AgentStreamEvent::SubagentCompleted {
                conversation_id: run.conversation_id,
                message_id: run.message_id,
                subagent_id: output.id.clone(),
                role: output.role.clone(),
                task: output.task.clone(),
                output: output.output.clone(),
                input_tokens: output.input_tokens,
                output_tokens: output.output_tokens,
                model: config.model.clone(),
            },
        )
        .await
}

async fn publish_subagent_failed(
    realtime: Option<&AgentRealtime>,
    run: &ActiveAgentRun,
    expert: &AgentExpert,
    error: &str,
) -> AiResult<()> {
    AgentRuntime::default()
        .publish(
            realtime,
            &AgentStreamEvent::SubagentFailed {
                conversation_id: run.conversation_id,
                message_id: run.message_id,
                subagent_id: expert.id.clone(),
                error: error.to_string(),
            },
        )
        .await
}