Build a RAG System with Accurate Source Citations

PDFs -> split by page -> Cohere embed -> Postgres (pgvector + full-text index)
Question -> vector + full-text search, fused by rank -> Cohere rerank -> top 5 chunks
Claude (documents with citations enabled) -> answer + cited passages -> range check

npm install @anthropic-ai/sdk cohere-ai pg @langchain/textsplitters @langchain/community @langchain/core@latest pdf-parse
npm install -D typescript tsx @types/node @types/pg

// src/ingest.ts
import { PDFLoader } from "@langchain/community/document_loaders/fs/pdf";
import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters";
import { CohereClientV2 } from "cohere-ai";
import pg from "pg";

const cohere = new CohereClientV2({ token: process.env.COHERE_API_KEY });
export const pool = new pg.Pool({ connectionString: process.env.DATABASE_URL });

export async function setup() {
  await pool.query("CREATE EXTENSION IF NOT EXISTS vector");
  await pool.query(`CREATE TABLE IF NOT EXISTS chunks (
    id bigserial PRIMARY KEY, doc_name text NOT NULL, page int, content text NOT NULL,
    embedding vector(1024) NOT NULL,
    tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', content)) STORED)`);
  await pool.query("CREATE INDEX IF NOT EXISTS chunks_tsv ON chunks USING gin (tsv)");
  await pool.query("CREATE INDEX IF NOT EXISTS chunks_emb ON chunks USING hnsw (embedding vector_cosine_ops)");
}

export async function embed(texts: string[], inputType: "search_document" | "search_query") {
  const res = await cohere.embed({
    model: "embed-english-v3.0", texts, inputType, embeddingTypes: ["float"],
  });
  return res.embeddings.float!;
}

export async function ingestPdf(path: string, docName: string) {
  const pages = await new PDFLoader(path, { splitPages: true }).load();
  const splitter = new RecursiveCharacterTextSplitter({ chunkSize: 1000, chunkOverlap: 100 });
  const chunks = await splitter.splitDocuments(pages);   // each keeps metadata.loc.pageNumber
  for (let i = 0; i < chunks.length; i += 96) {          // Cohere embed takes at most 96 texts per call
    const batch = chunks.slice(i, i + 96);
    const vectors = await embed(batch.map((c) => c.pageContent), "search_document");
    for (const [j, c] of batch.entries()) {
      await pool.query(
        "INSERT INTO chunks (doc_name, page, content, embedding) VALUES ($1, $2, $3, $4::vector)",
        [docName, c.metadata.loc?.pageNumber ?? null, c.pageContent, JSON.stringify(vectors[j])]);
    }
  }
}

// src/retrieve.ts
import { CohereClientV2 } from "cohere-ai";
import { embed, pool } from "./ingest.js";

const cohere = new CohereClientV2({ token: process.env.COHERE_API_KEY });
export interface Chunk { id: string; docName: string; page: number | null; content: string }

export async function retrieve(query: string, k = 20, topN = 5): Promise<Chunk[]> {
  const [qv] = await embed([query], "search_query");
  // Reciprocal rank fusion of vector search and Postgres full-text search (ts_rank, not true BM25)
  const { rows } = await pool.query(
    `WITH v AS (
       SELECT id, row_number() OVER (ORDER BY embedding <=> $1::vector) AS r
       FROM chunks ORDER BY embedding <=> $1::vector LIMIT $3),
     t AS (
       SELECT id, row_number() OVER (ORDER BY ts_rank_cd(tsv, q) DESC) AS r
       FROM chunks, websearch_to_tsquery('english', $2) q WHERE tsv @@ q
       ORDER BY ts_rank_cd(tsv, q) DESC LIMIT $3)
     SELECT c.id, c.doc_name AS "docName", c.page, c.content
     FROM chunks c LEFT JOIN v ON v.id = c.id LEFT JOIN t ON t.id = c.id
     WHERE v.id IS NOT NULL OR t.id IS NOT NULL
     ORDER BY coalesce(1.0 / (60 + v.r), 0) + coalesce(1.0 / (60 + t.r), 0) DESC LIMIT $3`,
    [JSON.stringify(qv), query, k]);
  const reranked = await cohere.rerank({
    model: "rerank-v3.5", query, documents: rows.map((r) => r.content), topN,
  });
  return reranked.results.map((r) => rows[r.index] as Chunk);
}

// src/cite.ts
import Anthropic from "@anthropic-ai/sdk";
import type { Chunk } from "./retrieve.js";

const client = new Anthropic();   // reads ANTHROPIC_API_KEY

export interface Citation { n: number; chunk: Chunk; citedText: string; verified: boolean }

export async function answer(question: string, chunks: Chunk[]) {
  const response = await client.messages.create({
    model: "claude-opus-5-5",
    max_tokens: 1024,
    system: "Answer only from the provided documents. If they do not cover the question, say so.",
    messages: [{
      role: "user",
      content: [
        ...chunks.map((c) => ({
          type: "document" as const,
          source: { type: "text" as const, media_type: "text/plain" as const, data: c.content },
          title: `${c.docName}, p. ${c.page ?? "?"}`,
          citations: { enabled: true },
        })),
        { type: "text" as const, text: question },
      ],
    }],
  });

  const citations: Citation[] = [];
  let text = "";
  for (const block of response.content) {
    if (block.type !== "text") continue;
    text += block.text;
    for (const c of block.citations ?? []) {
      if (c.type !== "char_location") continue;
      const chunk = chunks[c.document_index];
      let n = citations.findIndex((x) => x.chunk === chunk && x.citedText === c.cited_text) + 1;
      if (!n) {
        // The API returns the quoted span; check it really is inside our stored chunk.
        const verified = chunk.content.slice(c.start_char_index, c.end_char_index) === c.cited_text;
        citations.push({ n: citations.length + 1, chunk, citedText: c.cited_text, verified });
        n = citations.length;
      }
      text += ` [${n}]`;
    }
  }
  return { text, citations };
}

// src/main.ts
import { setup, ingestPdf, pool } from "./ingest.js";
import { retrieve } from "./retrieve.js";
import { answer } from "./cite.js";

await setup();
await ingestPdf("contracts/msa.pdf", "Master Services Agreement");
const chunks = await retrieve("When can the customer terminate for convenience?");
const { text, citations } = await answer("When can the customer terminate for convenience?", chunks);
console.log(text);
for (const c of citations) {
  console.log(`[${c.n}] ${c.chunk.docName}, p. ${c.chunk.page}${c.verified ? "" : "  (UNVERIFIED)"}: "${c.citedText}"`);
}
await pool.end();