pgvector_examples.py

"""
pgvector Embeddings - Code Examples
Azure Developer Associate (AI-200) Study Material

Examples for:
- Connecting to Azure PostgreSQL with pgvector
- Generating embeddings with Azure OpenAI
- Storing and querying vector embeddings
- Batch embedding pipelines

Prerequisites:
    pip install psycopg[binary] pgvector openai azure-identity numpy tenacity
"""

import os
import time
from typing import List

import numpy as np
import psycopg
from pgvector.psycopg import register_vector
from openai import AzureOpenAI
from azure.identity import DefaultAzureCredential
from tenacity import retry, stop_after_attempt, wait_exponential


# ============================================================================
# Section 1: Setup and Connection
# ============================================================================
# Connect to Azure Database for PostgreSQL Flexible Server with pgvector.
# pgvector must be enabled on the server: CREATE EXTENSION IF NOT EXISTS vector;

def get_connection() -> psycopg.Connection:
    """Connect to Azure PostgreSQL and register the pgvector type."""
    conn = psycopg.connect(
        host=os.environ["PGHOST"],           # e.g. myserver.postgres.database.azure.com
        port=os.environ.get("PGPORT", 5432),
        dbname=os.environ["PGDATABASE"],
        user=os.environ["PGUSER"],
        password=os.environ["PGPASSWORD"],
        sslmode="require",                   # Always require SSL for Azure
    )
    # Register the vector type so psycopg can encode/decode numpy arrays
    # as PostgreSQL vector values automatically.
    register_vector(conn)
    return conn


def create_documents_table(conn: psycopg.Connection, dimensions: int = 1536):
    """Create a table with a vector column for storing embeddings.

    Args:
        conn: Active psycopg connection with pgvector registered.
        dimensions: Embedding dimensions (1536 for text-embedding-3-small default,
                    3072 for text-embedding-3-large).
    """
    conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
    conn.execute(f"""
        CREATE TABLE IF NOT EXISTS documents (
            id          BIGSERIAL PRIMARY KEY,
            content     TEXT NOT NULL,
            metadata    JSONB DEFAULT '{{}}',
            embedding   vector({dimensions}),
            created_at  TIMESTAMPTZ DEFAULT now()
        )
    """)
    # Create an IVFFlat index for approximate nearest-neighbor search.
    # lists = sqrt(row_count) is a reasonable starting point; tune for your data.
    conn.execute("""
        CREATE INDEX IF NOT EXISTS idx_documents_embedding
        ON documents
        USING ivfflat (embedding vector_cosine_ops)
        WITH (lists = 100)
    """)
    # HNSW index alternative -- better recall, more memory, slower builds:
    # CREATE INDEX IF NOT EXISTS idx_documents_embedding_hnsw
    # ON documents
    # USING hnsw (embedding vector_cosine_ops)
    # WITH (m = 16, ef_construction = 64);
    conn.commit()


# ============================================================================
# Section 2: Generate Embeddings with Azure OpenAI
# ============================================================================
# Uses text-embedding-3-small (default 1536 dims) or text-embedding-3-large.
# The `dimensions` parameter lets you request reduced-dimension output.

def get_openai_client() -> AzureOpenAI:
    """Create an Azure OpenAI client from environment variables."""
    return AzureOpenAI(
        api_key=os.environ["AZURE_OPENAI_API_KEY"],
        api_version=os.environ.get("AZURE_OPENAI_API_VERSION", "2024-06-01"),
        azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
        # e.g. https://myresource.openai.azure.com/
    )


def generate_embedding(client: AzureOpenAI, text: str,
                        model: str = "text-embedding-3-small",
                        dimensions: int = 1536) -> np.ndarray:
    """Generate a single embedding vector for the given text.

    Args:
        client: Azure OpenAI client.
        text: Input text to embed (max ~8191 tokens for embedding-3 models).
        model: Deployment name of the embedding model.
        dimensions: Output dimensions (256, 512, 1024, 1536, or 3072).

    Returns:
        numpy array of shape (dimensions,).
    """
    response = client.embeddings.create(
        input=text,
        model=model,
        dimensions=dimensions,
    )
    return np.array(response.data[0].embedding, dtype=np.float32)


