Refactor train-umap/apply-umap into a ragnar umap train/ragnar umap apply Thor subcommand group with new specs and updated string references.

Author: cpetersen

Gemfile

@@ -3,3 +3,4 @@
 source "https://rubygems.org"
 
 gemspec
+#gem "thor-interactive", path: "../thor-interactive"

lib/ragnar/cli.rb

@@ -1,3 +1,4 @@
+require_relative "cli/umap"
 require_relative "cli_visualization"
 require_relative "config"
 require "thor/interactive"
@@ -116,83 +117,8 @@ def index(path)
       end
     end
 
-    desc "train-umap", "Train UMAP model on existing embeddings"
-    option :db_path, type: :string, desc: "Path to Lance database (default from config)"
-    option :n_components, type: :numeric, default: 50, desc: "Number of dimensions for reduction"
-    option :n_neighbors, type: :numeric, default: 15, desc: "Number of neighbors for UMAP"
-    option :min_dist, type: :numeric, default: 0.1, desc: "Minimum distance for UMAP"
-    option :model_path, type: :string, desc: "Path to save UMAP model"
-    def train_umap
-      say "Training UMAP model on embeddings...", :green
-
-      config = Config.instance
-      # Use model_path from options if provided, otherwise use config models_dir
-      model_path = if options[:model_path]
-        options[:model_path]
-      else
-        File.join(config.models_dir, "umap_model.bin")
-      end
-      
-      processor = UmapProcessor.new(
-        db_path: options[:db_path] || config.database_path,
-        model_path: model_path
-      )
-
-      begin
-        stats = processor.train(
-          n_components: options[:n_components] || 50,
-          n_neighbors: options[:n_neighbors] || 15,
-          min_dist: options[:min_dist] || 0.1
-        )
-
-        say "\nUMAP training complete!", :green
-        say "Embeddings processed: #{stats[:embeddings_count]}"
-        say "Original dimensions: #{stats[:original_dims]}"
-        say "Reduced dimensions: #{stats[:reduced_dims]}"
-        say "Model saved to: #{processor.model_path}"
-      rescue => e
-        say "Error during UMAP training: #{e.message}", :red
-        exit 1
-      end
-    end
-
-    desc "apply-umap", "Apply trained UMAP model to reduce embedding dimensions"
-    option :db_path, type: :string, desc: "Path to Lance database (default from config)"
-    option :model_path, type: :string, desc: "Path to UMAP model"
-    option :batch_size, type: :numeric, default: 100, desc: "Batch size for processing"
-    def apply_umap
-      config = Config.instance
-      model_path = if options[:model_path]
-        options[:model_path]
-      else
-        File.join(config.models_dir, "umap_model.bin")
-      end
-
-      unless File.exist?(model_path)
-        say "Error: UMAP model not found at: #{model_path}", :red
-        say "Please run 'train-umap' first to create a model.", :yellow
-        exit 1
-      end
-
-      say "Applying UMAP model to embeddings...", :green
-
-      processor = UmapProcessor.new(
-        db_path: options[:db_path] || config.database_path,
-        model_path: model_path
-      )
-
-      begin
-        stats = processor.apply(batch_size: options[:batch_size] || 100)
-
-        say "\nUMAP application complete!", :green
-        say "Embeddings processed: #{stats[:processed]}"
-        say "Already processed: #{stats[:skipped]}"
-        say "Errors: #{stats[:errors]}" if stats[:errors] > 0
-      rescue => e
-        say "Error applying UMAP: #{e.message}", :red
-        exit 1
-      end
-    end
+    desc "umap SUBCOMMAND", "UMAP dimensionality reduction commands"
+    subcommand "umap", CLI::Umap
 
     desc "topics", "Extract and display topics from indexed documents"
     option :db_path, type: :string, desc: "Path to Lance database (default from config)"

