leann 0.1.0

data/lib/leann/builder.rb

@@ -0,0 +1,317 @@
+# frozen_string_literal: true
+
+require "json"
+require "fileutils"
+require "securerandom"
+require "time"
+
+module Leann
+  # Builds a new Leann index
+  #
+  # @example DSL style
+  #   Leann.build("my_index") do
+  #     add "First document"
+  #     add "Second document", source: "manual"
+  #     add_file "README.md"
+  #   end
+  #
+  # @example Programmatic style
+  #   builder = Leann::Builder.new("my_index")
+  #   builder.add("First document")
+  #   builder.add("Second document")
+  #   builder.save
+  #
+  class Builder
+    # @return [String] Index name
+    attr_reader :name
+
+    # @return [String] Index path
+    attr_reader :path
+
+    # @return [Array<Hash>] Documents to be indexed
+    attr_reader :documents
+
+    # @param name [String] Index name
+    # @param embedding [Symbol] Embedding provider (:ruby_llm, :openai, :ollama, :fastembed)
+    # @param model [String, nil] Embedding model name
+    # @param path [String, nil] Custom path for index
+    # @param force [Boolean] Overwrite existing index
+    def initialize(name, embedding: nil, model: nil, path: nil, force: false, **_options)
+      @name = name
+      @path = resolve_path(name, path)
+      @embedding_provider = embedding || Leann.configuration.embedding_provider
+      @embedding_model = model || Leann.configuration.embedding_model_for(@embedding_provider)
+      @force = force
+      @documents = []
+
+      check_existing_index unless force
+    end
+
+    # Add a text document
+    #
+    # @param text [String] Document text
+    # @param metadata [Hash] Additional metadata (passed as keyword arguments)
+    # @return [self]
+    #
+    # @example
+    #   builder.add("Hello world")
+    #   builder.add("Document with metadata", source: "file.txt", chapter: 1)
+    def add(text, **metadata)
+      raise ArgumentError, "Text cannot be nil" if text.nil?
+      raise ArgumentError, "Text cannot be empty" if text.to_s.strip.empty?
+
+      doc = {
+        id: metadata.delete(:id) || generate_id,
+        text: text.to_s.strip,
+        metadata: metadata
+      }
+
+      @documents << doc
+      self
+    end
+
+    # Add document (alias for add)
+    alias << add
+
+    # Add content from a file
+    #
+    # @param file_path [String] Path to file
+    # @param metadata [Hash] Additional metadata
+    # @return [self]
+    #
+    # @example
+    #   builder.add_file("README.md")
+    #   builder.add_file("docs/guide.txt", category: "documentation")
+    def add_file(file_path, **metadata)
+      raise ArgumentError, "File not found: #{file_path}" unless File.exist?(file_path)
+
+      content = File.read(file_path)
+      file_metadata = {
+        source: file_path,
+        filename: File.basename(file_path),
+        extension: File.extname(file_path)
+      }.merge(metadata)
+
+      add(content, **file_metadata)
+    end
+
+    # Add all files from a directory
+    #
+    # @param directory [String] Directory path
+    # @param pattern [String] Glob pattern (default: "**/*")
+    # @param extensions [Array<String>, nil] Filter by extensions (e.g., [".md", ".txt"])
+    # @param metadata [Hash] Additional metadata for all files
+    # @return [self]
+    #
+    # @example
+    #   builder.add_directory("docs/")
+    #   builder.add_directory("src/", extensions: [".rb", ".py"])
+    def add_directory(directory, pattern: "**/*", extensions: nil, **metadata)
+      raise ArgumentError, "Directory not found: #{directory}" unless Dir.exist?(directory)
+
+      full_pattern = File.join(directory, pattern)
+      Dir.glob(full_pattern).each do |file_path|
+        next unless File.file?(file_path)
+        next if extensions && !extensions.include?(File.extname(file_path))
+
+        add_file(file_path, **metadata)
+      end
+
+      self
+    end
+
+    # Add multiple documents at once
+    #
+    # @param docs [Array<String>, Array<Hash>] Documents to add
+    # @return [self]
+    #
+    # @example
+    #   builder.add_all(["Doc 1", "Doc 2", "Doc 3"])
+    #   builder.add_all([
+    #     { text: "Doc 1", source: "a" },
+    #     { text: "Doc 2", source: "b" }
+    #   ])
+    def add_all(docs)
+      docs.each do |doc|
+        case doc
+        when String
+          add(doc)
+        when Hash
+          text = doc.delete(:text) || doc.delete("text")
+          add(text, **doc.transform_keys(&:to_sym))
+        else
+          raise ArgumentError, "Invalid document type: #{doc.class}"
+        end
+      end
+
+      self
+    end
+
+    # Get number of documents added
+    # @return [Integer]
+    def count
+      @documents.size
+    end
+    alias size count
+
+    # Check if any documents have been added
+    # @return [Boolean]
+    def empty?
+      @documents.empty?
+    end
+
+    # Build and save the index
+    # @return [Index] The built index
+    def save
+      raise EmptyIndexError if empty?
+
+      puts "Building index '#{name}' with #{count} documents..."
+
+      # Create directory if needed
+      FileUtils.mkdir_p(File.dirname(path))
+
+      # Delete existing if force mode
+      Index.delete(path) if @force && Index.exists?(path)
+
+      # Compute embeddings
+      embeddings = compute_embeddings
+
+      # Save passages
+      save_passages
+
+      # Build and save graph
+      save_graph(embeddings)
+
+      # Save metadata
+      save_metadata(embeddings)
+
+      puts "Index '#{name}' created successfully!"
+
+      Index.open(path)
+    end
+    alias build save
+
+    private
+
+    def resolve_path(name, custom_path)
+      if custom_path
+        custom_path.end_with?(Index::INDEX_EXTENSION) ? custom_path : "#{custom_path}#{Index::INDEX_EXTENSION}"
+      else
+        "#{name}#{Index::INDEX_EXTENSION}"
+      end
+    end
+
+    def check_existing_index
+      raise IndexExistsError, name if Index.exists?(path)
+    end
+
+    def generate_id
+      SecureRandom.uuid
+    end
+
+    def compute_embeddings
+      texts = @documents.map { |d| d[:text] }
+      embedding_provider.compute(texts)
+    end
+
+    def embedding_provider
+      @_embedding_provider ||= load_embedding_provider
+    end
+
+    def load_embedding_provider
+      require_relative "embedding/base"
+
+      case @embedding_provider
+      when :ruby_llm
+        require_relative "embedding/ruby_llm"
+        Embedding::RubyLLM.new(model: @embedding_model)
+      when :openai
+        require_relative "embedding/openai"
+        Embedding::OpenAI.new(model: @embedding_model)
+      when :ollama
+        require_relative "embedding/ollama"
+        Embedding::Ollama.new(model: @embedding_model)
+      when :fastembed
+        require_relative "embedding/fastembed"
+        Embedding::FastEmbed.new(model: @embedding_model)
+      else
+        raise ConfigurationError, "Unknown embedding provider: #{@embedding_provider}"
+      end
+    end
+
+    def save_passages
+      passages_file = "#{path}#{Index::PASSAGES_SUFFIX}"
+      offsets_file = "#{path}#{Index::OFFSETS_SUFFIX}"
+
+      offsets = {}
+
+      File.open(passages_file, "w") do |f|
+        @documents.each do |doc|
+          offsets[doc[:id]] = f.pos
+          f.puts(JSON.generate(doc))
+        end
+      end
+
+      File.write(offsets_file, JSON.generate(offsets))
+    end
+
+    def save_graph(embeddings)
+      ids = @documents.map { |d| d[:id] }
+
+      require_relative "backend/leann_graph"
+
+      graph = Backend::LeannGraph.new(
+        dimensions: embeddings.first.size,
+        m: Leann.configuration.hnsw_m,
+        ef_construction: Leann.configuration.hnsw_ef_construction
+      )
+
+      graph.build(ids, embeddings)
+      graph.save(path)
+
+      report_storage_savings(embeddings)
+    end
+
+    def report_storage_savings(embeddings)
+      embedding_size = embeddings.first.size * 4  # float32
+      total_embedding_bytes = embeddings.length * embedding_size
+
+      graph_file = "#{path}.graph.bin"
+      actual_size = File.exist?(graph_file) ? File.size(graph_file) : 0
+
+      savings = ((total_embedding_bytes - actual_size).to_f / total_embedding_bytes * 100).round(1)
+      puts "Storage savings: #{savings}% (#{format_bytes(total_embedding_bytes)} → #{format_bytes(actual_size)})"
+    end
+
+    def format_bytes(bytes)
+      if bytes < 1024
+        "#{bytes} B"
+      elsif bytes < 1024 * 1024
+        "#{(bytes / 1024.0).round(1)} KB"
+      else
+        "#{(bytes / (1024.0 * 1024)).round(2)} MB"
+      end
+    end
+
+    def save_metadata(embeddings)
+      meta_file = "#{path}#{Index::META_SUFFIX}"
+
+      metadata = {
+        version: "1.0",
+        name: name,
+        backend: "leann",
+        embedding_provider: @embedding_provider.to_s,
+        embedding_model: @embedding_model,
+        dimensions: embeddings.first&.size || 0,
+        document_count: @documents.size,
+        created_at: Time.now.utc.iso8601,
+        config: {
+          hnsw_m: Leann.configuration.hnsw_m,
+          hnsw_ef_construction: Leann.configuration.hnsw_ef_construction
+        }
+      }
+
+      File.write(meta_file, JSON.pretty_generate(metadata))
+    end
+  end
+end

