Creating AI Applications with PGVector and PostgreSQL: A Hands-On Introduction

Understanding PostgreSQL and pgvector

PostgreSQL, also known as Postgres, is an open-source object-relational database system with over 30 years of active development. Known for reliability, a wide range of features, and performance, Postgres is a popular choice for any application, from small projects to large enterprise software.

pgvector is a Postgres extension that adds vector similarity search capabilities to the database. In AI, vector similarity search is crucial for many applications, including recommendation systems, natural language processing, image recognition, and more. pgvector allows you to both store and query high-dimensional vectors efficiently within your Postgres database.

The Power of Embeddings in AI

Embeddings are at the heart of many modern AI applications. They are vector representations of data that capture semantic meaning and relationships. Embeddings measure how related text strings are, allowing for more nuanced comparisons than simple text matching.

An embedding is essentially just a list (vector) of floating-point numbers. The distance between two vectors measures their relatedness. Small distances suggest high relatedness, while large distances suggest low relatedness. This makes embeddings incredibly useful for tasks like semantic search, recommendation systems, and content-based filtering.

Why Vector Databases Matter

Vector databases, like Postgres with pgvector, offer significant advantages in AI applications:

1. Efficient Similarity Search: They excel at finding “similar” items, which is often more valuable than exact matches in AI applications.

2. Handling High-Dimensional Data: They’re optimized for storing and querying high-dimensional vectors produced by machine learning models.

3. Scalability: Specialized indexing methods allow for quick queries even with millions of vectors.

4. Semantic Understanding: They enable context-aware searches, understanding the meaning behind queries.

These capabilities make vector databases ideal for various AI applications, including:

• Recommendation systems
• Semantic search engines
• Content-based filtering
• Image and audio similarity search
• Natural language processing tasks

By augmenting Postgres with pgvector, developers can leverage these powerful capabilities while still benefiting from the robustness and extensive feature set of a mature relational database system.

Starting an AI Project with Batteries Included

Batteries Included offers a streamlined approach to working with AI projects, complete with integrated database support. We’ll be using Jupyter Notebooks to run Python code and interact with both the OpenAI API as well as a Postgres database.

Jupyter Notebooks provide an interactive environment where you can write and execute code, visualize data, and document your work all in one place. This makes them ideal for exploratory data analysis, prototyping AI models, and sharing results with others.

Let’s get started!

export DATABASE_URL=$(bi postgres access-info mycluster myusername --localhost)

!pip install psycopg2-binary
!pip install openai
!pip install pgvector

import psycopg2
import os

conn = psycopg2.connect(os.environ['DATABASE_URL'])
cur = conn.cursor()

import openai
import numpy as np
import psycopg2
from pgvector.psycopg2 import register_vector
from openai import OpenAI

# Set up OpenAI API
client = OpenAI(api_key=os.environ['OPENAI_KEY'])

# Install pgvector
cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
conn.commit()

register_vector(conn)

# Create a table to store documents and their embeddings
cur.execute("""
CREATE TABLE IF NOT EXISTS documents (
    id SERIAL PRIMARY KEY,
    content TEXT,
    embedding VECTOR(1536)
);
""")

conn.commit()

# Sample documents
documents = [
    "Foxes are great night-time predators",
    "Foxes can make over 40 different sounds",
    "Foxes make use of the earth's magnetic field to hunt",
    "Baby foxes are unable to see, walk or thermoregulate when they are born",
]

for doc in documents:
    # Generate embedding using OpenAI
    response = client.embeddings.create(input=doc, model="text-embedding-ada-002")
    embedding = response.data[0].embedding

    # Insert document + embedding into the database
    cur.execute("INSERT INTO documents (content, embedding) VALUES (%s, %s);", (doc, embedding))
    conn.commit()

Leveraging pgvector

Now that we have our environment set up and data stored, let’s explore how we can use pgvector in our AI applications.

1. Indexing: Let’s start by creating an index on the embedding column to speed up similarity searches:

cur.execute("CREATE INDEX ON documents USING ivfflat (embedding vector_l2_ops);")
conn.commit()

2. Similarity Search: When a query comes in, let’s generate an embedding for the query and use pgvector’s similarity search capabilities to find the nearest row in your database:

def find_similar_document(query):
    # Generate embedding for the query
    response = client.embeddings.create(input=query, model="text-embedding-ada-002")
    query_embedding = response.data[0].embedding

    # Perform similarity search
    cur.execute("""
    SELECT content, embedding <-> %s::vector AS distance
    FROM documents
    ORDER BY distance
    LIMIT 1;
    """, (query_embedding,))

    return cur.fetchone()

# Example usage
similar_doc = find_similar_document("I was wondering, how many sounds can foxes make?")
if similar_doc:
    doc, distance = similar_doc
    print(f"Distance: {distance:.4f}, Content: {doc}")
else:
    print("No similar document found.")

This function will return the most similar document to the given query, along with its distance score.

For example, performing a vector similarity search with the string:

I was wondering, how many sounds can foxes make?

Yields:

Distance: 0.4075, Content: Foxes can make over 40 different sounds

Wrapping Up

And that’s really it! As your projects grow, you can effortlessly manage millions of vectors while maintaining swift query times, thanks to Postgres and pgvector’s efficient indexing. Whether you’re developing a cutting-edge recommendation system or diving into novel AI applications, this combination offers a flexible and powerful platform to realize your ideas.