OpenAI-like helpers for storing an querying embedding vectors utilizing a postgres database with pgvector extension

!pgstore.md

PGStore

OpenAI-like helpers for storing an querying embedding vectors utilizing a postgres database with pgvector extension.

Also included is the docker-compose which instantiates a postgres database with the necessary extension.

Usage

import pathlib
import openai
import pgstore
vector_store = pgstore.PGVectorStore(
  client=openai.Client(
    BASE_URL=os.getenv('OPENAI_BASE_URL'),
    API_KEY=os.getenv('OPENAI_API_KEY'),
  ),
  database_url=os.getenv('DATABASE_URL')
)
vector_store.create('docstore', {'model': 'nomic-embed-text'})
vector_store.upload('docstore', 'README.md', pathlib.Path('README.md').read_text(), {'tag': 'test'})

for file in vector_store.search('docstore', search='Some text in README'):
  print(file)
for file in vector_store.search('docstore', filter=dict(tag='test')):
  print(file)

docker-compose.yaml

services:
  pgvector-postgres:
    image: pgvector/pgvector:pg17
    ports:
    - 5435:5432
    environment:
    - POSTGRES_PASSWORD
    volumes:
    - pgvector-postgres-data:/var/lib/postgresql/data
volumes:
  pgvector-postgres-data:

pgstore.py

''' OpenAI-like helpers for storing embedding vectors in a postgres database with the pgvector extension.
'''
import json
import openai
import psycopg2
import contextlib
import traceback

class PGFileStore:
  def __init__(self, *, database_url: str):
    self.pgconn = psycopg2.connect(database_url)
    
  @contextlib.contextmanager
  def _cursor(self):
    try:
      with self.pgconn.cursor() as cur:
        yield cur
    except:
      self.pgconn.rollback()
      raise
    else:
      self.pgconn.commit()

  def migrate(self):
    try:
      with self._cursor() as cur:
        cur.execute('create table file (id bigserial primary key, filename varchar, content varchar);')
    except:
      traceback.print_exc()

  def list(self):
    with self._cursor() as cur:
      cur.execute(f'select id from file;')
      for id in cur:
        yield PGFileStoreFile(file_store=self, id=id)

  def upload(self, filename: str, content: str, **_):
    with self._cursor() as cur:
      cur.execute(f'insert into file (filename, content) values (%s, %s) returning id;', (filename, content,))
      id, = cur.fetchone()
      return id

  def filename(self, id: str, **_):
    with self._cursor() as cur:
      cur.execute(f'select filename from file where id = %s;', (id,))
      filename, = cur.fetchone()
      return filename

  def content(self, id: str, **_):
    with self._cursor() as cur:
      cur.execute(f'select content from file where id = %s;', (id,))
      content, = cur.fetchone()
      return content

  def delete(self, id: str, **_):
    with self._cursor() as cur:
      cur.execute(f'delete from file where id = %s;', (id,))

class PGFileStoreFile:
  def __init__(self, *, file_store: PGFileStore, id):
    self.file_store = file_store
    self.id = id

  def __repr__(self):
    return f"PGFileStoreFile(id={self.id})"

  @property
  def filename(self):
    return self.file_store.filename(self.id)

  @property
  def content(self):
    return self.file_store.content(self.id)

  def delete(self):
    return self.file_store.delete(self.id)

class PGVectorStore:
  def __init__(self, *, client: openai.OpenAI, database_url: str):
    self.client = client
    self.file_store = PGFileStore(database_url=database_url)
    self.pgconn = psycopg2.connect(database_url)
    
  def migrate(self):
    try:
      with self._cursor() as cur:
        cur.execute('create extension if not exists vector;')
        cur.execute('create table vector_store (id bigserial primary key, name varchar unique, opts jsonb);')
    except:
      traceback.print_exc()

  @contextlib.contextmanager
  def _cursor(self):
    try:
      with self.pgconn.cursor() as cur:
        yield cur
    except:
      self.pgconn.rollback()
      raise
    else:
      self.pgconn.commit()

  def create(self, name: str, opts: dict, **_):
    response = self.client.embeddings.create(**opts, input=['test'])
    dimensions = len(response.data[0].embedding)
    with self._cursor() as cur:
      cur.execute(f'insert into vector_store (name, opts) values (%s, %s) returning id;', (name, json.dumps(opts)))
      id, = cur.fetchone()
      cur.execute(f'create table {json.dumps(f"vector_store_{id}")} (id bigserial primary key, file bigint references file (id) on delete cascade, attributes jsonb, embedding vector({dimensions}));')
      return id

  def upload(self, store: str, filename: str, content: str, attributes: dict):
    with self._cursor() as cur:
      cur.execute('select id, opts from vector_store where name = %s;', (store,))
      id, opts = cur.fetchone()
      file_id = self.file_store.upload(filename, content)
      response = self.client.embeddings.create(**opts, input=[content])
      embedding = response.data[0].embedding
      cur.execute(f'insert into {json.dumps(f"vector_store_{id}")} (file, attributes, embedding) values (%s, %s, %s) returning id;', (file_id, json.dumps(attributes), embedding))
      id, = cur.fetchone()
      return id
  
  def search(self, store: str, *, search=None, filters=None):
    with self._cursor() as cur:
      cur.execute('select id, opts from vector_store where name = %s;', (store,))
      id, opts = cur.fetchone()
      if search is not None:
        model = opts.pop('model')
        response = self.client.embeddings.create(model=model, input=[search])
        embedding = response.data[0].embedding
        if filters is not None:
          cur.execute(f'select file from {json.dumps(f"vector_store_{id}")} where attributes @> %s order by embedding <-> %s::vector;', (json.dumps(embedding), json.dumps(filters)))
        else:
          cur.execute(f'select file from {json.dumps(f"vector_store_{id}")} order by embedding <-> %s::vector;', (json.dumps(embedding),))
      elif filters is not None:
        cur.execute(f'select file from {json.dumps(f"vector_store_{id}")} where attributes @> %s;', (json.dumps(filters),))
      else:
        cur.execute(f'select file from {json.dumps(f"vector_store_{id}")};')
      record = cur.fetchone()
      while record is not None:
        file, = record
        yield PGFileStoreFile(file_store=self.file_store, id=file)
        record = cur.fetchone()

  def list(self, store: str, **_):
    yield from self.search(store)

  def delete_file(self, store: str, file_id: str, prune=False):
    with self._cursor() as cur:
      cur.execute('select id from vector_store where name = %s;', (store,))
      id, = cur.fetchone()
      cur.execute(f'delete from {json.dumps(f"vector_store_{id}")} where file = %s;', (file_id,))
      if prune: cur.execute(f"delete from file where file = %s;", (file_id,))
      return file_id

  def delete(self, store: str, prune=False):
    with self._cursor() as cur:
      cur.execute('select id from vector_store where name = %s;', (store,))
      id, = cur.fetchone()
      if prune: cur.execute(f"delete from file as f where f.id in (delete from {json.dumps(f"vector_store_{id}")} vsi returning vsi.file);")
      cur.execute('delete from vector_store where id = %s;', (id,))
      cur.execute(f'drop table {json.dumps(f"vector_store_{id}")};')
      return id