1use clap::{Parser, Subcommand};
2
3pub const DEFAULT_CHAT_MODEL: &str = "llama3.2-3b-instruct";
11
12pub 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 #[arg(short, long, action = clap::ArgAction::Count, global = true)]
25 pub verbose: u8,
26}
27
28#[derive(Subcommand, Debug, PartialEq)]
29pub enum Commands {
30 Model {
32 #[command(subcommand)]
33 action: ModelCommands,
34 },
35
36 Generate {
38 prompt: Option<String>,
40
41 #[arg(short, long, default_value = DEFAULT_GENERATE_MODEL)]
42 model: String,
43
44 #[arg(long)]
46 model_path: Option<String>,
47
48 #[arg(short = 'n', long, default_value_t = 100)]
50 max_tokens: usize,
51
52 #[arg(short, long, default_value_t = 0.7)]
54 temperature: f32,
55
56 #[arg(long)]
58 top_k: Option<usize>,
59
60 #[arg(long)]
62 top_p: Option<f32>,
63
64 #[arg(long)]
66 min_p: Option<f32>,
67
68 #[arg(long, default_value_t = 1.1)]
70 repetition_penalty: f32,
71
72 #[arg(long)]
74 greedy: bool,
75
76 #[arg(long)]
78 gpu: bool,
79
80 #[arg(long)]
82 no_stream: bool,
83
84 #[arg(short, long)]
86 quiet: bool,
87 },
88
89 Summarize {
91 #[arg(short, long)]
93 input: Option<String>,
94
95 #[arg(short, long, default_value = "distilbart-cnn")]
97 model: String,
98
99 #[arg(long)]
102 model_path: Option<String>,
103
104 #[arg(long)]
106 min_length: Option<usize>,
107
108 #[arg(long)]
110 max_length: Option<usize>,
111
112 #[arg(long)]
114 num_beams: Option<usize>,
115
116 #[arg(long)]
118 length_penalty: Option<f32>,
119
120 #[arg(long)]
122 no_repeat_ngram: Option<usize>,
123
124 #[arg(long)]
126 greedy: bool,
127
128 #[arg(long)]
130 no_stream: bool,
131
132 #[arg(long)]
134 gpu: bool,
135
136 #[arg(short, long)]
138 quiet: bool,
139 },
140
141 Translate {
143 #[arg(short, long)]
145 input: Option<String>,
146
147 #[arg(short, long, default_value = "flan-t5-base")]
149 model: String,
150
151 #[arg(long)]
154 model_path: Option<String>,
155
156 #[arg(long)]
158 src: Option<String>,
159
160 #[arg(long)]
162 dst: Option<String>,
163
164 #[arg(long)]
166 max_length: Option<usize>,
167
168 #[arg(long)]
170 num_beams: Option<usize>,
171
172 #[arg(long)]
174 length_penalty: Option<f32>,
175
176 #[arg(long)]
178 no_repeat_ngram: Option<usize>,
179
180 #[arg(long)]
182 greedy: bool,
183
184 #[arg(long)]
186 no_stream: bool,
187
188 #[arg(long)]
190 gpu: bool,
191
192 #[arg(short, long)]
194 quiet: bool,
195 },
196
197 Inspect {
199 path: String,
201 },
202
203 Embed {
205 input: Option<String>,
207
208 #[arg(short, long, default_value = "minilm-l6-v2")]
210 model: String,
211
212 #[arg(long)]
214 model_path: Option<String>,
215
216 #[arg(long, default_value = "raw")]
218 format: String,
219
220 #[arg(long)]
222 normalize: bool,
223
224 #[arg(long, default_value = "mean")]
231 pooling: String,
232
233 #[arg(long)]
235 gpu: bool,
236
237 #[arg(short, long)]
239 quiet: bool,
240 },
241
242 Transcribe {
244 file: String,
246
247 #[arg(short, long, default_value = "whisper-small")]
249 model: String,
250
251 #[arg(long)]
254 model_path: Option<String>,
255
256 #[arg(short, long)]
258 language: Option<String>,
259
260 #[arg(long)]
262 translate: bool,
263
264 #[arg(short, long)]
266 timestamps: bool,
267
268 #[arg(long)]
270 max_tokens: Option<usize>,
271
272 #[arg(long)]
274 no_stream: bool,
275
276 #[arg(long)]
278 gpu: bool,
279
280 #[arg(short, long)]
282 quiet: bool,
283 },
284
285 Classify {
287 input: Vec<String>,
289
290 #[arg(short, long, default_value = "distilbert-sentiment")]
292 model: String,
293
294 #[arg(long, value_name = "PATH")]
296 model_path: Option<String>,
297
298 #[arg(long, value_name = "LABELS")]
301 labels: Option<String>,
302
303 #[arg(long, default_value = "5")]
305 top_k: usize,
306
307 #[arg(long)]
309 threshold: Option<f32>,
310
311 #[arg(long)]
313 max_length: Option<usize>,
314
315 #[arg(long)]
317 batch_size: Option<usize>,
318
319 #[arg(long)]
321 multi_label: bool,
322
323 #[arg(short, long, default_value = "text")]
325 format: String,
326
327 #[arg(long)]
329 gpu: bool,
330
331 #[arg(long)]
333 dtype: Option<String>,
334
335 #[arg(short, long)]
337 quiet: bool,
338 },
339
340 Rerank {
342 query: String,
344
345 documents: Vec<String>,
347
348 #[arg(short, long, default_value = "minilm-l6-v2-cross-encoder")]
349 model: String,
350
351 #[arg(long)]
353 model_path: Option<String>,
354
355 #[arg(short = 'k', long)]
357 top_k: Option<usize>,
358
359 #[arg(short, long, default_value = "text")]
361 format: String,
362
363 #[arg(long)]
365 gpu: bool,
366
367 #[arg(short, long)]
369 quiet: bool,
370 },
371
372 Chat {
374 #[arg(short, long, default_value = DEFAULT_CHAT_MODEL)]
375 model: String,
376
377 #[arg(long)]
379 model_path: Option<String>,
380
381 #[arg(short, long)]
383 system: Option<String>,
384
385 #[arg(short, long, default_value_t = 0.7)]
387 temperature: f32,
388
389 #[arg(short = 'n', long, default_value_t = 512)]
391 max_tokens: usize,
392
393 #[arg(long)]
395 gpu: bool,
396
397 #[arg(short, long)]
399 quiet: bool,
400 },
401
402 Index {
404 #[command(subcommand)]
405 action: IndexCommands,
406 },
407
408 Search {
410 index_path: String,
412
413 query: String,
415
416 #[arg(short = 'k', long, default_value_t = 10)]
418 top_k: usize,
419
420 #[arg(long, default_value = "hybrid")]
422 mode: String,
423
424 #[arg(short, long, default_value = "minilm-l6-v2")]
426 model: String,
427
428 #[arg(long, default_value = None)]
432 rerank_model: Option<String>,
433
434 #[arg(short, long, default_value = "text")]
436 format: String,
437
438 #[arg(long)]
440 gpu: bool,
441
442 #[arg(short, long)]
444 quiet: bool,
445 },
446
447 Similarity {
449 text1: String,
451
452 text2: String,
454
455 #[arg(short, long, default_value = "minilm-l6-v2")]
457 model: String,
458
459 #[arg(long)]
461 gpu: bool,
462
463 #[arg(short, long)]
465 quiet: bool,
466 },
467}
468
469#[derive(Subcommand, Debug, PartialEq)]
470pub enum ModelCommands {
471 List {
473 #[arg(short, long)]
475 arch: Option<String>,
476
477 #[arg(short, long)]
479 task: Option<String>,
480
481 #[arg(short, long)]
483 downloaded: bool,
484 },
485
486 Download {
488 name: String,
489
490 #[arg(long)]
491 gguf: bool,
492
493 #[arg(short, long)]
494 quiet: bool,
495 },
496
497 Remove { name: String },
499
500 Info { name: String },
502
503 Search { query: String },
505}
506
507#[derive(Subcommand, Debug, PartialEq)]
508pub enum InspectCommands {
509 Model {
511 path: String,
513 },
514}
515
516#[derive(Subcommand, Debug, PartialEq)]
517pub enum IndexCommands {
518 Create {
520 output: String,
522
523 #[arg(conflicts_with = "from_chunks")]
525 inputs: Vec<String>,
526
527 #[arg(long, conflicts_with = "inputs")]
529 from_chunks: Option<String>,
530
531 #[arg(long, default_value_t = 512)]
538 chunk_size: usize,
539
540 #[arg(long, default_value_t = 100)]
542 chunk_overlap: usize,
543
544 #[arg(short, long, default_value = "minilm-l6-v2")]
546 model: String,
547
548 #[arg(long)]
550 gpu: bool,
551
552 #[arg(short, long)]
554 quiet: bool,
555 },
556
557 Add {
559 index_path: String,
561
562 inputs: Vec<String>,
564
565 #[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 #[arg(long)]
579 gpu: bool,
580
581 #[arg(short, long)]
583 quiet: bool,
584 },
585
586 Info {
588 index_path: String,
590 },
591}
592
593pub 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 #[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 #[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 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 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 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}