openrouter-batch 0.1.2

//! CLI that exercises the crate against the live OpenRouter API:
//!
//! ```text
//! cargo run --release -- structured --prompt "What is 2+2? State the number."
//! cargo run --release -- embeddings --input "first text" --input "second text"
//! cargo run --release -- batch
//! cargo run --release -- poll --batch-id <batch-id>
//! ```
//!
//! Key resolution order: `--api-key`, then `OPENROUTER_API_KEY` in the
//! environment, then a `.env` file in the current directory. The key is
//! passed to the transport and never printed.

use std::time::Duration;

use anyhow::{Context, bail};
use clap::{Parser, Subcommand};
use openrouter_batch::{
    BatchObject, BatchStatus, BatchTransport, EmbeddingTransport, HttpTransport, Message,
    StructuredParams, embedding_vectors, submit_structured,
};
use serde_json::json;

/// Demo strict-mode schema shared by the structured and batch commands.
fn demo_schema() -> serde_json::Value {
    json!({
        "type": "object",
        "properties": {
            "answer": {"type": "string"},
            "confidence": {"type": "integer", "minimum": 0, "maximum": 100}
        },
        "required": ["answer", "confidence"],
        "additionalProperties": false
    })
}

fn demo_params(max_tokens: u32) -> StructuredParams {
    StructuredParams {
        schema: demo_schema(),
        schema_name: "cli_demo".to_string(),
        max_tokens,
        disable_reasoning: false,
        require_parameters: false,
        reasoning_max_tokens: None,
    }
}

#[derive(Parser)]
#[command(
    name = "openrouter-cli",
    about = "Live smoke tests for openrouter-batch-client"
)]
struct Cli {
    #[command(subcommand)]
    command: Command,

    /// OpenRouter bearer token. Falls back to OPENROUTER_API_KEY, then .env.
    #[arg(long, global = true)]
    api_key: Option<String>,

    /// API root. Falls back to OPENROUTER_API_ROOT, then .env, then the public endpoint.
    #[arg(long, global = true)]
    api_root: Option<String>,

    /// Total client request timeout in seconds.
    #[arg(long, global = true, default_value_t = 300)]
    timeout: u64,
}

#[derive(Subcommand)]
enum Command {
    /// One synchronous structured chat completion.
    Structured {
        #[arg(long, default_value = "z-ai/glm-5.3-flash")]
        model: String,
        #[arg(
            long,
            default_value = "What is 2+2? State the number and your confidence."
        )]
        prompt: String,
        #[arg(long, default_value_t = 500)]
        max_tokens: u32,
    },
    /// One synchronous embeddings call, decoded back into input-ordered vectors.
    Embeddings {
        #[arg(long, default_value = "perplexity/pplx-embed-v1-0.6b")]
        model: String,
        /// Input texts; repeat the flag for an ordered array.
        #[arg(long = "input", required = true)]
        inputs: Vec<String>,
    },
    /// Submit a two-item structured batch, poll to a terminal status, print results.
    Batch {
        #[arg(long, default_value = "z-ai/glm-5.3-flash:batch")]
        model: String,
        #[arg(long, default_value_t = 5)]
        poll_interval_seconds: u64,
        /// Stop polling after this many seconds even if the batch is unfinished.
        #[arg(long, default_value_t = 600)]
        poll_timeout_seconds: u64,
    },
    /// Poll an existing batch id to a terminal status and print results.
    /// Use this to resume watching a batch after a poll timeout; the batch
    /// keeps processing server-side and results stay retrievable for 30 days.
    Poll {
        /// Batch id returned when the batch was submitted.
        #[arg(long)]
        batch_id: String,
        #[arg(long, default_value_t = 5)]
        poll_interval_seconds: u64,
        /// Stop polling after this many seconds even if the batch is unfinished.
        #[arg(long, default_value_t = 600)]
        poll_timeout_seconds: u64,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let cli = Cli::parse();
    let env = dotenv_map();
    let api_key = cli
        .api_key
        .or_else(|| std::env::var("OPENROUTER_API_KEY").ok())
        .or_else(|| env.get("OPENROUTER_API_KEY").cloned())
        .filter(|key| !key.is_empty())
        .context("no API key: pass --api-key, set OPENROUTER_API_KEY, or add a .env file")?;
    let api_root = cli
        .api_root
        .or_else(|| std::env::var("OPENROUTER_API_ROOT").ok())
        .or_else(|| env.get("OPENROUTER_API_ROOT").cloned())
        .filter(|root| !root.is_empty())
        .unwrap_or_else(|| "https://openrouter.ai/api".to_string());
    let transport =
        HttpTransport::with_timeout(&api_root, &api_key, Duration::from_secs(cli.timeout))
            .context("transport setup")?;

    match cli.command {
        Command::Structured {
            model,
            prompt,
            max_tokens,
        } => {
            let response = transport
                .chat_completion(
                    &model,
                    &demo_params(max_tokens),
                    vec![Message::user(prompt)],
                )
                .await?;
            println!("provider_request_id: {:?}", response.provider_request_id);
            println!("content: {}", response.content()?);
            println!(
                "raw response:\n{}",
                serde_json::to_string_pretty(&response.raw)?
            );
        }
        Command::Embeddings { model, inputs } => {
            let raw = transport
                .embed(&model, &inputs, &uuid::Uuid::new_v4().to_string())
                .await?;
            // Prove the full round trip: decode the raw payload back into
            // input-ordered vectors, exactly as a consumer would.
            let vectors = embedding_vectors(&raw.raw, inputs.len()).context("decode")?;
            println!("provider_request_id: {:?}", raw.provider_request_id);
            println!(
                "vectors: {} dimensions: {}",
                vectors.len(),
                vectors[0].len()
            );
            println!(
                "vector[0][0..4]: {:?}",
                &vectors[0][..vectors[0].len().min(4)]
            );
        }
        Command::Batch {
            model,
            poll_interval_seconds,
            poll_timeout_seconds,
        } => {
            run_batch(
                &transport,
                &model,
                poll_interval_seconds,
                poll_timeout_seconds,
            )
            .await?;
        }
        Command::Poll {
            batch_id,
            poll_interval_seconds,
            poll_timeout_seconds,
        } => {
            let batch = poll_to_terminal(
                &transport,
                &batch_id,
                poll_interval_seconds,
                poll_timeout_seconds,
                None,
            )
            .await?;
            print_batch(batch);
        }
    }
    Ok(())
}

