ragnar-cli 0.1.0.pre.1

data/lib/ragnar/database.rb

@@ -0,0 +1,267 @@
+module Ragnar
+  class Database
+    attr_reader :db_path, :table_name
+    
+    def initialize(db_path, table_name: "documents")
+      @db_path = db_path
+      @table_name = table_name
+      ensure_database_exists
+    end
+    
+    def add_documents(documents)
+      return if documents.empty?
+      
+      # Convert documents to Lance-compatible format
+      data = documents.map do |doc|
+        {
+          id: doc[:id],
+          chunk_text: doc[:chunk_text],
+          file_path: doc[:file_path],
+          chunk_index: doc[:chunk_index],
+          embedding: doc[:embedding],
+          metadata: doc[:metadata].to_json
+        }
+      end
+      
+      # Define schema for the table with vector type
+      embedding_size = documents.first[:embedding].size
+      schema = {
+        id: :string,
+        chunk_text: :string,
+        file_path: :string,
+        chunk_index: :int64,
+        embedding: { type: "vector", dimension: embedding_size },
+        metadata: :string
+      }
+      
+      # Use the new open_or_create method from Lancelot
+      # This automatically handles both creating new and opening existing datasets
+      dataset = Lancelot::Dataset.open_or_create(@db_path, schema: schema)
+      dataset.add_documents(data)
+    end
+    
+    def get_embeddings(limit: nil, offset: 0)
+      return [] unless dataset_exists?
+      
+      dataset = Lancelot::Dataset.open(@db_path)
+      
+      # Get all documents or a subset
+      docs = if limit && offset > 0
+        # Get limit + offset items, then drop offset
+        dataset.first(limit + offset).drop(offset)
+      elsif limit
+        dataset.first(limit)
+      else
+        dataset.to_a.drop(offset)
+      end
+      
+      docs.map do |doc|
+        {
+          id: doc[:id],
+          embedding: doc[:embedding],
+          reduced_embedding: doc[:reduced_embedding]
+        }
+      end
+    end
+    
+    def update_reduced_embeddings(updates)
+      return if updates.empty?
+      
+      dataset = Lancelot::Dataset.open(@db_path)
+      
+      # Get all existing documents and safely extract their data
+      all_docs = dataset.to_a.map do |doc|
+        # Safely extract fields we know about
+        {
+          id: doc[:id],
+          content: doc[:content],
+          chunk_text: doc[:chunk_text],
+          file_path: doc[:file_path],
+          chunk_index: doc[:chunk_index],
+          embedding: doc[:embedding],
+          metadata: doc[:metadata],
+          reduced_embedding: doc[:reduced_embedding]
+        }
+      end
+      
+      # Create a map for quick lookup
+      update_map = updates.each_with_object({}) do |update, map|
+        map[update[:id]] = update[:reduced_embedding]
+      end
+      
+      # Update documents with reduced embeddings
+      updated_docs = all_docs.map do |doc|
+        if update_map[doc[:id]]
+          doc.merge(reduced_embedding: update_map[doc[:id]])
+        else
+          doc
+        end
+      end
+      
+      # Need to recreate the dataset with updated data
+      # First, backup the schema including the new reduced_embedding field
+      embedding_size = all_docs.first[:embedding].size
+      reduced_size = updates.first[:reduced_embedding].size
+      
+      schema = {
+        id: :string,
+        chunk_text: :string,
+        file_path: :string,
+        chunk_index: :int64,
+        embedding: { type: "vector", dimension: embedding_size },
+        reduced_embedding: { type: "vector", dimension: reduced_size },
+        metadata: :string
+      }
+      
+      # Remove old dataset and create new one with updated data
+      FileUtils.rm_rf(@db_path)
+      # Use open_or_create which will create since we just deleted the path
+      dataset = Lancelot::Dataset.open_or_create(@db_path, schema: schema)
+      dataset.add_documents(updated_docs)
+    end
+    
+    def search_similar(embedding, k: 10, use_reduced: false)
+      return [] unless dataset_exists?
+      
+      dataset = Lancelot::Dataset.open(@db_path)
+      
+      embedding_field = use_reduced ? :reduced_embedding : :embedding
+      
+      # Perform vector search
+      results = dataset.vector_search(
+        embedding.to_a, 
+        column: embedding_field,
+        limit: k
+      )
+      
+      results.map do |row|
+        {
+          id: row[:id],
+          chunk_text: row[:chunk_text],
+          file_path: row[:file_path],
+          chunk_index: row[:chunk_index],
+          distance: row[:_distance],
+          metadata: JSON.parse(row[:metadata] || "{}")
+        }
+      end
+    end
+    
+    def count
+      return 0 unless dataset_exists?
+      
+      dataset = Lancelot::Dataset.open(@db_path)
+      dataset.to_a.size
+    end
+    
+    def get_stats
+      unless dataset_exists?
+        return {
+          document_count: 0,
+          total_documents: 0,
+          unique_files: 0,
+          total_chunks: 0,
+          with_embeddings: 0,
+          with_reduced_embeddings: 0,
+          total_size_mb: 0.0
+        }
+      end
+      
+      dataset = Lancelot::Dataset.open(@db_path)
+      
+      # Get all documents
+      all_docs = dataset.to_a
+      
+      stats = {
+        document_count: all_docs.size,  # Add for compatibility with specs
+        total_documents: all_docs.size,
+        total_chunks: all_docs.size,
+        unique_files: all_docs.map { |d| d[:file_path] }.uniq.size,
+        with_embeddings: 0,
+        with_reduced_embeddings: 0,
+        avg_chunk_size: 0,
+        total_size_mb: 0,  # Add for CLI stats command
+        embedding_dims: nil,
+        reduced_dims: nil
+      }
+      
+      chunk_sizes = []
+      total_bytes = 0
+      
+      all_docs.each do |doc|
+        if doc[:embedding] && !doc[:embedding].empty?
+          stats[:with_embeddings] += 1
+          stats[:embedding_dims] ||= doc[:embedding].size
+        end
+        
+        if doc[:reduced_embedding] && !doc[:reduced_embedding].empty?
+          stats[:with_reduced_embeddings] += 1
+          stats[:reduced_dims] ||= doc[:reduced_embedding].size
+        end
+        
+        if doc[:chunk_text]
+          chunk_size = doc[:chunk_text].size
+          chunk_sizes << chunk_size
+          total_bytes += chunk_size
+        end
+      end
+      
+      stats[:avg_chunk_size] = (chunk_sizes.sum.to_f / chunk_sizes.size).round if chunk_sizes.any?
+      stats[:total_size_mb] = (total_bytes / 1024.0 / 1024.0).round(2)
+      
+      stats
+    end
+    
+    def get_all_documents_with_embeddings(limit: nil)
+      return [] unless dataset_exists?
+      
+      dataset = Lancelot::Dataset.open(@db_path)
+      all_docs = limit ? dataset.first(limit) : dataset.to_a
+      
+      all_docs.select { |doc| doc[:embedding] && !doc[:embedding].empty? }
+    end
+    
+    def full_text_search(query, limit: 10)
+      return [] unless dataset_exists?
+      
+      dataset = Lancelot::Dataset.open(@db_path)
+      
+      # Use Lancelot's full-text search
+      results = dataset.full_text_search(
+        query,
+        columns: [:chunk_text],
+        limit: limit
+      )
+      
+      results.map do |row|
+        {
+          id: row[:id],
+          chunk_text: row[:chunk_text],
+          file_path: row[:file_path],
+          chunk_index: row[:chunk_index],
+          metadata: JSON.parse(row[:metadata] || "{}")
+        }
+      end
+    end
+    
+    def dataset_exists?
+      return false unless File.exist?(@db_path)
+      
+      begin
+        Lancelot::Dataset.open(@db_path)
+        true
+      rescue
+        false
+      end
+    end
+    
+    private
+    
+    def ensure_database_exists
+      # Don't create directory - Lance will handle this
+    end
+    
+    def table_exists?
+      dataset_exists?
+    end
+  end
+end

