Docs Q&A with citations

"""Ask your docs: embeddings, retrieval, and a streamed answer that cites the sections it used.

pip install openai
RUNCRATE_API_KEY=... python docs_qa_citations.py ./docs "How do I make a job retry on a new machine?"
The index (index.json) is cached by content hash: a rerun embeds only the sections that changed.
"""
import hashlib, json, math, os, re, sys, time
from pathlib import Path

from openai import OpenAI

client = OpenAI(base_url="https://api.runcrate.ai/v1", api_key=os.environ["RUNCRATE_API_KEY"], max_retries=4)
EMBED = "BAAI/bge-m3"                    # multilingual, 8K-token inputs
ANSWER = "deepseek-ai/DeepSeek-V4.1-Flash"  # on the free tier
TOP_K, MIN_SCORE = 6, 0.45               # below MIN_SCORE the docs don't cover it
INDEX = Path("index.json")
costs = {"embed_index": 0.0, "embed_query": 0.0, "answer": 0.0}

def sections(root: Path):
    """Split every .md/.mdx file at its ## headings, about 1,500 characters a section."""
    for path in sorted(root.rglob("*.md*")):
        text = re.sub(r"^---.*?---", "", path.read_text(errors="ignore"), flags=re.S)  # front matter
        text = re.sub(r"^(import|export) .*$", "", text, flags=re.M)                      # MDX imports
        for block in re.split(r"\n(?=## )", text):
            title = (re.match(r"## (.+)", block) or re.match(r"#? ?(.+)", block.strip() or "x")).group(1).strip()
            for i in range(0, len(block), 1500):
                chunk = block[i:i + 1500].strip()
                if len(chunk) > 80:
                    yield {"source": f"{path.relative_to(root)}#{title}", "text": chunk}

def embed(texts: list[str], bucket: str) -> list[list[float]]:
    out = []
    for i in range(0, len(texts), 64):  # 64 inputs a request
        raw = client.embeddings.with_raw_response.create(model=EMBED, input=texts[i:i + 64])
        costs[bucket] += float(raw.headers.get("x-llm-cost", 0))
        out += [d.embedding for d in raw.parse().data]
    return out

def build_index(root: Path) -> list[dict]:
    old = {c["hash"]: c for c in json.loads(INDEX.read_text())} if INDEX.exists() else {}
    chunks = list(sections(root))
    for c in chunks:
        c["hash"] = hashlib.sha256(f"{EMBED}\n{c['text']}".encode()).hexdigest()
    todo = [c for c in chunks if c["hash"] not in old]
    for c, v in zip(todo, embed([c["text"] for c in todo], "embed_index")):
        c["vector"] = v
    for c in chunks:
        c.setdefault("vector", old.get(c["hash"], {}).get("vector"))
    INDEX.write_text(json.dumps(chunks))
    print(f"index: {len(chunks)} sections, {len(todo)} embedded now, {len(chunks) - len(todo)} reused")
    return chunks

def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    return dot / (math.sqrt(sum(x * x for x in a)) * math.sqrt(sum(y * y for y in b)))