lib/ragnar/cli/umap.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module Ragnar
+  class CLI < Thor
+    class Umap < Thor
+      desc "train", "Train UMAP model on existing embeddings"
+      option :db_path, type: :string, desc: "Path to Lance database (default from config)"
+      option :n_components, type: :numeric, default: 50, desc: "Number of dimensions for reduction"
+      option :n_neighbors, type: :numeric, default: 15, desc: "Number of neighbors for UMAP"
+      option :min_dist, type: :numeric, default: 0.1, desc: "Minimum distance for UMAP"
+      option :model_path, type: :string, desc: "Path to save UMAP model"
+      def train
+        say "Training UMAP model on embeddings...", :green
+
+        config = Config.instance
+        model_path = if options[:model_path]
+          options[:model_path]
+        else
+          File.join(config.models_dir, "umap_model.bin")
+        end
+
+        processor = UmapProcessor.new(
+          db_path: options[:db_path] || config.database_path,
+          model_path: model_path
+        )
+
+        begin
+          stats = processor.train(
+            n_components: options[:n_components] || 50,
+            n_neighbors: options[:n_neighbors] || 15,
+            min_dist: options[:min_dist] || 0.1
+          )
+
+          say "\nUMAP training complete!", :green
+          say "Embeddings processed: #{stats[:embeddings_count]}"
+          say "Original dimensions: #{stats[:original_dims]}"
+          say "Reduced dimensions: #{stats[:reduced_dims]}"
+          say "Model saved to: #{processor.model_path}"
+        rescue => e
+          say "Error during UMAP training: #{e.message}", :red
+          exit 1
+        end
+      end
+
+      desc "apply", "Apply trained UMAP model to reduce embedding dimensions"
+      option :db_path, type: :string, desc: "Path to Lance database (default from config)"
+      option :model_path, type: :string, desc: "Path to UMAP model"
+      option :batch_size, type: :numeric, default: 100, desc: "Batch size for processing"
+      def apply
+        config = Config.instance
+        model_path = if options[:model_path]
+          options[:model_path]
+        else
+          File.join(config.models_dir, "umap_model.bin")
+        end
+
+        unless File.exist?(model_path)
+          say "Error: UMAP model not found at: #{model_path}", :red
+          say "Please run 'ragnar umap train' first to create a model.", :yellow
+          exit 1
+        end
+
+        say "Applying UMAP model to embeddings...", :green
+
+        processor = UmapProcessor.new(
+          db_path: options[:db_path] || config.database_path,
+          model_path: model_path
+        )
+
+        begin
+          stats = processor.apply(batch_size: options[:batch_size] || 100)
+
+          say "\nUMAP application complete!", :green
+          say "Embeddings processed: #{stats[:processed]}"
+          say "Already processed: #{stats[:skipped]}"
+          say "Errors: #{stats[:errors]}" if stats[:errors] > 0
+        rescue => e
+          say "Error applying UMAP: #{e.message}", :red
+          exit 1
+        end
+      end
+    end
+  end
+end

lib/ragnar/indexer.rb

@@ -402,7 +402,7 @@ def suggest_umap_training
       puts "Training a UMAP model can significantly improve query performance."
       puts ""
       puts "To train UMAP, run:"
-      puts "  ragnar train-umap"
+      puts "  ragnar umap train"
       puts ""
       puts "This will:"
       puts "  • Reduce embedding dimensions from #{stats[:embedding_dims]} to 50-64"
@@ -427,7 +427,7 @@ def check_umap_retraining_needed?
         if current_doc_count > model_doc_count * 2
           puts "\n⚠️  UMAP model may be outdated"
           puts "Model was trained on #{model_doc_count} documents, now have #{current_doc_count}"
-          puts "Consider retraining with: ragnar train-umap"
+          puts "Consider retraining with: ragnar umap train"
           return true
         end
       rescue => e

lib/ragnar/umap_processor.rb