data/lib/ragnar/embedder.rb

@@ -0,0 +1,137 @@
+module Ragnar
+  class Embedder
+    attr_reader :model, :model_name
+    
+    def initialize(model_name: Ragnar::DEFAULT_EMBEDDING_MODEL)
+      @model_name = model_name
+      @model = load_model(model_name)
+    end
+    
+    def embed_text(text)
+      return nil if text.nil? || text.empty? || (text.respond_to?(:strip) && text.strip.empty?)
+      
+      # Use Candle to generate embeddings
+      # The embedding method returns a tensor, we need to convert to array
+      embedding = @model.embedding(text)
+      
+      # Convert tensor to array - Candle tensors need double to_a
+      # First to_a gives [tensor], second to_a on the tensor gives the float array
+      if embedding.respond_to?(:to_a)
+        result = embedding.to_a
+        if result.is_a?(Array) && result.first.respond_to?(:to_a)
+          result.first.to_a
+        else
+          result
+        end
+      else
+        embedding
+      end
+    rescue => e
+      puts "Error generating embedding: #{e.message}"
+      nil
+    end
+    
+    def embed_batch(texts, show_progress: true)
+      embeddings = []
+      
+      if show_progress
+        progressbar = TTY::ProgressBar.new(
+          "Generating embeddings [:bar] :percent :current/:total",
+          total: texts.size,
+          bar_format: :block,
+          width: 30
+        )
+      end
+      
+      texts.each do |text|
+        embedding = embed_text(text)
+        embeddings << embedding
+        progressbar.advance if show_progress
+      end
+      
+      embeddings
+    end
+    
+    def embed_chunks(chunks, show_progress: true)
+      texts = chunks.map do |chunk|
+        if chunk.is_a?(Hash)
+          chunk[:text] || chunk["text"]
+        else
+          chunk.to_s
+        end
+      end
+      
+      embed_batch(texts, show_progress: show_progress)
+    end
+    
+    private
+    
+    def load_model(model_name)
+      # Initialize Candle embedding model using the new standardized from_pretrained method
+      begin
+        # Try to load the model using from_pretrained
+        Candle::EmbeddingModel.from_pretrained(model_name)
+      rescue => e
+        puts "Warning: Could not load model #{model_name}, falling back to default"
+        puts "Error: #{e.message}"
+        
+        # Fall back to default model
+        begin
+          Candle::EmbeddingModel.from_pretrained("jinaai/jina-embeddings-v2-base-en")
+        rescue => fallback_error
+          puts "Error loading fallback model: #{fallback_error.message}"
+          # Last resort: try the old initialization method for backwards compatibility
+          Candle::EmbeddingModel.new
+        end
+      end
+    end
+    
+    def self.available_models
+      # List of commonly used embedding models
+      # This could be expanded or made dynamic
+      [
+        "BAAI/bge-small-en-v1.5",
+        "BAAI/bge-base-en-v1.5",
+        "BAAI/bge-large-en-v1.5",
+        "sentence-transformers/all-MiniLM-L6-v2",
+        "sentence-transformers/all-mpnet-base-v2",
+        "thenlper/gte-small",
+        "thenlper/gte-base",
+        "thenlper/gte-large"
+      ]
+    end
+    
+    def self.model_info(model_name)
+      # Provide information about embedding models
+      info = {
+        "BAAI/bge-small-en-v1.5" => {
+          dimensions: 384,
+          max_tokens: 512,
+          description: "Small, fast, good quality embeddings"
+        },
+        "BAAI/bge-base-en-v1.5" => {
+          dimensions: 768,
+          max_tokens: 512,
+          description: "Balanced size and quality"
+        },
+        "BAAI/bge-large-en-v1.5" => {
+          dimensions: 1024,
+          max_tokens: 512,
+          description: "Large, highest quality embeddings"
+        },
+        "sentence-transformers/all-MiniLM-L6-v2" => {
+          dimensions: 384,
+          max_tokens: 256,
+          description: "Fast, lightweight model"
+        },
+        "sentence-transformers/all-mpnet-base-v2" => {
+          dimensions: 768,
+          max_tokens: 384,
+          description: "High quality general purpose embeddings"
+        }
+      }
+      
+      info[model_name] || { description: "Model information not available" }
+    end
+  end
+end

