pgvector Setup

SELECT extname, extversion FROM pg_extension WHERE extname = 'vector';

CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding VECTOR(1536)  -- OpenAI ada-002 dimensions
);

INSERT INTO documents (content, embedding)
VALUES (
    'Rivestack is a managed PostgreSQL service',
    '[0.1, 0.2, 0.3, ...]'  -- your embedding vector
);

import psycopg2
import openai

# Generate embedding
response = openai.embeddings.create(
    model="text-embedding-ada-002",
    input="Rivestack is a managed PostgreSQL service"
)
embedding = response.data[0].embedding

# Store in Rivestack
conn = psycopg2.connect("postgresql://user:pass@host:5432/mydb?sslmode=require")
cur = conn.cursor()
cur.execute(
    "INSERT INTO documents (content, embedding) VALUES (%s, %s)",
    ("Rivestack is a managed PostgreSQL service", str(embedding))
)
conn.commit()

SELECT id, content, embedding <=> '[0.1, 0.2, 0.3, ...]' AS distance
FROM documents
ORDER BY embedding <=> '[0.1, 0.2, 0.3, ...]'
LIMIT 5;

CREATE INDEX ON documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 200);

-- Per query — works on every port, including the pooled :6432 (see warning):
BEGIN;
SET LOCAL hnsw.ef_search = 100;
SELECT id FROM documents ORDER BY embedding <=> $1 LIMIT 10;
COMMIT;

-- Or set a persistent default for the whole database (applies to new connections):
ALTER DATABASE appdb SET hnsw.ef_search = 100;

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

SET ivfflat.probes = 10;  -- default is 1, higher = better recall

import psycopg2
import openai

def search(query: str, top_k: int = 5):
    # 1. Generate embedding for the query
    response = openai.embeddings.create(
        model="text-embedding-ada-002",
        input=query
    )
    query_embedding = response.data[0].embedding

    # 2. Search Rivestack for similar documents
    conn = psycopg2.connect("postgresql://user:pass@host:5432/mydb?sslmode=require")
    cur = conn.cursor()
    cur.execute("""
        SELECT content, embedding <=> %s::vector AS distance
        FROM documents
        ORDER BY embedding <=> %s::vector
        LIMIT %s
    """, (str(query_embedding), str(query_embedding), top_k))

    results = cur.fetchall()
    conn.close()
    return results

# 3. Use results as context for your LLM
results = search("How does Rivestack handle backups?")
context = "\n".join([row[0] for row in results])