Vector search on Azure Database for PostgreSQL with the Azure AI extension

This month Microsoft announced the public preview of the Azure AI extension on Azure Database for PostgreSQL Flexible Server which simplifies the integration of various Azure AI Services, including Azure OpenAI and Azure AI Language, with your data. Through the azure_ai extension, you can generate, store, index, and query embeddings using SQL code.

In this article, you will learn how to:

Enable the azure_ai extension.
Generate vector embeddings for text through Azure OpenAI within your database.
Create a table for storing and searching for similar vectors.

For a quick introduction to the concepts of vector similarity search and vector embeddings, you can read my post Building a vector similarity search app with Azure AI Vision and PostgreSQL.

Prerequisites

To proceed with this tutorial, ensure that you have the following prerequisites installed and configured:

An Azure subscription: If you don’t have an active subscription, you can sign up for either an Azure free account or an Azure for Students account.
An Azure OpenAI resource with the model text-embedding-ada-002 deployed: You can request access to the Azure OpenAI service by filling out the application form.
An Azure Database for PostgreSQL Flexible Server instance and a database: If you don’t have a resource, you can use the Azure portal for creating one.
Python 3.10.x

Activate the Azure AI extension

Before installing an extension in Azure Database for PostgreSQL Flexible Server, you should add it in the list of the allowed extensions.

Using the Azure portal, navigate to your Azure Database for PostgreSQL resource.
On the left pane, select Server parameters.
Search for the azure.etensions parameter and add the AZURE_AI and VECTOR extensions to the allowlist.

It is recommended to install both the azure_ai and the pgvector (vector) extensions because they are commonly used together.

Then, you can install the extension, by connecting to your database and running the CREATE EXTENSION command from the psql command prompt:

1
2
CREATE EXTENSION azure_ai;
CREATE EXTENSION vector;

Configure the Azure AI extension

The azure_ai extension includes functions that allow you to call the Azure OpenAI service and generate embeddings within your database. In order to use this functionality, you should first specify the endpoint and the key of your Azure OpenAI resource.

The azure_ai.set_setting() function lets you set the necessary configuration values for Azure AI services. Execute the following commands from the psql command prompt to specify the endpoint and the key of your Azure OpenAI resource:

1
2
SELECT azure_ai.set_setting('azure_openai.endpoint','{endpoint}');
SELECT azure_ai.set_setting('azure_openai.subscription_key', '{api-key}');

Get started with the Azure AI extension

The code provided illustrates the core features of the Azure AI extension. Assume that you want to create a database to store blog post details, such as titles and abstracts. Your goal is to generate vector embeddings from the text data provided to search for similar articles more effectively.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
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
import os
from dotenv import load_dotenv
import psycopg2

load_dotenv()
host = os.getenv("POSTGRESQL_HOST")
database_name = os.getenv("POSTGRESQL_DB_NAME")
username = os.getenv("POSTGRESQL_USERNAME")
password = os.getenv("POSTGRESQL_PASSWORD")
sslmode = "require"
table_name = "blogposts"

conn_string = f"host={host} user={username} dbname={database_name} password={password} sslmode={sslmode}"
conn = psycopg2.connect(conn_string)
print("Connection created successfully!")

# Create a cursor object
cursor = conn.cursor()

# Create a table
cursor.execute(f"DROP TABLE IF EXISTS {table_name};")
cursor.execute(
    f"CREATE TABLE {table_name} ("
    "id int PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,"
    "title text,"
    "abstract text,"
    "vector VECTOR(1536) GENERATED ALWAYS AS ("
    "azure_openai.create_embeddings('text-embedding-ada-002', title || abstract) ::vector"
    ") STORED"
    ");"
)
print("Finished creating table")

