RAG in TypeScript: Embeddings, pgvector, Hybrid Search, and Citations

export interface Chunk {
  source: string; // URL or path, used for citations
  title: string; // "Billing > Refunds"
  content: string;
}
 
const MAX_CHARS = 1500; // roughly 300 to 400 tokens of English
 
export function chunkMarkdown(source: string, markdown: string): Chunk[] {
  const chunks: Chunk[] = [];
  const headings: string[] = [];
  let buffer: string[] = [];
 
  const flush = () => {
    const content = buffer.join("\n\n").trim();
    if (content) chunks.push({ source, title: headings.filter(Boolean).join(" > "), content });
    buffer = [];
  };
 
  for (const block of markdown.split(/\n{2,}/)) {
    const heading = /^(#{1,3})\s+(.+)$/.exec(block.trim());
    if (heading) {
      flush();
      const level = heading[1]!.length;
      headings.length = level; // drop deeper headings from the previous section
      headings[level - 1] = heading[2]!.trim();
      continue;
    }
    if (buffer.join("\n\n").length + block.length > MAX_CHARS) flush();
    buffer.push(block.trim());
  }
  flush();
  return chunks;
}

const API = "https://api.voyageai.com/v1";
const MODEL = "voyage-4"; // 1024 dimensions by default
const BATCH = 128;
 
async function post<T>(path: string, body: unknown): Promise<T> {
  const res = await fetch(`${API}${path}`, {
    method: "POST",
    headers: { "content-type": "application/json", authorization: `Bearer ${process.env.VOYAGE_API_KEY}` },
    body: JSON.stringify(body),
  });
  if (!res.ok) throw new Error(`Voyage ${path} failed: ${res.status} ${await res.text()}`);
  return (await res.json()) as T;
}
 
export async function embed(texts: string[], inputType: "document" | "query"): Promise<number[][]> {
  const vectors: number[][] = [];
  for (let i = 0; i < texts.length; i += BATCH) {
    const { data } = await post<{ data: { embedding: number[]; index: number }[] }>("/embeddings", {
      input: texts.slice(i, i + BATCH),
      model: MODEL,
      input_type: inputType,
    });
    vectors.push(...data.sort((a, b) => a.index - b.index).map((d) => d.embedding));
  }
  return vectors;
}
 
export async function rerank(query: string, documents: string[], topK: number): Promise<{ index: number; score: number }[]> {
  const { data } = await post<{ data: { index: number; relevance_score: number }[] }>("/rerank", {
    query,
    documents,
    model: "rerank-3",
    top_k: topK,
  });
  return data.map((d) => ({ index: d.index, score: d.relevance_score }));
}

CREATE EXTENSION IF NOT EXISTS vector;
 
CREATE TABLE chunks (
  id         bigserial PRIMARY KEY,
  source     text NOT NULL,
  title      text NOT NULL,
  content    text NOT NULL,
  embedding  vector(1024) NOT NULL,
  -- Keyword search over title and content, kept up to date by Postgres.
  tsv        tsvector GENERATED ALWAYS AS (to_tsvector('english', title || ' ' || content)) STORED
);
 
CREATE INDEX chunks_embedding_idx ON chunks USING hnsw (embedding vector_cosine_ops);
CREATE INDEX chunks_tsv_idx ON chunks USING gin (tsv);

import pg from "pg";
 
export const pool = new pg.Pool({ connectionString: process.env.DATABASE_URL });
 
// pgvector accepts the text form "[0.1,0.2,...]".
export const toVector = (v: number[]) => JSON.stringify(v);

import { readFile, readdir } from "node:fs/promises";
import path from "node:path";
import { chunkMarkdown } from "./lib/chunk.ts";
import { pool, toVector } from "./lib/db.ts";
import { embed } from "./lib/voyage.ts";
 
const dir = process.argv[2] ?? "docs";
 
for (const file of (await readdir(dir)).filter((f) => f.endsWith(".md"))) {
  const chunks = chunkMarkdown(file, await readFile(path.join(dir, file), "utf8"));
  // Embedding the title with the content helps short chunks match.
  const vectors = await embed(chunks.map((c) => `${c.title}\n\n${c.content}`), "document");
 
  const client = await pool.connect();
  try {
    await client.query("BEGIN");
    await client.query("DELETE FROM chunks WHERE source = $1", [file]);
    for (const [i, c] of chunks.entries()) {
      await client.query("INSERT INTO chunks (source, title, content, embedding) VALUES ($1, $2, $3, $4)", [
        c.source, c.title, c.content, toVector(vectors[i]!),
      ]);
    }
    await client.query("COMMIT");
    console.log(`${file}: ${chunks.length} chunks`);
  } catch (error) {
    await client.query("ROLLBACK");
    throw error;
  } finally {
    client.release();
  }
}
await pool.end();

import { pool, toVector } from "./db.ts";
import { embed, rerank } from "./voyage.ts";
 
export interface Hit {
  id: number;
  source: string;
  title: string;
  content: string;
}
 
const CANDIDATES = 20;
 
export function reciprocalRankFusion(lists: Hit[][], k = 60): Hit[] {
  const scores = new Map<number, { hit: Hit; score: number }>();
  for (const list of lists) {
    list.forEach((hit, rank) => {
      const entry = scores.get(hit.id) ?? { hit, score: 0 };
      entry.score += 1 / (k + rank + 1);
      scores.set(hit.id, entry);
    });
  }
  return [...scores.values()].sort((a, b) => b.score - a.score).map((e) => e.hit);
}
 
export async function search(question: string, topK = 5): Promise<Hit[]> {
  const [queryVector] = await embed([question], "query");
 
  const [byMeaning, byKeyword] = await Promise.all([
    // <=> is cosine distance; ORDER BY distance uses the HNSW index.
    pool.query<Hit>(
      "SELECT id, source, title, content FROM chunks ORDER BY embedding <=> $1 LIMIT $2",
      [toVector(queryVector!), CANDIDATES],
    ),
    pool.query<Hit>(
      `SELECT id, source, title, content FROM chunks
       WHERE tsv @@ websearch_to_tsquery('english', $1)
       ORDER BY ts_rank_cd(tsv, websearch_to_tsquery('english', $1)) DESC LIMIT $2`,
      [question, CANDIDATES],
    ),
  ]);
 
  const fused = reciprocalRankFusion([byMeaning.rows, byKeyword.rows]).slice(0, CANDIDATES);
  if (fused.length === 0) return [];
 
  const ranked = await rerank(question, fused.map((h) => `${h.title}\n\n${h.content}`), topK);
  return ranked.map((r) => fused[r.index]!);
}

import Anthropic from "@anthropic-ai/sdk";
import { pool } from "./lib/db.ts";
import { search } from "./lib/search.ts";
 
const client = new Anthropic();
const question = process.argv[2] ?? "How long do refunds take?";
 
const hits = await search(question);
 
const message = await client.messages.create({
  model: "claude-opus-5-5",
  max_tokens: 4096,
  system:
    "Answer questions about Acme Billing using only the search results provided. If they do not contain the answer, say so plainly and suggest contacting support.",
  messages: [
    {
      role: "user",
      content: [
        ...hits.map((h): Anthropic.SearchResultBlockParam => ({
          type: "search_result",
          source: h.source,
          title: h.title,
          content: [{ type: "text", text: h.content }],
          citations: { enabled: true },
        })),
        { type: "text", text: question },
      ],
    },
  ],
});
 
// Print the answer with numbered references.
const sources: string[] = [];
let output = "";
for (const block of message.content) {
  if (block.type !== "text") continue;
  output += block.text;
  for (const c of block.citations ?? []) {
    if (c.type !== "search_result_location") continue;
    const ref = `${c.source}: ${c.title ?? ""}`;
    if (!sources.includes(ref)) sources.push(ref);
    output += ` [${sources.indexOf(ref) + 1}]`;
  }
}
console.log(output);
sources.forEach((s, i) => console.log(`[${i + 1}] ${s}`));
await pool.end();

import { pool } from "./lib/db.ts";
import { search } from "./lib/search.ts";
 
const cases = [
  { question: "How long do refunds take?", expected: "refunds.md" },
  { question: "What does E_CARD_DECLINED mean?", expected: "errors.md" },
  { question: "Can I pay by bank transfer?", expected: "payments.md" },
];
 
let found = 0;
for (const { question, expected } of cases) {
  const hits = await search(question, 5);
  const ok = hits.some((h) => h.source === expected);
  if (ok) found++;
  console.log(ok ? "✓" : "✗", question, "→", hits.map((h) => h.source).join(", "));
}
console.log(`recall@5: ${((found / cases.length) * 100).toFixed(0)}%`);
await pool.end();