leann 0.1.0

data/lib/leann/rails/storage/active_record_backend.rb

@@ -0,0 +1,332 @@
+# frozen_string_literal: true
+
+module Leann
+  module Rails
+    # ActiveRecord-based storage backend for LEANN graphs
+    #
+    # Stores graph structure (neighbor lists) in the passages table,
+    # avoiding the need for separate binary files.
+    #
+    class ActiveRecordBackend
+      attr_reader :index, :dimensions, :m, :ef_construction
+
+      # @param index [Leann::Rails::Index] The index record
+      def initialize(index)
+        @index = index
+        @dimensions = index.dimensions
+        config = index.config || {}
+        @m = config["hnsw_m"] || Leann.configuration.hnsw_m
+        @ef_construction = config["hnsw_ef_construction"] || Leann.configuration.hnsw_ef_construction
+        @entry_point_id = nil
+        @max_level = 0
+      end
+
+      # Build the graph from documents and embeddings
+      #
+      # @param documents [Array<Hash>] Documents with :id, :text, :metadata
+      # @param embeddings [Array<Array<Float>>] Embedding vectors
+      def build(documents, embeddings)
+        return if documents.empty?
+
+        puts "Building LEANN graph with #{documents.size} nodes (M=#{@m})..."
+
+        # Build in-memory graph first using the core algorithm
+        graph = build_graph(documents, embeddings)
+
+        # Store passages with neighbor information
+        store_passages(documents, graph)
+
+        # Store graph metadata in the index
+        update_index_metadata(graph)
+
+        puts "Graph built and stored in database: #{documents.size} passages"
+      end
+
+      # Search the graph
+      #
+      # @param query_embedding [Array<Float>] Query vector
+      # @param embedding_provider [Leann::Embedding::Base] Provider for recomputing embeddings
+      # @param passages [Hash] Passage texts by ID (for embedding recomputation)
+      # @param limit [Integer] Number of results
+      # @return [Array<Array>] Array of [id, score] pairs
+      def search(query_embedding, embedding_provider:, passages:, limit:)
+        return [] if @index.passages.empty?
+
+        # Load graph metadata
+        config = @index.config || {}
+        entry_point_id = config["entry_point_id"]
+        max_level = config["max_level"] || 0
+
+        return [] unless entry_point_id
+
+        # Perform HNSW search with on-the-fly embedding recomputation
+        ef_search = [limit * 2, 50].max
+        search_hnsw(query_embedding, entry_point_id, max_level, ef_search, embedding_provider, passages, limit)
+      end
+
+      private
+
+      def build_graph(documents, embeddings)
+        # Simple HNSW-like graph construction
+        graph = {}
+        levels = {}
+        entry_point = nil
+        max_level = 0
+
+        documents.each_with_index do |doc, idx|
+          id = doc[:id]
+          embedding = embeddings[idx]
+          level = random_level
+
+          graph[id] = { embedding: embedding, neighbors: Array.new(level + 1) { [] } }
+          levels[id] = level
+
+          if entry_point.nil?
+            entry_point = id
+            max_level = level
+          else
+            # Insert into graph
+            insert_node(graph, id, embedding, entry_point, max_level, level)
+
+            if level > max_level
+              max_level = level
+              entry_point = id
+            end
+          end
+        end
+
+        { graph: graph, entry_point: entry_point, max_level: max_level }
+      end
+
+      def random_level
+        level = 0
+        ml = 1.0 / Math.log(@m)
+        while rand < Math.exp(-level / ml) && level < 16
+          level += 1
+        end
+        level
+      end
+
+      def insert_node(graph, new_id, new_embedding, entry_point, max_level, node_level)
+        current = entry_point
+
+        # Traverse from top to node's level
+        (max_level).downto(node_level + 1) do |level|
+          current = greedy_search_level(graph, new_embedding, current, level)
+        end
+
+        # Insert at each level
+        node_level.downto(0) do |level|
+          neighbors = search_level(graph, new_embedding, current, level, @ef_construction)
+
+          # Select M best neighbors
+          selected = select_neighbors(graph, new_embedding, neighbors, @m)
+
+          graph[new_id][:neighbors][level] = selected
+
+          # Add bidirectional connections
+          selected.each do |neighbor_id|
+            neighbor_neighbors = graph[neighbor_id][:neighbors][level] || []
+            neighbor_neighbors << new_id
+
+            # Prune if too many
+            if neighbor_neighbors.size > @m * 2
+              neighbor_embedding = graph[neighbor_id][:embedding]
+              graph[neighbor_id][:neighbors][level] = select_neighbors(
+                graph, neighbor_embedding, neighbor_neighbors, @m * 2
+              )
+            else
+              graph[neighbor_id][:neighbors][level] = neighbor_neighbors
+            end
+          end
+
+          current = selected.first if selected.any?
+        end
+      end
+
+      def greedy_search_level(graph, query, entry, level)
+        current = entry
+        current_dist = cosine_distance(query, graph[current][:embedding])
+
+        loop do
+          changed = false
+          neighbors = graph[current][:neighbors][level] || []
+
+          neighbors.each do |neighbor|
+            next unless graph[neighbor]
+
+            dist = cosine_distance(query, graph[neighbor][:embedding])
+            if dist < current_dist
+              current = neighbor
+              current_dist = dist
+              changed = true
+            end
+          end
+
+          break unless changed
+        end
+
+        current
+      end
+
+      def search_level(graph, query, entry, level, ef)
+        visited = Set.new([entry])
+        candidates = [[cosine_distance(query, graph[entry][:embedding]), entry]]
+        results = candidates.dup
+
+        while candidates.any?
+          candidates.sort_by!(&:first)
+          current_dist, current = candidates.shift
+
+          break if results.any? && current_dist > results.last.first
+
+          neighbors = graph[current][:neighbors][level] || []
+          neighbors.each do |neighbor|
+            next if visited.include?(neighbor)
+            next unless graph[neighbor]
+
+            visited << neighbor
+            dist = cosine_distance(query, graph[neighbor][:embedding])
+
+            if results.size < ef || dist < results.last.first
+              candidates << [dist, neighbor]
+              results << [dist, neighbor]
+              results.sort_by!(&:first)
+              results.pop if results.size > ef
+            end
+          end
+        end
+
+        results.map(&:last)
+      end
+
+      def select_neighbors(graph, query, candidates, m)
+        return candidates if candidates.size <= m
+
+        scored = candidates.map do |id|
+          [cosine_distance(query, graph[id][:embedding]), id]
+        end
+
+        scored.sort_by(&:first).first(m).map(&:last)
+      end
+
+      def cosine_distance(a, b)
+        dot = 0.0
+        norm_a = 0.0
+        norm_b = 0.0
+
+        a.each_with_index do |val, i|
+          dot += val * b[i]
+          norm_a += val * val
+          norm_b += b[i] * b[i]
+        end
+
+        similarity = dot / (Math.sqrt(norm_a) * Math.sqrt(norm_b) + 1e-10)
+        1.0 - similarity
+      end
+
+      def store_passages(documents, graph_data)
+        graph = graph_data[:graph]
+
+        # Bulk insert passages
+        passage_records = documents.map do |doc|
+          node = graph[doc[:id]]
+          # Store only level-0 neighbors (most important for search)
+          neighbors = node[:neighbors][0] || []
+
+          {
+            leann_index_id: @index.id,
+            external_id: doc[:id],
+            text: doc[:text],
+            metadata: doc[:metadata] || {},
+            neighbors: neighbors,
+            created_at: Time.current,
+            updated_at: Time.current
+          }
+        end
+
+        Passage.insert_all(passage_records)
+      end
+
+      def update_index_metadata(graph_data)
+        @index.update!(
+          config: (@index.config || {}).merge(
+            "entry_point_id" => graph_data[:entry_point],
+            "max_level" => graph_data[:max_level],
+            "hnsw_m" => @m,
+            "hnsw_ef_construction" => @ef_construction
+          )
+        )
+      end
+
+      def search_hnsw(query_embedding, entry_point_id, max_level, ef, embedding_provider, passages, limit)
+        # Load passages with neighbors
+        passage_map = @index.passages.index_by(&:external_id)
+
+        return [] if passage_map.empty?
+
+        # Get entry point
+        entry = passage_map[entry_point_id]
+        return [] unless entry
+
+        # Cache for computed embeddings
+        embedding_cache = {}
+
+        get_embedding = lambda do |id|
+          return embedding_cache[id] if embedding_cache[id]
+
+          text = passages[id] || passage_map[id]&.text
+          return nil unless text
+
+          embedding_cache[id] = embedding_provider.compute([text]).first
+        end
+
+        entry_embedding = get_embedding.call(entry_point_id)
+        return [] unless entry_embedding
+
+        # Simple greedy search at level 0 (most passages only have level 0)
+        visited = Set.new([entry_point_id])
+        current_dist = cosine_distance(query_embedding, entry_embedding)
+        candidates = [[current_dist, entry_point_id]]
+        results = candidates.dup
+
+        while candidates.any?
+          candidates.sort_by!(&:first)
+          dist, current_id = candidates.shift
+
+          break if results.size >= ef && dist > results.last.first
+
+          # Get neighbors from database
+          current_passage = passage_map[current_id]
+          next unless current_passage
+
+          neighbors = current_passage.neighbor_ids
+
+          neighbors.each do |neighbor_id|
+            next if visited.include?(neighbor_id)
+
+            visited << neighbor_id
+
+            neighbor_embedding = get_embedding.call(neighbor_id)
+            next unless neighbor_embedding
+
+            neighbor_dist = cosine_distance(query_embedding, neighbor_embedding)
+
+            if results.size < ef || neighbor_dist < results.last.first
+              candidates << [neighbor_dist, neighbor_id]
+              results << [neighbor_dist, neighbor_id]
+              results.sort_by!(&:first)
+              results.pop if results.size > ef
+            end
+          end
+        end
+
+        # Convert distances to similarity scores
+        results.first(limit).map do |dist, id|
+          score = 1.0 - dist
+          [id, score]
+        end
+      end
+    end
+  end
+end