# Insert data
data = [
    (
        "Image Vector Similarity Search with Azure Computer Vision and PostgreSQL",
        "Understand how vector similarity search works and build an image retrieval "
        "system using Computer Vision and Azure Cosmos DB for PostgreSQL."
    ),
    (
        "Extract text from images with Azure Computer Vision 4.0 Read OCR",
        "This post will take you through the newest Read OCR API of Azure "
        "Computer Vision, which is used for extracting text from images."
    ),
    (
        "Automate customer reviews processing with Form Recognizer and Azure OpenAI",
        "In this article, you will build a document processing solution using "
        "Form Recognizer, Logic Apps, and Azure OpenAI."
    ),
    (
        "Automate document processing with Form Recognizer and Logic Apps",
        "In this article, you will build an end-to-end document processing "
        "solution using Form Recognizer, Logic Apps, and Power BI."
    ),
    (
        "Explore Azure Computer Vision 4.0 (Florence model)",
        "In this post, you will explore the latest features of Azure Computer "
        "Vision and create a basic image analysis app."
    )
]

cursor.executemany(f"INSERT INTO {table_name} (title, abstract) VALUES (%s,%s)", data)
conn.commit()

# Print the data inserted on the table
cursor.execute(f"SELECT * from {table_name}")
rows = cursor.fetchall()
print("Data:")
print("*" * 10)
for row in rows:
    print(f"Post {row[0]}: {row[1]}")
    print(f"Summary: {row[2]}")
    print(f"Vector {row[3][0:15]}...]")
    print("*" * 10)

# Search for "How to extract text from documents using Azure services"
search_query = "How to extract text from documents using Azure services"
cursor.execute(
    f"SELECT title, abstract FROM {table_name} ORDER BY vector <#> "
    f"azure_openai.create_embeddings('text-embedding-ada-002', '{search_query}')::vector"
)

# Display the results
rows = cursor.fetchall()
print(f"Search results for: \"{search_query}\"")
print("*" * 10)
for row in rows:
    print(f"Title: {row[0]}")
    print(f"Summary: {row[1]}")
    print("*" * 10)

conn.commit()
cursor.close()
conn.close()

In the following sections, we’ll delve into creating a table to store vector embeddings, generating embeddings with Azure OpenAI, and searching for similar vectors.

Create a table to store vectors

The pgvector extension introduces a data type called VECTOR for indicating that a table column holds vector embeddings. When creating the column, it’s essential to specify the dimension of the vectors. In our scenario, the text-embedding-ada-002 generates 1536-dimensional vectors.

Additionally, we want to specify that the vector column is automatically populated with data upon the insertion of a new record into the table. The vector embeddings are generated by concatenating the title and abstract fields and calling the text-embedding-ada-002 model.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
cursor.execute(
    f"CREATE TABLE {table_name} ("
    "id int PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,"
    "title text,"
    "abstract text,"
    "vector VECTOR(1536) GENERATED ALWAYS AS ("
    "azure_openai.create_embeddings('text-embedding-ada-002', title || abstract) ::vector"
    ") STORED"
    ");"
)

Insert data into the table

We will insert information (title and abstract) about 5 blog posts into the table. Notice that the vector embeddings are created automatically.

1
cursor.executemany(f"INSERT INTO {table_name} (title, abstract) VALUES (%s,%s)", data)

Search for similar vectors

After populating the table with vector data, you can search for blog posts that are most similar to a text prompt. To calculate similarity and retrieve images, you will use SELECT statements and the built-it vector operators of the PostgreSQL database. To calculate the vector embedding of the search terms, you will use the azure_ai extension.

The following query computes the negative inner product (<#>) between the vector of the search_query and the vectors stored in the table, sorts the results by the calculated distance, and returns the title and the abstract of the posts.

1
2
3
4
cursor.execute(
    f"SELECT title, abstract FROM {table_name} ORDER BY vector <#> "
    f"azure_openai.create_embeddings('text-embedding-ada-002', '{search_query}')::vector"
)

Summary and next steps

In this post, we explored the basics of the Azure AI extension in Azure Database for PostgreSQL Flexible Server. Here are some resources to learn more about vector similarity search in PostgreSQL: