Module: Legion::Data::Vector

Extended by:
Logging::Helper
Defined in:
lib/legion/data/vector.rb

Class Method Summary collapse

Methods included from Logging::Helper

handle_exception

Class Method Details

.available?Boolean

Returns:

  • (Boolean)


11
12
13
14
15
16
17
18
19
# File 'lib/legion/data/vector.rb', line 11

def available?
  return false unless Legion::Data.connection
  return false unless Legion::Data.connection.adapter_scheme == :postgres

  Legion::Data.connection.fetch("SELECT 1 FROM pg_extension WHERE extname = 'vector'").any?
rescue StandardError => e
  handle_exception(e, level: :warn, handled: true, operation: :vector_available?)
  false
end

.cosine_search(table:, column:, query_vector:, limit: 10, min_similarity: 0.0) ⇒ Object



32
33
34
35
36
37
38
39
40
41
42
43
44
45
# File 'lib/legion/data/vector.rb', line 32

def cosine_search(table:, column:, query_vector:, limit: 10, min_similarity: 0.0)
  return [] unless available?

  log.debug "Vector cosine_search: table=#{table} column=#{column} limit=#{limit}"
  vec_literal = vector_literal(query_vector)
  ds = Legion::Data.connection[table]
                   .select_all
                   .select_append(Sequel.lit("1 - (#{column} <=> ?)", vec_literal).as(:similarity))
                   .order(Sequel.lit("#{column} <=> ?", vec_literal))
                   .limit(limit)

  ds = ds.where(Sequel.lit("1 - (#{column} <=> ?) >= ?", vec_literal, min_similarity)) if min_similarity.positive?
  ds.all
end

.ensure_extension!Object



21
22
23
24
25
26
27
28
29
30
# File 'lib/legion/data/vector.rb', line 21

def ensure_extension!
  return false unless Legion::Data.connection&.adapter_scheme == :postgres

  Legion::Data.connection.run('CREATE EXTENSION IF NOT EXISTS vector')
  log.info 'pgvector extension enabled'
  true
rescue StandardError => e
  handle_exception(e, level: :warn, handled: true, operation: :ensure_vector_extension)
  false
end

.l2_search(table:, column:, query_vector:, limit: 10) ⇒ Object



47
48
49
50
51
52
53
54
55
56
57
58
# File 'lib/legion/data/vector.rb', line 47

def l2_search(table:, column:, query_vector:, limit: 10)
  return [] unless available?

  log.debug "Vector l2_search: table=#{table} column=#{column} limit=#{limit}"
  vec_literal = vector_literal(query_vector)
  Legion::Data.connection[table]
              .select_all
              .select_append(Sequel.lit("#{column} <-> ?", vec_literal).as(:distance))
              .order(Sequel.lit("#{column} <-> ?", vec_literal))
              .limit(limit)
              .all
end