aether_cli/generate_command/
mod.rs

1use crate::output::OutputFormat;
2use futures::StreamExt;
3use llm::catalog::{ReasoningEffortError, validate_reasoning_effort};
4use llm::parser::ModelProviderParser;
5use llm::{ChatMessage, Context, LlmError, LlmResponse, ModelSettings, ReasoningEffort, StreamingModelProvider};
6use std::io::{Read, stdin};
7use std::process::ExitCode;
8use thiserror::Error;
9
10#[derive(clap::Args)]
11pub struct GenerateArgs {
12    /// Model to call, as `provider:model` (e.g. `anthropic:claude-sonnet-4-5`).
13    #[arg(long)]
14    pub model: String,
15
16    /// Prompt text to send.
17    #[arg(long, conflicts_with = "prompt_file", required_unless_present = "prompt_file")]
18    pub prompt: Option<String>,
19
20    /// Prompt to send from a file path, or `-` to read from stdin.
21    #[arg(long, value_name = "PATH_OR_DASH", conflicts_with = "prompt", required_unless_present = "prompt")]
22    pub prompt_file: Option<String>,
23
24    /// Optional system prompt.
25    #[arg(long)]
26    pub system: Option<String>,
27
28    /// Sampling controls (temperature, top-p, max-tokens) for the model call.
29    #[command(flatten)]
30    pub model_settings: ModelSettingsArgs,
31
32    /// Reasoning effort for models that support extended thinking
33    #[arg(long)]
34    pub reasoning_effort: Option<ReasoningEffort>,
35
36    /// Output format. `text` prints the raw response; `json`/`pretty` wrap it as
37    /// `{ "text": <response>, "model": <model> }`.
38    #[arg(long, default_value = "text")]
39    pub output: OutputFormat,
40}
41
42#[derive(Debug, Clone, Default, clap::Args)]
43pub struct ModelSettingsArgs {
44    /// Sampling temperature. Lower is more deterministic (e.g. `0` for grading).
45    #[arg(long)]
46    pub temperature: Option<f32>,
47    /// Nucleus sampling: the probability mass to sample from.
48    #[arg(long)]
49    pub top_p: Option<f32>,
50    /// Upper bound on the number of tokens generated in the response.
51    #[arg(long)]
52    pub max_tokens: Option<u32>,
53}
54
55impl From<ModelSettingsArgs> for ModelSettings {
56    fn from(args: ModelSettingsArgs) -> Self {
57        ModelSettings { temperature: args.temperature, top_p: args.top_p, max_tokens: args.max_tokens }
58    }
59}
60
61#[derive(Debug, Error)]
62pub enum GenerateCommandError {
63    #[error("provide exactly one of --prompt or --prompt-file")]
64    PromptSource,
65
66    #[error("failed to read prompt from {path}: {source}")]
67    ReadPrompt { path: String, source: std::io::Error },
68
69    #[error("invalid reasoning effort: {0}")]
70    ReasoningEffort(#[from] ReasoningEffortError),
71
72    #[error("failed to initialize model `{model}`: {source}")]
73    Model { model: String, source: LlmError },
74
75    #[error("model stream error: {0}")]
76    Stream(LlmError),
77}
78
79/// Call a model with a single prompt and print its response. The judge is a special case of this:
80/// the caller supplies a grading prompt and parses the structured verdict out of the response.
81pub async fn run(args: GenerateArgs) -> Result<ExitCode, GenerateCommandError> {
82    let prompt = resolve_prompt(args.prompt.as_deref(), args.prompt_file.as_deref())?;
83    validate_reasoning_effort(&args.model, args.reasoning_effort)?;
84    let (provider, _) = ModelProviderParser::default()
85        .parse(&args.model)
86        .await
87        .map_err(|source| GenerateCommandError::Model { model: args.model.clone(), source })?;
88
89    let messages = {
90        let mut messages = Vec::new();
91        if let Some(system) = args.system.as_deref() {
92            messages.push(ChatMessage::system(system));
93        }
94        messages.push(ChatMessage::user(&prompt));
95        messages
96    };
97
98    let mut context = Context::new(messages, vec![]);
99    context.set_model_settings(args.model_settings.clone().into());
100    context.set_reasoning_effort(args.reasoning_effort.unwrap_or_default());
101
102    let mut stream = provider.stream_response(&context);
103    let mut text = String::new();
104    while let Some(result) = stream.next().await {
105        match result {
106            Ok(LlmResponse::Text { chunk }) => text.push_str(&chunk),
107            Err(error) => return Err(GenerateCommandError::Stream(error)),
108            _ => {}
109        }
110    }
111    println!("{}", format_output(&text, &args.model, args.output));
112    Ok(ExitCode::SUCCESS)
113}
114
115fn format_output(text: &str, model: &str, format: OutputFormat) -> String {
116    match format {
117        OutputFormat::Text => text.to_string(),
118        OutputFormat::Json => serde_json::json!({ "text": text, "model": model }).to_string(),
119        OutputFormat::Pretty => serde_json::to_string_pretty(&serde_json::json!({ "text": text, "model": model }))
120            .expect("generate response serializes to JSON"),
121    }
122}
123
124fn resolve_prompt(prompt: Option<&str>, prompt_file: Option<&str>) -> Result<String, GenerateCommandError> {
125    match (prompt, prompt_file) {
126        (Some(prompt), None) => Ok(prompt.to_string()),
127        (None, Some(prompt_file)) => read_prompt_file(prompt_file),
128        _ => Err(GenerateCommandError::PromptSource),
129    }
130}
131
132fn read_prompt_file(source: &str) -> Result<String, GenerateCommandError> {
133    if source == "-" {
134        let mut prompt = String::new();
135        stdin()
136            .read_to_string(&mut prompt)
137            .map_err(|error| GenerateCommandError::ReadPrompt { path: "-".to_string(), source: error })?;
138        return Ok(prompt);
139    }
140    std::fs::read_to_string(source)
141        .map_err(|error| GenerateCommandError::ReadPrompt { path: source.to_string(), source: error })
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use clap::Parser;
148
149    #[test]
150    fn format_output_wraps_json_and_passes_text_through() {
151        assert_eq!(format_output("hi", "anthropic:m", OutputFormat::Text), "hi");
152
153        let json: serde_json::Value =
154            serde_json::from_str(&format_output("hi", "anthropic:m", OutputFormat::Json)).unwrap();
155        assert_eq!(json["text"], "hi");
156        assert_eq!(json["model"], "anthropic:m");
157    }
158
159    #[tokio::test]
160    async fn run_rejects_reasoning_effort_unsupported_by_model_before_initializing_provider() {
161        let args = GenerateArgs {
162            model: "anthropic:claude-opus-4-6".to_string(),
163            prompt: Some("the prompt".to_string()),
164            prompt_file: None,
165            system: None,
166            model_settings: ModelSettingsArgs::default(),
167            reasoning_effort: Some(ReasoningEffort::Xhigh),
168            output: OutputFormat::Text,
169        };
170
171        let error = run(args).await.unwrap_err();
172
173        assert!(matches!(error, GenerateCommandError::ReasoningEffort(_)));
174    }
175
176    #[tokio::test]
177    async fn run_errors_on_unknown_provider() {
178        let dir = tempfile::tempdir().unwrap();
179        let prompt_path = dir.path().join("prompt.txt");
180        std::fs::write(&prompt_path, "the prompt").unwrap();
181        let args = GenerateArgs {
182            model: "definitely-not-a-provider:nope".to_string(),
183            prompt: None,
184            prompt_file: Some(prompt_path.to_string_lossy().into_owned()),
185            system: None,
186            model_settings: ModelSettingsArgs::default(),
187            reasoning_effort: None,
188            output: OutputFormat::Json,
189        };
190
191        let error = run(args).await.unwrap_err();
192
193        assert!(matches!(error, GenerateCommandError::Model { .. }), "got: {error:?}");
194    }
195
196    #[derive(clap::Parser)]
197    struct TestCli {
198        #[command(flatten)]
199        args: GenerateArgs,
200    }
201
202    fn parse_args(argv: &[&str]) -> GenerateArgs {
203        TestCli::try_parse_from(argv).unwrap().args
204    }
205
206    #[test]
207    fn model_settings_and_reasoning_flags_parse_convert_and_validate() {
208        let set = parse_args(&[
209            "gen",
210            "--model",
211            "anthropic:m",
212            "--prompt",
213            "hi",
214            "--temperature",
215            "0",
216            "--top-p",
217            "0.5",
218            "--max-tokens",
219            "64",
220            "--reasoning-effort",
221            "high",
222        ]);
223        assert_eq!(
224            ModelSettings::from(set.model_settings),
225            ModelSettings { temperature: Some(0.0), top_p: Some(0.5), max_tokens: Some(64) }
226        );
227        assert_eq!(set.reasoning_effort, Some(ReasoningEffort::High));
228
229        let absent = parse_args(&["gen", "--model", "anthropic:m", "--prompt", "hi"]);
230        assert!(ModelSettings::from(absent.model_settings).is_empty());
231        assert_eq!(absent.reasoning_effort, None);
232
233        let bad =
234            TestCli::try_parse_from(["gen", "--model", "anthropic:m", "--prompt", "hi", "--reasoning-effort", "nope"]);
235        assert!(bad.is_err());
236    }
237
238    #[test]
239    fn resolve_prompt_uses_inline_prompt() {
240        assert_eq!(resolve_prompt(Some("say hi"), None).unwrap(), "say hi");
241    }
242
243    #[test]
244    fn resolve_prompt_reads_a_file() {
245        let dir = tempfile::tempdir().unwrap();
246        let path = dir.path().join("prompt.txt");
247        std::fs::write(&path, "graded prompt").unwrap();
248
249        assert_eq!(resolve_prompt(None, Some(&path.to_string_lossy())).unwrap(), "graded prompt");
250    }
251
252    #[test]
253    fn resolve_prompt_reports_missing_file() {
254        let error = resolve_prompt(None, Some("/nonexistent/prompt.txt")).unwrap_err();
255
256        assert!(matches!(error, GenerateCommandError::ReadPrompt { .. }), "got: {error:?}");
257    }
258
259    #[test]
260    fn resolve_prompt_rejects_missing_prompt_source() {
261        let error = resolve_prompt(None, None).unwrap_err();
262
263        assert!(matches!(error, GenerateCommandError::PromptSource), "got: {error:?}");
264    }
265
266    #[test]
267    fn resolve_prompt_rejects_multiple_prompt_sources() {
268        let error = resolve_prompt(Some("inline"), Some("prompt.txt")).unwrap_err();
269
270        assert!(matches!(error, GenerateCommandError::PromptSource), "got: {error:?}");
271    }
272}
aether_cli/generate_command/mod.rs

aether_cli/generate_command/
mod.rs
