RAG: a complete overview

// Split a document into pieces of up to ~800 characters along paragraph borders.
// Every piece keeps its source and heading — the answer will cite them.
export type Chunk = { source: string; heading: string; text: string };

export function chunk(source: string, markdown: string, maxChars = 800): Chunk[] {
  const chunks: Chunk[] = [];
  let heading = '';
  let buffer: string[] = [];

  const flush = () => {
    const text = buffer.join('\n\n').trim();
    if (text) chunks.push({ source, heading, text });
    buffer = [];
  };

  for (const block of markdown.split(/\n{2,}/)) {
    if (block.startsWith('#')) {
      flush(); // a new heading always starts a new piece
      heading = block.replace(/^#+\s*/, '');
      continue;
    }
    if (buffer.join('\n\n').length + block.length > maxChars) flush();
    buffer.push(block);
  }
  flush();
  return chunks;
}

// Indexing: every piece gets a vector and goes into PostgreSQL with pgvector.
// CREATE TABLE chunks (id bigserial PRIMARY KEY, source text, heading text,
//                      text text, embedding vector(384));
import pg from 'pg';
import { chunk } from './chunk.ts';

const db = new pg.Pool({ connectionString: process.env.DATABASE_URL });

// Any OpenAI-compatible embeddings API — a cloud provider or a local model
export async function embed(texts: string[]): Promise<number[][]> {
  const res = await fetch(`${process.env.EMBEDDINGS_URL}/v1/embeddings`, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      Authorization: `Bearer ${process.env.EMBEDDINGS_KEY}`,
    },
    body: JSON.stringify({ model: process.env.EMBEDDINGS_MODEL, input: texts }),
  });
  if (!res.ok) throw new Error(`embeddings: HTTP ${res.status}`);
  const { data } = (await res.json()) as { data: { embedding: number[] }[] };
  return data.map((d) => d.embedding);
}

export async function indexDocument(source: string, markdown: string): Promise<number> {
  const pieces = chunk(source, markdown);
  const vectors = await embed(pieces.map((p) => `${p.heading}\n${p.text}`));
  // One statement: the old pieces of the document are replaced atomically
  await db.query(
    `WITH old AS (DELETE FROM chunks WHERE source = $1)
     INSERT INTO chunks (source, heading, text, embedding)
     SELECT $1, h, t, e::vector FROM unnest($2::text[], $3::text[], $4::text[]) AS u(h, t, e)`,
    [source, pieces.map((p) => p.heading), pieces.map((p) => p.text), vectors.map((v) => JSON.stringify(v))],
  );
  return pieces.length;
}

// The answer: find the closest pieces, give them to the model numbered,
// require citations — and do not call the model at all if nothing was found
import pg from 'pg';
import { embed } from './index.ts';

const db = new pg.Pool({ connectionString: process.env.DATABASE_URL });
const API = process.env.ANTHROPIC_BASE_URL ?? 'https://api.anthropic.com';

export async function ask(question: string) {
  const [q] = await embed([question]);
  const { rows } = await db.query(
    `SELECT source, heading, text, 1 - (embedding <=> $1) AS score
       FROM chunks ORDER BY embedding <=> $1 LIMIT 4`,
    [JSON.stringify(q)],
  );
  // The threshold is chosen on your own questions; below it — not an answer
  const found = rows.filter((r) => r.score > 0.5);
  if (found.length === 0) return { answer: null, sources: [] }; // hand over to a person

  const context = found.map((r, i) => `[${i + 1}] ${r.heading}\n${r.text}`).join('\n\n');
  const res = await fetch(`${API}/v1/messages`, {
    method: 'POST',
    headers: {
      'content-type': 'application/json',
      'x-api-key': process.env.ANTHROPIC_API_KEY ?? '',
      'anthropic-version': '2023-06-01',
    },
    body: JSON.stringify({
      model: 'claude-sonnet-5',
      max_tokens: 600,
      system:
        'Answer only from the sources below and cite them as [1], [2]. ' +
        'If the sources do not contain the answer, say so.\n\n' + context,
      messages: [{ role: 'user', content: question }],
    }),
  });
  if (!res.ok) throw new Error(`model: HTTP ${res.status}`);
  const data = (await res.json()) as { content: { type: string; text: string }[] };
  return { answer: data.content[0].text, sources: found.map((r) => r.source) };
}