data/lib/leann/configuration.rb

@@ -0,0 +1,148 @@
+# frozen_string_literal: true
+
+module Leann
+  # Global configuration for Leann
+  #
+  # @example With RubyLLM (recommended)
+  #   # If RubyLLM is present, LEANN uses it automatically
+  #   # Just configure RubyLLM as usual:
+  #   RubyLLM.configure do |config|
+  #     config.openai_api_key = ENV["OPENAI_API_KEY"]
+  #   end
+  #
+  # @example Manual configuration
+  #   Leann.configure do |config|
+  #     config.embedding_provider = :openai
+  #     config.openai_api_key = ENV["OPENAI_API_KEY"]
+  #   end
+  #
+  class Configuration
+    # Embedding provider (:ruby_llm, :openai, :ollama, :fastembed)
+    # Defaults to :ruby_llm if RubyLLM gem is available, otherwise :openai
+    # @return [Symbol]
+    attr_accessor :embedding_provider
+
+    # OpenAI API key (only needed if not using RubyLLM)
+    # @return [String, nil]
+    attr_accessor :openai_api_key
+
+    # OpenAI base URL (for custom endpoints)
+    # @return [String, nil]
+    attr_accessor :openai_base_url
+
+    # Ollama host URL
+    # @return [String]
+    attr_accessor :ollama_host
+
+    # Default embedding model
+    # @return [String]
+    attr_accessor :default_embedding_model
+
+    # Index storage directory
+    # @return [String]
+    attr_accessor :index_directory
+
+    # HNSW M parameter (graph connectivity)
+    # @return [Integer]
+    attr_accessor :hnsw_m
+
+    # HNSW ef_construction parameter
+    # @return [Integer]
+    attr_accessor :hnsw_ef_construction
+
+    # Default chunk size for text splitting
+    # @return [Integer]
+    attr_accessor :chunk_size
+
+    # Default chunk overlap
+    # @return [Integer]
+    attr_accessor :chunk_overlap
+
+    def initialize
+      # Default to RubyLLM if available, otherwise OpenAI
+      @embedding_provider = ruby_llm_available? ? :ruby_llm : :openai
+
+      @openai_api_key = ENV["OPENAI_API_KEY"]
+      @openai_base_url = ENV["OPENAI_BASE_URL"]
+      @ollama_host = ENV.fetch("OLLAMA_HOST", "http://localhost:11434")
+      @default_embedding_model = nil  # Let provider choose default
+
+      @index_directory = ".leann"
+      @hnsw_m = 32
+      @hnsw_ef_construction = 200
+
+      @chunk_size = 512
+      @chunk_overlap = 64
+    end
+
+    # Check if RubyLLM gem is available
+    # @return [Boolean]
+    def ruby_llm_available?
+      defined?(::RubyLLM) || gem_available?("ruby_llm")
+    end
+
+    # Check if FastEmbed gem is available
+    # @return [Boolean]
+    def fastembed_available?
+      defined?(::Fastembed) || gem_available?("fastembed")
+    end
+
+    # Validate configuration
+    # @raise [ConfigurationError] if configuration is invalid
+    def validate!
+      case embedding_provider
+      when :ruby_llm
+        unless ruby_llm_available?
+          raise ConfigurationError, "RubyLLM gem is required. Add 'ruby_llm' to your Gemfile."
+        end
+      when :openai
+        raise ConfigurationError, "OpenAI API key is required" if openai_api_key.nil? || openai_api_key.empty?
+      when :ollama
+        # Ollama doesn't require API key, just needs to be running
+      when :fastembed
+        unless fastembed_available?
+          raise ConfigurationError, "FastEmbed gem is required. Add 'fastembed' to your Gemfile."
+        end
+      else
+        raise ConfigurationError, "Unknown embedding provider: #{embedding_provider}"
+      end
+
+      true
+    end
+
+    # Get embedding model for a provider
+    # @return [String, nil]
+    def embedding_model_for(provider)
+      # Return custom model if explicitly set
+      return @default_embedding_model if @custom_embedding_model
+
+      # Provider-specific defaults
+      case provider
+      when :ruby_llm
+        nil  # RubyLLM uses its own configured default
+      when :openai
+        "text-embedding-3-small"
+      when :ollama
+        "nomic-embed-text"
+      when :fastembed
+        "BAAI/bge-small-en-v1.5"
+      else
+        @default_embedding_model
+      end
+    end
+
+    def default_embedding_model=(value)
+      @default_embedding_model = value
+      @custom_embedding_model = true
+    end
+
+    private
+
+    def gem_available?(name)
+      Gem::Specification.find_by_name(name)
+      true
+    rescue Gem::LoadError
+      false
+    end
+  end
+end