data/lib/ragnar/indexer.rb

@@ -0,0 +1,234 @@
+require 'parsekit'
+
+module Ragnar
+  class Indexer
+    attr_reader :database, :chunker, :embedder
+
+    def initialize(db_path: Ragnar::DEFAULT_DB_PATH,
+                   chunk_size: Ragnar::DEFAULT_CHUNK_SIZE,
+                   chunk_overlap: Ragnar::DEFAULT_CHUNK_OVERLAP,
+                   embedding_model: Ragnar::DEFAULT_EMBEDDING_MODEL,
+                   show_progress: true)
+      @database = Database.new(db_path)
+      @chunker = Chunker.new(chunk_size: chunk_size, chunk_overlap: chunk_overlap)
+      @embedder = Embedder.new(model_name: embedding_model)
+      @show_progress = show_progress
+    end
+
+    def index_path(path)
+      stats = {
+        files_processed: 0,
+        chunks_created: 0,
+        errors: 0
+      }
+
+      files = collect_files(path)
+
+      if files.empty?
+        puts "No text files found at path: #{path}"
+        return stats
+      end
+
+      puts "Found #{files.size} file(s) to process" if @show_progress
+
+      file_progress = if @show_progress
+        TTY::ProgressBar.new(
+          "Processing [:bar] :percent :current/:total - :filename",
+          total: files.size,
+          bar_format: :block,
+          width: 30,
+          clear: true
+        )
+      else
+        nil
+      end
+
+      files.each do |file_path|
+        begin
+          if file_progress
+            # Update the progress bar with current filename
+            filename = File.basename(file_path)
+            filename = filename[0..27] + "..." if filename.length > 30
+            file_progress.advance(0, filename: filename)
+          end
+
+          process_file(file_path, stats, file_progress)
+          stats[:files_processed] += 1
+        rescue => e
+          if file_progress
+            file_progress.log "Error: #{File.basename(file_path)} - #{e.message}"
+          else
+            puts "Error processing #{File.basename(file_path)}: #{e.message}" if @show_progress
+          end
+          stats[:errors] += 1
+        ensure
+          file_progress&.advance
+        end
+      end
+
+      stats
+    end
+
+    def index_text(text, metadata = {})
+      chunks = @chunker.chunk_text(text, metadata)
+      process_chunks(chunks, metadata[:file_path] || "inline_text")
+    end
+    
+    # Convenience methods for compatibility
+    def index_files(files)
+      stats = {
+        files_processed: 0,
+        chunks_created: 0,
+        errors: 0
+      }
+      
+      files.each do |file|
+        next unless File.exist?(file)
+        process_file(file, stats)
+        stats[:files_processed] += 1
+      end
+      
+      stats
+    end
+    
+    def index_directory(dir_path)
+      index_path(dir_path)
+    end
+
+    private
+
+    def collect_files(path)
+      if File.file?(path)
+        [path]
+      elsif File.directory?(path)
+        # Now we support many more file types through parser-core
+        pattern = "*.{txt,md,markdown,text,pdf,docx,doc,xlsx,xls,pptx,ppt,csv,json,xml,html,htm,rb,py,js,rs,go,java,cpp,c,h}"
+        Dir.glob(File.join(path, "**", pattern))
+      else
+        []
+      end
+    end
+
+    def process_file(file_path, stats, progress_bar = nil)
+      # Extract text using parser-core
+      begin
+        text = extract_text_from_file(file_path)
+
+        if text.nil? || text.strip.empty?
+          progress_bar.log("  Skipped: #{File.basename(file_path)} (empty or unsupported)") if progress_bar
+          return
+        end
+
+        # Create metadata
+        metadata = {
+          file_path: file_path,
+          file_name: File.basename(file_path),
+          file_type: File.extname(file_path).downcase[1..-1] || 'unknown'
+        }
+
+        # Chunk the extracted text
+        chunks = @chunker.chunk_text(text, metadata)
+
+        if chunks.empty?
+          progress_bar.log("  Skipped: #{File.basename(file_path)} (text too short)") if progress_bar
+          return
+        end
+
+        # Process chunks and create documents
+        chunk_count = process_chunks(chunks, file_path, progress_bar)
+        stats[:chunks_created] += chunk_count
+      rescue => e
+        if progress_bar
+          progress_bar.log("  Error processing file: #{e.message}")
+          progress_bar.log("  Backtrace: #{e.backtrace.first}")
+        end
+        raise e
+      end
+    end
+
+    def process_chunks(chunks, file_path, progress_bar = nil)
+      return 0 if chunks.empty?
+
+      # Extract texts for embedding
+      texts = chunks.map { |c| c[:text] }
+
+      # Generate embeddings (silently)
+      embeddings = @embedder.embed_batch(texts, show_progress: false)
+
+      # Prepare documents for database
+      documents = []
+      chunks.each_with_index do |chunk, idx|
+        embedding = embeddings[idx]
+        next unless embedding  # Skip if embedding failed
+
+        doc = {
+          id: SecureRandom.uuid,
+          chunk_text: chunk[:text],
+          file_path: file_path,
+          chunk_index: chunk[:index],
+          embedding: embedding,
+          metadata: chunk[:metadata] || {}
+        }
+
+        # Note: No need to add reduced_embedding field anymore!
+        # Lancelot now supports optional fields after our fix
+
+        documents << doc
+      end
+
+      # Store in database
+      if documents.any?
+        @database.add_documents(documents)
+        # Successfully stored chunks (silent to preserve progress bar)
+      end
+
+      documents.size
+    end
+
+    def extract_text_from_file(file_path)
+      # Use parser-core to extract text from various file formats
+      begin
+        ParseKit.parse_file(file_path)
+      rescue => e
+        # If parser-core fails, try reading as plain text for known text formats
+        ext = File.extname(file_path).downcase
+        if %w[.txt .md .markdown .text .log .rb .py .js .rs .go .java .cpp .c .h].include?(ext)
+          File.read(file_path, encoding: 'UTF-8')
+        else
+          raise e
+        end
+      end
+    end
+
+    def self.supported_extensions
+      # Extended list of supported formats through parser-core
+      %w[.txt .md .markdown .text .log .csv .json .xml .html .htm
+         .pdf .docx .doc .xlsx .xls .pptx .ppt
+         .rb .py .js .rs .go .java .cpp .c .h]
+    end
+
+    def self.is_text_file?(file_path)
+      # Check by extension
+      ext = File.extname(file_path).downcase
+      return true if supported_extensions.include?(ext)
+
+      # Check if file appears to be text
+      begin
+        # Read first 8KB to check if it's text
+        sample = File.read(file_path, 8192, mode: 'rb')
+        return false if sample.nil?
+
+        # Check for binary content
+        null_count = sample.count("\x00")
+        return false if null_count > 0
+
+        # Check if mostly printable ASCII
+        printable = sample.count("\t\n\r\x20-\x7E")
+        ratio = printable.to_f / sample.size
+        ratio > 0.9
+      rescue
+        false
+      end
+    end
+  end
+end