def ask(chunks: list[dict], question: str) -> dict:
    t0 = time.time()
    q = embed([question], "embed_query")[0]
    hits = sorted(({"score": cosine(q, c["vector"]), **c} for c in chunks), key=lambda h: -h["score"])[:TOP_K]
    retrieval_s = time.time() - t0
    if hits[0]["score"] < MIN_SCORE:
        print("The docs don't cover this.")
        return {"question": question, "answer": None, "best_score": hits[0]["score"], "retrieval_seconds": retrieval_s}

    context = "\n\n".join(f"[{i + 1}] ({h['source']})\n{h['text']}" for i, h in enumerate(hits))
    stream = client.chat.completions.create(
        model=ANSWER, stream=True, stream_options={"include_usage": True}, temperature=0,
        messages=[
            {"role": "system", "content": "Answer only from the numbered sources. Cite every claim like [2]. If the sources do not answer it, say so."},
            {"role": "user", "content": f"Sources:\n{context}\n\nQuestion: {question}"},
        ],
    )
    answer, first = "", None
    for chunk in stream:
        if chunk.usage:
            costs["answer"] += float(getattr(chunk.usage, "cost", 0) or 0)
        for ch in chunk.choices:
            if ch.delta.content:
                first = first or time.time() - t0
                answer += ch.delta.content
                print(ch.delta.content, end="", flush=True)
    cited = sorted({int(n) for n in re.findall(r"\[(\d+)\]", answer) if 0 < int(n) <= len(hits)})
    print("\n\nSources:")
    for n in cited:
        print(f"  [{n}] {hits[n - 1]['source']}  (similarity {hits[n - 1]['score']:.2f})")
    return {"question": question, "answer": answer, "cited": [hits[n - 1]["source"] for n in cited],
            "retrieval_seconds": round(retrieval_s, 2), "first_token_seconds": round(first or 0, 2), "total_seconds": round(time.time() - t0, 2)}

if __name__ == "__main__":
    t0 = time.time()
    chunks = build_index(Path(sys.argv[1]))
    index_s = time.time() - t0
    results = [ask(chunks, q) for q in sys.argv[2:]]
    Path("run.json").write_text(json.dumps({"index_seconds": round(index_s, 2), "sections": len(chunks), "costs": costs,
                                            "total_usd": sum(costs.values()), "questions": results}, indent=2))
    print(f"\nindex {index_s:.1f}s, costs {json.dumps(costs)}")

// Ask your docs: embeddings, retrieval, and a streamed answer that cites the sections it used.
//
// npm install openai
// RUNCRATE_API_KEY=... npx tsx docs-qa-citations.ts ./docs "How do I make a job retry on a new machine?"
// The index (index.json) is cached by content hash: a rerun embeds only the sections that changed.
import { createHash } from "node:crypto";
import { existsSync, readdirSync, readFileSync, writeFileSync } from "node:fs";
import { join, relative } from "node:path";
import OpenAI from "openai";

const client = new OpenAI({ baseURL: "https://api.runcrate.ai/v1", apiKey: process.env.RUNCRATE_API_KEY, maxRetries: 4 });
const EMBED = "BAAI/bge-m3"; // multilingual, 8K-token inputs
const ANSWER = "deepseek-ai/DeepSeek-V4.1-Flash"; // on the free tier
const TOP_K = 6;
const MIN_SCORE = 0.45; // below this the docs don't cover it
const INDEX = "index.json";
const costs = { embed_index: 0, embed_query: 0, answer: 0 };

type Chunk = { source: string; text: string; hash: string; vector?: number[] };

