RAG with Postgres + pgvector

CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS pg_trgm;

CREATE TABLE documents (
  id        BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
  content   TEXT NOT NULL,
  metadata  JSONB DEFAULT '{}',
  embedding VECTOR(1536) NOT NULL
);

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

import { OpenAI } from 'openai';

const openai = new OpenAI();

async function ingest(raw: string) {
  const { data } = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: raw,
  });

  await sql`
    INSERT INTO documents (content, embedding)
    VALUES (${raw}, ${JSON.stringify(data[0].embedding)}::vector)
  `;
}

async function search(query: string, k = 5) {
  const { data } = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: query,
  });

  const results = await sql`
    SELECT
      id,
      content,
      metadata,
      1 - (embedding <=> ${JSON.stringify(data[0].embedding)}::vector) AS similarity
    FROM documents
    ORDER BY embedding <=> ${JSON.stringify(data[0].embedding)}::vector
    LIMIT ${k}
  `;

  return results;
}

function buildPrompt(query: string, docs: { content: string }[]) {
  const context = docs.map((d, i) => `[doc ${i + 1}] ${d.content}`).join('\n\n');

  return `You are a helpful assistant. Answer using ONLY the context below.
If the answer is not in the context, say "I don't know."

Context:
${context}

Question: ${query}`;
}