def generate_embeddings_batch(client: AzureOpenAI, texts: List[str],
                              model: str = "text-embedding-3-small",
                              dimensions: int = 1536) -> List[np.ndarray]:
    """Generate embeddings for multiple texts in a single API call.

    Azure OpenAI supports batching up to 2048 texts per request.
    For large batches, chunk into groups of ~100 to stay within token limits.

    Args:
        client: Azure OpenAI client.
        texts: List of input texts.
        model: Deployment name.
        dimensions: Output vector dimensions.

    Returns:
        List of numpy arrays, one per input text (order preserved).
    """
    response = client.embeddings.create(
        input=texts,
        model=model,
        dimensions=dimensions,
    )
    # The API returns results sorted by index, but sort explicitly to be safe.
    sorted_data = sorted(response.data, key=lambda x: x.index)
    return [np.array(item.embedding, dtype=np.float32) for item in sorted_data]


# ============================================================================
# Section 3: Store Embeddings
# ============================================================================

def insert_single_document(conn: psycopg.Connection, content: str,
                           embedding: np.ndarray, metadata: dict | None = None):
    """Insert a single document with its embedding.

    pgvector's psycopg adapter automatically converts numpy arrays to the
    PostgreSQL vector format when registered via register_vector().
    """
    conn.execute(
        """
        INSERT INTO documents (content, embedding, metadata)
        VALUES (%s, %s, %s)
        """,
        (content, embedding, psycopg.types.json.Json(metadata or {})),
    )
    conn.commit()


def insert_documents_batch(conn: psycopg.Connection,
                           documents: List[dict]):
    """Insert multiple documents using executemany for moderate batch sizes.

    Args:
        documents: List of dicts with keys: content, embedding, metadata.
    """
    with conn.cursor() as cur:
        cur.executemany(
            """
            INSERT INTO documents (content, embedding, metadata)
            VALUES (%(content)s, %(embedding)s, %(metadata)s)
            """,
            [
                {
                    "content": doc["content"],
                    "embedding": doc["embedding"],
                    "metadata": psycopg.types.json.Json(doc.get("metadata", {})),
                }
                for doc in documents
            ],
        )
    conn.commit()


def bulk_insert_with_copy(conn: psycopg.Connection,
                          documents: List[dict]):
    """Bulk insert using PostgreSQL COPY for maximum throughput.

    COPY is significantly faster than INSERT for large batches (10k+ rows).
    """
    with conn.cursor() as cur:
        with cur.copy(
            "COPY documents (content, embedding, metadata) FROM STDIN"
        ) as copy:
            for doc in documents:
                # Convert numpy array to pgvector literal: '[0.1,0.2,...]'
                vec_literal = "[" + ",".join(
                    str(float(v)) for v in doc["embedding"]
                ) + "]"
                metadata_str = psycopg.types.json.Json(
                    doc.get("metadata", {})
                ).dumps(doc.get("metadata", {}))
                copy.write_row((doc["content"], vec_literal, metadata_str))
    conn.commit()


# ============================================================================
# Section 4: Similarity Search
# ============================================================================
# pgvector supports three distance operators:
#   <=>  Cosine distance   (1 - cosine_similarity; lower = more similar)
#   <->  L2 distance       (Euclidean; lower = more similar)
#   <#>  Inner product      (negative inner product; lower = more similar)

def search_by_cosine(conn: psycopg.Connection, query_embedding: np.ndarray,
                     limit: int = 5) -> list:
    """Find the most similar documents using cosine distance.

    This is the most common choice for text embeddings because it is
    insensitive to vector magnitude (focuses on direction).
    """
    results = conn.execute(
        """
        SELECT id, content, metadata,
               1 - (embedding <=> %s) AS similarity
        FROM documents
        ORDER BY embedding <=> %s
        LIMIT %s
        """,
        (query_embedding, query_embedding, limit),
    ).fetchall()
    return [
        {"id": r[0], "content": r[1], "metadata": r[2], "similarity": r[3]}
        for r in results
    ]


def search_by_l2(conn: psycopg.Connection, query_embedding: np.ndarray,
                 limit: int = 5) -> list:
    """Find the most similar documents using L2 (Euclidean) distance.

    L2 distance is magnitude-sensitive. Good when embeddings are normalized
    or when you want absolute distance rather than angle similarity.
    """
    results = conn.execute(
        """
        SELECT id, content, metadata,
               embedding <-> %s AS l2_distance
        FROM documents
        ORDER BY embedding <-> %s
        LIMIT %s
        """,
        (query_embedding, query_embedding, limit),
    ).fetchall()
    return [
        {"id": r[0], "content": r[1], "metadata": r[2], "l2_distance": r[3]}
        for r in results
    ]


