Class: Prescient::Pgvector::Store

Inherits:
Object
  • Object
show all
Defined in:
lib/prescient/pgvector.rb

Overview

PostgreSQL pgvector integration.

The integration accepts a PG-compatible connection object, so applications choose and manage their own PostgreSQL driver and connection lifecycle. Stores provider embeddings and performs nearest-neighbor searches.

Constant Summary collapse

METRICS =

Returns Supported pgvector distance operators.

Returns:

  • (Hash<Symbol, String>) —

    Supported pgvector distance operators

{
  cosine:        '<=>',
  euclidean:     '<->',
  inner_product: '<#>',
}.freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(connection:, dimensions:, table: 'prescient_embeddings') ⇒ Store

Returns a new instance of Store.

Parameters:

  • connection (Object) —

    PG-compatible object responding to exec and exec_params

  • dimensions (Integer) —

    Required dimensions for every embedding

  • table (String, Symbol) (defaults to: 'prescient_embeddings') —

    Safe PostgreSQL table identifier



28
29
30
31
32
 
# File 'lib/prescient/pgvector.rb', line 28

def initialize(connection:, dimensions:, table: 'prescient_embeddings')
  @connection = connection
  @dimensions = validate_dimensions(dimensions)
  @table = validate_table(table)
end

Instance Attribute Details

#dimensions ⇒ Integer (readonly)

Returns Required vector dimensions.

Returns:

  • (Integer) —

    Required vector dimensions



19
20
21
 
# File 'lib/prescient/pgvector.rb', line 19

def dimensions
  @dimensions
end

#table ⇒ String (readonly)

Returns Embeddings table name.

Returns:

  • (String) —

    Embeddings table name



22
23
24
 
# File 'lib/prescient/pgvector.rb', line 22

def table
  @table
end

Instance Method Details

#create_index!(metric: :cosine) ⇒ void

This method returns an undefined value.

Create an HNSW index for the selected distance metric.

Parameters:

  • metric (Symbol) (defaults to: :cosine) —

    :cosine, :euclidean, or :inner_product



58
59
60
61
62
63
64
 
# File 'lib/prescient/pgvector.rb', line 58

def create_index!(metric: :cosine)
  metric_name = validate_metric(metric)
  @connection.exec(<<~SQL)
    CREATE INDEX IF NOT EXISTS #{table}_#{metric_name}_embedding_idx
    ON #{table} USING hnsw (embedding #{metric_operator_class(metric_name)})
  SQL
end

#install! ⇒ void

This method returns an undefined value.

Create the pgvector extension and the embeddings table.



37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
 
# File 'lib/prescient/pgvector.rb', line 37

def install!
  @connection.exec('CREATE EXTENSION IF NOT EXISTS vector')
  @connection.exec(<<~SQL)
    CREATE TABLE IF NOT EXISTS #{table} (
      id text PRIMARY KEY,
      provider text NOT NULL,
      model text NOT NULL,
      dimensions integer NOT NULL CHECK (dimensions = #{dimensions}),
      embedding vector(#{dimensions}) NOT NULL,
      content text,
      metadata jsonb NOT NULL DEFAULT '{}'::jsonb,
      created_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP,
      updated_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP
    )
  SQL
end

#search(embedding:, limit: 10, metric: :cosine, provider: nil, model: nil) ⇒ Array<Hash>

Find the nearest stored embeddings.

Parameters:

  • embedding (Array<Numeric>) —

    Query vector

  • limit (Integer) (defaults to: 10) —

    Maximum result count

  • metric (Symbol) (defaults to: :cosine) —

    Distance metric

  • provider (String, Symbol, nil) (defaults to: nil) —

    Optional provider filter

  • model (String, nil) (defaults to: nil) —

    Optional model filter

Returns:

  • (Array<Hash>) —

    Records ordered by ascending distance



97
98
99
100
101
102
103
104
105
 
# File 'lib/prescient/pgvector.rb', line 97

def search(embedding:, limit: 10, metric: :cosine, provider: nil, model: nil)
  vector = serialize_vector(embedding)
  limit = validate_limit(limit)
  metric = validate_metric(metric)
  filters, parameters = search_filters(provider, model)
  result = @connection.exec_params(search_query(metric, filters), [vector, limit, *parameters])

  result.map { |row| record_from(row) }
end

#upsert(id:, embedding:, provider:, model:, content: nil, metadata: {}) ⇒ Hash

Insert or replace an embedding record.

Returns:

  • (Hash) —

    Stored record metadata



69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
 
# File 'lib/prescient/pgvector.rb', line 69

def upsert(id:, embedding:, provider:, model:, content: nil, metadata: {})
  vector = serialize_vector(embedding)
  parameters = [id.to_s, provider.to_s, model.to_s, dimensions, vector, content, JSON.generate(metadata)]
  result = @connection.exec_params(<<~SQL, parameters)
    INSERT INTO #{table} (id, provider, model, dimensions, embedding, content, metadata)
    VALUES ($1, $2, $3, $4, $5::vector, $6, $7::jsonb)
    ON CONFLICT (id) DO UPDATE SET
      provider = EXCLUDED.provider,
      model = EXCLUDED.model,
      dimensions = EXCLUDED.dimensions,
      embedding = EXCLUDED.embedding,
      content = EXCLUDED.content,
      metadata = EXCLUDED.metadata,
      updated_at = CURRENT_TIMESTAMP
    RETURNING id, provider, model, dimensions, content, metadata
  SQL

  record_from(result.first)
end