data/lib/leann/embedding/base.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module Leann
+  module Embedding
+    # Base class for embedding providers
+    #
+    # Subclasses must implement:
+    # - #compute(texts) -> Array<Array<Float>>
+    #
+    class Base
+      # @return [String] Model name
+      attr_reader :model
+
+      # @return [Integer, nil] Embedding dimensions
+      attr_reader :dimensions
+
+      # @param model [String] Embedding model name
+      def initialize(model:)
+        @model = model
+        @dimensions = nil
+      end
+
+      # Compute embeddings for a list of texts
+      #
+      # @param texts [Array<String>] Texts to embed
+      # @return [Array<Array<Float>>] Embeddings (one per text)
+      # @raise [NotImplementedError] if not overridden
+      def compute(texts)
+        raise NotImplementedError, "Subclasses must implement #compute"
+      end
+
+      # Compute embedding for a single text
+      #
+      # @param text [String]
+      # @return [Array<Float>]
+      def compute_one(text)
+        compute([text]).first
+      end
+
+      protected
+
+      # Normalize embedding to unit length (L2 normalization)
+      # @param embedding [Array<Float>]
+      # @return [Array<Float>]
+      def normalize(embedding)
+        norm = Math.sqrt(embedding.sum { |x| x * x })
+        return embedding if norm.zero?
+
+        embedding.map { |x| x / norm }
+      end
+
+      # Batch processing helper
+      # @param items [Array]
+      # @param batch_size [Integer]
+      # @yield [Array] Each batch
+      def in_batches(items, batch_size)
+        items.each_slice(batch_size) do |batch|
+          yield batch
+        end
+      end
+    end
+  end
+end

