aether_cli/generate_command/
mod.rs

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

aether_cli/generate_command/
mod.rs
