Module: Sequel::Plugins::Pgvector::ClassMethods

Defined in:
lib/sequel/plugins/pgvector/class_methods.rb

Instance Attribute Summary collapse

Instance Method Summary collapse

Instance Attribute Details

#vector_columns ⇒ Object

Returns the value of attribute vector_columns.



5
6
7
 
# File 'lib/sequel/plugins/pgvector/class_methods.rb', line 5

def vector_columns
  @vector_columns
end

Instance Method Details

#nearest_neighbors(column, value, distance:, threshold: nil) ⇒ Object 



7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
 
# File 'lib/sequel/plugins/pgvector/class_methods.rb', line 7

def nearest_neighbors(column, value, distance:, threshold: nil)
    value = ::Pgvector.encode(value) unless value.is_a?(String)
    quoted_column = dataset.quote_identifier(column)
    distance = distance.to_s

    operator =
        case distance
        when "inner_product"
        "<#>"
        when "cosine"
        "<=>"
        when "euclidean"
        "<->"
        else
        raise ArgumentError, "Invalid distance: #{distance}"
        end

    order = "#{quoted_column} #{operator} ?"

    neighbor_distance =
        if distance == "inner_product"
        "(#{order}) * -1"
        else
        order
        end

    query = select_append(Sequel.lit("#{neighbor_distance} AS neighbor_distance", value))
        .exclude(column => nil)
        .order(Sequel.lit(order, value))

    # Apply the WHERE condition only if threshold is provided
    query = query.where(Sequel.lit("(#{neighbor_distance}) < ?", value, threshold)) if threshold

    query
end