ballerinax/ai.pgvector Ballerina library

Overview

Pgvector is a PostgreSQL extension that introduces a vector data type and similarity search capabilities for working with embeddings.

The Ballerina pgvector module provides an API for integrating with the pgvector extension for PostgreSQL. Its implementation allows it to be used as a Ballerina AI ai:VectorStore, enabling users to store, retrieve, and search high-dimensional vectors.

Setup guide

You need a running PostgreSQL instance with the pgvector extension enabled. For that you can use the official pgvector Docker image.

docker run --name pgvector-db \
  -e POSTGRES_PASSWORD=mypassword \
  -e POSTGRES_DB=vector_db \
  -p 5432:5432 \
  -d pgvector/pgvector:pg17

To enable the pgvector extension, connect to the database and execute the following query.

CREATE EXTENSION IF NOT EXISTS vector;

Quick Start

To use the pgvector vector store in your Ballerina project, modify the .bal file as follows.

Step 1: Import the module

import ballerina/ai;
import ballerinax/ai.pgvector;

Step 2: Initialize the Pgvector vector store

ai:VectorStore vectorStore = check new(
   host,
   user,
   password,
   database,
   tableName,
   configs = {
      vectorDimension: 1536
   }
);

Step 3: Invoke the operations

ai:Error? result = vectorStore.add(
   [
      {
         id: "1",
         embedding: [1.0, 2.0, 3.0],
         chunk: {
               'type: "text", 
               content: "This is a chunk"
         }
      }
   ]
);

ai:VectorMatch[]|ai:Error matches = vectorStore.query({
   embedding: [1.0, 2.0, 3.0],
   filters: {
      // optional metadata filters
   }
});