samelhousseini avatar

postgres-ai

Azure PostgreSQL AI Integration - Vector search, RAG pipelines, and agent memory with pgvector and a

by samelhousseini|Open Source

PostgreSQL AI Integration Skill

Folder Contents

FileTypeDescription
SKILL.mdDocumentationMain skill documentation with Azure PostgreSQL setup, pgvector configuration, and troubleshooting guide
PRD.mdDocumentationProduct Requirements Document for the skill
.env.sampleConfigurationSample environment variables for PostgreSQL and Azure OpenAI
requirements.txtDependenciesPython package dependencies (psycopg2, pgvector, langchain-postgres, langgraph)
scripts/
scripts/__init__.pyModulePackage initializer with exports
scripts/postgres_client.pyClientCore PostgreSQL client with pgvector support, connection management, and vector operations
scripts/vector_search.pySearchVector similarity search with HNSW/IVFFlat indexes, cosine/L2/inner-product distance operators
scripts/rag_pipeline.pyPipelineRAG pipeline with LangChain PGVector integration and retriever configuration
scripts/agent_memory.pyMemoryLangGraph checkpointing (PostgresSaver) and long-term memory (PostgresStore) for agent state

CRITICAL: No Mock Functionality

ALL implementations must be real and fully connected to Azure PostgreSQL and Azure OpenAI services.

  • NO mock database connections
  • NO fake embeddings or search results
  • NO simulated agent memory
  • NO placeholder vectors or retrieval results
  • NO hardcoded similarity scores

Everything must connect to real Azure PostgreSQL with pgvector/azure_ai and return real results.

If any functionality cannot be implemented with real connections (e.g., missing credentials, extensions not enabled), STOP and confirm with the user before proceeding.


Overview

Azure Database for PostgreSQL Flexible Server provides powerful AI capabilities through extensions:

  • pgvector: Vector similarity search with up to 16,000 dimensions
  • azure_ai: In-database Azure OpenAI calls for embeddings and completions

This skill covers vector search fundamentals, RAG pipeline implementation, and persistent agent memory.


Azure PostgreSQL Setup & Troubleshooting Guide

This section documents the complete setup process and common issues encountered when configuring Azure PostgreSQL for AI workloads.

Step 1: Create Azure PostgreSQL Flexible Server

# Create resource group
az group create --name rg-postgres-ai --location eastus

# Create PostgreSQL Flexible Server
az postgres flexible-server create \
  --name your-server-name \
  --resource-group rg-postgres-ai \
  --location eastus \
  --admin-user samer \
  --admin-password "YourSecurePassword123!" \
  --sku-name Standard_B2s \
  --tier Burstable \
  --storage-size 32 \
  --version 16 \
  --public-access 0.0.0.0 \
  --yes

Step 2: Configure Network Access (Make Server Publicly Accessible)

By default, Azure PostgreSQL may block external connections. To enable public access:

Option A: Azure Portal

  1. Navigate to your PostgreSQL server in Azure Portal
  2. Go to Settings > Networking
  3. Under Public access, select Allow public access from any Azure service within Azure to this server
  4. Under Firewall rules, click + Add current client IP address
  5. Optionally add 0.0.0.0 to 255.255.255.255 to allow all IPs (not recommended for production)
  6. Click Save

Option B: Azure CLI

# Allow public access
az postgres flexible-server update \
  --name your-server-name \
  --resource-group rg-postgres-ai \
  --public-access Enabled

# Add firewall rule for your IP
az postgres flexible-server firewall-rule create \
  --name allow-my-ip \
  --resource-group rg-postgres-ai \
  --server-name your-server-name \
  --start-ip-address YOUR_IP \
  --end-ip-address YOUR_IP

# Or allow all IPs (development only!)
az postgres flexible-server firewall-rule create \
  --name allow-all \
  --resource-group rg-postgres-ai \
  --server-name your-server-name \
  --start-ip-address 0.0.0.0 \
  --end-ip-address 255.255.255.255

Step 3: Configure pgvector Extension

The vector extension must be explicitly allowed before it can be created.

Option A: Azure Portal

  1. Navigate to your PostgreSQL server
  2. Go to Settings > Server parameters
  3. Search for azure.extensions
  4. Add VECTOR to the comma-separated list of allowed extensions
  5. Click Save (server may restart)

Option B: Azure CLI