data/lib/leann/rails.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+require_relative "../leann"
+
+module Leann
+  module Rails
+    autoload :Index, "leann/rails/active_record/index"
+    autoload :Passage, "leann/rails/active_record/passage"
+    autoload :ActiveRecordBackend, "leann/rails/storage/active_record_backend"
+    autoload :Builder, "leann/rails/builder"
+    autoload :Searcher, "leann/rails/searcher"
+
+    class << self
+      # Build a new index stored in the database
+      #
+      # @param name [String] Index name (unique identifier)
+      # @param options [Hash] Options for building
+      # @option options [Symbol] :embedding (:openai) Embedding provider
+      # @option options [String] :model Embedding model name
+      #
+      # @example
+      #   Leann::Rails.build("products") do
+      #     add "Red running shoes for athletes", category: "shoes"
+      #     add "Blue denim jeans, slim fit", category: "pants"
+      #   end
+      #
+      # @return [Leann::Rails::Index] The built index record
+      def build(name, **options, &block)
+        builder = Builder.new(name, **options)
+        builder.instance_eval(&block) if block_given?
+        builder.save
+      end
+
+      # Search an existing database index
+      #
+      # @param name [String] Index name
+      # @param query [String] Search query
+      # @param limit [Integer] Maximum results
+      # @param threshold [Float] Minimum similarity score
+      # @param filters [Hash] Metadata filters
+      #
+      # @example
+      #   results = Leann::Rails.search("products", "comfortable shoes")
+      #
+      # @return [Leann::SearchResults]
+      def search(name, query, limit: 5, threshold: nil, filters: nil)
+        index = Index.find_by!(name: name)
+        searcher = Searcher.new(index)
+        searcher.search(query, limit: limit, threshold: threshold, filters: filters)
+      end
+
+      # Open an existing index
+      #
+      # @param name [String] Index name
+      # @return [Leann::Rails::Index]
+      def open(name)
+        Index.find_by!(name: name)
+      end
+
+      # Check if an index exists
+      #
+      # @param name [String] Index name
+      # @return [Boolean]
+      def exists?(name)
+        Index.exists?(name: name)
+      end
+
+      # Delete an index and all its passages
+      #
+      # @param name [String] Index name
+      # @return [Boolean]
+      def delete(name)
+        index = Index.find_by(name: name)
+        return false unless index
+
+        index.destroy
+        true
+      end
+
+      # List all indexes
+      #
+      # @return [Array<String>]
+      def list
+        Index.pluck(:name).sort
+      end
+    end
+  end
+end
+
+require "leann/rails/railtie" if defined?(::Rails::Railtie)