@@ -160,7 +160,7 @@ def train(n_components: Ragnar::DEFAULT_REDUCED_DIMENSIONS, n_neighbors: 15, min
           error_msg += "  • There are duplicate or corrupted embeddings\n\n"
           error_msg += "Suggested solutions:\n"
           error_msg += "  1. Try with more conservative parameters:\n"
-          error_msg += "     ragnar train-umap --n-components 10 --n-neighbors 5\n\n"
+          error_msg += "     ragnar umap train --n-components 10 --n-neighbors 5\n\n"
           error_msg += "  2. Re-index your documents to regenerate embeddings:\n"
           error_msg += "     ragnar index <path> --force\n\n"
           error_msg += "  3. Check your embedding model configuration\n\n"
@@ -173,7 +173,7 @@ def train(n_components: Ragnar::DEFAULT_REDUCED_DIMENSIONS, n_neighbors: 15, min
           error_msg += "Error: #{e.message}\n\n"
           error_msg += "This may be due to incompatible parameters or data issues.\n"
           error_msg += "Try using more conservative parameters:\n"
-          error_msg += "  ragnar train-umap --n-components 10 --n-neighbors 5\n"
+          error_msg += "  ragnar umap train --n-components 10 --n-neighbors 5\n"
         end
 
         raise RuntimeError, error_msg

lib/ragnar/umap_transform_service.rb

@@ -160,7 +160,7 @@ def load_model!
       return if @umap_model
 
       unless File.exist?(@model_path)
-        raise "UMAP model not found at #{@model_path}. Please train a model first using 'ragnar train-umap'."
+        raise "UMAP model not found at #{@model_path}. Please train a model first using 'ragnar umap train'."
       end
 
       @umap_model = ClusterKit::Dimensionality::UMAP.load_model(@model_path)

ragnar.gemspec

@@ -36,7 +36,7 @@ Gem::Specification.new do |spec|
   spec.add_dependency "baran", "~> 0.2"
   spec.add_dependency "parsekit", "~> 0.1", ">= 0.1.2"
   spec.add_dependency "tty-progressbar", "~> 0.18"
-  spec.add_dependency "thor-interactive", "~> 0.1.0.pre.4"
+  spec.add_dependency "thor-interactive", "~> 0.1.0.pre.5"
 
   # Development dependencies
   spec.add_development_dependency "rake", "~> 13.0"

spec/unit/cli_umap_spec.rb

@@ -0,0 +1,132 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe Ragnar::CLI::Umap do
+  let(:config) do
+    instance_double(
+      Ragnar::Config,
+      database_path: "./ragnar_database",
+      models_dir: "/tmp/ragnar_test_models"
+    )
+  end
+  let(:model_path) { File.join(config.models_dir, "umap_model.bin") }
+
+  before do
+    allow(Ragnar::Config).to receive(:instance).and_return(config)
+  end
+
+  describe "train" do
+    let(:processor) { instance_double(Ragnar::UmapProcessor, model_path: model_path) }
+
+    before do
+      allow(Ragnar::UmapProcessor).to receive(:new).and_return(processor)
+    end
+
+    it "trains UMAP model with default options" do
+      stats = { embeddings_count: 100, original_dims: 768, reduced_dims: 50 }
+      allow(processor).to receive(:train).and_return(stats)
+
+      expect {
+        Ragnar::CLI.start(["umap", "train"])
+      }.not_to raise_error
+    end
+
+    it "passes custom options to the processor" do
+      stats = { embeddings_count: 100, original_dims: 768, reduced_dims: 10 }
+      expect(Ragnar::UmapProcessor).to receive(:new).with(
+        db_path: "./ragnar_database",
+        model_path: model_path
+      ).and_return(processor)
+
+      expect(processor).to receive(:train).with(
+        n_components: 10,
+        n_neighbors: 5,
+        min_dist: 0.1
+      ).and_return(stats)
+
+      Ragnar::CLI.start(["umap", "train", "--n-components", "10", "--n-neighbors", "5"])
+    end
+
+    it "uses custom model path when provided" do
+      custom_path = "/custom/path/model.bin"
+      stats = { embeddings_count: 100, original_dims: 768, reduced_dims: 50 }
+
+      expect(Ragnar::UmapProcessor).to receive(:new).with(
+        db_path: "./ragnar_database",
+        model_path: custom_path
+      ).and_return(processor)
+
+      allow(processor).to receive(:train).and_return(stats)
+
+      Ragnar::CLI.start(["umap", "train", "--model-path", custom_path])
+    end
+
+    it "uses custom db path when provided" do
+      custom_db = "/custom/db/path"
+      stats = { embeddings_count: 100, original_dims: 768, reduced_dims: 50 }
+
+      expect(Ragnar::UmapProcessor).to receive(:new).with(
+        db_path: custom_db,
+        model_path: model_path
+      ).and_return(processor)
+
+      allow(processor).to receive(:train).and_return(stats)
+
+      Ragnar::CLI.start(["umap", "train", "--db-path", custom_db])
+    end
+
+    it "handles training errors gracefully" do
+      allow(processor).to receive(:train).and_raise("No embeddings found")
+
+      expect {
+        Ragnar::CLI.start(["umap", "train"])
+      }.to raise_error(SystemExit)
+    end
+  end
+
+  describe "apply" do
+    let(:processor) { instance_double(Ragnar::UmapProcessor) }
+
+    before do
+      allow(Ragnar::UmapProcessor).to receive(:new).and_return(processor)
+      allow(File).to receive(:exist?).and_call_original
+    end
+
+    it "applies UMAP model when model exists" do
+      allow(File).to receive(:exist?).with(model_path).and_return(true)
+      stats = { processed: 50, skipped: 0, errors: 0 }
+      allow(processor).to receive(:apply).and_return(stats)
+
+      expect {
+        Ragnar::CLI.start(["umap", "apply"])
+      }.not_to raise_error
+    end
+
+    it "passes batch size option to the processor" do
+      allow(File).to receive(:exist?).with(model_path).and_return(true)
+      stats = { processed: 50, skipped: 0, errors: 0 }
+
+      expect(processor).to receive(:apply).with(batch_size: 200).and_return(stats)
+
+      Ragnar::CLI.start(["umap", "apply", "--batch-size", "200"])
+    end
+
+    it "exits with error when model does not exist" do
+      allow(File).to receive(:exist?).with(model_path).and_return(false)
+
+      expect {
+        Ragnar::CLI.start(["umap", "apply"])
+      }.to raise_error(SystemExit)
+    end
+
+    it "handles apply errors gracefully" do
+      allow(File).to receive(:exist?).with(model_path).and_return(true)
+      allow(processor).to receive(:apply).and_raise("Database error")
+
+      expect {
+        Ragnar::CLI.start(["umap", "apply"])
+      }.to raise_error(SystemExit)
+    end
+  end
+end

spec/unit/umap_transform_service_spec.rb

@@ -336,7 +336,7 @@
         it "raises error with helpful message" do
           expect {
             service.send(:load_model!)
-          }.to raise_error(/UMAP model not found.*train-umap/)
+          }.to raise_error(/UMAP model not found.*umap train/)
         end
       end
     end