az postgres flexible-server parameter set \
  --resource-group rg-postgres-ai \
  --server-name your-server-name \
  --name azure.extensions \
  --value "vector,azure_ai,pg_cron"

Then create the extension in your database:

CREATE EXTENSION IF NOT EXISTS vector;

Step 4: Authentication Troubleshooting

We encountered several authentication issues. Here's what we learned:

Issue 1: Microsoft Entra ID (Azure AD) Token Authentication

Azure PostgreSQL supports two authentication methods:

  1. Password authentication - Traditional username/password
  2. Microsoft Entra ID authentication - Uses JWT tokens

Symptoms of Entra ID issues:

FATAL: password authentication failed for user "email@domain.com"

If using Entra ID tokens:

  • Tokens expire quickly (typically 1 hour)
  • Username must be your full email: user@domain.com
  • Password is the JWT access token (very long string starting with eyJ...)
  • Get a fresh token:
    az account get-access-token --resource-type oss-rdbms --query accessToken -o tsv
    

Common mistake: Using an expired JWT token. Tokens have an exp claim that determines expiration.

Issue 2: Wrong Username Format

We tried multiple username formats before finding the correct one:

Username TriedResult
samer.elhousseini@microsoft.comFailed with JWT token (token expired)
samerFailed with password "samer"
samer@pgsql-testFailed
pgadminFailed
adminFailed
samerSUCCESS with password Mobility12#

Key insight: The admin username is the one set during server creation (--admin-user), NOT your Azure email address.

Issue 3: Password vs Token Confusion

The .env file had two password-related variables that caused confusion:

# WRONG: Using Entra ID email with password
PGUSER=samer.elhousseini@microsoft.com
PGPASSWORD=eyJ0eXAiOiJKV1Q...  # Expired JWT token

# CORRECT: Using admin credentials set during server creation
PGUSER=samer
PGPASSWORD=Mobility12#

Resolution: Use the PostgreSQL admin credentials (username and password) that were specified when creating the server, not Entra ID credentials.

Step 5: Verify Connection

Test your connection before proceeding:

import psycopg2

conn = psycopg2.connect(
    host="your-server.postgres.database.azure.com",
    database="postgres",
    user="samer",          # Admin username from server creation
    password="Mobility12#", # Admin password from server creation
    port=5432,
    sslmode="require"      # SSL is mandatory for Azure PostgreSQL
)
print("Connected successfully!")
conn.close()

Step 6: HNSW Index Dimension Limit

Issue encountered:

ERROR: column cannot have more than 2000 dimensions for hnsw index

Azure PostgreSQL's pgvector implementation limits HNSW indexes to 2000 dimensions. However, text-embedding-3-large produces 3072 dimensions by default.

Solution: Request reduced dimensions from the embedding API:

from openai import AzureOpenAI

client = AzureOpenAI(...)

response = client.embeddings.create(
    model="text-embedding-3-large",
    input="Your text here",
    dimensions=1536  # Explicitly request 1536 dimensions
)

Alternative: Use IVFFlat index instead (no dimension limit, but requires training data):

CREATE INDEX ON documents USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

Quick Reference: Common Errors and Solutions

ErrorCauseSolution
password authentication failedWrong credentials or expired tokenUse admin username/password, not Entra ID
connection refusedFirewall blockingAdd your IP to firewall rules in Azure Portal
extension "vector" is not allow-listedExtension not enabledAdd VECTOR to azure.extensions parameter
cannot have more than 2000 dimensionsHNSW dimension limitUse dimensions=1536 in embedding API call
SSL requiredMissing sslmodeAdd sslmode=require to connection string

Final Working Configuration

# .env file
PGHOST=pgsql-test.postgres.database.azure.com
PGUSER=samer
PGPORT=5432
PGDATABASE=postgres
PGPASSWORD=Mobility12#

# Azure OpenAI for embeddings
AZURE_OPENAI_ENDPOINT=https://dev-aoai-swedencentral.openai.azure.com
AZURE_OPENAI_API_KEY=your-api-key
AZURE_OPENAI_EMBEDDING_DEPLOYMENT=text-embedding-3-large

Building Blocks

ScriptPurpose
postgres_client.pyCore PostgreSQL client with pgvector support
vector_search.pyVector similarity search operations
rag_pipeline.pyRAG with LangChain and PostgreSQL
agent_memory.pyLangGraph checkpointing and memory stores

Environment Variables

