1use clap::error::ErrorKind;
7use clap::{ArgGroup, Args, Parser, Subcommand};
8use nemo_fabric_core::{
9 ResolveContext, RunRequest, RunStatus, doctor_plan, resolve_diagnostic_plan_from_config,
10 resolve_run_plan_from_config, run_plan,
11};
12
13use crate::examples;
14use crate::presets;
15use crate::scaffold::{self, Language};
16
17#[derive(Debug, Parser)]
18#[command(name = "nemo-fabric")]
19#[command(about = "Experiment with NeMo Fabric presets and maintained examples")]
20#[command(arg_required_else_help = true)]
21#[command(version)]
22struct Cli {
23 #[command(subcommand)]
24 command: Command,
25}
26
27#[derive(Debug, Subcommand)]
28enum Command {
29 Preset {
31 #[command(subcommand)]
32 command: PresetCommand,
33 },
34 Example {
36 #[command(subcommand)]
37 command: ExampleCommand,
38 },
39 Plan(Selector),
41 Doctor(Selector),
43 Run {
45 #[command(flatten)]
46 selector: Selector,
47 #[arg(long, default_value = "")]
49 input: String,
50 },
51 Version,
53}
54
55#[derive(Debug, Subcommand)]
56enum PresetCommand {
57 List,
59 Show {
61 name: String,
63 },
64}
65
66#[derive(Debug, Subcommand)]
67enum ExampleCommand {
68 List,
70 Show {
72 name: String,
74 },
75 Init {
77 name: String,
79 destination: Option<std::path::PathBuf>,
81 #[arg(long, value_enum, default_value_t = Language::Python)]
83 language: Language,
84 #[arg(long)]
86 variant: Option<String>,
87 },
88}
89
90#[derive(Debug, Args)]
91#[command(group(
92 ArgGroup::new("source")
93 .required(true)
94 .multiple(false)
95 .args(["preset", "example"])
96))]
97struct Selector {
98 #[arg(long, value_name = "NAME")]
100 preset: Option<String>,
101 #[arg(long, value_name = "NAME")]
103 example: Option<String>,
104 #[arg(
106 long,
107 value_name = "NAME",
108 requires = "example",
109 conflicts_with = "preset"
110 )]
111 variant: Option<String>,
112 #[arg(
114 long,
115 value_name = "MODEL",
116 conflicts_with = "example",
117 value_parser = clap::builder::NonEmptyStringValueParser::new()
118 )]
119 model: Option<String>,
120 #[arg(
122 long,
123 value_name = "FLOAT",
124 conflicts_with = "example",
125 value_parser = parse_temperature
126 )]
127 temperature: Option<f64>,
128}
129
130fn parse_temperature(value: &str) -> Result<f64, String> {
131 let temperature = value
132 .parse::<f64>()
133 .map_err(|error| format!("invalid temperature: {error}"))?;
134 if !temperature.is_finite() {
135 return Err("temperature must be finite".to_string());
136 }
137 Ok(temperature)
138}
139
140enum SelectedSource {
141 Preset(presets::SelectedPreset),
142 Example(examples::SelectedExample),
143}
144
145impl SelectedSource {
146 fn config(&self) -> &nemo_fabric_core::FabricConfig {
147 match self {
148 Self::Preset(selected) => &selected.config,
149 Self::Example(selected) => &selected.config,
150 }
151 }
152
153 fn base_dir(&self) -> &std::path::Path {
154 match self {
155 Self::Preset(selected) => selected.base_dir(),
156 Self::Example(selected) => selected.base_dir(),
157 }
158 }
159}
160
161pub fn execute_from<I, T>(args: I) -> Result<(), String>
163where
164 I: IntoIterator<Item = T>,
165 T: Into<std::ffi::OsString> + Clone,
166{
167 let cli = match Cli::try_parse_from(args) {
168 Ok(cli) => cli,
169 Err(error)
170 if matches!(
171 error.kind(),
172 ErrorKind::DisplayHelp | ErrorKind::DisplayVersion
173 ) =>
174 {
175 error.print().map_err(|error| error.to_string())?;
176 return Ok(());
177 }
178 Err(error) => return Err(error.to_string()),
179 };
180 run(cli).map_err(|error| error.to_string())
181}
182
183fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
184 match cli.command {
185 Command::Preset {
186 command: PresetCommand::List,
187 } => {
188 for preset in presets::all() {
189 println!("{:<12} {}", preset.name, preset.description);
190 }
191 }
192 Command::Preset {
193 command: PresetCommand::Show { name },
194 } => {
195 let preset = find_preset(&name)?;
196 println!("{}\n {}", preset.name, preset.description);
197 if preset.required_env.is_empty() {
198 println!(" credentials: none");
199 } else {
200 println!(" required environment: {}", preset.required_env.join(", "));
201 }
202 }
203 Command::Example {
204 command: ExampleCommand::List,
205 } => {
206 for example in examples::all() {
207 println!("{:<12} {}", example.name, example.description);
208 }
209 }
210 Command::Example {
211 command: ExampleCommand::Show { name },
212 } => {
213 let example = find_example(&name)?;
214 println!("{}\n {}", example.name, example.description);
215 println!(" default variant: {}", example.default_variant);
216 println!(" variants: {}", example.variants.join(", "));
217 }
218 Command::Example {
219 command:
220 ExampleCommand::Init {
221 name,
222 destination,
223 language,
224 variant,
225 },
226 } => {
227 let example = find_example(&name)?;
228 let destination = destination.unwrap_or_else(|| std::path::PathBuf::from(&name));
229 let destination = scaffold::init(example, variant.as_deref(), language, &destination)?;
230 println!("created {}", destination.display());
231 }
232 Command::Plan(selector) => {
233 let selected = select_source(&selector)?;
234 let plan = resolve_run_plan_from_config(
235 selected.config().clone(),
236 ResolveContext::new(selected.base_dir()),
237 )?;
238 println!("{}", serde_json::to_string_pretty(&plan)?);
239 }
240 Command::Doctor(selector) => {
241 let selected = select_source(&selector)?;
242 let plan = resolve_diagnostic_plan_from_config(
243 selected.config().clone(),
244 ResolveContext::new(selected.base_dir()),
245 )?;
246 println!("{}", serde_json::to_string_pretty(&doctor_plan(&plan))?);
247 }
248 Command::Run { selector, input } => {
249 let selected = select_source(&selector)?;
250 let plan = resolve_run_plan_from_config(
251 selected.config().clone(),
252 ResolveContext::new(selected.base_dir()),
253 )?;
254 let result = match run_plan(&plan, RunRequest::text(input)) {
255 Ok(result) => result,
256 Err(error) => {
257 println!(
258 "{}",
259 serde_json::to_string_pretty(&serde_json::json!({
260 "status": "failed",
261 "error": {
262 "code": "runtime_error",
263 "message": "run failed before a normalized result was available",
264 },
265 }))?
266 );
267 return Err(error.into());
268 }
269 };
270 println!("{}", serde_json::to_string_pretty(&result)?);
271 require_successful_run(result.status)?;
272 }
273 Command::Version => println!("{}", nemo_fabric_core::version()),
274 }
275 Ok(())
276}
277
278fn require_successful_run(status: RunStatus) -> Result<(), String> {
279 match status {
280 RunStatus::Succeeded => Ok(()),
281 RunStatus::Failed => Err("run completed with status failed".to_string()),
282 RunStatus::Cancelled => Err("run completed with status cancelled".to_string()),
283 }
284}
285
286fn select_source(selector: &Selector) -> Result<SelectedSource, Box<dyn std::error::Error>> {
287 if let Some(name) = &selector.preset {
288 let mut selected = find_preset(name)?.stage()?;
289 apply_model_overrides(&mut selected.config, selector)?;
290 return Ok(SelectedSource::Preset(selected));
291 }
292 if let Some(name) = &selector.example {
293 return Ok(SelectedSource::Example(
294 find_example(name)?.select(selector.variant.as_deref())?,
295 ));
296 }
297 unreachable!("Clap requires exactly one source selector")
298}
299
300fn apply_model_overrides(
301 config: &mut nemo_fabric_core::FabricConfig,
302 selector: &Selector,
303) -> Result<(), String> {
304 if selector.model.is_none() && selector.temperature.is_none() {
305 return Ok(());
306 }
307 let default_model = config
308 .models
309 .get_mut("default")
310 .ok_or_else(|| "the selected preset does not define a default model".to_string())?;
311 if let Some(model) = &selector.model {
312 default_model.model.clone_from(model);
313 }
314 if let Some(temperature) = selector.temperature {
315 default_model.temperature = Some(temperature);
316 }
317 Ok(())
318}
319
320fn find_preset(name: &str) -> Result<presets::Preset, String> {
321 presets::find(name).ok_or_else(|| {
322 format!(
323 "unknown preset {name:?}; available: {}",
324 presets::all()
325 .iter()
326 .map(|preset| preset.name)
327 .collect::<Vec<_>>()
328 .join(", ")
329 )
330 })
331}
332
333fn find_example(name: &str) -> Result<examples::Example, String> {
334 examples::find(name).ok_or_else(|| {
335 format!(
336 "unknown example {name:?}; available: {}",
337 examples::all()
338 .iter()
339 .map(|example| example.name)
340 .collect::<Vec<_>>()
341 .join(", ")
342 )
343 })
344}
345
346#[cfg(test)]
347mod tests {
348 use clap::Parser;
349
350 use super::*;
351
352 #[test]
353 fn missing_subcommand_is_an_error() {
354 assert!(execute_from(["nemo-fabric"]).is_err());
355 }
356
357 #[test]
358 fn explicit_help_is_successful() {
359 assert!(execute_from(["nemo-fabric", "--help"]).is_ok());
360 }
361
362 #[test]
363 fn failed_and_cancelled_runs_are_errors() {
364 assert!(require_successful_run(RunStatus::Succeeded).is_ok());
365 assert!(require_successful_run(RunStatus::Failed).is_err());
366 assert!(require_successful_run(RunStatus::Cancelled).is_err());
367 }
368
369 #[test]
370 fn parses_preset_plan() {
371 let cli = Cli::try_parse_from(["nemo-fabric", "plan", "--preset", "scripted"])
372 .expect("parse plan");
373 assert!(matches!(cli.command, Command::Plan(_)));
374 }
375
376 #[test]
377 fn preset_model_overrides_preserve_provider_settings() {
378 let selector = Selector {
379 preset: Some("hermes".to_string()),
380 example: None,
381 variant: None,
382 model: Some("meta/llama-3.3-70b-instruct".to_string()),
383 temperature: Some(0.2),
384 };
385 let selected = select_source(&selector).expect("select preset");
386 let model = selected.config().models.get("default").expect("model");
387
388 assert_eq!(model.provider, "nvidia");
389 assert_eq!(model.model, "meta/llama-3.3-70b-instruct");
390 assert_eq!(model.temperature, Some(0.2));
391 assert_eq!(model.api_key_env.as_deref(), Some("NVIDIA_API_KEY"));
392 }
393
394 #[test]
395 fn model_overrides_require_a_preset_with_a_default_model() {
396 assert!(
397 Cli::try_parse_from([
398 "nemo-fabric",
399 "plan",
400 "--example",
401 "code-review",
402 "--model",
403 "other-model",
404 ])
405 .is_err()
406 );
407
408 let selector = Selector {
409 preset: Some("scripted".to_string()),
410 example: None,
411 variant: None,
412 model: Some("other-model".to_string()),
413 temperature: None,
414 };
415 assert_eq!(
416 select_source(&selector)
417 .err()
418 .expect("scripted has no default model")
419 .to_string(),
420 "the selected preset does not define a default model"
421 );
422
423 assert!(
424 Cli::try_parse_from([
425 "nemo-fabric",
426 "plan",
427 "--preset",
428 "hermes",
429 "--temperature",
430 "NaN",
431 ])
432 .is_err()
433 );
434 }
435
436 #[test]
437 fn parses_example_variant() {
438 let cli = Cli::try_parse_from([
439 "nemo-fabric",
440 "run",
441 "--example",
442 "code-review",
443 "--variant",
444 "hermes",
445 ])
446 .expect("parse example run");
447 assert!(matches!(cli.command, Command::Run { .. }));
448 }
449
450 #[test]
451 fn selectors_are_mutually_exclusive() {
452 assert!(
453 Cli::try_parse_from([
454 "nemo-fabric",
455 "plan",
456 "--preset",
457 "scripted",
458 "--example",
459 "code-review",
460 ])
461 .is_err()
462 );
463 assert!(
464 Cli::try_parse_from([
465 "nemo-fabric",
466 "plan",
467 "--preset",
468 "scripted",
469 "--variant",
470 "hermes",
471 ])
472 .is_err()
473 );
474 }
475
476 #[test]
477 fn plan_requires_a_preset() {
478 assert!(Cli::try_parse_from(["nemo-fabric", "plan"]).is_err());
479 }
480}