data/lib/leann/ruby_llm/search.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require_relative "../../leann"
+
+# Only require ruby_llm if not already defined (allows mocking in tests)
+unless defined?(::RubyLLM::Tool)
+  begin
+    require "ruby_llm"
+  rescue LoadError
+    raise LoadError, "RubyLLM is required for Leann::RubyLLM::Search. Add 'ruby_llm' to your Gemfile."
+  end
+end
+
+module Leann
+  module RubyLLM
+    # A RubyLLM tool for searching LEANN indexes
+    #
+    # @example Basic usage
+    #   chat = ::RubyLLM.chat(model: "gpt-4o")
+    #     .with_tool(Leann::RubyLLM::Search.new("my_index"))
+    #
+    #   chat.ask("What does LEANN do?")
+    #   # => LLM searches the index and generates an answer
+    #
+    # @example Multiple indexes
+    #   docs_search = Leann::RubyLLM::Search.new("docs", name: "search_docs")
+    #   code_search = Leann::RubyLLM::Search.new("codebase", name: "search_code")
+    #
+    #   chat = ::RubyLLM.chat(model: "gpt-4o")
+    #     .with_tools(docs_search, code_search)
+    #
+    class Search < ::RubyLLM::Tool
+      # @param index_name [String] Name of the LEANN index to search
+      # @param name [String] Tool name (defaults to "leann_search")
+      # @param limit [Integer] Default number of results (default: 5)
+      def initialize(index_name, name: "leann_search", limit: 5)
+        @index_name = index_name
+        @default_limit = limit
+        @tool_name = name
+        super()
+      end
+
+      def name
+        @tool_name
+      end
+
+      def description
+        "Searches the '#{@index_name}' knowledge base for relevant documents. " \
+          "Use this to find information before answering questions."
+      end
+
+      def params
+        ::RubyLLM::Schema.new do
+          string :query,
+                 description: "The search query to find relevant documents",
+                 required: true
+          integer :limit,
+                  description: "Maximum number of results to return (default: #{@default_limit})",
+                  required: false
+        end
+      end
+
+      def execute(query:, limit: nil)
+        limit ||= @default_limit
+        results = Leann.search(@index_name, query, limit: limit)
+
+        if results.empty?
+          { found: false, message: "No relevant documents found for: #{query}" }
+        else
+          {
+            found: true,
+            count: results.size,
+            documents: results.map do |r|
+              {
+                text: r.text,
+                score: r.score.round(3),
+                metadata: r.metadata
+              }
+            end
+          }
+        end
+      rescue Leann::IndexNotFoundError
+        { error: "Index '#{@index_name}' not found" }
+      rescue StandardError => e
+        { error: e.message }
+      end
+    end
+  end
+end