# PostgreSQL Connection
PGHOST=your-server.postgres.database.azure.com
PGDATABASE=your-database
PGUSER=your-username
PGPASSWORD=your-password
PGPORT=5432

# Azure OpenAI (for embeddings)
AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com/
AZURE_OPENAI_API_KEY=your-api-key
AZURE_OPENAI_EMBEDDING_DEPLOYMENT=text-embedding-ada-002
AZURE_OPENAI_CHAT_DEPLOYMENT=gpt-4o
AZURE_OPENAI_API_VERSION=2024-02-01

Core Capabilities

1. Vector Search with pgvector

# Create vector table (max 2000 dimensions for HNSW index in Azure PostgreSQL)
CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding vector(1536)  # Use dimensions=1536 with text-embedding-3-large
);

# Create HNSW index (best for OpenAI embeddings)
CREATE INDEX ON documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

# Similarity search
SELECT content, 1 - (embedding <=> query_vector) as similarity
FROM documents
ORDER BY embedding <=> query_vector
LIMIT 5;

2. Index Selection Guide

ScenarioRecommended Index
Fast build, limited memoryIVFFlat
Best query performanceHNSW
Very large datasets (100M+)DiskANN
Empty table at startHNSW (no training needed)

3. azure_ai Extension

Generate embeddings directly in SQL:

-- Configure Azure OpenAI
SELECT azure_ai.set_setting('azure_openai.endpoint', 'https://...');
SELECT azure_ai.set_setting('azure_openai.subscription_key', 'key');

-- Generate embeddings
UPDATE documents
SET embedding = azure_openai.create_embeddings(
    'text-embedding-ada-002',
    content
)::vector
WHERE embedding IS NULL;

4. RAG Pipeline with LangChain

from langchain_postgres import PGVector
from langchain_openai import AzureOpenAIEmbeddings

vector_store = PGVector(
    collection_name="documents",
    connection=DATABASE_URL,
    embeddings=embeddings,
)

retriever = vector_store.as_retriever(search_kwargs={"k": 3})

5. Agent Memory with LangGraph

from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.store.postgres import PostgresStore

# Session checkpointing
with PostgresSaver.from_conn_string(DATABASE_URL) as checkpointer:
    checkpointer.setup()
    graph = builder.compile(checkpointer=checkpointer)

# Long-term memory
with PostgresStore.from_conn_string(DATABASE_URL) as store:
    store.setup()
    store.put(namespace=("users", "user123"), key="preferences", value={...})

Performance Tuning

-- IVFFlat: increase probes for better recall
SET ivfflat.probes = 10;  -- Default: 1

-- HNSW: increase ef_search for better recall
SET hnsw.ef_search = 100;  -- Default: 40

-- Parallel workers for large scans
SET max_parallel_workers_per_gather = 4;

Distance Operators

OperatorDistance TypeUse Case
<->L2 (Euclidean)Image similarity
<#>Inner productNormalized vectors
<=>CosineText embeddings (OpenAI)

Common Patterns

Hybrid Search (Vector + Metadata)

SELECT content, 1 - (embedding <=> query_vec) as similarity
FROM documents
WHERE metadata->>'category' = 'tech'
ORDER BY embedding <=> query_vec
LIMIT 5;

Batch Embedding Generation

# Process in batches of 100
for batch in chunks(texts, 100):
    embeddings = openai.embeddings.create(model="...", input=batch)
    # Insert to database

Lessons Learned

HNSW Dimension Limit

Azure PostgreSQL limits HNSW indexes to 2000 dimensions. Use dimensions=1536 parameter with text-embedding-3-large:

response = openai.embeddings.create(
    model="text-embedding-3-large",
    input=text,
    dimensions=1536  # Truncate to fit HNSW limit
)

Connection String Format

Azure PostgreSQL requires SSL:

postgresql://user:password@host:5432/database?sslmode=require

pgvector Registration

Always register pgvector with psycopg2:

from pgvector.psycopg2 import register_vector
conn = psycopg2.connect(DATABASE_URL)
register_vector(conn)

Extension Enablement

Extensions must be enabled in Azure portal first:

az postgres flexible-server parameter set \
  --name azure.extensions \
  --value "vector,azure_ai"

LangGraph Checkpoint Tables

PostgresSaver.setup() creates required tables automatically. Call once before first use.

Dependencies

psycopg2-binary
asyncpg
pgvector
python-dotenv
openai
langchain
langchain-openai
langchain-postgres
langgraph
langgraph-checkpoint-postgres

References