hyperhive/swarmctl/src/inference.rs

//! `swarmctl agent inference` — the model, effort and subagent models the
//! swarm sets for one agent, through swarm-controller's
//! `/api/agents/{name}/inference` routes.
//!
//! The agent picks a change up on its next read of its key, within ten
//! seconds, and its next turn runs on it. A field the swarm does not set
//! leaves the agent on its own value.
//!
//! The bodies come from `swarm-controller-api`, which the controller imports
//! too.

use anyhow::{Context as _, Result, bail};
use swarm_controller_api::inference::{
    AgentInferenceView, EffectiveInference, Effort, InferenceSettings,
};

use crate::agent::{parse_ident, send};
use crate::remote::Target;

/// What `agent inference set` was asked to change.
pub(crate) struct SetFlags {
    pub model: Option<String>,
    pub effort: Option<String>,
    pub subagent_models: Option<Vec<String>>,
}

impl SetFlags {
    /// `current` with each given flag written over it.
    fn merge(self, current: InferenceSettings) -> Result<InferenceSettings> {
        if self.model.is_none() && self.effort.is_none() && self.subagent_models.is_none() {
            bail!("nothing to change: pass --model, --effort or --subagent-models");
        }
        let effort = self.effort.as_deref().map(parse_effort).transpose()?;
        Ok(InferenceSettings {
            model: self.model.or(current.model),
            effort: effort.or(current.effort),
            subagent_models: self.subagent_models.or(current.subagent_models),
            max_subagents: current.max_subagents,
        })
    }
}

/// The level a `--effort` flag names. The controller refuses an unknown one
/// too; this fails before the round trip.
fn parse_effort(value: &str) -> Result<Effort> {
    serde_json::from_value(serde_json::Value::String(value.to_owned())).map_err(|_| {
        anyhow::anyhow!("invalid effort {value:?}: expected low, medium, high, xhigh or max")
    })
}

fn runtime() -> Result<tokio::runtime::Runtime> {
    tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("starting a tokio runtime for the controller request")
}

/// The agent's settings path. The name is checked first, so it cannot
/// reshape the URI.
fn member(agent: &str) -> Result<String> {
    let agent = parse_ident(agent, "agent name")?;
    Ok(format!("/api/agents/{agent}/inference"))
}

/// Run `swarmctl agent inference get`.
pub(crate) fn get(target: &Target, agent: &str) -> Result<()> {
    let uri = member(agent)?;
    let view: AgentInferenceView = runtime()?.block_on(send(
        target,
        hyper::Method::GET,
        &uri,
        None::<&()>,
        "inference-get",
    ))?;
    print!(
        "{}{}",
        render_configured(view.configured.as_ref()),
        render_effective(view.effective.as_ref())
    );
    Ok(())
}

/// Run `swarmctl agent inference set`: read the agent's settings, write the
/// flags over them, and store the result.
pub(crate) fn set(target: &Target, agent: &str, flags: SetFlags) -> Result<()> {
    let uri = member(agent)?;
    let rt = runtime()?;
    let current: AgentInferenceView = rt.block_on(send(
        target,
        hyper::Method::GET,
        &uri,
        None::<&()>,
        "inference-get",
    ))?;
    let wanted = flags.merge(current.configured.unwrap_or_default())?;
    let stored: InferenceSettings = rt.block_on(send(
        target,
        hyper::Method::PUT,
        &uri,
        Some(&wanted),
        "inference-set",
    ))?;
    print!("{}", render(&stored));
    Ok(())
}

/// Run `swarmctl agent inference clear`.
pub(crate) fn clear(target: &Target, agent: &str) -> Result<()> {
    let uri = member(agent)?;
    runtime()?.block_on(send::<(), ()>(
        target,
        hyper::Method::DELETE,
        &uri,
        None,
        "inference-clear",
    ))?;
    println!("cleared; {agent} goes back to its own model and effort");
    Ok(())
}

/// The swarm's setting, or one line saying it has none.
fn render_configured(s: Option<&InferenceSettings>) -> String {
    s.map_or_else(
        || "no swarm setting; the agent runs on its own model and effort\n".to_owned(),
        render,
    )
}

/// One line per field, naming the unset ones.
fn render(s: &InferenceSettings) -> String {
    const UNSET: &str = "(not set; the agent's own)";
    format!(
        "model: {}\neffort: {}\nsubagent_models: {}\nmax_subagents: {}\n",
        s.model.as_deref().unwrap_or(UNSET),
        s.effort.map_or(UNSET, Effort::as_str),
        s.subagent_models
            .as_ref()
            .map_or_else(|| UNSET.to_owned(), |m| m.join(", ")),
        s.max_subagents
            .map_or_else(|| "(not set; no cap)".to_owned(), |n| n.to_string())
    )
}

