Sha256: d9b5a64a0273aa8eef68f599f7fe59bcb65534f0be51dde34d4e30127a97d350

Contents?: true

Size: 1.89 KB

Versions: 2

Compression:

Stored size: 1.89 KB

Contents

# frozen_string_literal: true

module Boxcars
  module VectorStore
    module InMemory
      class Search
        include VectorStore

        def initialize(params)
          validate_params(params[:vector_documents])
          @vector_documents = params[:vector_documents]
        end

        def call(query_vector:, count: 1)
          raise ::Boxcars::ArgumentError, 'query_vector is empty' if query_vector.empty?

          search(query_vector, count)
        end

        private

        attr_reader :vector_documents

        def validate_params(vector_documents)
          return if valid_vector_store?(vector_documents)

          raise ::Boxcars::ArgumentError, "vector_documents is not valid"
        end

        def valid_vector_store?(vector_documents)
          vector_documents && vector_documents[:type] == :in_memory &&
            vector_documents[:vector_store].is_a?(Array) &&
            vector_documents[:vector_store].all? do |doc|
              doc.is_a?(Boxcars::VectorStore::Document)
            end
        end

        def search(query_vector, num_neighbors)
          results = vector_documents[:vector_store].map do |doc|
            {
              document: doc,
              similarity: cosine_similarity(query_vector, doc.embedding)
            }
          end
          results.sort_by { |result| -result[:similarity] }
                 .first(num_neighbors)
        rescue StandardError => e
          raise_argument_error("Error searching for #{query_vector}: #{e.message}")
          raise_error
        end

        def cosine_similarity(vector1, vector2)
          dot_product = vector1.zip(vector2).reduce(0) { |sum, (a, b)| sum + (a * b) }
          magnitude1 = Math.sqrt(vector1.reduce(0) { |sum, a| sum + (a**2) })
          magnitude2 = Math.sqrt(vector2.reduce(0) { |sum, b| sum + (b**2) })
          dot_product / (magnitude1 * magnitude2)
        end
      end
    end
  end
end

Version data entries

2 entries across 2 versions & 1 rubygems

Version Path
boxcars-0.2.13 lib/boxcars/vector_store/in_memory/search.rb
boxcars-0.2.12 lib/boxcars/vector_store/in_memory/search.rb