async fn run_batch(
    transport: &HttpTransport,
    model: &str,
    poll_interval_seconds: u64,
    poll_timeout_seconds: u64,
) -> anyhow::Result<()> {
    let items: Vec<(String, String)> = vec![
        (
            "req-1".to_string(),
            "What is 2+2? State the number and your confidence.".to_string(),
        ),
        (
            "req-2".to_string(),
            "What is 3*7? State the number and your confidence.".to_string(),
        ),
    ];
    let submitted = submit_structured(
        transport,
        model,
        &uuid::Uuid::new_v4().to_string(),
        &demo_params(500),
        items
            .iter()
            .map(|(id, prompt)| (id.clone(), vec![Message::user(prompt.clone())]))
            .collect(),
    )
    .await?;
    println!("submitted: {} status: {:?}", submitted.id, submitted.status);

    let batch = poll_to_terminal(
        transport,
        &submitted.id,
        poll_interval_seconds,
        poll_timeout_seconds,
        Some(submitted.status),
    )
    .await?;
    print_batch(batch);
    Ok(())
}

/// Poll a batch id until a terminal status or the deadline. Prints each
/// status transition. `last_status` seeds the already-known status so an
/// unchanged status is not printed again.
async fn poll_to_terminal(
    transport: &HttpTransport,
    batch_id: &str,
    poll_interval_seconds: u64,
    poll_timeout_seconds: u64,
    last_status: Option<BatchStatus>,
) -> anyhow::Result<BatchObject> {
    let deadline = std::time::Instant::now() + Duration::from_secs(poll_timeout_seconds);
    let mut last_status = last_status;
    let mut consecutive_not_found = 0u32;
    loop {
        if std::time::Instant::now() >= deadline {
            let seen = last_status
                .map(|status| format!("{status:?}"))
                .unwrap_or_else(|| "unobserved".to_string());
            bail!("poll timeout: batch {} still {}", batch_id, seen);
        }
        tokio::time::sleep(Duration::from_secs(poll_interval_seconds)).await;
        match transport.get(batch_id).await {
            // The submit → first-poll window can 404 (eventual consistency);
            // treat it as not-yet-visible rather than gone.
            Err(error) if error.not_found() => {
                consecutive_not_found += 1;
                if consecutive_not_found > 12 {
                    return Err(error).context("batch never became visible");
                }
            }
            Err(error) => return Err(error.into()),
            Ok(batch) => {
                consecutive_not_found = 0;
                if last_status != Some(batch.status) {
                    println!("status: {:?}", batch.status);
                    last_status = Some(batch.status);
                }
                if batch.status.is_terminal() {
                    return Ok(batch);
                }
            }
        }
    }
}

/// Print the batch summary and each decoded result line.
fn print_batch(batch: BatchObject) {
    println!("request_counts: {:?}", batch.request_counts);
    println!("usage: {:?}", batch.usage);
    for line in batch.results.unwrap_or_default() {
        match (line.custom_id(), line.chat_content()) {
            (Some(id), Ok(content)) => println!("{id}: {content}"),
            (Some(id), Err(error)) => println!("{id}: per-line error: {error}"),
            (None, _) => println!("undecodable line: {}", line.raw()),
        }
    }
}

/// Minimal `.env` reader for KEY=VALUE lines in the current directory.
/// Values are returned to the caller only; they are never printed.
fn dotenv_map() -> std::collections::HashMap<String, String> {
    let mut map = std::collections::HashMap::new();
    let Ok(contents) = std::fs::read_to_string(".env") else {
        return map;
    };
    for line in contents.lines() {
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        if let Some((key, value)) = line.split_once('=') {
            map.entry(key.trim().to_string())
                .or_insert_with(|| value.trim().to_string());
        }
    }
    map
}