/// What the agent runs on, each value with where it comes from.
fn render_effective(e: Option<&EffectiveInference>) -> String {
    e.map_or_else(
        || {
            "running: not reported since swarm-controller started; the agent reports within a \
             minute\n"
                .to_owned()
        },
        |e| {
            format!(
                "running model: {} ({})\nrunning effort: {} ({})\n",
                e.model,
                e.model_source.as_str(),
                e.effort,
                e.effort_source.as_str()
            )
        },
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    fn flags() -> SetFlags {
        SetFlags {
            model: None,
            effort: None,
            subagent_models: None,
        }
    }

    /// A flag replaces its field and every other field keeps what is stored.
    #[test]
    fn set_merges_into_the_stored_value() {
        let current: InferenceSettings =
            serde_json::from_str(r#"{"model":"sonnet","subagent_models":["haiku"]}"#).unwrap();
        let merged = SetFlags {
            effort: Some("high".to_owned()),
            ..flags()
        }
        .merge(current)
        .unwrap();
        assert_eq!(
            serde_json::to_string(&merged).unwrap(),
            r#"{"model":"sonnet","effort":"high","subagent_models":["haiku"]}"#
        );
        let merged = SetFlags {
            model: Some("opus".to_owned()),
            ..flags()
        }
        .merge(merged)
        .unwrap();
        assert_eq!(merged.model.as_deref(), Some("opus"));
        assert_eq!(merged.effort, Some(Effort::High));
    }

    /// `set` has no flag for the cap, and a stored one survives it.
    #[test]
    fn set_keeps_the_stored_max_subagents() {
        let current: InferenceSettings = serde_json::from_str(r#"{"max_subagents":3}"#).unwrap();
        let merged = SetFlags {
            model: Some("opus".to_owned()),
            ..flags()
        }
        .merge(current)
        .unwrap();
        assert_eq!(merged.max_subagents, Some(3));
    }

    /// An effort the controller would answer 422 for fails before the request.
    #[test]
    fn set_with_an_unknown_effort_is_refused() {
        let flags = |effort: &str| SetFlags {
            effort: Some(effort.to_owned()),
            ..flags()
        };
        let err = flags("extreme")
            .merge(InferenceSettings::default())
            .unwrap_err();
        assert!(err.to_string().contains("invalid effort \"extreme\""));
        assert_eq!(
            flags("xhigh")
                .merge(InferenceSettings::default())
                .unwrap()
                .effort,
            Some(Effort::Xhigh)
        );
    }
    #[test]
    fn set_with_no_flag_is_refused() {
        assert!(flags().merge(InferenceSettings::default()).is_err());
    }

    #[test]
    fn get_shows_the_running_values_with_their_sources() {
        let view: AgentInferenceView = serde_json::from_str(
            r#"{"configured":{"model":"opus"},"effective":{"model":"opus","effort":"medium",
                "model_source":"swarm","effort_source":"local_file"}}"#,
        )
        .unwrap();
        assert_eq!(
            view.configured.as_ref().and_then(|s| s.model.as_deref()),
            Some("opus")
        );
        assert_eq!(
            render_effective(view.effective.as_ref()),
            "running model: opus (swarm)\nrunning effort: medium (local_file)\n"
        );
        let unknown: AgentInferenceView =
            serde_json::from_str(r#"{"configured":null,"effective":null}"#).unwrap();
        assert!(render_configured(unknown.configured.as_ref()).starts_with("no swarm setting"));
        assert!(render_effective(unknown.effective.as_ref()).starts_with("running: not reported"));
        // A stored setting with no fields still lists every field as unset.
        let empty: AgentInferenceView =
            serde_json::from_str(r#"{"configured":{},"effective":null}"#).unwrap();
        assert_eq!(
            render_configured(empty.configured.as_ref()).lines().count(),
            4
        );
    }

    #[test]
    fn an_empty_value_renders_every_field_as_unset() {
        let row = render(&serde_json::from_str("{}").unwrap());
        assert_eq!(row.lines().count(), 4);
        assert!(row.ends_with("max_subagents: (not set; no cap)\n"), "{row}");
        assert!(
            row.lines()
                .take(3)
                .all(|l| l.ends_with("(not set; the agent's own)")),
            "{row}"
        );
    }
}