Skip to main content

kjarni_cli/
lib.rs

1use clap::{Parser, Subcommand};
2
3/// Default model for `kjarni chat`.
4///
5/// Must be a CLI name that `ModelType::from_cli_name` resolves. Both this and
6/// [`DEFAULT_GENERATE_MODEL`] previously named models that do not exist
7/// (`llama-3.2-8b-instruct`, `llama-3.2-1b`), so the bare `kjarni chat` and
8/// `kjarni generate` commands failed for everyone. `tests::default_models_exist`
9/// now fails the build if either drifts from the registry again.
10pub const DEFAULT_CHAT_MODEL: &str = "llama3.2-3b-instruct";
11
12/// Default model for `kjarni generate`.
13pub const DEFAULT_GENERATE_MODEL: &str = "llama3.2-1b-instruct";
14
15#[derive(Parser)]
16#[command(name = "kjarni")]
17#[command(about = "Kjarni: The SQLite of AI", long_about = None)]
18#[command(version)]
19pub struct Cli {
20    #[command(subcommand)]
21    pub command: Commands,
22
23    /// Verbosity level (-v, -vv, -vvv)
24    #[arg(short, long, action = clap::ArgAction::Count, global = true)]
25    pub verbose: u8,
26}
27
28#[derive(Subcommand, Debug, PartialEq)]
29pub enum Commands {
30    /// Manage models (list, download, info)
31    Model {
32        #[command(subcommand)]
33        action: ModelCommands,
34    },
35
36    /// Generate text from a prompt
37    Generate {
38        /// The prompt (or file path, or stdin if not provided)
39        prompt: Option<String>,
40
41        #[arg(short, long, default_value = DEFAULT_GENERATE_MODEL)]
42        model: String,
43
44        /// Load weights from a local file or directory instead of the registry
45        #[arg(long)]
46        model_path: Option<String>,
47
48        /// Maximum tokens to generate
49        #[arg(short = 'n', long, default_value_t = 100)]
50        max_tokens: usize,
51
52        /// Sampling temperature (0.0 = greedy, higher = more random)
53        #[arg(short, long, default_value_t = 0.7)]
54        temperature: f32,
55
56        /// Top-K sampling (limits to K most likely tokens)
57        #[arg(long)]
58        top_k: Option<usize>,
59
60        /// Top-P (nucleus) sampling threshold
61        #[arg(long)]
62        top_p: Option<f32>,
63
64        /// Min-P sampling threshold  
65        #[arg(long)]
66        min_p: Option<f32>,
67
68        /// Repetition penalty (1.0 = no penalty)
69        #[arg(long, default_value_t = 1.1)]
70        repetition_penalty: f32,
71
72        /// Use greedy decoding (ignores temperature)
73        #[arg(long)]
74        greedy: bool,
75
76        /// Use GPU
77        #[arg(long)]
78        gpu: bool,
79
80        /// Disable streaming output
81        #[arg(long)]
82        no_stream: bool,
83
84        /// Suppress status messages
85        #[arg(short, long)]
86        quiet: bool,
87    },
88
89    /// Summarize text
90    Summarize {
91        /// Input text (or read from stdin)
92        #[arg(short, long)]
93        input: Option<String>,
94
95        /// Model to use
96        #[arg(short, long, default_value = "distilbart-cnn")]
97        model: String,
98
99        /// Path to local model
100        /// Load weights from a local file or directory instead of the registry
101        #[arg(long)]
102        model_path: Option<String>,
103
104        /// Minimum summary length
105        #[arg(long)]
106        min_length: Option<usize>,
107
108        /// Maximum summary length
109        #[arg(long)]
110        max_length: Option<usize>,
111
112        /// Number of beams for beam search
113        #[arg(long)]
114        num_beams: Option<usize>,
115
116        /// Length penalty for beam search (< 1 shorter, > 1 longer)
117        #[arg(long)]
118        length_penalty: Option<f32>,
119
120        /// Block repeated n-grams of this size
121        #[arg(long)]
122        no_repeat_ngram: Option<usize>,
123
124        /// Use greedy decoding (deterministic, fastest)
125        #[arg(long)]
126        greedy: bool,
127
128        /// Disable streaming output
129        #[arg(long)]
130        no_stream: bool,
131
132        /// Use GPU
133        #[arg(long)]
134        gpu: bool,
135
136        /// Suppress progress messages
137        #[arg(short, long)]
138        quiet: bool,
139    },
140
141    /// Translate text between languages
142    Translate {
143        /// Input text (or read from stdin)
144        #[arg(short, long)]
145        input: Option<String>,
146
147        /// Model to use
148        #[arg(short, long, default_value = "flan-t5-base")]
149        model: String,
150
151        /// Path to local model
152        /// Load weights from a local file or directory instead of the registry
153        #[arg(long)]
154        model_path: Option<String>,
155
156        /// Source language (e.g., en, de, fr)
157        #[arg(long)]
158        src: Option<String>,
159
160        /// Target language (e.g., en, de, fr)
161        #[arg(long)]
162        dst: Option<String>,
163
164        /// Maximum output length
165        #[arg(long)]
166        max_length: Option<usize>,
167
168        /// Number of beams for beam search
169        #[arg(long)]
170        num_beams: Option<usize>,
171
172        /// Length penalty for beam search (< 1 shorter, > 1 longer)
173        #[arg(long)]
174        length_penalty: Option<f32>,
175
176        /// Block repeated n-grams of this size
177        #[arg(long)]
178        no_repeat_ngram: Option<usize>,
179
180        /// Use greedy decoding (deterministic, fastest)
181        #[arg(long)]
182        greedy: bool,
183
184        /// Disable streaming output
185        #[arg(long)]
186        no_stream: bool,
187
188        /// Use GPU
189        #[arg(long)]
190        gpu: bool,
191
192        /// Suppress progress messages
193        #[arg(short, long)]
194        quiet: bool,
195    },
196
197    /// Show a model's metadata, config and tensor layout
198    Inspect {
199        /// Path to a .gguf, a .safetensors directory, or a name in ~/.cache/kjarni
200        path: String,
201    },
202
203    /// Generate embeddings for text
204    Embed {
205        /// Input text, file path, or stdin if not provided
206        input: Option<String>,
207
208        /// Embedding model
209        #[arg(short, long, default_value = "minilm-l6-v2")]
210        model: String,
211
212        /// Load weights from a local file or directory instead of the registry
213        #[arg(long)]
214        model_path: Option<String>,
215
216        /// Output format: raw, json
217        #[arg(long, default_value = "raw")]
218        format: String,
219
220        /// Scale each vector to unit length
221        #[arg(long)]
222        normalize: bool,
223
224        /// Pooling strategy: mean, cls, max
225        ///
226        /// Mean matches what sentence-transformers does for these models, and what the
227        /// library bindings and presets already default to. Pooling changes the vector
228        /// itself, so a different choice here does not produce a slightly different
229        /// answer, it produces one that cannot be compared against the others.
230        #[arg(long, default_value = "mean")]
231        pooling: String,
232
233        /// Run on the GPU
234        #[arg(long)]
235        gpu: bool,
236
237        /// Suppress status messages
238        #[arg(short, long)]
239        quiet: bool,
240    },
241
242    /// Transcribe audio to text
243    Transcribe {
244        /// Path to audio file (wav, mp3, flac, ogg)
245        file: String,
246
247        /// Model to use (whisper-small, whisper-large-v3)
248        #[arg(short, long, default_value = "whisper-small")]
249        model: String,
250
251        /// Path to local model directory (not yet implemented)
252        /// Load weights from a local file or directory instead of the registry
253        #[arg(long)]
254        model_path: Option<String>,
255
256        /// Language code (e.g., en, fr, de). Omit for auto-detect.
257        #[arg(short, long)]
258        language: Option<String>,
259
260        /// Translate to English instead of transcribing
261        #[arg(long)]
262        translate: bool,
263
264        /// Include timestamps in output
265        #[arg(short, long)]
266        timestamps: bool,
267
268        /// Maximum tokens per 30-second chunk
269        #[arg(long)]
270        max_tokens: Option<usize>,
271
272        /// Disable streaming (wait for full result)
273        #[arg(long)]
274        no_stream: bool,
275
276        /// Use GPU acceleration
277        #[arg(long)]
278        gpu: bool,
279
280        /// Suppress progress output
281        #[arg(short, long)]
282        quiet: bool,
283    },
284
285    /// Classify text using a classification model
286    Classify {
287        /// Input text(s) to classify. Use - for stdin.
288        input: Vec<String>,
289
290        /// Model name from registry
291        #[arg(short, long, default_value = "distilbert-sentiment")]
292        model: String,
293
294        /// Load model from local path instead of registry
295        #[arg(long, value_name = "PATH")]
296        model_path: Option<String>,
297
298        /// Custom labels (comma-separated, order must match model output)
299        /// Example: --labels "negative,positive" or --labels "neikvætt,jákvætt"
300        #[arg(long, value_name = "LABELS")]
301        labels: Option<String>,
302
303        /// Return top K predictions
304        #[arg(long, default_value = "5")]
305        top_k: usize,
306
307        /// Minimum confidence threshold (0.0-1.0)
308        #[arg(long)]
309        threshold: Option<f32>,
310
311        /// Maximum sequence length (truncates longer inputs)
312        #[arg(long)]
313        max_length: Option<usize>,
314
315        /// Batch size for inference
316        #[arg(long)]
317        batch_size: Option<usize>,
318
319        /// Use multi-label classification (sigmoid instead of softmax)
320        #[arg(long)]
321        multi_label: bool,
322
323        /// Output format: json, jsonl, text
324        #[arg(short, long, default_value = "text")]
325        format: String,
326
327        /// Run on GPU
328        #[arg(long)]
329        gpu: bool,
330
331        /// Model precision: f32, f16, bf16
332        #[arg(long)]
333        dtype: Option<String>,
334
335        /// Suppress progress output
336        #[arg(short, long)]
337        quiet: bool,
338    },
339
340    /// Rerank documents by relevance to a query
341    Rerank {
342        /// The query to rank against
343        query: String,
344
345        /// Documents to rerank (or read from stdin, one per line)
346        documents: Vec<String>,
347
348        #[arg(short, long, default_value = "minilm-l6-v2-cross-encoder")]
349        model: String,
350
351        /// Load weights from a local file or directory instead of the registry
352        #[arg(long)]
353        model_path: Option<String>,
354
355        /// Return only top K results
356        #[arg(short = 'k', long)]
357        top_k: Option<usize>,
358
359        /// Output format: json, jsonl, text, docs
360        #[arg(short, long, default_value = "text")]
361        format: String,
362
363        /// Run on the GPU
364        #[arg(long)]
365        gpu: bool,
366
367        /// Suppress progress output
368        #[arg(short, long)]
369        quiet: bool,
370    },
371
372    /// Interactive chat mode
373    Chat {
374        #[arg(short, long, default_value = DEFAULT_CHAT_MODEL)]
375        model: String,
376
377        /// Load weights from a local file or directory instead of the registry
378        #[arg(long)]
379        model_path: Option<String>,
380
381        /// System prompt to set assistant behavior
382        #[arg(short, long)]
383        system: Option<String>,
384
385        /// Sampling temperature
386        #[arg(short, long, default_value_t = 0.7)]
387        temperature: f32,
388
389        /// Max tokens per response
390        #[arg(short = 'n', long, default_value_t = 512)]
391        max_tokens: usize,
392
393        /// Run on the GPU
394        #[arg(long)]
395        gpu: bool,
396
397        /// Suppress progress output
398        #[arg(short, long)]
399        quiet: bool,
400    },
401
402    /// Create or manage search indexes
403    Index {
404        #[command(subcommand)]
405        action: IndexCommands,
406    },
407
408    /// Search an index
409    Search {
410        /// Path to the index file
411        index_path: String,
412
413        /// Search query
414        query: String,
415
416        /// Number of results to return
417        #[arg(short = 'k', long, default_value_t = 10)]
418        top_k: usize,
419
420        /// Search mode: hybrid, semantic, keyword
421        #[arg(long, default_value = "hybrid")]
422        mode: String,
423
424        /// Encoder model for semantic search
425        #[arg(short, long, default_value = "minilm-l6-v2")]
426        model: String,
427
428        /// Reranking model (optional)
429        /// Use a cross-encoder model to rerank initial results
430        /// Example: --rerank-model "ms-marco-minilm"
431        #[arg(long, default_value = None)]
432        rerank_model: Option<String>,
433
434        /// Output format: json, jsonl, text
435        #[arg(short, long, default_value = "text")]
436        format: String,
437
438        /// Run on the GPU
439        #[arg(long)]
440        gpu: bool,
441
442        /// Suppress progress output
443        #[arg(short, long)]
444        quiet: bool,
445    },
446
447    /// Compute similarity between two texts
448    Similarity {
449        /// First text (or file path)
450        text1: String,
451
452        /// Second text (or file path)
453        text2: String,
454
455        /// Encoder model
456        #[arg(short, long, default_value = "minilm-l6-v2")]
457        model: String,
458
459        /// Run on the GPU
460        #[arg(long)]
461        gpu: bool,
462
463        /// Suppress progress output
464        #[arg(short, long)]
465        quiet: bool,
466    },
467}
468
469#[derive(Subcommand, Debug, PartialEq)]
470pub enum ModelCommands {
471    /// List all available models
472    List {
473        /// Filter by architecture (e.g., llama, bert, t5)
474        #[arg(short, long)]
475        arch: Option<String>,
476
477        /// Filter by task (e.g., chat, embedding, classification, summarization)
478        #[arg(short, long)]
479        task: Option<String>,
480
481        /// Show only downloaded models
482        #[arg(short, long)]
483        downloaded: bool,
484    },
485
486    /// Download a model
487    Download {
488        name: String,
489
490        #[arg(long)]
491        gguf: bool,
492
493        #[arg(short, long)]
494        quiet: bool,
495    },
496
497    /// Remove a downloaded model
498    Remove { name: String },
499
500    /// Show detailed info about a model
501    Info { name: String },
502
503    /// Search for models by name or description
504    Search { query: String },
505}
506
507#[derive(Subcommand, Debug, PartialEq)]
508pub enum InspectCommands {
509    /// Inspect a model file
510    Model {
511        /// Path to the model file
512        path: String,
513    },
514}
515
516#[derive(Subcommand, Debug, PartialEq)]
517pub enum IndexCommands {
518    /// Create a new index from documents
519    Create {
520        /// Output index file path
521        output: String,
522
523        /// Input files/directories (uses built-in chunking)
524        #[arg(conflicts_with = "from_chunks")]
525        inputs: Vec<String>,
526
527        /// Pre-chunked JSONL file (bypass chunking)
528        #[arg(long, conflicts_with = "inputs")]
529        from_chunks: Option<String>,
530
531        /// Chunk size in characters
532        ///
533        /// The default fits inside the encoder's window. minilm-l6-v2 reads 256
534        /// tokens, roughly 900 characters, so the previous default of 1000 left
535        /// about nine chunks in ten longer than the model reads, with the tail
536        /// dropped silently. The C# Indexer has always used 512; this matches it.
537        #[arg(long, default_value_t = 512)]
538        chunk_size: usize,
539
540        /// Chunk overlap in characters
541        #[arg(long, default_value_t = 100)]
542        chunk_overlap: usize,
543
544        /// Encoder model for embeddings
545        #[arg(short, long, default_value = "minilm-l6-v2")]
546        model: String,
547
548        /// Run on the GPU
549        #[arg(long)]
550        gpu: bool,
551
552        /// Suppress progress output
553        #[arg(short, long)]
554        quiet: bool,
555    },
556
557    /// Add documents to an existing index
558    Add {
559        /// Index file path
560        index_path: String,
561
562        /// Input files or directories to add
563        inputs: Vec<String>,
564
565        /// Chunk size in characters
566        ///
567        /// Matches `index create`; see the note there.
568        #[arg(long, default_value_t = 512)]
569        chunk_size: usize,
570
571        #[arg(long, default_value_t = 100)]
572        chunk_overlap: usize,
573
574        #[arg(short, long, default_value = "minilm-l6-v2")]
575        model: String,
576
577        /// Run on the GPU
578        #[arg(long)]
579        gpu: bool,
580
581        /// Suppress progress output
582        #[arg(short, long)]
583        quiet: bool,
584    },
585
586    /// Show index info
587    Info {
588        /// Index file path
589        index_path: String,
590    },
591}
592
593/// Convert verbosity count to log level string
594pub fn verbosity_to_log_level(verbose: u8) -> &'static str {
595    match verbose {
596        0 => "warn",
597        1 => "info",
598        2 => "debug",
599        _ => "trace",
600    }
601}
602
603#[cfg(test)]
604mod tests {
605    use super::*;
606    use clap::Parser;
607
608    /// Every model name hardcoded outside the registry must actually resolve.
609    ///
610    /// This exists because three separate lists had drifted from the registry at
611    /// once: both CLI defaults, and four of the seven entries in
612    /// `kjarni::chat::suggest_chat_models`. Each was presumably correct when
613    /// written and went stale when models gained an `-instruct` suffix, with
614    /// nothing to notice. A name that does not resolve is a user-facing failure
615    /// on the very first command someone types, so it should break the build
616    /// instead.
617    #[test]
618    fn hardcoded_model_names_resolve() {
619        use kjarni::ModelType;
620
621        let mut bad = Vec::new();
622
623        for name in [DEFAULT_CHAT_MODEL, DEFAULT_GENERATE_MODEL] {
624            if ModelType::from_cli_name(name).is_none() {
625                bad.push(format!("CLI default '{name}'"));
626            }
627        }
628
629        for name in kjarni::chat::suggested_models() {
630            if ModelType::from_cli_name(name).is_none() {
631                bad.push(format!("chat suggestion '{name}'"));
632            }
633        }
634
635        assert!(
636            bad.is_empty(),
637            "these names are not in the model registry, so anyone who uses them \
638             gets an error:\n  {}\n\
639             Run `kjarni model list` for the valid names.",
640            bad.join("\n  ")
641        );
642    }
643
644    /// The chat default must be usable for chat, not merely a real model.
645    ///
646    /// `kjarni chat` with no arguments should work; defaulting to an embedding
647    /// model would resolve fine here and still fail at runtime.
648    #[test]
649    fn default_chat_model_is_a_chat_model() {
650        assert!(
651            kjarni::chat::is_chat_model(DEFAULT_CHAT_MODEL).is_ok(),
652            "DEFAULT_CHAT_MODEL '{DEFAULT_CHAT_MODEL}' cannot be used for chat"
653        );
654    }
655    fn parse_args(args: &[&str]) -> Result<Cli, clap::Error> {
656        let mut full_args = vec!["kjarni"];
657        full_args.extend(args);
658        Cli::try_parse_from(full_args)
659    }
660
661    #[test]
662    fn test_verbosity_to_log_level_zero() {
663        assert_eq!(verbosity_to_log_level(0), "warn");
664    }
665
666    #[test]
667    fn test_verbosity_to_log_level_one() {
668        assert_eq!(verbosity_to_log_level(1), "info");
669    }
670
671    #[test]
672    fn test_verbosity_to_log_level_two() {
673        assert_eq!(verbosity_to_log_level(2), "debug");
674    }
675
676    #[test]
677    fn test_verbosity_to_log_level_three() {
678        assert_eq!(verbosity_to_log_level(3), "trace");
679    }
680
681    #[test]
682    fn test_verbosity_to_log_level_high() {
683        assert_eq!(verbosity_to_log_level(10), "trace");
684        assert_eq!(verbosity_to_log_level(255), "trace");
685    }
686
687    #[test]
688    fn test_generate_minimal() {
689        let cli = parse_args(&["generate"]).unwrap();
690
691        match cli.command {
692            Commands::Generate {
693                prompt,
694                model,
695                max_tokens,
696                temperature,
697                greedy,
698                gpu,
699                quiet,
700                ..
701            } => {
702                assert!(prompt.is_none());
703                assert_eq!(model, DEFAULT_GENERATE_MODEL);
704                assert_eq!(max_tokens, 100);
705                assert!((temperature - 0.7).abs() < 0.001);
706                assert!(!greedy);
707                assert!(!gpu);
708                assert!(!quiet);
709            }
710            _ => panic!("Expected Generate command"),
711        }
712    }
713
714    #[test]
715    fn test_generate_with_prompt() {
716        let cli = parse_args(&["generate", "Hello world"]).unwrap();
717
718        match cli.command {
719            Commands::Generate { prompt, .. } => {
720                assert_eq!(prompt, Some("Hello world".to_string()));
721            }
722            _ => panic!("Expected Generate command"),
723        }
724    }
725
726    #[test]
727    fn test_generate_with_model() {
728        let cli = parse_args(&["generate", "-m", "phi3.5-mini"]).unwrap();
729
730        match cli.command {
731            Commands::Generate { model, .. } => {
732                assert_eq!(model, "phi3.5-mini");
733            }
734            _ => panic!("Expected Generate command"),
735        }
736    }
737
738    #[test]
739    fn test_generate_with_max_tokens() {
740        let cli = parse_args(&["generate", "-n", "500"]).unwrap();
741
742        match cli.command {
743            Commands::Generate { max_tokens, .. } => {
744                assert_eq!(max_tokens, 500);
745            }
746            _ => panic!("Expected Generate command"),
747        }
748    }
749
750    #[test]
751    fn test_generate_with_temperature() {
752        let cli = parse_args(&["generate", "-t", "1.5"]).unwrap();
753
754        match cli.command {
755            Commands::Generate { temperature, .. } => {
756                assert!((temperature - 1.5).abs() < 0.001);
757            }
758            _ => panic!("Expected Generate command"),
759        }
760    }
761
762    #[test]
763    fn test_generate_with_sampling_params() {
764        let cli = parse_args(&[
765            "generate", "--top-k", "50", "--top-p", "0.9", "--min-p", "0.05",
766        ])
767        .unwrap();
768
769        match cli.command {
770            Commands::Generate {
771                top_k,
772                top_p,
773                min_p,
774                ..
775            } => {
776                assert_eq!(top_k, Some(50));
777                assert_eq!(top_p, Some(0.9));
778                assert_eq!(min_p, Some(0.05));
779            }
780            _ => panic!("Expected Generate command"),
781        }
782    }
783
784    #[test]
785    fn test_generate_with_greedy() {
786        let cli = parse_args(&["generate", "--greedy"]).unwrap();
787
788        match cli.command {
789            Commands::Generate { greedy, .. } => {
790                assert!(greedy);
791            }
792            _ => panic!("Expected Generate command"),
793        }
794    }
795
796    #[test]
797    fn test_generate_with_gpu() {
798        let cli = parse_args(&["generate", "--gpu"]).unwrap();
799
800        match cli.command {
801            Commands::Generate { gpu, .. } => {
802                assert!(gpu);
803            }
804            _ => panic!("Expected Generate command"),
805        }
806    }
807
808    #[test]
809    fn test_generate_with_no_stream() {
810        let cli = parse_args(&["generate", "--no-stream"]).unwrap();
811
812        match cli.command {
813            Commands::Generate { no_stream, .. } => {
814                assert!(no_stream);
815            }
816            _ => panic!("Expected Generate command"),
817        }
818    }
819
820    #[test]
821    fn test_generate_with_quiet() {
822        let cli = parse_args(&["generate", "-q"]).unwrap();
823
824        match cli.command {
825            Commands::Generate { quiet, .. } => {
826                assert!(quiet);
827            }
828            _ => panic!("Expected Generate command"),
829        }
830    }
831
832    #[test]
833    fn test_generate_all_options() {
834        let cli = parse_args(&[
835            "generate",
836            "test prompt",
837            "-m",
838            "llama3.2-3b-instruct",
839            "-n",
840            "256",
841            "-t",
842            "0.8",
843            "--top-k",
844            "40",
845            "--top-p",
846            "0.95",
847            "--repetition-penalty",
848            "1.2",
849            "--greedy",
850            "--gpu",
851            "--no-stream",
852            "-q",
853        ])
854        .unwrap();
855
856        match cli.command {
857            Commands::Generate {
858                prompt,
859                model,
860                max_tokens,
861                temperature,
862                top_k,
863                top_p,
864                repetition_penalty,
865                greedy,
866                gpu,
867                no_stream,
868                quiet,
869                ..
870            } => {
871                assert_eq!(prompt, Some("test prompt".to_string()));
872                assert_eq!(model, "llama3.2-3b-instruct");
873                assert_eq!(max_tokens, 256);
874                assert!((temperature - 0.8).abs() < 0.001);
875                assert_eq!(top_k, Some(40));
876                assert_eq!(top_p, Some(0.95));
877                assert!((repetition_penalty - 1.2).abs() < 0.001);
878                assert!(greedy);
879                assert!(gpu);
880                assert!(no_stream);
881                assert!(quiet);
882            }
883            _ => panic!("Expected Generate command"),
884        }
885    }
886
887    #[test]
888    fn test_chat_defaults() {
889        let cli = parse_args(&["chat"]).unwrap();
890
891        match cli.command {
892            Commands::Chat {
893                model,
894                system,
895                temperature,
896                max_tokens,
897                gpu,
898                quiet,
899                ..
900            } => {
901                assert_eq!(model, DEFAULT_CHAT_MODEL);
902                assert!(system.is_none());
903                assert!((temperature - 0.7).abs() < 0.001);
904                assert_eq!(max_tokens, 512);
905                assert!(!gpu);
906                assert!(!quiet);
907            }
908            _ => panic!("Expected Chat command"),
909        }
910    }
911
912    #[test]
913    fn test_chat_with_system_prompt() {
914        let cli = parse_args(&["chat", "-s", "You are a helpful assistant"]).unwrap();
915
916        match cli.command {
917            Commands::Chat { system, .. } => {
918                assert_eq!(system, Some("You are a helpful assistant".to_string()));
919            }
920            _ => panic!("Expected Chat command"),
921        }
922    }
923
924    #[test]
925    fn test_chat_with_model() {
926        let cli = parse_args(&["chat", "-m", "phi3.5-mini"]).unwrap();
927
928        match cli.command {
929            Commands::Chat { model, .. } => {
930                assert_eq!(model, "phi3.5-mini");
931            }
932            _ => panic!("Expected Chat command"),
933        }
934    }
935
936    #[test]
937    fn test_classify_defaults() {
938        let cli = parse_args(&["classify"]).unwrap();
939
940        match cli.command {
941            Commands::Classify {
942                input,
943                model,
944                top_k,
945                format,
946                multi_label,
947                gpu,
948                quiet,
949                ..
950            } => {
951                assert!(input.is_empty());
952                assert_eq!(model, "distilbert-sentiment");
953                assert_eq!(top_k, 5);
954                assert_eq!(format, "text");
955                assert!(!multi_label);
956                assert!(!gpu);
957                assert!(!quiet);
958            }
959            _ => panic!("Expected Classify command"),
960        }
961    }
962
963    #[test]
964    fn test_classify_with_input() {
965        let cli = parse_args(&["classify", "This is great!"]).unwrap();
966
967        match cli.command {
968            Commands::Classify { input, .. } => {
969                assert_eq!(input, vec!["This is great!".to_string()]);
970            }
971            _ => panic!("Expected Classify command"),
972        }
973    }
974
975    #[test]
976    fn test_classify_with_multiple_inputs() {
977        let cli = parse_args(&["classify", "text one", "text two", "text three"]).unwrap();
978
979        match cli.command {
980            Commands::Classify { input, .. } => {
981                assert_eq!(input.len(), 3);
982                assert_eq!(input[0], "text one");
983                assert_eq!(input[1], "text two");
984                assert_eq!(input[2], "text three");
985            }
986            _ => panic!("Expected Classify command"),
987        }
988    }
989
990    #[test]
991    fn test_classify_with_labels() {
992        let cli = parse_args(&["classify", "--labels", "bad,good"]).unwrap();
993
994        match cli.command {
995            Commands::Classify { labels, .. } => {
996                assert_eq!(labels, Some("bad,good".to_string()));
997            }
998            _ => panic!("Expected Classify command"),
999        }
1000    }
1001
1002    #[test]
1003    fn test_classify_with_multi_label() {
1004        let cli = parse_args(&["classify", "--multi-label"]).unwrap();
1005
1006        match cli.command {
1007            Commands::Classify { multi_label, .. } => {
1008                assert!(multi_label);
1009            }
1010            _ => panic!("Expected Classify command"),
1011        }
1012    }
1013
1014    #[test]
1015    fn test_search_minimal() {
1016        let cli = parse_args(&["search", "./index", "my query"]).unwrap();
1017
1018        match cli.command {
1019            Commands::Search {
1020                index_path,
1021                query,
1022                top_k,
1023                mode,
1024                model,
1025                format,
1026                ..
1027            } => {
1028                assert_eq!(index_path, "./index");
1029                assert_eq!(query, "my query");
1030                assert_eq!(top_k, 10);
1031                assert_eq!(mode, "hybrid");
1032                assert_eq!(model, "minilm-l6-v2");
1033                assert_eq!(format, "text");
1034            }
1035            _ => panic!("Expected Search command"),
1036        }
1037    }
1038
1039    #[test]
1040    fn test_search_with_top_k() {
1041        let cli = parse_args(&["search", "./index", "query", "-k", "20"]).unwrap();
1042
1043        match cli.command {
1044            Commands::Search { top_k, .. } => {
1045                assert_eq!(top_k, 20);
1046            }
1047            _ => panic!("Expected Search command"),
1048        }
1049    }
1050
1051    #[test]
1052    fn test_search_with_mode() {
1053        let cli = parse_args(&["search", "./index", "query", "--mode", "semantic"]).unwrap();
1054
1055        match cli.command {
1056            Commands::Search { mode, .. } => {
1057                assert_eq!(mode, "semantic");
1058            }
1059            _ => panic!("Expected Search command"),
1060        }
1061    }
1062
1063    #[test]
1064    fn test_search_with_rerank_model() {
1065        let cli = parse_args(&[
1066            "search",
1067            "./index",
1068            "query",
1069            "--rerank-model",
1070            "ms-marco-minilm",
1071        ])
1072        .unwrap();
1073
1074        match cli.command {
1075            Commands::Search { rerank_model, .. } => {
1076                assert_eq!(rerank_model, Some("ms-marco-minilm".to_string()));
1077            }
1078            _ => panic!("Expected Search command"),
1079        }
1080    }
1081
1082    #[test]
1083    fn test_similarity_minimal() {
1084        let cli = parse_args(&["similarity", "text one", "text two"]).unwrap();
1085
1086        match cli.command {
1087            Commands::Similarity {
1088                text1,
1089                text2,
1090                model,
1091                gpu,
1092                quiet,
1093            } => {
1094                assert_eq!(text1, "text one");
1095                assert_eq!(text2, "text two");
1096                assert_eq!(model, "minilm-l6-v2");
1097                assert!(!gpu);
1098                assert!(!quiet);
1099            }
1100            _ => panic!("Expected Similarity command"),
1101        }
1102    }
1103
1104    #[test]
1105    fn test_model_list_defaults() {
1106        let cli = parse_args(&["model", "list"]).unwrap();
1107
1108        match cli.command {
1109            Commands::Model {
1110                action:
1111                    ModelCommands::List {
1112                        arch,
1113                        task,
1114                        downloaded,
1115                    },
1116            } => {
1117                assert!(arch.is_none());
1118                assert!(task.is_none());
1119                assert!(!downloaded);
1120            }
1121            _ => panic!("Expected Model List command"),
1122        }
1123    }
1124
1125    #[test]
1126    fn test_model_list_with_filters() {
1127        let cli = parse_args(&[
1128            "model",
1129            "list",
1130            "--arch",
1131            "bert",
1132            "--task",
1133            "embedding",
1134            "--downloaded",
1135        ])
1136        .unwrap();
1137
1138        match cli.command {
1139            Commands::Model {
1140                action:
1141                    ModelCommands::List {
1142                        arch,
1143                        task,
1144                        downloaded,
1145                    },
1146            } => {
1147                assert_eq!(arch, Some("bert".to_string()));
1148                assert_eq!(task, Some("embedding".to_string()));
1149                assert!(downloaded);
1150            }
1151            _ => panic!("Expected Model List command"),
1152        }
1153    }
1154
1155    #[test]
1156    fn test_model_download() {
1157        let cli = parse_args(&["model", "download", "minilm-l6-v2"]).unwrap();
1158
1159        match cli.command {
1160            Commands::Model {
1161                action: ModelCommands::Download { name, gguf, quiet },
1162            } => {
1163                assert_eq!(name, "minilm-l6-v2");
1164                assert!(!gguf);
1165                assert!(!quiet);
1166            }
1167            _ => panic!("Expected Model Download command"),
1168        }
1169    }
1170
1171    #[test]
1172    fn test_model_download_gguf() {
1173        let cli = parse_args(&["model", "download", "llama3.2-1b", "--gguf"]).unwrap();
1174
1175        match cli.command {
1176            Commands::Model {
1177                action: ModelCommands::Download { name, gguf, .. },
1178            } => {
1179                assert_eq!(name, "llama3.2-1b");
1180                assert!(gguf);
1181            }
1182            _ => panic!("Expected Model Download command"),
1183        }
1184    }
1185
1186    #[test]
1187    fn test_model_info() {
1188        let cli = parse_args(&["model", "info", "phi3.5-mini"]).unwrap();
1189
1190        match cli.command {
1191            Commands::Model {
1192                action: ModelCommands::Info { name },
1193            } => {
1194                assert_eq!(name, "phi3.5-mini");
1195            }
1196            _ => panic!("Expected Model Info command"),
1197        }
1198    }
1199
1200    #[test]
1201    fn test_model_remove() {
1202        let cli = parse_args(&["model", "remove", "old-model"]).unwrap();
1203
1204        match cli.command {
1205            Commands::Model {
1206                action: ModelCommands::Remove { name },
1207            } => {
1208                assert_eq!(name, "old-model");
1209            }
1210            _ => panic!("Expected Model Remove command"),
1211        }
1212    }
1213
1214    #[test]
1215    fn test_model_search() {
1216        let cli = parse_args(&["model", "search", "llama"]).unwrap();
1217
1218        match cli.command {
1219            Commands::Model {
1220                action: ModelCommands::Search { query },
1221            } => {
1222                assert_eq!(query, "llama");
1223            }
1224            _ => panic!("Expected Model Search command"),
1225        }
1226    }
1227
1228    #[test]
1229    fn test_index_create_minimal() {
1230        let cli = parse_args(&["index", "create", "output.idx"]).unwrap();
1231
1232        match cli.command {
1233            Commands::Index {
1234                action:
1235                    IndexCommands::Create {
1236                        output,
1237                        inputs,
1238                        chunk_size,
1239                        chunk_overlap,
1240                        model,
1241                        ..
1242                    },
1243            } => {
1244                assert_eq!(output, "output.idx");
1245                assert!(inputs.is_empty());
1246                assert_eq!(chunk_size, 512);
1247                assert_eq!(chunk_overlap, 100);
1248                assert_eq!(model, "minilm-l6-v2");
1249            }
1250            _ => panic!("Expected Index Create command"),
1251        }
1252    }
1253
1254    #[test]
1255    fn test_index_create_with_inputs() {
1256        let cli = parse_args(&[
1257            "index",
1258            "create",
1259            "out.idx",
1260            "file1.txt",
1261            "file2.txt",
1262            "dir/",
1263        ])
1264        .unwrap();
1265
1266        match cli.command {
1267            Commands::Index {
1268                action: IndexCommands::Create { inputs, .. },
1269            } => {
1270                assert_eq!(inputs.len(), 3);
1271                assert_eq!(inputs[0], "file1.txt");
1272                assert_eq!(inputs[1], "file2.txt");
1273                assert_eq!(inputs[2], "dir/");
1274            }
1275            _ => panic!("Expected Index Create command"),
1276        }
1277    }
1278
1279    #[test]
1280    fn test_index_create_with_options() {
1281        let cli = parse_args(&[
1282            "index",
1283            "create",
1284            "out.idx",
1285            "--chunk-size",
1286            "500",
1287            "--chunk-overlap",
1288            "100",
1289            "-m",
1290            "nomic-embed-text",
1291            "--gpu",
1292            "-q",
1293        ])
1294        .unwrap();
1295
1296        match cli.command {
1297            Commands::Index {
1298                action:
1299                    IndexCommands::Create {
1300                        chunk_size,
1301                        chunk_overlap,
1302                        model,
1303                        gpu,
1304                        quiet,
1305                        ..
1306                    },
1307            } => {
1308                assert_eq!(chunk_size, 500);
1309                assert_eq!(chunk_overlap, 100);
1310                assert_eq!(model, "nomic-embed-text");
1311                assert!(gpu);
1312                assert!(quiet);
1313            }
1314            _ => panic!("Expected Index Create command"),
1315        }
1316    }
1317
1318    #[test]
1319    fn test_index_add() {
1320        let cli = parse_args(&["index", "add", "existing.idx", "newfile.txt"]).unwrap();
1321
1322        match cli.command {
1323            Commands::Index {
1324                action:
1325                    IndexCommands::Add {
1326                        index_path, inputs, ..
1327                    },
1328            } => {
1329                assert_eq!(index_path, "existing.idx");
1330                assert_eq!(inputs, vec!["newfile.txt".to_string()]);
1331            }
1332            _ => panic!("Expected Index Add command"),
1333        }
1334    }
1335
1336    #[test]
1337    fn test_index_info() {
1338        let cli = parse_args(&["index", "info", "my.idx"]).unwrap();
1339
1340        match cli.command {
1341            Commands::Index {
1342                action: IndexCommands::Info { index_path },
1343            } => {
1344                assert_eq!(index_path, "my.idx");
1345            }
1346            _ => panic!("Expected Index Info command"),
1347        }
1348    }
1349
1350    #[test]
1351    fn test_verbose_zero() {
1352        let cli = parse_args(&["generate"]).unwrap();
1353        assert_eq!(cli.verbose, 0);
1354    }
1355
1356    #[test]
1357    fn test_verbose_one() {
1358        let cli = parse_args(&["-v", "generate"]).unwrap();
1359        assert_eq!(cli.verbose, 1);
1360    }
1361
1362    #[test]
1363    fn test_verbose_two() {
1364        let cli = parse_args(&["-vv", "generate"]).unwrap();
1365        assert_eq!(cli.verbose, 2);
1366    }
1367
1368    #[test]
1369    fn test_verbose_three() {
1370        let cli = parse_args(&["-vvv", "generate"]).unwrap();
1371        assert_eq!(cli.verbose, 3);
1372    }
1373
1374    #[test]
1375    fn test_verbose_long_form() {
1376        let cli = parse_args(&["--verbose", "--verbose", "generate"]).unwrap();
1377        assert_eq!(cli.verbose, 2);
1378    }
1379
1380    #[test]
1381    fn test_verbose_after_command() {
1382        // Global flag can come after command
1383        let cli = parse_args(&["generate", "-v"]).unwrap();
1384        assert_eq!(cli.verbose, 1);
1385    }
1386
1387    #[test]
1388    fn test_missing_command() {
1389        let result = parse_args(&[]);
1390        assert!(result.is_err());
1391    }
1392
1393    #[test]
1394    fn test_unknown_command() {
1395        let result = parse_args(&["unknown"]);
1396        assert!(result.is_err());
1397    }
1398
1399    #[test]
1400    fn test_missing_required_arg() {
1401        // transcribe requires a file
1402        let result = parse_args(&["transcribe"]);
1403        assert!(result.is_err());
1404    }
1405
1406    #[test]
1407    fn test_invalid_number() {
1408        let result = parse_args(&["generate", "-n", "not_a_number"]);
1409        assert!(result.is_err());
1410    }
1411
1412    #[test]
1413    fn test_invalid_float() {
1414        let result = parse_args(&["generate", "-t", "not_a_float"]);
1415        assert!(result.is_err());
1416    }
1417
1418    #[test]
1419    fn test_rerank_minimal() {
1420        let cli = parse_args(&["rerank", "my query"]).unwrap();
1421
1422        match cli.command {
1423            Commands::Rerank {
1424                query,
1425                documents,
1426                model,
1427                top_k,
1428                format,
1429                ..
1430            } => {
1431                assert_eq!(query, "my query");
1432                assert!(documents.is_empty());
1433                assert_eq!(model, "minilm-l6-v2-cross-encoder");
1434                assert!(top_k.is_none());
1435                assert_eq!(format, "text");
1436            }
1437            _ => panic!("Expected Rerank command"),
1438        }
1439    }
1440
1441    #[test]
1442    fn test_rerank_with_documents() {
1443        let cli = parse_args(&["rerank", "query", "doc1", "doc2", "doc3"]).unwrap();
1444
1445        match cli.command {
1446            Commands::Rerank {
1447                query, documents, ..
1448            } => {
1449                assert_eq!(query, "query");
1450                assert_eq!(documents.len(), 3);
1451            }
1452            _ => panic!("Expected Rerank command"),
1453        }
1454    }
1455
1456    #[test]
1457    fn test_rerank_with_top_k() {
1458        let cli = parse_args(&["rerank", "query", "-k", "5"]).unwrap();
1459
1460        match cli.command {
1461            Commands::Rerank { top_k, .. } => {
1462                assert_eq!(top_k, Some(5));
1463            }
1464            _ => panic!("Expected Rerank command"),
1465        }
1466    }
1467
1468    #[test]
1469    fn test_summarize_defaults() {
1470        let cli = parse_args(&["summarize"]).unwrap();
1471
1472        match cli.command {
1473            Commands::Summarize {
1474                input,
1475                model,
1476                min_length,
1477                max_length,
1478                num_beams,
1479                ..
1480            } => {
1481                assert!(input.is_none());
1482                assert_eq!(model, "distilbart-cnn");
1483                assert!(min_length.is_none());
1484                assert!(max_length.is_none());
1485                assert!(num_beams.is_none());
1486            }
1487            _ => panic!("Expected Summarize command"),
1488        }
1489    }
1490
1491    #[test]
1492    fn test_summarize_with_options() {
1493        let cli = parse_args(&[
1494            "summarize",
1495            "--input",
1496            "input.txt",
1497            "--min-length",
1498            "50",
1499            "--max-length",
1500            "200",
1501            "--num-beams",
1502            "4",
1503            "--length-penalty",
1504            "1.5",
1505        ])
1506        .unwrap();
1507
1508        match cli.command {
1509            Commands::Summarize {
1510                input,
1511                min_length,
1512                max_length,
1513                num_beams,
1514                length_penalty,
1515                ..
1516            } => {
1517                assert_eq!(input, Some("input.txt".to_string()));
1518                assert_eq!(min_length, Some(50));
1519                assert_eq!(max_length, Some(200));
1520                assert_eq!(num_beams, Some(4));
1521                assert_eq!(length_penalty, Some(1.5));
1522            }
1523            _ => panic!("Expected Summarize command"),
1524        }
1525    }
1526
1527    #[test]
1528    fn test_translate_defaults() {
1529        let cli = parse_args(&["translate"]).unwrap();
1530
1531        match cli.command {
1532            Commands::Translate {
1533                input,
1534                model,
1535                src,
1536                dst,
1537                ..
1538            } => {
1539                assert!(input.is_none());
1540                assert_eq!(model, "flan-t5-base");
1541                assert!(src.is_none());
1542                assert!(dst.is_none());
1543            }
1544            _ => panic!("Expected Translate command"),
1545        }
1546    }
1547
1548    #[test]
1549    fn test_translate_with_languages() {
1550        let cli = parse_args(&["translate", "--src", "en", "--dst", "is"]).unwrap();
1551
1552        match cli.command {
1553            Commands::Translate { src, dst, .. } => {
1554                assert_eq!(src, Some("en".to_string()));
1555                assert_eq!(dst, Some("is".to_string()));
1556            }
1557            _ => panic!("Expected Translate command"),
1558        }
1559    }
1560    #[test]
1561    fn test_transcribe_minimal() {
1562        let cli = parse_args(&["transcribe", "audio.wav"]).unwrap();
1563
1564        match cli.command {
1565            Commands::Transcribe {
1566                file,
1567                model,
1568                language,
1569                ..
1570            } => {
1571                assert_eq!(file, "audio.wav");
1572                assert_eq!(model, "whisper-small");
1573                assert!(language.is_none());
1574            }
1575            _ => panic!("Expected Transcribe command"),
1576        }
1577    }
1578
1579    #[test]
1580    fn test_classify_with_model() {
1581        let cli = parse_args(&["classify", "i hate mondays", "--model", "toxic-bert"]).unwrap();
1582        match cli.command {
1583            Commands::Classify { input, model, .. } => {
1584                assert_eq!(input, vec!["i hate mondays".to_string()]);
1585                assert_eq!(model, "toxic-bert");
1586            }
1587            _ => panic!("Expected Classify command"),
1588        }
1589    }
1590
1591    #[test]
1592    fn test_transcribe_with_options() {
1593        let cli = parse_args(&[
1594            "transcribe",
1595            "audio.mp3",
1596            "-m",
1597            "whisper-large-v3",
1598            "--language",
1599            "is",
1600        ])
1601        .unwrap();
1602
1603        match cli.command {
1604            Commands::Transcribe {
1605                file,
1606                model,
1607                language,
1608                ..
1609            } => {
1610                assert_eq!(file, "audio.mp3");
1611                assert_eq!(model, "whisper-large-v3");
1612                assert_eq!(language, Some("is".to_string()));
1613            }
1614            _ => panic!("Expected Transcribe command"),
1615        }
1616    }
1617}
1618
1619mod send_sync_tests {
1620    use kjarni::{
1621        Classifier, Embedder, Indexer, Reranker, Searcher, chat::Chat, generator::Generator,
1622    };
1623    // Compile time verificatio
1624    const _: () = {
1625        const fn assert_send<T: Send>() {}
1626        const fn assert_sync<T: Sync>() {}
1627        assert_send::<Embedder>();
1628        assert_sync::<Embedder>();
1629
1630        assert_send::<Indexer>();
1631        assert_sync::<Indexer>();
1632
1633        assert_send::<Searcher>();
1634        assert_sync::<Searcher>();
1635
1636        assert_send::<Reranker>();
1637        assert_sync::<Reranker>();
1638
1639        assert_send::<Generator>();
1640        assert_sync::<Generator>();
1641
1642        assert_send::<Chat>();
1643        assert_sync::<Chat>();
1644
1645        assert_send::<Classifier>();
1646        assert_sync::<Classifier>();
1647    };
1648}