1use std::path::PathBuf;
2
3use clap::{ArgGroup, Args, Parser, Subcommand, ValueEnum};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
6pub enum OutputFormat {
7 Cli,
8 Json,
9}
10
11#[derive(Debug, Parser)]
12#[command(name = "kbolt", version, about = "local-first retrieval engine")]
13pub struct Cli {
14 #[arg(short = 's', long = "space", value_name = "name")]
15 pub space: Option<String>,
16
17 #[arg(
18 short = 'f',
19 long = "format",
20 value_enum,
21 default_value_t = OutputFormat::Cli
22 )]
23 pub format: OutputFormat,
24
25 #[command(subcommand)]
26 pub command: Command,
27}
28
29#[derive(Debug, Subcommand)]
30pub enum Command {
31 Doctor,
32 Setup(SetupArgs),
33 Local(LocalArgs),
34 Space(SpaceArgs),
35 Collection(CollectionArgs),
36 Ignore(IgnoreArgs),
37 Models(ModelsArgs),
38 Eval(EvalArgs),
39 Schedule(ScheduleArgs),
40 Mcp,
41 Search(SearchArgs),
42 Update(UpdateArgs),
43 Status,
44 Ls(LsArgs),
45 Get(GetArgs),
46 MultiGet(MultiGetArgs),
47}
48
49#[derive(Debug, Args)]
50pub struct SpaceArgs {
51 #[command(subcommand)]
52 pub command: SpaceCommand,
53}
54
55#[derive(Debug, Args)]
56pub struct SetupArgs {
57 #[command(subcommand)]
58 pub command: SetupCommand,
59}
60
61#[derive(Debug, Args)]
62pub struct LocalArgs {
63 #[command(subcommand)]
64 pub command: LocalCommand,
65}
66
67#[derive(Debug, Args)]
68pub struct CollectionArgs {
69 #[command(subcommand)]
70 pub command: CollectionCommand,
71}
72
73#[derive(Debug, Args)]
74pub struct IgnoreArgs {
75 #[command(subcommand)]
76 pub command: IgnoreCommand,
77}
78
79#[derive(Debug, Args)]
80pub struct ModelsArgs {
81 #[command(subcommand)]
82 pub command: ModelsCommand,
83}
84
85#[derive(Debug, Args)]
86pub struct EvalArgs {
87 #[command(subcommand)]
88 pub command: EvalCommand,
89}
90
91#[derive(Debug, Args, PartialEq, Eq)]
92pub struct EvalImportArgs {
93 #[command(subcommand)]
94 pub dataset: EvalImportCommand,
95}
96
97#[derive(Debug, Args, PartialEq, Eq)]
98pub struct EvalRunArgs {
99 #[arg(long, value_name = "path")]
100 pub file: Option<PathBuf>,
101}
102
103#[derive(Debug, Args)]
104pub struct ScheduleArgs {
105 #[command(subcommand)]
106 pub command: ScheduleCommand,
107}
108
109#[derive(Debug, Args, PartialEq, Eq)]
110pub struct UpdateArgs {
111 #[arg(long = "collection", value_delimiter = ',')]
112 pub collections: Vec<String>,
113 #[arg(long)]
114 pub no_embed: bool,
115 #[arg(long)]
116 pub dry_run: bool,
117 #[arg(long)]
118 pub verbose: bool,
119}
120
121#[derive(Debug, Args, PartialEq, Eq)]
122pub struct LsArgs {
123 pub collection: String,
124 pub prefix: Option<String>,
125 #[arg(long)]
126 pub all: bool,
127}
128
129#[derive(Debug, Args, PartialEq, Eq)]
130pub struct GetArgs {
131 pub identifier: String,
132 #[arg(long)]
133 pub offset: Option<usize>,
134 #[arg(long)]
135 pub limit: Option<usize>,
136}
137
138#[derive(Debug, Args, PartialEq, Eq)]
139pub struct MultiGetArgs {
140 #[arg(value_delimiter = ',')]
141 pub locators: Vec<String>,
142 #[arg(long, default_value_t = 20)]
143 pub max_files: usize,
144 #[arg(long, default_value_t = 51_200)]
145 pub max_bytes: usize,
146}
147
148#[derive(Debug, Args, PartialEq)]
149pub struct SearchArgs {
150 pub query: String,
151 #[arg(long = "collection", value_delimiter = ',')]
152 pub collections: Vec<String>,
153 #[arg(long, default_value_t = 10)]
154 pub limit: usize,
155 #[arg(long, default_value_t = 0.0)]
156 pub min_score: f32,
157 #[arg(long)]
158 pub deep: bool,
159 #[arg(long)]
160 pub keyword: bool,
161 #[arg(long)]
162 pub semantic: bool,
163 #[arg(long, conflicts_with = "rerank")]
164 pub no_rerank: bool,
165 #[arg(long, conflicts_with = "no_rerank")]
166 pub rerank: bool,
167 #[arg(long)]
168 pub debug: bool,
169}
170
171#[derive(Debug, Subcommand, PartialEq, Eq)]
172pub enum SpaceCommand {
173 Add {
174 name: String,
175 #[arg(long)]
176 description: Option<String>,
177 #[arg(long)]
178 strict: bool,
179 dirs: Vec<PathBuf>,
180 },
181 Describe {
182 name: String,
183 text: String,
184 },
185 Rename {
186 old: String,
187 new: String,
188 },
189 Remove {
190 name: String,
191 },
192 Current,
193 Default {
194 name: Option<String>,
195 },
196 List,
197 Info {
198 name: String,
199 },
200}
201
202#[derive(Debug, Subcommand, PartialEq, Eq)]
203pub enum SetupCommand {
204 Local,
205}
206
207#[derive(Debug, Subcommand, PartialEq, Eq)]
208pub enum LocalCommand {
209 Status,
210 Start,
211 Stop,
212 Enable { feature: LocalFeature },
213}
214
215#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
216pub enum LocalFeature {
217 Deep,
218}
219
220#[derive(Debug, Subcommand, PartialEq, Eq)]
221pub enum CollectionCommand {
222 Add {
223 path: PathBuf,
224 #[arg(long)]
225 name: Option<String>,
226 #[arg(long)]
227 description: Option<String>,
228 #[arg(long, value_delimiter = ',')]
229 extensions: Option<Vec<String>>,
230 #[arg(long)]
231 no_index: bool,
232 },
233 List,
234 Info {
235 name: String,
236 },
237 Describe {
238 name: String,
239 text: String,
240 },
241 Rename {
242 old: String,
243 new: String,
244 },
245 Remove {
246 name: String,
247 },
248}
249
250#[derive(Debug, Subcommand, PartialEq, Eq)]
251pub enum IgnoreCommand {
252 Show { collection: String },
253 Add { collection: String, pattern: String },
254 Remove { collection: String, pattern: String },
255 Edit { collection: String },
256 List,
257}
258
259#[derive(Debug, Subcommand, PartialEq, Eq)]
260pub enum ModelsCommand {
261 List,
262}
263
264#[derive(Debug, Subcommand, PartialEq, Eq)]
265pub enum EvalCommand {
266 Run(EvalRunArgs),
267 Import(EvalImportArgs),
268}
269
270#[derive(Debug, Subcommand, PartialEq, Eq)]
271pub enum EvalImportCommand {
272 #[command(
273 about = "import a canonical BEIR dataset from an extracted directory",
274 long_about = "Import a canonical BEIR dataset from an extracted directory.\n\nExpected source layout:\n corpus.jsonl\n queries.jsonl\n qrels/test.tsv\n\nThis command always imports the test split."
275 )]
276 Beir(EvalImportBeirArgs),
277}
278
279#[derive(Debug, Args, PartialEq, Eq)]
280pub struct EvalImportBeirArgs {
281 #[arg(long, value_name = "name")]
282 pub dataset: String,
283 #[arg(long, value_name = "dir")]
284 pub source: PathBuf,
285 #[arg(long, value_name = "dir")]
286 pub output: PathBuf,
287 #[arg(long, value_name = "name")]
288 pub collection: Option<String>,
289}
290
291#[derive(Debug, Subcommand, PartialEq, Eq)]
292pub enum ScheduleCommand {
293 Add(ScheduleAddArgs),
294 Status,
295 Remove(ScheduleRemoveArgs),
296}
297
298#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
299pub enum ScheduleDayArg {
300 Mon,
301 Tue,
302 Wed,
303 Thu,
304 Fri,
305 Sat,
306 Sun,
307}
308
309#[derive(Debug, Args, PartialEq, Eq)]
310#[command(group(
311 ArgGroup::new("trigger")
312 .required(true)
313 .args(["every", "at"])
314))]
315pub struct ScheduleAddArgs {
316 #[arg(long, conflicts_with = "at")]
317 pub every: Option<String>,
318 #[arg(long, conflicts_with = "every")]
319 pub at: Option<String>,
320 #[arg(long = "on", value_delimiter = ',', requires = "at", value_enum)]
321 pub on: Vec<ScheduleDayArg>,
322 #[arg(long)]
323 pub space: Option<String>,
324 #[arg(long = "collection", requires = "space")]
325 pub collections: Vec<String>,
326}
327
328#[derive(Debug, Args, PartialEq, Eq)]
329#[command(group(
330 ArgGroup::new("selector")
331 .required(true)
332 .args(["id", "all", "space"])
333))]
334pub struct ScheduleRemoveArgs {
335 pub id: Option<String>,
336 #[arg(long, conflicts_with_all = ["id", "space", "collections"])]
337 pub all: bool,
338 #[arg(long, conflicts_with = "id")]
339 pub space: Option<String>,
340 #[arg(long = "collection", requires = "space", conflicts_with = "id")]
341 pub collections: Vec<String>,
342}
343
344#[cfg(test)]
345mod tests {
346 use std::path::PathBuf;
347
348 use clap::Parser;
349
350 use super::{
351 Cli, CollectionCommand, Command, EvalCommand, EvalImportArgs, EvalImportBeirArgs,
352 EvalImportCommand, EvalRunArgs, GetArgs, LocalCommand, LocalFeature, MultiGetArgs,
353 OutputFormat, ScheduleAddArgs, ScheduleCommand, ScheduleDayArg, ScheduleRemoveArgs,
354 SearchArgs, SetupCommand, SpaceCommand, UpdateArgs,
355 };
356
357 fn parse<const N: usize>(args: [&str; N]) -> Cli {
358 Cli::try_parse_from(args).expect("parse cli")
359 }
360
361 #[test]
362 fn parses_output_format_variants() {
363 let parsed = parse(["kbolt", "status"]);
364 assert_eq!(parsed.format, OutputFormat::Cli);
365 let parsed = parse(["kbolt", "--format", "json", "status"]);
366 assert_eq!(parsed.format, OutputFormat::Json);
367 }
368
369 #[test]
370 fn parses_doctor_command() {
371 let parsed = parse(["kbolt", "doctor"]);
372 assert!(matches!(parsed.command, Command::Doctor));
373 }
374
375 #[test]
376 fn parses_setup_local_command() {
377 let parsed = parse(["kbolt", "setup", "local"]);
378 assert!(matches!(
379 parsed.command,
380 Command::Setup(args) if args.command == SetupCommand::Local
381 ));
382 }
383
384 #[test]
385 fn parses_local_enable_deep_command() {
386 let parsed = parse(["kbolt", "local", "enable", "deep"]);
387 assert!(matches!(
388 parsed.command,
389 Command::Local(args)
390 if args.command == LocalCommand::Enable {
391 feature: LocalFeature::Deep
392 }
393 ));
394 }
395
396 #[test]
397 fn parses_global_space_override() {
398 let parsed = parse(["kbolt", "--space", "work", "space", "current"]);
399 assert_eq!(parsed.space.as_deref(), Some("work"));
400 assert!(matches!(
401 parsed.command,
402 Command::Space(space) if space.command == SpaceCommand::Current
403 ));
404 }
405
406 #[test]
407 fn parses_collection_add_with_options() {
408 let parsed = parse([
409 "kbolt",
410 "collection",
411 "add",
412 "/tmp/work-api",
413 "--name",
414 "api",
415 "--description",
416 "api docs",
417 "--extensions",
418 "rs,md",
419 "--no-index",
420 ]);
421 assert_eq!(parsed.space, None);
422
423 assert!(matches!(
424 parsed.command,
425 Command::Collection(collection)
426 if collection.command
427 == CollectionCommand::Add {
428 path: PathBuf::from("/tmp/work-api"),
429 name: Some("api".to_string()),
430 description: Some("api docs".to_string()),
431 extensions: Some(vec!["rs".to_string(), "md".to_string()]),
432 no_index: true
433 }
434 ));
435 }
436
437 #[test]
438 fn parses_update_with_defaults() {
439 let parsed = parse(["kbolt", "update"]);
440 assert_eq!(parsed.space, None);
441 assert!(matches!(
442 parsed.command,
443 Command::Update(UpdateArgs {
444 collections,
445 no_embed: false,
446 dry_run: false,
447 verbose: false,
448 }) if collections.is_empty()
449 ));
450 }
451
452 #[test]
453 fn parses_update_with_flags() {
454 let parsed = parse([
455 "kbolt",
456 "--space",
457 "work",
458 "update",
459 "--collection",
460 "api,wiki",
461 "--no-embed",
462 "--dry-run",
463 "--verbose",
464 ]);
465 assert_eq!(parsed.space.as_deref(), Some("work"));
466 assert!(matches!(
467 parsed.command,
468 Command::Update(UpdateArgs {
469 collections,
470 no_embed: true,
471 dry_run: true,
472 verbose: true,
473 }) if collections == vec!["api".to_string(), "wiki".to_string()]
474 ));
475 }
476
477 #[test]
478 fn parses_get_with_options() {
479 let parsed = parse(["kbolt", "get", "api/src/lib.rs"]);
480 assert_eq!(parsed.space, None);
481 assert!(matches!(
482 parsed.command,
483 Command::Get(GetArgs {
484 identifier,
485 offset: None,
486 limit: None,
487 }) if identifier == "api/src/lib.rs"
488 ));
489
490 let parsed = parse([
491 "kbolt", "--space", "work", "get", "#abc123", "--offset", "10", "--limit", "25",
492 ]);
493 assert_eq!(parsed.space.as_deref(), Some("work"));
494 assert!(matches!(
495 parsed.command,
496 Command::Get(GetArgs {
497 identifier,
498 offset: Some(10),
499 limit: Some(25),
500 }) if identifier == "#abc123"
501 ));
502 }
503
504 #[test]
505 fn parses_multi_get_with_options() {
506 let parsed = parse(["kbolt", "multi-get", "api/a.md,#abc123"]);
507 assert_eq!(parsed.space, None);
508 assert!(matches!(
509 parsed.command,
510 Command::MultiGet(MultiGetArgs {
511 locators,
512 max_files: 20,
513 max_bytes: 51_200,
514 }) if locators == vec!["api/a.md".to_string(), "#abc123".to_string()]
515 ));
516
517 let parsed = parse([
518 "kbolt",
519 "--space",
520 "work",
521 "multi-get",
522 "api/a.md,api/b.md",
523 "--max-files",
524 "5",
525 "--max-bytes",
526 "1024",
527 ]);
528 assert_eq!(parsed.space.as_deref(), Some("work"));
529 assert!(matches!(
530 parsed.command,
531 Command::MultiGet(MultiGetArgs {
532 locators,
533 max_files: 5,
534 max_bytes: 1024,
535 }) if locators == vec!["api/a.md".to_string(), "api/b.md".to_string()]
536 ));
537 }
538
539 #[test]
540 fn parses_search_with_defaults_and_flags() {
541 let parsed = parse(["kbolt", "search", "alpha"]);
542 assert_eq!(parsed.space, None);
543 assert!(matches!(
544 parsed.command,
545 Command::Search(SearchArgs {
546 query,
547 collections,
548 limit: 10,
549 min_score,
550 deep: false,
551 keyword: false,
552 semantic: false,
553 no_rerank: false,
554 rerank: false,
555 debug: false,
556 }) if query == "alpha" && collections.is_empty() && min_score == 0.0
557 ));
558
559 let parsed = parse([
560 "kbolt",
561 "--space",
562 "work",
563 "search",
564 "alpha beta",
565 "--collection",
566 "api,wiki",
567 "--limit",
568 "7",
569 "--min-score",
570 "0.25",
571 "--keyword",
572 "--no-rerank",
573 "--debug",
574 ]);
575 assert_eq!(parsed.space.as_deref(), Some("work"));
576 assert!(matches!(
577 parsed.command,
578 Command::Search(SearchArgs {
579 query,
580 collections,
581 limit: 7,
582 min_score,
583 deep: false,
584 keyword: true,
585 semantic: false,
586 no_rerank: true,
587 rerank: false,
588 debug: true,
589 }) if query == "alpha beta"
590 && collections == vec!["api".to_string(), "wiki".to_string()]
591 && min_score == 0.25
592 ));
593 }
594
595 #[test]
596 fn parses_search_rerank_opt_in_flag() {
597 let parsed = parse(["kbolt", "search", "alpha", "--rerank"]);
598 assert!(matches!(
599 parsed.command,
600 Command::Search(SearchArgs {
601 rerank: true,
602 no_rerank: false,
603 ..
604 })
605 ));
606 }
607
608 #[test]
609 fn parses_schedule_add_interval_and_weekly_variants() {
610 let parsed = parse(["kbolt", "schedule", "add", "--every", "30m"]);
611 assert!(matches!(
612 parsed.command,
613 Command::Schedule(schedule)
614 if schedule.command
615 == ScheduleCommand::Add(ScheduleAddArgs {
616 every: Some("30m".to_string()),
617 at: None,
618 on: vec![],
619 space: None,
620 collections: vec![],
621 })
622 ));
623
624 let parsed = parse([
625 "kbolt",
626 "schedule",
627 "add",
628 "--at",
629 "3pm",
630 "--on",
631 "mon,fri",
632 "--space",
633 "work",
634 "--collection",
635 "api",
636 "--collection",
637 "docs",
638 ]);
639 assert!(matches!(
640 parsed.command,
641 Command::Schedule(schedule)
642 if schedule.command
643 == ScheduleCommand::Add(ScheduleAddArgs {
644 every: None,
645 at: Some("3pm".to_string()),
646 on: vec![ScheduleDayArg::Mon, ScheduleDayArg::Fri],
647 space: Some("work".to_string()),
648 collections: vec!["api".to_string(), "docs".to_string()],
649 })
650 ));
651 }
652
653 #[test]
654 fn parses_schedule_remove_selectors() {
655 let parsed = parse(["kbolt", "schedule", "remove", "s2"]);
656 assert!(matches!(
657 parsed.command,
658 Command::Schedule(schedule)
659 if schedule.command
660 == ScheduleCommand::Remove(ScheduleRemoveArgs {
661 id: Some("s2".to_string()),
662 all: false,
663 space: None,
664 collections: vec![],
665 })
666 ));
667
668 let parsed = parse([
669 "kbolt",
670 "schedule",
671 "remove",
672 "--space",
673 "work",
674 "--collection",
675 "api",
676 ]);
677 assert!(matches!(
678 parsed.command,
679 Command::Schedule(schedule)
680 if schedule.command
681 == ScheduleCommand::Remove(ScheduleRemoveArgs {
682 id: None,
683 all: false,
684 space: Some("work".to_string()),
685 collections: vec!["api".to_string()],
686 })
687 ));
688 }
689
690 #[test]
691 fn parses_eval_run_with_optional_manifest_path() {
692 let parsed = parse(["kbolt", "eval", "run"]);
693 assert!(matches!(
694 parsed.command,
695 Command::Eval(eval) if eval.command == EvalCommand::Run(EvalRunArgs { file: None })
696 ));
697
698 let parsed = parse(["kbolt", "eval", "run", "--file", "/tmp/scifact.toml"]);
699 assert!(matches!(
700 parsed.command,
701 Command::Eval(eval)
702 if eval.command
703 == EvalCommand::Run(EvalRunArgs {
704 file: Some(PathBuf::from("/tmp/scifact.toml"))
705 })
706 ));
707 }
708
709 #[test]
710 fn parses_eval_import_beir_with_required_paths() {
711 let parsed = parse([
712 "kbolt",
713 "eval",
714 "import",
715 "beir",
716 "--dataset",
717 "fiqa",
718 "--source",
719 "/tmp/fiqa-source",
720 "--output",
721 "/tmp/fiqa-bench",
722 ]);
723
724 let Command::Eval(eval) = parsed.command else {
725 panic!("expected eval command");
726 };
727 assert_eq!(
728 eval.command,
729 EvalCommand::Import(EvalImportArgs {
730 dataset: EvalImportCommand::Beir(EvalImportBeirArgs {
731 dataset: "fiqa".to_string(),
732 source: PathBuf::from("/tmp/fiqa-source"),
733 output: PathBuf::from("/tmp/fiqa-bench"),
734 collection: None,
735 })
736 })
737 );
738 }
739
740 #[test]
741 fn parses_eval_import_beir_with_collection_override() {
742 let parsed = parse([
743 "kbolt",
744 "eval",
745 "import",
746 "beir",
747 "--dataset",
748 "fiqa",
749 "--source",
750 "/tmp/fiqa-source",
751 "--output",
752 "/tmp/fiqa-bench",
753 "--collection",
754 "finance",
755 ]);
756
757 let Command::Eval(eval) = parsed.command else {
758 panic!("expected eval command");
759 };
760 assert_eq!(
761 eval.command,
762 EvalCommand::Import(EvalImportArgs {
763 dataset: EvalImportCommand::Beir(EvalImportBeirArgs {
764 dataset: "fiqa".to_string(),
765 source: PathBuf::from("/tmp/fiqa-source"),
766 output: PathBuf::from("/tmp/fiqa-bench"),
767 collection: Some("finance".to_string()),
768 })
769 })
770 );
771 }
772}