Module: IndexUtil::VectorIndex

Defined in:
lib/index_util/vector_index.rb

Class Method Summary collapse

Class Method Details

.delete(db, fragment_id) ⇒ Object 



43
44
45
46
47
 
# File 'lib/index_util/vector_index.rb', line 43

def delete(db, fragment_id)
  db[:fragment_vectors].where(fragment_id: Integer(fragment_id)).delete
rescue Sequel::DatabaseError
  nil
end

.ensure!(db, dimensions:) ⇒ Object 



21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
 
# File 'lib/index_util/vector_index.rb', line 21

def ensure!(db, dimensions:)
  dimensions = Integer(dimensions)
  stored_dimensions = IndexMetadata.get(db, :embedding_dimensions)
  if stored_dimensions && stored_dimensions.to_i != dimensions
    raise Error, "Embedding dimensions changed from #{stored_dimensions} to #{dimensions}; rebuild the index with --index"
  end

  db.run <<~SQL
    CREATE VIRTUAL TABLE IF NOT EXISTS fragment_vectors USING vec0(
      fragment_id integer primary key,
      embedding float[#{dimensions}]
    )
  SQL
  IndexMetadata.set(db, :embedding_dimensions, dimensions)
rescue Sequel::DatabaseError => e
  raise Error, "Failed to create sqlite-vec table: #{e.message}"
end

.insert(db, fragment_id, vector) ⇒ Object 



39
40
41
 
# File 'lib/index_util/vector_index.rb', line 39

def insert(db, fragment_id, vector)
  db["INSERT INTO fragment_vectors(fragment_id, embedding) VALUES (?, ?)", Integer(fragment_id), FragmentEmbedding.pack(vector)].insert
end

.load!(db) ⇒ Object 



7
8
9
10
11
12
13
14
15
16
17
18
19
 
# File 'lib/index_util/vector_index.rb', line 7

def load!(db)
  require "sqlite_vec"
  db.synchronize do |connection|
    connection.enable_load_extension(true)
    SqliteVec.load(connection)
    connection.enable_load_extension(false)
  end
  db.get(Sequel.lit("vec_version()"))
rescue LoadError => e
  raise Error, "sqlite-vec gem is required by index_util: #{e.message}"
rescue StandardError => e
  raise Error, "sqlite-vec support is required by index_util: #{e.message}"
end

.search(db, query_embedding, limit:) ⇒ Object 



49
50
51
52
53
54
55
56
57
58
59
 
# File 'lib/index_util/vector_index.rb', line 49

def search(db, query_embedding, limit:)
  db.fetch(<<~SQL, FragmentEmbedding.pack(query_embedding), Integer(limit)).all
    SELECT fragments.id, fragments.fragment_document, fragments.content, fragment_vectors.distance AS distance
    FROM fragment_vectors
    JOIN fragments ON fragments.id = fragment_vectors.fragment_id
    WHERE embedding MATCH ? AND k = ?
    ORDER BY distance
  SQL
rescue Sequel::DatabaseError
  []
end