
postgres-ai
Azure PostgreSQL AI Integration - Vector search, RAG pipelines, and agent memory with pgvector and a
PostgreSQL AI Integration Skill
Folder Contents
| File | Type | Description |
|---|---|---|
SKILL.md | Documentation | Main skill documentation with Azure PostgreSQL setup, pgvector configuration, and troubleshooting guide |
PRD.md | Documentation | Product Requirements Document for the skill |
.env.sample | Configuration | Sample environment variables for PostgreSQL and Azure OpenAI |
requirements.txt | Dependencies | Python package dependencies (psycopg2, pgvector, langchain-postgres, langgraph) |
| scripts/ | ||
scripts/__init__.py | Module | Package initializer with exports |
scripts/postgres_client.py | Client | Core PostgreSQL client with pgvector support, connection management, and vector operations |
scripts/vector_search.py | Search | Vector similarity search with HNSW/IVFFlat indexes, cosine/L2/inner-product distance operators |
scripts/rag_pipeline.py | Pipeline | RAG pipeline with LangChain PGVector integration and retriever configuration |
scripts/agent_memory.py | Memory | LangGraph 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
- Navigate to your PostgreSQL server in Azure Portal
- Go to Settings > Networking
- Under Public access, select Allow public access from any Azure service within Azure to this server
- Under Firewall rules, click + Add current client IP address
- Optionally add
0.0.0.0to255.255.255.255to allow all IPs (not recommended for production) - 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
- Navigate to your PostgreSQL server
- Go to Settings > Server parameters
- Search for
azure.extensions - Add
VECTORto the comma-separated list of allowed extensions - 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:
- Password authentication - Traditional username/password
- 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 Tried | Result |
|---|---|
samer.elhousseini@microsoft.com | Failed with JWT token (token expired) |
samer | Failed with password "samer" |
samer@pgsql-test | Failed |
pgadmin | Failed |
admin | Failed |
samer | SUCCESS 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
| Error | Cause | Solution |
|---|---|---|
password authentication failed | Wrong credentials or expired token | Use admin username/password, not Entra ID |
connection refused | Firewall blocking | Add your IP to firewall rules in Azure Portal |
extension "vector" is not allow-listed | Extension not enabled | Add VECTOR to azure.extensions parameter |
cannot have more than 2000 dimensions | HNSW dimension limit | Use dimensions=1536 in embedding API call |
SSL required | Missing sslmode | Add 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
| Script | Purpose |
|---|---|
postgres_client.py | Core PostgreSQL client with pgvector support |
vector_search.py | Vector similarity search operations |
rag_pipeline.py | RAG with LangChain and PostgreSQL |
agent_memory.py | LangGraph 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
| Scenario | Recommended Index |
|---|---|
| Fast build, limited memory | IVFFlat |
| Best query performance | HNSW |
| Very large datasets (100M+) | DiskANN |
| Empty table at start | HNSW (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
| Operator | Distance Type | Use Case |
|---|---|---|
<-> | L2 (Euclidean) | Image similarity |
<#> | Inner product | Normalized vectors |
<=> | Cosine | Text 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