data/lib/leann/embedding/fastembed.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+require_relative "base"
+
+module Leann
+  module Embedding
+    # FastEmbed provider for local embeddings
+    #
+    # Uses ONNX Runtime for fast, local embedding generation without
+    # requiring an API key or external service.
+    #
+    # @example
+    #   provider = Leann::Embedding::FastEmbed.new(model: "BAAI/bge-small-en-v1.5")
+    #   embeddings = provider.compute(["Hello", "World"])
+    #
+    class FastEmbed < Base
+      MAX_BATCH_SIZE = 64
+
+      # Supported models with their dimensions
+      MODELS = {
+        "BAAI/bge-small-en-v1.5" => 384,
+        "BAAI/bge-base-en-v1.5" => 768,
+        "intfloat/multilingual-e5-small" => 384,
+        "nomic-ai/nomic-embed-text-v1.5" => 768
+      }.freeze
+
+      DEFAULT_MODEL = "BAAI/bge-small-en-v1.5"
+
+      # @param model [String] FastEmbed model name
+      # @param cache_dir [String, nil] Model cache directory
+      # @param threads [Integer, nil] Number of ONNX threads
+      def initialize(model: nil, cache_dir: nil, threads: nil)
+        model ||= DEFAULT_MODEL
+        super(model: model)
+
+        @cache_dir = cache_dir || ENV["FASTEMBED_CACHE_PATH"]
+        @threads = threads
+        @client = nil
+
+        check_gem!
+      end
+
+      # Compute embeddings for texts
+      #
+      # @param texts [Array<String>]
+      # @return [Array<Array<Float>>]
+      def compute(texts)
+        return [] if texts.empty?
+
+        all_embeddings = []
+
+        in_batches(texts, MAX_BATCH_SIZE) do |batch|
+          batch_embeddings = compute_batch(batch)
+          all_embeddings.concat(batch_embeddings)
+          print "." # Progress indicator
+        end
+
+        puts " Done! (#{all_embeddings.size} embeddings)" unless texts.size < MAX_BATCH_SIZE
+
+        # FastEmbed returns normalized vectors by default
+        all_embeddings
+      end
+
+      # Get dimensions for the configured model
+      # @return [Integer]
+      def dimensions
+        @dimensions ||= MODELS[model] || detect_dimensions
+      end
+
+      private
+
+      def check_gem!
+        unless defined?(::Fastembed)
+          raise ConfigurationError, <<~MSG
+            FastEmbed gem is required for local embeddings.
+
+            Add to your Gemfile:
+              gem 'fastembed'
+
+            Or install directly:
+              gem install fastembed
+          MSG
+        end
+      end
+
+      def client
+        @client ||= begin
+          options = { model_name: model }
+          options[:cache_dir] = @cache_dir if @cache_dir
+          options[:threads] = @threads if @threads
+
+          ::Fastembed::TextEmbedding.new(**options)
+        end
+      end
+
+      def compute_batch(texts)
+        # FastEmbed returns a lazy enumerator, convert to array
+        client.embed(texts, batch_size: texts.size).to_a
+      rescue ::Fastembed::Error => e
+        raise EmbeddingError.new(
+          "FastEmbed error: #{e.message}",
+          provider: :fastembed,
+          original_error: e
+        )
+      rescue StandardError => e
+        raise EmbeddingError.new(
+          "FastEmbed error: #{e.message}",
+          provider: :fastembed,
+          original_error: e
+        )
+      end
+
+      def detect_dimensions
+        # Compute a single embedding to detect dimensions
+        sample = client.embed(["test"], batch_size: 1).first
+        sample.size
+      end
+    end
+  end
+end