RAG App on Supabase pgvector with Cited Answers and Evals

create extension if not exists vector with schema extensions;

create table doc_chunks (
  id bigint primary key generated always as identity,
  doc_path text not null,
  doc_title text not null,
  chunk_index int not null,
  content text not null,
  content_hash text not null,
  embedding extensions.vector(1024) not null,
  unique (doc_path, chunk_index)
);

alter table doc_chunks enable row level security;

create index on doc_chunks using hnsw (embedding vector_cosine_ops);

create or replace function match_doc_chunks (
  query_embedding extensions.vector(1024),
  match_threshold float,
  match_count int
)
returns table (id bigint, doc_path text, doc_title text, content text, similarity float)
language sql stable
as $$
  select c.id, c.doc_path, c.doc_title, c.content,
         1 - (c.embedding <=> query_embedding) as similarity
  from doc_chunks c
  where c.embedding <=> query_embedding < 1 - match_threshold
  order by c.embedding <=> query_embedding asc
  limit least(match_count, 20);
$$;

// src/embed.ts
export function requireEnv(name: string): string {
  const value = process.env[name];
  if (!value) throw new Error(`Missing environment variable ${name}`);
  return value;
}

export async function embed(texts: string[], inputType: "document" | "query"): Promise<number[][]> {
  const res = await fetch(requireEnv("EMBEDDINGS_URL"), {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${requireEnv("VOYAGE_API_KEY")}`,
    },
    body: JSON.stringify({ input: texts, model: requireEnv("EMBEDDING_MODEL"), input_type: inputType }),
  });
  if (!res.ok) throw new Error(`Embedding request failed with ${res.status}`);
  const json = (await res.json()) as { data: { embedding: number[]; index: number }[] };
  return json.data.sort((a, b) => a.index - b.index).map((item) => item.embedding);
}

// src/chunk.ts
export function chunkMarkdown(text: string, maxChars = 1500, overlap = 200): string[] {
  const sections = text.split(/\n(?=#{1,3} )/);
  const chunks: string[] = [];
  for (const section of sections) {
    for (let start = 0; start < section.length; start += maxChars - overlap) {
      chunks.push(section.slice(start, start + maxChars).trim());
      if (start + maxChars >= section.length) break;
    }
  }
  return chunks.filter((chunk) => chunk.length > 0);
}

// src/ingest.ts (core of the per-document loop)
import { createHash } from "node:crypto";
import { createClient } from "@supabase/supabase-js";
import { chunkMarkdown } from "./chunk";
import { embed, requireEnv } from "./embed";

const supabase = createClient(requireEnv("SUPABASE_URL"), requireEnv("SUPABASE_SECRET_KEY"));

export async function ingestDocument(docPath: string, docTitle: string, text: string): Promise<void> {
  const chunks = chunkMarkdown(text);
  const hashes = chunks.map((chunk) => createHash("sha256").update(chunk).digest("hex"));

  const { data: existing, error: readError } = await supabase
    .from("doc_chunks").select("chunk_index, content_hash").eq("doc_path", docPath);
  if (readError) throw new Error(`Read failed for ${docPath}: ${readError.message}`);
  const known = new Map((existing ?? []).map((row) => [row.chunk_index, row.content_hash]));

  const changed = chunks
    .map((content, index) => ({ content, index, hash: hashes[index] }))
    .filter((chunk) => known.get(chunk.index) !== chunk.hash);

  if (changed.length > 0) {
    const vectors = await embed(changed.map((chunk) => chunk.content), "document");
    const rows = changed.map((chunk, i) => ({
      doc_path: docPath, doc_title: docTitle, chunk_index: chunk.index,
      content: chunk.content, content_hash: chunk.hash, embedding: JSON.stringify(vectors[i]),
    }));
    const { error } = await supabase.from("doc_chunks").upsert(rows, { onConflict: "doc_path,chunk_index" });
    if (error) throw new Error(`Upsert failed for ${docPath}: ${error.message}`);
  }

  const { error: pruneError } = await supabase
    .from("doc_chunks").delete().eq("doc_path", docPath).gte("chunk_index", chunks.length);
  if (pruneError) throw new Error(`Prune failed for ${docPath}: ${pruneError.message}`);
}

// src/answer.ts
import Anthropic from "@anthropic-ai/sdk";
import { createClient } from "@supabase/supabase-js";
import { embed, requireEnv } from "./embed";

const supabase = createClient(requireEnv("SUPABASE_URL"), requireEnv("SUPABASE_SECRET_KEY"));
const client = new Anthropic(); // reads ANTHROPIC_API_KEY

interface Match { doc_path: string; doc_title: string; content: string; similarity: number }

export async function answer(question: string) {
  const [queryEmbedding] = await embed([question], "query");
  const { data, error } = await supabase.rpc("match_doc_chunks", {
    query_embedding: queryEmbedding,
    match_threshold: 0.5,
    match_count: 6,
  });
  if (error) throw new Error(`Retrieval failed: ${error.message}`);
  const matches = (data ?? []) as Match[];
  if (matches.length === 0) return { text: "I couldn't find this in the docs.", sources: [] as string[] };

  const response = await client.messages.create({
    model: "claude-sonnet-5-5",
    max_tokens: 1024,
    system: "Answer only from the search results. If they don't answer the question, say so plainly.",
    messages: [{
      role: "user",
      content: [
        ...matches.map((match) => ({
          type: "search_result" as const,
          source: match.doc_path,
          title: match.doc_title,
          content: [{ type: "text" as const, text: match.content }],
          citations: { enabled: true },
        })),
        { type: "text" as const, text: question },
      ],
    }],
  });

  const textBlocks = response.content.filter((block) => block.type === "text");
  const sources = textBlocks.flatMap((block) =>
    (block.citations ?? []).flatMap((citation) =>
      citation.type === "search_result_location" ? [citation.source] : [],
    ),
  );
  return { text: textBlocks.map((block) => block.text).join(""), sources: [...new Set(sources)] };
}

// evals/run.ts
import { answer } from "../src/answer";

const cases: { question: string; expectSource: string | null }[] = [
  { question: "How many days do customers have to request a refund?", expectSource: "policies/refunds.md" },
  { question: "Who approves discounts above 20 percent?", expectSource: "sales/discount-approval.md" },
  { question: "Do we offer weekend support?", expectSource: "support/hours.md" },
  { question: "What is the CEO's home address?", expectSource: null },
];

let passed = 0;
for (const testCase of cases) {
  const result = await answer(testCase.question);
  const ok = testCase.expectSource === null
    ? result.sources.length === 0
    : result.sources.includes(testCase.expectSource);
  if (ok) passed++;
  console.log(ok ? "PASS" : "FAIL", testCase.question, "->", result.sources.join(", ") || "(no sources)");
}
console.log(`${passed}/${cases.length} passed`);
if (passed < cases.length) process.exitCode = 1;