megacode-agent-sdk 0.32.0

//! Pip: a text-message assistant that runs errands. Dana texts "book me a
//! train to Leeds on Friday"; Pip's server starts one harness session for the
//! errand and gives the model two tools written here in Rust.
//!
//! ```text
//! cargo run -p megacode-agent-sdk --example pip -- \
//!     sdk/rust/examples/pip.mock.toml /path/to/harness
//! ```
//!
//! The first argument is a mock file, the model's replies scripted, so the
//! example runs with no API key; leave it out to use a real model. The second
//! names the harness binary; without it the SDK reads
//! `MEGACODE_AGENT_SDK_BINARY`, and only then downloads the release.

use std::sync::{Arc, Mutex};

use megacode_agent_sdk::{Client, Decision, Error, Event, Options, Tool, ToolContext, ToolOutput};
use serde::Deserialize;
use serde_json::json;

const SYSTEM_PROMPT: &str = "You are Pip, a personal assistant that runs errands for Dana. \
    Look things up with search_trains, then pay with pay. Be brief.";

#[derive(Deserialize)]
struct Search {
    to: String,
    day: String,
}

#[derive(Deserialize)]
struct Pay {
    merchant: String,
    amount_pence: u64,
}

/// A session that is not a coding agent: no built-in tools, Pip's own system
/// prompt, nothing read from the machine, and two tools run by this program.
pub fn pip(mock: Option<&str>) -> Options {
    let search = Tool::new(
        "search_trains",
        "Find trains to a destination on a day.",
        json!({
            "type": "object",
            "properties": {"to": {"type": "string"}, "day": {"type": "string"}},
            "required": ["to", "day"],
        }),
        |args: Search, _ctx: ToolContext| async move {
            Ok::<_, String>(ToolOutput::text(format!(
                "{} to {}: 9:10 for 42 pounds, 11:40 for 55 pounds",
                args.day, args.to
            )))
        },
    );
    // `ask` makes harness send an approval request before every call.
    let pay = Tool::new(
        "pay",
        "Pay a merchant with Dana's card.",
        json!({
            "type": "object",
            "properties": {"merchant": {"type": "string"}, "amount_pence": {"type": "integer"}},
            "required": ["merchant", "amount_pence"],
        }),
        |args: Pay, _ctx: ToolContext| async move {
            Ok::<_, String>(ToolOutput::text(format!(
                "paid {} pence to {}",
                args.amount_pence, args.merchant
            )))
        },
    )
    .ask(true);
    let options = Options::plain(SYSTEM_PROMPT).tool(search).tool(pay);
    match mock {
        Some(file) => options.mock(file),
        None => options,
    }
}

/// What the errand sent Dana and what the session did.
pub struct Errand {
    pub texts: Vec<String>,
    pub events: Vec<Event>,
}

/// Runs the errand on `options`. `dana_replies` is what Dana texts back to
/// the one question; a real server would wait for her message here.
pub async fn errand(options: Options, dana_replies: &str) -> Result<Errand, Error> {
    let texts = Arc::new(Mutex::new(Vec::new()));
    let asked = texts.clone();
    let yes = dana_replies.trim().eq_ignore_ascii_case("yes");
    // The callback answers the approval request, so the stream never shows it.
    let options = options.on_approval(move |request| {
        let asked = asked.clone();
        async move {
            // The gate asks before the arguments are checked, so the amount
            // may be missing or not a whole number: that is a no.
            let Some(pence) = request.input["amount_pence"].as_u64() else {
                return Decision::Deny {
                    reason: Some("The amount to pay is missing or not a number.".into()),
                };
            };
            let amount = if pence % 100 == 0 {
                format!("£{}", pence / 100)
            } else {
                format!("£{}.{:02}", pence / 100, pence % 100)
            };
            asked.lock().unwrap().push(format!("OK to pay {amount}?"));
            if yes {
                Decision::AllowOnce
            } else {
                Decision::Deny {
                    reason: Some("Dana said no.".into()),
                }
            }
        }
    });
    let client = Client::start(options).await?;
    let mut turn = client.send("book me a train to Leeds on Friday").await?;
    let mut events = Vec::new();
    while let Some(event) = turn.next().await {
        events.push(event?);
    }
    client.close().await;
    let texts = texts.lock().unwrap().clone();
    Ok(Errand { texts, events })
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<(), Error> {
    let mut args = std::env::args().skip(1);
    let mock = args.next();
    let mut options = pip(mock.as_deref());
    if let Some(binary) = args.next() {
        options = options.binary(binary);
    }
    let done = errand(options, "yes").await?;
    for text in &done.texts {
        println!("to Dana: {text}");
    }
    for event in &done.events {
        if let Event::Result(result) = event {
            println!("{}", result.result.as_deref().unwrap_or("(no answer)"));
        }
    }
    Ok(())
}