data/lib/leann/search_result.rb

@@ -0,0 +1,195 @@
+# frozen_string_literal: true
+
+module Leann
+  # Represents a single search result
+  #
+  # @example
+  #   results = Leann.search("my_index", "query")
+  #   results.each do |result|
+  #     puts result.text
+  #     puts result.score
+  #     puts result.metadata[:source]
+  #   end
+  #
+  class SearchResult
+    # @return [String] Document ID
+    attr_reader :id
+
+    # @return [String] Document text
+    attr_reader :text
+
+    # @return [Float] Similarity score (higher is better)
+    attr_reader :score
+
+    # @return [Hash] Document metadata
+    attr_reader :metadata
+
+    # @param id [String]
+    # @param text [String]
+    # @param score [Float]
+    # @param metadata [Hash]
+    def initialize(id:, text:, score:, metadata: {})
+      @id = id
+      @text = text
+      @score = score.to_f
+      @metadata = metadata.transform_keys(&:to_sym)
+    end
+
+    # Truncate text to a maximum length
+    # @param max_length [Integer]
+    # @param omission [String]
+    # @return [String]
+    def truncated_text(max_length = 100, omission: "...")
+      return text if text.length <= max_length
+
+      text[0, max_length - omission.length] + omission
+    end
+
+    # Human-readable string representation
+    # @return [String]
+    def to_s
+      "[#{format("%.3f", score)}] #{truncated_text(80)}"
+    end
+
+    # Detailed inspection
+    # @return [String]
+    def inspect
+      "#<Leann::SearchResult id=#{id.inspect} score=#{format("%.4f", score)} text=#{truncated_text(50).inspect}>"
+    end
+
+    # Convert to hash
+    # @return [Hash]
+    def to_h
+      {
+        id: id,
+        text: text,
+        score: score,
+        metadata: metadata
+      }
+    end
+
+    # Compare by score (for sorting)
+    # @param other [SearchResult]
+    # @return [Integer]
+    def <=>(other)
+      other.score <=> score # Descending order
+    end
+
+    # Check equality
+    # @param other [SearchResult]
+    # @return [Boolean]
+    def ==(other)
+      return false unless other.is_a?(SearchResult)
+
+      id == other.id && text == other.text && score == other.score
+    end
+    alias eql? ==
+
+    # Hash code for use as hash key
+    # @return [Integer]
+    def hash
+      [id, text, score].hash
+    end
+  end
+
+  # Collection of search results with utility methods
+  class SearchResults
+    include Enumerable
+
+    # @return [Array<SearchResult>]
+    attr_reader :results
+
+    # @return [String] Original query
+    attr_reader :query
+
+    # @return [Float] Search duration in seconds
+    attr_reader :duration
+
+    # @param results [Array<SearchResult>]
+    # @param query [String]
+    # @param duration [Float]
+    def initialize(results, query: nil, duration: nil)
+      @results = results
+      @query = query
+      @duration = duration
+    end
+
+    # Iterate over results
+    def each(&block)
+      results.each(&block)
+    end
+
+    # Number of results
+    # @return [Integer]
+    def size
+      results.size
+    end
+    alias length size
+    alias count size
+
+    # Check if empty
+    # @return [Boolean]
+    def empty?
+      results.empty?
+    end
+
+    # Get first result
+    # @return [SearchResult, nil]
+    def first
+      results.first
+    end
+
+    # Get top n results
+    # @param n [Integer]
+    # @return [Array<SearchResult>]
+    def top(n)
+      results.first(n)
+    end
+
+    # Get result by index
+    # @param index [Integer]
+    # @return [SearchResult, nil]
+    def [](index)
+      results[index]
+    end
+
+    # Filter results by minimum score
+    # @param min_score [Float]
+    # @return [SearchResults]
+    def above(min_score)
+      filtered = results.select { |r| r.score >= min_score }
+      SearchResults.new(filtered, query: query, duration: duration)
+    end
+
+    # Get all texts
+    # @return [Array<String>]
+    def texts
+      results.map(&:text)
+    end
+
+    # Join all texts
+    # @param separator [String]
+    # @return [String]
+    def combined_text(separator: "\n\n")
+      texts.join(separator)
+    end
+
+    # Pretty print results
+    # @return [String]
+    def to_s
+      lines = ["Search results for: #{query.inspect}"]
+      lines << "Found #{size} results in #{format("%.3f", duration || 0)}s"
+      lines << "-" * 60
+      results.each_with_index do |r, i|
+        lines << "#{i + 1}. #{r}"
+      end
+      lines.join("\n")
+    end
+
+    # Convert to array of hashes
+    # @return [Array<Hash>]
+    def to_a
+      results.map(&:to_h)
+    end
+  end
+end