/** Split every .md/.mdx file at its ## headings, about 1,500 characters a section. */
function sections(root: string): Chunk[] {
  const files = (readdirSync(root, { recursive: true }) as string[]).filter((f) => /\.mdx?$/.test(f)).sort();
  const out: Chunk[] = [];
  for (const file of files) {
    const text = readFileSync(join(root, file), "utf8")
      .replace(/^---[\s\S]*?---/, "") // front matter
      .replace(/^(import|export) .*$/gm, ""); // MDX imports
    for (const block of text.split(/\n(?=## )/)) {
      const title = (block.match(/^## (.+)/)?.[1] ?? block.trim().split("\n")[0].replace(/^#\s*/, "")).trim();
      for (let i = 0; i < block.length; i += 1500) {
        const chunk = block.slice(i, i + 1500).trim();
        if (chunk.length > 80) {
          out.push({ source: `${relative(root, join(root, file))}#${title}`, text: chunk, hash: createHash("sha256").update(`${EMBED}\n${chunk}`).digest("hex") });
        }
      }
    }
  }
  return out;
}

async function embed(texts: string[], bucket: keyof typeof costs): Promise<number[][]> {
  const out: number[][] = [];
  for (let i = 0; i < texts.length; i += 64) {
    // 64 inputs a request
    const { data, response } = await client.embeddings.create({ model: EMBED, input: texts.slice(i, i + 64) }).withResponse();
    costs[bucket] += Number(response.headers.get("x-llm-cost") ?? 0);
    out.push(...data.data.map((d) => d.embedding));
  }
  return out;
}

async function buildIndex(root: string): Promise<Chunk[]> {
  const old = new Map<string, Chunk>(existsSync(INDEX) ? (JSON.parse(readFileSync(INDEX, "utf8")) as Chunk[]).map((c) => [c.hash, c]) : []);
  const chunks = sections(root);
  const todo = chunks.filter((c) => !old.has(c.hash));
  const vectors = await embed(todo.map((c) => c.text), "embed_index");
  todo.forEach((c, i) => (c.vector = vectors[i]));
  for (const c of chunks) c.vector ??= old.get(c.hash)?.vector;
  writeFileSync(INDEX, JSON.stringify(chunks));
  console.log(`index: ${chunks.length} sections, ${todo.length} embedded now, ${chunks.length - todo.length} reused`);
  return chunks;
}

function cosine(a: number[], b: number[]): number {
  let dot = 0, na = 0, nb = 0;
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i];
    na += a[i] * a[i];
    nb += b[i] * b[i];
  }
  return dot / Math.sqrt(na * nb);
}

async function ask(chunks: Chunk[], question: string) {
  const t0 = Date.now();
  const [q] = await embed([question], "embed_query");
  const hits = chunks.map((c) => ({ ...c, score: cosine(q, c.vector ?? []) })).sort((a, b) => b.score - a.score).slice(0, TOP_K);
  const retrievalSeconds = (Date.now() - t0) / 1000;
  if (hits[0].score < MIN_SCORE) {
    console.log("The docs don't cover this.");
    return { question, answer: null, best_score: hits[0].score, retrieval_seconds: retrievalSeconds };
  }
  const context = hits.map((h, i) => `[${i + 1}] (${h.source})\n${h.text}`).join("\n\n");
  const stream = await client.chat.completions.create({
    model: ANSWER,
    stream: true,
    stream_options: { include_usage: true },
    temperature: 0,
    messages: [
      { role: "system", content: "Answer only from the numbered sources. Cite every claim like [2]. If the sources do not answer it, say so." },
      { role: "user", content: `Sources:\n${context}\n\nQuestion: ${question}` },
    ],
  });
  let answer = "";
  let first = 0;
  for await (const chunk of stream) {
    if (chunk.usage) costs.answer += Number((chunk.usage as { cost?: number }).cost ?? 0);
    for (const ch of chunk.choices) {
      if (ch.delta.content) {
        first ||= (Date.now() - t0) / 1000;
        answer += ch.delta.content;
        process.stdout.write(ch.delta.content);
      }
    }
  }
  const cited = [...new Set([...answer.matchAll(/\[(\d+)\]/g)].map((m) => Number(m[1])))].filter((n) => n > 0 && n <= hits.length).sort((a, b) => a - b);
  console.log("\n\nSources:");
  for (const n of cited) console.log(`  [${n}] ${hits[n - 1].source}  (similarity ${hits[n - 1].score.toFixed(2)})`);
  return { question, answer, cited: cited.map((n) => hits[n - 1].source), retrieval_seconds: retrievalSeconds, first_token_seconds: first, total_seconds: (Date.now() - t0) / 1000 };
}

const [root, ...questions] = process.argv.slice(2);
const t0 = Date.now();
const chunks = await buildIndex(root);
const indexSeconds = (Date.now() - t0) / 1000;
const results = [];
for (const q of questions) results.push(await ask(chunks, q));
const total = costs.embed_index + costs.embed_query + costs.answer;
writeFileSync("run.json", JSON.stringify({ index_seconds: indexSeconds, sections: chunks.length, costs, total_usd: total, questions: results }, null, 2));
console.log(`\nindex ${indexSeconds.toFixed(1)}s, costs ${JSON.stringify(costs)}`);