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

aether_cli/generate_command/
mod.rs