def search_by_inner_product(conn: psycopg.Connection,
                            query_embedding: np.ndarray,
                            limit: int = 5) -> list:
    """Find the most similar documents using negative inner product.

    pgvector uses negative inner product (<#>) so that ORDER BY ASC gives
    highest similarity first. Best for normalized embeddings.
    """
    results = conn.execute(
        """
        SELECT id, content, metadata,
               (embedding <#> %s) * -1 AS inner_product
        FROM documents
        ORDER BY embedding <#> %s
        LIMIT %s
        """,
        (query_embedding, query_embedding, limit),
    ).fetchall()
    return [
        {"id": r[0], "content": r[1], "metadata": r[2], "inner_product": r[3]}
        for r in results
    ]


def search_with_filter(conn: psycopg.Connection, query_embedding: np.ndarray,
                       category: str, min_similarity: float = 0.7,
                       limit: int = 5) -> list:
    """Filtered similarity search: combine WHERE clauses with vector ordering.

    This pattern is very common in RAG: filter by metadata first, then rank
    remaining rows by vector similarity.
    """
    results = conn.execute(
        """
        SELECT id, content, metadata,
               1 - (embedding <=> %s) AS similarity
        FROM documents
        WHERE metadata->>'category' = %s
          AND 1 - (embedding <=> %s) >= %s
        ORDER BY embedding <=> %s
        LIMIT %s
        """,
        (query_embedding, category, query_embedding, min_similarity,
         query_embedding, limit),
    ).fetchall()
    return [
        {"id": r[0], "content": r[1], "metadata": r[2], "similarity": r[3]}
        for r in results
    ]


# ============================================================================
# Section 5: Embedding Pipeline Pattern
# ============================================================================
# End-to-end pipeline: chunk text, generate embeddings with rate limiting,
# bulk insert, and verify with spot-check queries.

def chunk_text(text: str, chunk_size: int = 500,
               overlap: int = 50) -> List[str]:
    """Split text into overlapping chunks by character count.

    For production, consider token-based chunking or sentence-boundary
    splitting (e.g., with langchain or tiktoken).
    """
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - overlap
    return chunks


@retry(
    stop=stop_after_attempt(5),
    wait=wait_exponential(multiplier=1, min=2, max=60),
)
def generate_embeddings_with_retry(client: AzureOpenAI, batch: List[str],
                                   model: str = "text-embedding-3-small",
                                   dimensions: int = 1536) -> List[np.ndarray]:
    """Generate embeddings with exponential backoff retry.

    Azure OpenAI has rate limits (tokens per minute). The tenacity @retry
    decorator handles 429 (rate limit) and transient 5xx errors automatically
    because the OpenAI SDK raises retryable exceptions.
    """
    return generate_embeddings_batch(client, batch, model, dimensions)


def run_embedding_pipeline(source_texts: List[dict],
                           batch_size: int = 50,
                           chunk_size: int = 500):
    """Full pipeline: chunk, embed, store, verify.

    Args:
        source_texts: List of dicts with 'content' and 'metadata' keys.
        batch_size: Number of chunks to embed per API call.
        chunk_size: Character count per chunk.
    """
    client = get_openai_client()
    conn = get_connection()

    # Step 1: Create table
    create_documents_table(conn, dimensions=1536)

    # Step 2: Chunk all documents
    all_chunks = []
    for doc in source_texts:
        chunks = chunk_text(doc["content"], chunk_size=chunk_size)
        for i, chunk in enumerate(chunks):
            all_chunks.append({
                "content": chunk,
                "metadata": {
                    **doc.get("metadata", {}),
                    "chunk_index": i,
                    "total_chunks": len(chunks),
                },
            })

    print(f"Processing {len(all_chunks)} chunks from {len(source_texts)} documents")

    # Step 3: Generate embeddings in batches with rate limiting
    for batch_start in range(0, len(all_chunks), batch_size):
        batch_end = min(batch_start + batch_size, len(all_chunks))
        batch = all_chunks[batch_start:batch_end]
        batch_texts = [chunk["content"] for chunk in batch]

        embeddings = generate_embeddings_with_retry(client, batch_texts)

        for chunk, embedding in zip(batch, embeddings):
            chunk["embedding"] = embedding

        # Step 4: Bulk insert each batch
        insert_documents_batch(conn, batch)

        print(f"  Inserted batch {batch_start}-{batch_end}")

        # Respect rate limits: pause between batches
        time.sleep(0.5)

    # Step 5: Verify with a spot-check query
    sample_text = all_chunks[0]["content"]
    sample_embedding = generate_embedding(client, sample_text)
    results = search_by_cosine(conn, sample_embedding, limit=3)

    print("\nSpot-check: top 3 results for first chunk")
    for r in results:
        print(f"  similarity={r['similarity']:.4f}  content={r['content'][:80]}...")

    conn.close()


