aether_cli/generate_command/
mod.rs

1use crate::output::OutputFormat;
2use futures::StreamExt;
3use llm::parser::ModelProviderParser;
4use llm::types::IsoString;
5use llm::{ChatMessage, ContentBlock, Context, LlmError, LlmResponse, 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    /// Output format. `text` prints the raw response; `json`/`pretty` wrap it as
29    /// `{ "text": <response>, "model": <model> }`.
30    #[arg(long, default_value = "text")]
31    pub output: OutputFormat,
32}
33
34#[derive(Debug, Error)]
35pub enum GenerateCommandError {
36    #[error("provide exactly one of --prompt or --prompt-file")]
37    PromptSource,
38
39    #[error("failed to read prompt from {path}: {source}")]
40    ReadPrompt { path: String, source: std::io::Error },
41
42    #[error("failed to initialize model `{model}`: {source}")]
43    Model { model: String, source: LlmError },
44
45    #[error("model stream error: {0}")]
46    Stream(LlmError),
47}
48
49/// Call a model with a single prompt and print its response. The judge is a special case of this:
50/// the caller supplies a grading prompt and parses the structured verdict out of the response.
51pub async fn run(args: GenerateArgs) -> Result<ExitCode, GenerateCommandError> {
52    let prompt = resolve_prompt(args.prompt.as_deref(), args.prompt_file.as_deref())?;
53    let (provider, _) = ModelProviderParser::default()
54        .parse(&args.model)
55        .await
56        .map_err(|source| GenerateCommandError::Model { model: args.model.clone(), source })?;
57
58    let messages = {
59        let mut messages = Vec::new();
60        if let Some(system) = args.system.as_deref() {
61            messages.push(ChatMessage::System { content: system.to_string(), timestamp: IsoString::now() });
62        }
63        messages.push(ChatMessage::User { content: vec![ContentBlock::text(&prompt)], timestamp: IsoString::now() });
64        messages
65    };
66
67    let mut stream = provider.stream_response(&Context::new(messages, vec![]));
68    let mut text = String::new();
69    while let Some(result) = stream.next().await {
70        match result {
71            Ok(LlmResponse::Text { chunk }) => text.push_str(&chunk),
72            Err(error) => return Err(GenerateCommandError::Stream(error)),
73            _ => {}
74        }
75    }
76    println!("{}", format_output(&text, &args.model, args.output));
77    Ok(ExitCode::SUCCESS)
78}
79
80fn format_output(text: &str, model: &str, format: OutputFormat) -> String {
81    match format {
82        OutputFormat::Text => text.to_string(),
83        OutputFormat::Json => serde_json::json!({ "text": text, "model": model }).to_string(),
84        OutputFormat::Pretty => serde_json::to_string_pretty(&serde_json::json!({ "text": text, "model": model }))
85            .expect("generate response serializes to JSON"),
86    }
87}
88
89fn resolve_prompt(prompt: Option<&str>, prompt_file: Option<&str>) -> Result<String, GenerateCommandError> {
90    match (prompt, prompt_file) {
91        (Some(prompt), None) => Ok(prompt.to_string()),
92        (None, Some(prompt_file)) => read_prompt_file(prompt_file),
93        _ => Err(GenerateCommandError::PromptSource),
94    }
95}
96
97fn read_prompt_file(source: &str) -> Result<String, GenerateCommandError> {
98    if source == "-" {
99        let mut prompt = String::new();
100        stdin()
101            .read_to_string(&mut prompt)
102            .map_err(|error| GenerateCommandError::ReadPrompt { path: "-".to_string(), source: error })?;
103        return Ok(prompt);
104    }
105    std::fs::read_to_string(source)
106        .map_err(|error| GenerateCommandError::ReadPrompt { path: source.to_string(), source: error })
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn format_output_wraps_json_and_passes_text_through() {
115        assert_eq!(format_output("hi", "anthropic:m", OutputFormat::Text), "hi");
116
117        let json: serde_json::Value =
118            serde_json::from_str(&format_output("hi", "anthropic:m", OutputFormat::Json)).unwrap();
119        assert_eq!(json["text"], "hi");
120        assert_eq!(json["model"], "anthropic:m");
121    }
122
123    #[tokio::test]
124    async fn run_errors_on_unknown_provider() {
125        let dir = tempfile::tempdir().unwrap();
126        let prompt_path = dir.path().join("prompt.txt");
127        std::fs::write(&prompt_path, "the prompt").unwrap();
128        let args = GenerateArgs {
129            model: "definitely-not-a-provider:nope".to_string(),
130            prompt: None,
131            prompt_file: Some(prompt_path.to_string_lossy().into_owned()),
132            system: None,
133            output: OutputFormat::Json,
134        };
135
136        let error = run(args).await.unwrap_err();
137
138        assert!(matches!(error, GenerateCommandError::Model { .. }), "got: {error:?}");
139    }
140
141    #[test]
142    fn resolve_prompt_uses_inline_prompt() {
143        assert_eq!(resolve_prompt(Some("say hi"), None).unwrap(), "say hi");
144    }
145
146    #[test]
147    fn resolve_prompt_reads_a_file() {
148        let dir = tempfile::tempdir().unwrap();
149        let path = dir.path().join("prompt.txt");
150        std::fs::write(&path, "graded prompt").unwrap();
151
152        assert_eq!(resolve_prompt(None, Some(&path.to_string_lossy())).unwrap(), "graded prompt");
153    }
154
155    #[test]
156    fn resolve_prompt_reports_missing_file() {
157        let error = resolve_prompt(None, Some("/nonexistent/prompt.txt")).unwrap_err();
158
159        assert!(matches!(error, GenerateCommandError::ReadPrompt { .. }), "got: {error:?}");
160    }
161
162    #[test]
163    fn resolve_prompt_rejects_missing_prompt_source() {
164        let error = resolve_prompt(None, None).unwrap_err();
165
166        assert!(matches!(error, GenerateCommandError::PromptSource), "got: {error:?}");
167    }
168
169    #[test]
170    fn resolve_prompt_rejects_multiple_prompt_sources() {
171        let error = resolve_prompt(Some("inline"), Some("prompt.txt")).unwrap_err();
172
173        assert!(matches!(error, GenerateCommandError::PromptSource), "got: {error:?}");
174    }
175}
aether_cli/generate_command/mod.rs

aether_cli/generate_command/
mod.rs