data/lib/ragnar/llm_manager.rb

@@ -0,0 +1,43 @@
+module Ragnar
+  # Singleton manager for LLM instances to avoid reloading models
+  class LLMManager
+    include Singleton
+    
+    def initialize
+      @llms = {}
+      @mutex = Mutex.new
+    end
+    
+    # Get or create an LLM instance
+    # @param model_id [String] The model identifier
+    # @param gguf_file [String, nil] Optional GGUF file for quantized models
+    # @return [Candle::LLM] The LLM instance
+    def get_llm(model_id: "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", 
+                gguf_file: "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf")
+      cache_key = "#{model_id}:#{gguf_file}"
+      
+      @mutex.synchronize do
+        @llms[cache_key] ||= begin
+          puts "Loading LLM: #{model_id}..." unless @llms.key?(cache_key)
+          if gguf_file
+            Candle::LLM.from_pretrained(model_id, gguf_file: gguf_file)
+          else
+            Candle::LLM.from_pretrained(model_id)
+          end
+        end
+      end
+    end
+    
+    # Clear all cached models (useful for memory management)
+    def clear_cache
+      @mutex.synchronize do
+        @llms.clear
+      end
+    end
+    
+    # Get the default LLM for the application
+    def default_llm
+      get_llm
+    end
+  end
+end