# ============================================================================
# Section 6: Entra ID Authentication
# ============================================================================
# Use Azure Entra ID (formerly Azure AD) instead of passwords.
# Requires the PostgreSQL user to be an Entra principal and the server
# to have Entra authentication enabled.

def get_connection_with_entra() -> psycopg.Connection:
    """Connect to Azure PostgreSQL using a Microsoft Entra ID access token.

    DefaultAzureCredential tries managed identity first (when running on Azure),
    then falls back to Azure CLI, environment variables, etc.
    """
    credential = DefaultAzureCredential()

    # Request a token for the Azure OSS RDBMS audience
    token = credential.get_token(
        "https://ossrdbms-aad.database.windows.net/.default"
    )

    conn = psycopg.connect(
        host=os.environ["PGHOST"],
        port=os.environ.get("PGPORT", 5432),
        dbname=os.environ["PGDATABASE"],
        user=os.environ["PGUSER"],       # Entra principal name
        password=token.token,            # Access token used as password
        sslmode="require",
    )
    register_vector(conn)
    return conn


def get_openai_client_with_managed_identity() -> AzureOpenAI:
    """Create an Azure OpenAI client using managed identity (no API key).

    This is the recommended approach when running on Azure compute
    (App Service, Container Apps, AKS, Functions, VMs with managed identity).
    """
    credential = DefaultAzureCredential()
    token = credential.get_token(
        "https://cognitiveservices.azure.com/.default"
    )

    return AzureOpenAI(
        api_key=token.token,   # Token used in place of API key
        api_version=os.environ.get("AZURE_OPENAI_API_VERSION", "2024-06-01"),
        azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    )


def entra_pipeline_example():
    """Demonstrate end-to-end flow using Entra ID for both PostgreSQL and OpenAI.

    No secrets required -- authentication uses managed identity tokens.
    """
    conn = get_connection_with_entra()
    client = get_openai_client_with_managed_identity()

    # Generate and store an embedding
    text = "Azure PostgreSQL Flexible Server supports pgvector for similarity search"
    embedding = generate_embedding(client, text)
    insert_single_document(conn, text, embedding, metadata={"source": "entra_example"})

    # Query it back
    query_embedding = generate_embedding(client, "vector database on Azure")
    results = search_by_cosine(conn, query_embedding, limit=3)
    for r in results:
        print(f"similarity={r['similarity']:.4f}  {r['content'][:80]}")

    conn.close()


# ============================================================================
# Usage Example
# ============================================================================

if __name__ == "__main__":
    # Set required environment variables before running:
    #   PGHOST, PGDATABASE, PGUSER, PGPASSWORD
    #   AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT

    sample_docs = [
        {
            "content": "Azure Database for PostgreSQL Flexible Server is a fully "
                       "managed database service. It supports the pgvector extension "
                       "for storing and querying vector embeddings, enabling semantic "
                       "search and RAG (Retrieval-Augmented Generation) patterns.",
            "metadata": {"category": "azure", "topic": "postgresql"},
        },
        {
            "content": "Azure OpenAI Service provides access to embedding models like "
                       "text-embedding-3-small and text-embedding-3-large. These models "
                       "convert text into dense vector representations that capture "
                       "semantic meaning, useful for search and recommendations.",
            "metadata": {"category": "azure", "topic": "openai"},
        },
    ]

    run_embedding_pipeline(sample_docs, batch_size=10, chunk_size=200)