Class: Ask::RAG::VectorStore::PGVector

Inherits:
VectorStore
  • Object
show all
Defined in:
lib/ask/rag/vector_store/pgvector.rb

Overview

PGVector-backed vector store using PostgreSQL + the pgvector extension.

Stores document text, metadata, and embeddings in a PostgreSQL table. Requires the pgvector gem and an ActiveRecord connection.

Examples:

Setup

# In your migration:
#   create_table :documents do |t|
#     t.text :content
#     t.jsonb :metadata
#     t.vector :embedding, limit: 1536
#   end
#   add_index :documents, :embedding, using: :ivfflat, opclass: :vector_cosine_ops

Usage

store = Ask::RAG::VectorStore::PGVector.new(
  table_name: :embeddings,
  embedding_column: :embedding
)
store.add(chunks, model: "text-embedding-3-small")
results = store.similarity_search("query", limit: 5)

Instance Method Summary collapse

Constructor Details

#initialize(table_name:, content_column: :content, metadata_column: :metadata, embedding_column: :embedding, model_class: nil) ⇒ PGVector

Returns a new instance of PGVector.

Parameters:

  • table_name (Symbol) —

    the ActiveRecord table name

  • content_column (Symbol) (defaults to: :content) —

    column storing text content (default: :content)

  • metadata_column (Symbol) (defaults to: :metadata) —

    column storing JSON metadata (default: :metadata)

  • embedding_column (Symbol) (defaults to: :embedding) —

    column storing the vector (default: :embedding)

  • model_class (Class, nil) (defaults to: nil) —

    ActiveRecord model class (default: inferred from table_name)



34
35
36
37
38
39
40
41
42
 
# File 'lib/ask/rag/vector_store/pgvector.rb', line 34

def initialize(table_name:, content_column: :content, metadata_column: :metadata,
               embedding_column: :embedding, model_class: nil)
  @table_name = table_name.to_s
  @content_column = content_column.to_s
  @metadata_column = metadata_column.to_s
  @embedding_column = embedding_column.to_s
  @model_class = model_class || infer_model_class
  @embedding_model = nil
end

Instance Method Details

#add(documents, model:, batch_size: 20) ⇒ Object 



44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
 
# File 'lib/ask/rag/vector_store/pgvector.rb', line 44

def add(documents, model:, batch_size: 20)
  ids = []
  @embedding_model = model

  documents.each_slice(batch_size) do |batch|
    texts = batch.map(&:content)
    vectors = embed_texts(texts, model)

    batch.each_with_index do |doc, idx|
      record = @model_class.create!(
        @content_column => doc.content,
        @metadata_column => doc.metadata,
        @embedding_column => vectors[idx]
      )
      ids << record.id.to_s
    end
  end

  ids
end

#clear ⇒ Object 



105
106
107
 
# File 'lib/ask/rag/vector_store/pgvector.rb', line 105

def clear
  @model_class.delete_all
end

#delete(ids) ⇒ Object 



101
102
103
 
# File 'lib/ask/rag/vector_store/pgvector.rb', line 101

def delete(ids)
  @model_class.where(id: ids).delete_all
end

#similarity_search(query, limit: 10, filter: nil, mmr: false, diversity_bonus: 0.3) ⇒ Object 



65
66
67
68
 
# File 'lib/ask/rag/vector_store/pgvector.rb', line 65

def similarity_search(query, limit: 10, filter: nil, mmr: false, diversity_bonus: 0.3)
  query_vector = embed_query(query)
  similarity_search_by_vector(query_vector, limit: limit, filter: filter)
end

#similarity_search_by_vector(vector, limit: 10, filter: nil) ⇒ Object 



70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
 
# File 'lib/ask/rag/vector_store/pgvector.rb', line 70

def similarity_search_by_vector(vector, limit: 10, filter: nil)
  vector_str = vector.is_a?(Array) ? "[#{vector.join(',')}]" : vector.to_s
  column = "#{@table_name}.#{@embedding_column}"
  model_class = @model_class

  scope = model_class
    .select("#{@table_name}.*, 1 - (#{column} <=> '#{vector_str}') AS score")
    .where("#{column} IS NOT NULL")

  if filter
    filter.each do |key, value|
      scope = scope.where("#{@metadata_column} @> ?", { key.to_s => value }.to_json)
    end
  end

  records = scope
    .order(Arel.sql("#{column} <=> '#{vector_str}'"))
    .limit(limit)

  records.map do |record|
    Ask::Document.new(
      content: record.send(@content_column),
      metadata: (record.send(@metadata_column) || {}).merge(
        score: record.score,
        db_id: record.id
      ),
      id: record.id.to_s
    )
  end
end

#size ⇒ Object 



109
110
111
 
# File 'lib/ask/rag/vector_store/pgvector.rb', line 109

def size
  @model_class.count
end