Module 5. Knowledge retrieval (RAG) for company documents

Lab 5. Bilingual HR and procedures assistant

Goal: build, run and measure a Hebrew and English assistant over about 200 pages of HR policies and procedures, with hybrid retrieval, cited answers and an eval script. Four files, each runnable on its own.

Corpus. Use a friendly client's non-sensitive procedures, or make a synthetic set: ask Claude to write 12 to 15 documents (leave, sick days, expenses, remote work, onboarding, safety, equipment, travel) in Hebrew and English, with a few deliberate conflicts and one superseded version. Save them as a mix of PDF, DOCX and HTML in docs/. Do not put real employee data in your lab.

Setup

mkdir hr-rag && cd hr-rag
python -m venv .venv && source .venv/bin/activate
pip install anthropic chromadb sentence-transformers rank-bm25 pypdf python-docx beautifulsoup4 python-dotenv
echo "ANTHROPIC_API_KEY=sk-ant-..." > .env
mkdir docs     # put your 200 pages here

The first run downloads the embedding model, which is large. Run pip and the first ingest on a good connection.

ingest.py

"""ingest.py - read PDF/DOCX/HTML, chunk, optionally add Claude-written context, embed, store in Chroma.

Usage:
    python ingest.py ./docs
    python ingest.py ./docs --contextualize
    python ingest.py ./docs --rebuild
"""
import argparse
import hashlib
import json
import re
from datetime import datetime, timezone
from pathlib import Path

import anthropic
import chromadb
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer

load_dotenv()

EMBED_MODEL = "BAAI/bge-m3"          # multilingual: Hebrew, Russian, English
COLLECTION = "hr_docs"
INDEX_DIR = "index"
MANIFEST = Path(INDEX_DIR) / "manifest.json"
CONTEXT_MODEL = "claude-haiku-5-5"   # check current model IDs at docs.claude.com
TARGET_CHARS = 1400                  # tune with your eval, not by feel
OVERLAP_CHARS = 200
MAX_DOC_CHARS = 300_000              # keep the whole-document prompt inside the model context


# ---------- 1. readers: file -> list of {page, heading, text} ----------
def read_pdf(path):
    from pypdf import PdfReader
    units = []
    for i, page in enumerate(PdfReader(str(path)).pages, start=1):
        # Check Hebrew output by eye: some PDFs extract reversed or with broken digits.
        units.append({"page": i, "heading": "", "text": page.extract_text() or ""})
    return units


def read_docx(path):
    import docx
    d = docx.Document(str(path))
    units, heading, buf = [], "", []
    for p in d.paragraphs:
        txt = p.text.strip()
        if not txt:
            continue
        if p.style.name.lower().startswith("heading"):
            if buf:
                units.append({"page": 0, "heading": heading, "text": "\n\n".join(buf)})
                buf = []
            heading = txt
        else:
            buf.append(txt)
    if buf:
        units.append({"page": 0, "heading": heading, "text": "\n\n".join(buf)})
    for t in d.tables:  # keep table rows as "cell | cell | cell" lines
        rows = [" | ".join(c.text.strip() for c in r.cells) for r in t.rows]
        units.append({"page": 0, "heading": heading, "text": "\n".join(rows)})
    return units


def read_html(path):
    from bs4 import BeautifulSoup
    soup = BeautifulSoup(path.read_text(encoding="utf-8", errors="ignore"), "html.parser")
    for tag in soup(["script", "style", "nav", "footer", "header"]):
        tag.decompose()
    units, heading, buf = [], "", []
    for el in soup.find_all(["h1", "h2", "h3", "p", "li"]):
        txt = el.get_text(" ", strip=True)
        if not txt:
            continue
        if el.name in ("h1", "h2", "h3"):
            if buf:
                units.append({"page": 0, "heading": heading, "text": "\n\n".join(buf)})
                buf = []
            heading = txt
        else:
            buf.append(txt)
    if buf:
        units.append({"page": 0, "heading": heading, "text": "\n\n".join(buf)})
    return units


READERS = {".pdf": read_pdf, ".docx": read_docx, ".html": read_html, ".htm": read_html}


# ---------- 2. chunking: paragraphs first, sentences if a paragraph is too long ----------
def split_text(text, target=TARGET_CHARS, overlap=OVERLAP_CHARS):
    text = re.sub(r"[ \t]+", " ", text)
    pieces = []
    for para in (p.strip() for p in re.split(r"\n\s*\n", text)):
        if not para:
            continue
        if len(para) <= target:
            pieces.append(para)
        else:
            pieces.extend(s.strip() for s in re.split(r"(?<=[.!?])\s+", para) if s.strip())
    chunks, cur = [], ""
    for piece in pieces:
        if cur and len(cur) + len(piece) + 1 > target:
            chunks.append(cur)
            tail = cur[-overlap:].split(" ", 1)[-1]   # drop a cut-off first word
            cur = tail + " " + piece
        else:
            cur = (cur + "\n" + piece) if cur else piece
    if cur.strip():
        chunks.append(cur)
    return chunks


def detect_lang(text):
    he = len(re.findall(r"[֐-׿]", text))
    ru = len(re.findall(r"[Ѐ-ӿ]", text))
    en = len(re.findall(r"[A-Za-z]", text))
    return max((("he", he), ("ru", ru), ("en", en)), key=lambda x: x[1])[0]


# ---------- 3. contextual retrieval: Claude writes a situating note per chunk ----------
def make_context(client, doc_text, chunk_text):
    try:
        resp = client.messages.create(
            model=CONTEXT_MODEL,
            max_tokens=150,
            messages=[{"role": "user", "content": [
                # Same document block for every chunk of this file: cache it.
                {"type": "text", "text": f"<document>\n{doc_text}\n</document>",
                 "cache_control": {"type": "ephemeral"}},
                {"type": "text", "text": (
                    "Here is a chunk from the document above:\n<chunk>\n" + chunk_text + "\n</chunk>\n"
                    "Write one or two short sentences that place this chunk in the whole document "
                    "(which policy, which section, which topic) so a search engine can find it. "
                    "Write in the same language as the chunk. Output only those sentences.")},
            ]}],
        )
        return resp.content[0].text.strip()
    except anthropic.APIError as e:
        print("  context failed, continuing without it:", e)
        return ""


# ---------- 4. main: only changed files are re-processed ----------
def sha256(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("folder")
    ap.add_argument("--contextualize", action="store_true")
    ap.add_argument("--rebuild", action="store_true")
    args = ap.parse_args()

    root = Path(args.folder)
    Path(INDEX_DIR).mkdir(exist_ok=True)
    db = chromadb.PersistentClient(path=INDEX_DIR)
    if args.rebuild:
        try:
            db.delete_collection(COLLECTION)
        except Exception:
            pass
        manifest = {}
    else:
        manifest = json.loads(MANIFEST.read_text()) if MANIFEST.exists() else {}
    col = db.get_or_create_collection(
        COLLECTION, metadata={"hnsw:space": "cosine", "embed_model": EMBED_MODEL})

    client = anthropic.Anthropic() if args.contextualize else None
    model = None
    seen = set()

    for path in sorted(p for p in root.rglob("*") if p.suffix.lower() in READERS):
        key = str(path.relative_to(root))
        seen.add(key)
        digest = sha256(path)
        if manifest.get(key) == digest:
            continue
        print("indexing", key)
        col.delete(where={"source": key})            # drop the old version of this file

        units = [u for u in READERS[path.suffix.lower()](path) if u["text"].strip()]
        doc_text = "\n\n".join(u["text"] for u in units)[:MAX_DOC_CHARS]
        mtime = datetime.fromtimestamp(path.stat().st_mtime, timezone.utc).date().isoformat()

        rows = []
        for u in units:
            for piece in split_text(u["text"]):
                ctx = make_context(client, doc_text, piece) if client else ""
                rows.append({"text": piece, "context": ctx, "page": u["page"], "heading": u["heading"]})
        if not rows:
            manifest[key] = digest
            continue

        if model is None:
            model = SentenceTransformer(EMBED_MODEL)
        to_embed = [(r["context"] + "\n" + r["text"]).strip() for r in rows]
        vecs = model.encode(to_embed, normalize_embeddings=True, batch_size=16).tolist()
        col.add(
            ids=[f"{key}#{n}" for n in range(len(rows))],
            embeddings=vecs,
            documents=[r["text"] for r in rows],
            metadatas=[{"source": key, "page": r["page"], "heading": r["heading"],
                        "lang": detect_lang(r["text"]), "context": r["context"],
                        "mtime": mtime, "sha": digest[:12]} for r in rows],
        )
        manifest[key] = digest

    for gone in set(manifest) - seen:                # file was deleted from the folder
        print("removing", gone)
        col.delete(where={"source": gone})
        del manifest[gone]

    MANIFEST.write_text(json.dumps(manifest, indent=2))
    print("chunks in index:", col.count())


if __name__ == "__main__":
    main()

retrieve.py

"""retrieve.py - hybrid retrieval: vectors (Chroma) + BM25, fused with reciprocal rank fusion.

Usage:
    python retrieve.py "כמה ימי מחלה מותר לצבור?"
"""
import json
import re
import sys
import unicodedata

import anthropic
import chromadb
from dotenv import load_dotenv
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

load_dotenv()

EMBED_MODEL = "BAAI/bge-m3"
COLLECTION = "hr_docs"
INDEX_DIR = "index"
RERANK_MODEL = "claude-haiku-5-5"    # check current model IDs at docs.claude.com

NIQQUD = re.compile(r"[֑-ׇ]")
TOKEN = re.compile(r"\w+", re.UNICODE)
HE_PREFIXES = "והבכלמש"              # vav, he, bet, kaf, lamed, mem, shin


def normalize(text):
    t = unicodedata.normalize("NFKC", text)
    t = NIQQUD.sub("", t)                                   # vowel marks
    t = t.replace("״", '"').replace("“", '"').replace("”", '"')   # gershayim, curly quotes
    t = t.replace("׳", "'").replace("’", "'")
    t = re.sub(r"(?<=\w)[\"'](?=\w)", "", t)                # בע"מ -> בעמ, דו"ח -> דוח
    return t.lower()


def tokenize(text):
    """Word tokens plus Hebrew prefix-stripped variants (over-generates on purpose)."""
    out = []
    for w in TOKEN.findall(normalize(text)):
        out.append(w)
        stem = w
        for _ in range(2):                                  # allow two stacked prefixes
            if len(stem) > 3 and stem[0] in HE_PREFIXES:
                stem = stem[1:]
                out.append(stem)
            else:
                break
    return out


def rrf(rankings, k=60):
    """Reciprocal rank fusion: robust, no score calibration needed."""
    score = {}
    for ranking in rankings:
        for rank, doc_id in enumerate(ranking):
            score[doc_id] = score.get(doc_id, 0.0) + 1.0 / (k + rank + 1)
    return sorted(score, key=score.get, reverse=True), score


def chroma_where(where):
    if not where:
        return None
    items = [{k: v} for k, v in where.items()]
    return items[0] if len(items) == 1 else {"$and": items}


def rerank_with_claude(query, hits):
    """Optional listwise rerank by a small Claude model. A dedicated reranker model is the alternative."""
    client = anthropic.Anthropic()
    listing = "\n\n".join(f"[{n}] {h['text'][:600]}" for n, h in enumerate(hits))
    msg = client.messages.create(
        model=RERANK_MODEL, max_tokens=200,
        messages=[{"role": "user", "content": (
            f"Question: {query}\n\nPassages:\n{listing}\n\n"
            "Return a JSON list of passage numbers, most useful for answering the question first. "
            "Output only the JSON list.")}])
    m = re.search(r"\[[\d,\s]*\]", msg.content[0].text)
    if not m:
        return hits
    order = [i for i in json.loads(m.group(0)) if isinstance(i, int) and 0 <= i < len(hits)]
    order += [i for i in range(len(hits)) if i not in order]
    seen, result = set(), []
    for i in order:
        if i not in seen:
            seen.add(i)
            result.append(hits[i])
    return result


class Retriever:
    def __init__(self):
        db = chromadb.PersistentClient(path=INDEX_DIR)
        self.col = db.get_collection(COLLECTION)
        if (self.col.metadata or {}).get("embed_model") != EMBED_MODEL:
            raise SystemExit("Index was built with a different embedding model. Re-run ingest.py --rebuild.")
        self.model = SentenceTransformer(EMBED_MODEL)
        data = self.col.get(include=["documents", "metadatas"])
        self.ids, self.docs, self.metas = data["ids"], data["documents"], data["metadatas"]
        self.pos = {i: n for n, i in enumerate(self.ids)}
        # BM25 sees the Claude-written context too, when ingest was run with --contextualize
        corpus = [tokenize(m.get("context", "") + " " + d) for d, m in zip(self.docs, self.metas)]
        self.bm25 = BM25Okapi(corpus)

    def _hit(self, doc_id, score):
        n = self.pos[doc_id]
        return {"id": doc_id, "text": self.docs[n], "meta": self.metas[n], "score": score}

    def search(self, query, k=5, candidates=30, mode="hybrid", where=None, rerank=False):
        """mode: 'hybrid' | 'vector' | 'bm25'. where: equality filter such as {'lang': 'he'}."""
        rankings = []
        if mode in ("hybrid", "vector"):
            qv = self.model.encode([query], normalize_embeddings=True).tolist()
            res = self.col.query(query_embeddings=qv, n_results=min(candidates, len(self.ids)),
                                 where=chroma_where(where))
            rankings.append(res["ids"][0])
        if mode in ("hybrid", "bm25"):
            scores = self.bm25.get_scores(tokenize(query))
            order = sorted(range(len(scores)), key=lambda n: -scores[n])
            ids = [self.ids[n] for n in order if scores[n] > 0
                   and all(self.metas[n].get(a) == b for a, b in (where or {}).items())]
            rankings.append(ids[:candidates])
        fused, score = rrf(rankings)
        pool = fused[: max(k, 20)] if rerank else fused[:k]
        hits = [self._hit(i, score[i]) for i in pool]
        if rerank:
            hits = rerank_with_claude(query, hits)
        return hits[:k]


if __name__ == "__main__":
    r = Retriever()
    for h in r.search(" ".join(sys.argv[1:]) or "annual leave"):
        m = h["meta"]
        print(f"{h['score']:.4f}  {m['source']} p.{m['page']}  {h['text'][:100]!r}")

answer.py

"""answer.py - retrieve, then let Claude answer with citations that point at the retrieved passages.

Usage:
    python answer.py "How do I request remote work days?"
"""
import sys

import anthropic
from dotenv import load_dotenv

from retrieve import Retriever

load_dotenv()

MODEL = "claude-sonnet-5-5"          # check current model IDs at docs.claude.com
client = anthropic.Anthropic()

SYSTEM = """You answer employee questions about HR policy and company procedures.
Rules:
- Use only the provided documents. If they do not contain the answer, say so plainly and suggest asking HR. Do not guess.
- Reply in the language of the question (Hebrew, English or Russian). Be short and concrete.
- If documents disagree or look outdated, say so and name both."""


def doc_blocks(hits):
    blocks = []
    for h in hits:
        m = h["meta"]
        title = m["source"] + (f" p.{m['page']}" if m.get("page") else "")
        blocks.append({
            "type": "document",
            "source": {"type": "text", "media_type": "text/plain", "data": h["text"]},
            "title": title,
            "context": f"Section: {m.get('heading') or '-'}. Last modified: {m.get('mtime', '?')}.",
            "citations": {"enabled": True},
        })
    return blocks


def ask(question, retriever, k=5, rerank=False, hits=None):
    hits = hits if hits is not None else retriever.search(question, k=k, rerank=rerank)
    resp = client.messages.create(
        model=MODEL, max_tokens=800, system=SYSTEM,
        messages=[{"role": "user", "content": doc_blocks(hits) + [{"type": "text", "text": question}]}],
    )
    parts, numbering, citations = [], {}, []
    for block in resp.content:
        if block.type != "text":
            continue
        parts.append(block.text)
        for c in (block.citations or []):
            idx = c.document_index
            if idx not in numbering:
                numbering[idx] = len(numbering) + 1
                m = hits[idx]["meta"]
                citations.append({"n": numbering[idx], "source": m["source"], "page": m.get("page"),
                                  "quote": c.cited_text, "chunk_id": hits[idx]["id"]})
            parts.append(f" [{numbering[idx]}]")
    return {"answer": "".join(parts).strip(), "citations": citations, "hits": hits}


if __name__ == "__main__":
    result = ask(" ".join(sys.argv[1:]), Retriever())
    print(result["answer"], "\n")
    for c in result["citations"]:
        page = f" p.{c['page']}" if c["page"] else ""
        print(f"[{c['n']}] {c['source']}{page}: {c['quote'][:120]!r}")

eval.py

"""eval.py - recall@k for retrieval, plus faithfulness and correctness judged by Claude.

qa.jsonl, one JSON object per line:
  {"q": "...", "ref": "reference answer", "source": "file.pdf", "page": 14, "lang": "he"}
For DOCX/HTML leave out "page". For a question the documents cannot answer, use "source": null.

Usage:
    python eval.py --qa qa.jsonl --k 5 --mode hybrid
    python eval.py --qa qa.jsonl --no-judge          # retrieval only, no API cost
"""
import argparse
import json
import re
from collections import defaultdict

import anthropic
from dotenv import load_dotenv

from answer import ask
from retrieve import Retriever

load_dotenv()

JUDGE_MODEL = "claude-sonnet-5-5"    # check current model IDs; ideally not the weakest model you ship

JUDGE_PROMPT = """You grade a document Q&A system. Be strict.

QUESTION: {q}

REFERENCE ANSWER (written by a human from the documents; may say the documents have no answer):
{ref}

PASSAGES THE SYSTEM WAS GIVEN:
{passages}

SYSTEM ANSWER:
{ans}

Return JSON only: {{"faithful": true|false, "correct": true|false, "unsupported_claims": ["..."]}}
faithful: every factual claim in the system answer is supported by the passages above.
correct: the system answer agrees with the reference answer in substance.
If the system says the documents do not contain the answer, faithful is true, and correct is true only if the reference also says there is no answer."""


def is_hit(hit, gold):
    m = hit["meta"]
    return m["source"] == gold["source"] and (gold.get("page") is None or m.get("page") == gold["page"])


def parse_json(text):
    m = re.search(r"\{.*\}", text, re.S)
    try:
        return json.loads(m.group(0)) if m else None
    except json.JSONDecodeError:
        return None


def judge(client, row, result):
    passages = "\n\n".join(f"[{n + 1}] {h['meta']['source']}: {h['text']}" for n, h in enumerate(result["hits"]))
    msg = client.messages.create(
        model=JUDGE_MODEL, max_tokens=400,
        messages=[{"role": "user", "content": JUDGE_PROMPT.format(
            q=row["q"], ref=row["ref"], passages=passages, ans=result["answer"])}])
    return parse_json(msg.content[0].text) or {"faithful": False, "correct": False,
                                               "unsupported_claims": ["judge returned invalid JSON"]}


def pct(values):
    return f"{100 * sum(values) / len(values):.0f}%" if values else "n/a"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--qa", default="qa.jsonl")
    ap.add_argument("--k", type=int, default=5)
    ap.add_argument("--mode", default="hybrid", choices=["hybrid", "vector", "bm25"])
    ap.add_argument("--rerank", action="store_true")
    ap.add_argument("--no-judge", action="store_true")
    args = ap.parse_args()

    rows = [json.loads(line) for line in open(args.qa, encoding="utf-8") if line.strip()]
    retriever, client = Retriever(), anthropic.Anthropic()
    stats = defaultdict(lambda: defaultdict(list))
    details = []

    for row in rows:
        lang = row.get("lang", "?")
        hits = retriever.search(row["q"], k=args.k, mode=args.mode, rerank=args.rerank)
        rec = {"q": row["q"], "lang": lang, "top": [f"{h['meta']['source']} p.{h['meta']['page']}" for h in hits]}
        if row.get("source"):
            rec["hit"] = any(is_hit(h, row) for h in hits)
            for g in ("all", lang):
                stats[g]["recall"].append(rec["hit"])
        if not args.no_judge:
            result = ask(row["q"], retriever, hits=hits)
            verdict = judge(client, row, result)
            rec.update(answer=result["answer"], faithful=bool(verdict.get("faithful")),
                       correct=bool(verdict.get("correct")), unsupported=verdict.get("unsupported_claims", []))
            for g in ("all", lang):
                stats[g]["faithful"].append(rec["faithful"])
                stats[g]["correct"].append(rec["correct"])
                if row.get("source"):
                    stats[g]["cited"].append(bool(result["citations"]))
        details.append(rec)

    print(f"\nmode={args.mode} k={args.k} rerank={args.rerank} questions={len(rows)}")
    print(f"{'group':8} {'recall@k':>9} {'faithful':>9} {'correct':>9} {'cited':>7}")
    for g in ["all"] + sorted(k for k in stats if k != "all"):
        s = stats[g]
        print(f"{g:8} {pct(s['recall']):>9} {pct(s['faithful']):>9} {pct(s['correct']):>9} {pct(s['cited']):>7}")

    print("\nRetrieval misses:")
    for d in details:
        if d.get("hit") is False:
            print(" -", d["q"], "->", d["top"][:3])
    print("\nUnfaithful answers:")
    for d in details:
        if d.get("faithful") is False:
            print(" -", d["q"], "->", d["unsupported"])

    with open("eval_report.json", "w", encoding="utf-8") as f:
        json.dump(details, f, ensure_ascii=False, indent=2)


if __name__ == "__main__":
    main()

qa.jsonl (30 lines in the real file; three shown, answers to be filled from your documents)

{"q": "כמה ימי חופשה שנתית מגיעים לעובד אחרי שלוש שנים?", "ref": "<answer from your leave policy>", "source": "leave_policy_he.pdf", "page": 3, "lang": "he"}
{"q": "How do I claim travel expenses after a client visit?", "ref": "<answer from your expenses procedure>", "source": "expenses_en.docx", "lang": "en"}
{"q": "מה מדיניות הבונוס על ביצועים ברבעון הרביעי?", "ref": "The documents do not say.", "source": null, "lang": "he"}

Steps

1. Run python ingest.py docs. Check the chunk count. Open the index and read ten random chunks, including three Hebrew ones from PDFs. Is the text readable and in the right order?
2. Run python retrieve.py "ימי מחלה" and then a query with an English term, such as a system name. Do the right files come back?
3. Run python answer.py "How many sick days can I carry over?" and the same question in Hebrew. Check that every claim has a citation and the language matches the question.
4. Write the 30 questions in qa.jsonl as described in 5.8. Do this from the documents, not from the retrieval output.
5. Run python eval.py --no-judge --mode vector, then --mode bm25, then --mode hybrid. Record recall@5 for each, in total and per language.
6. Run the full eval (python eval.py). Hand-check five judged answers. Fix what you find: chunking, tokenizer, prompt.
7. Re-ingest with --rebuild --contextualize and rerun. Then try --rerank. Fill the table below.
8. Edit one document, run ingest.py again, and confirm that only that file is re-processed and the old text is no longer returned.

Stretch: turn off the prefix variants in tokenize() and see how BM25-only recall changes on Hebrew questions. That one experiment is your best client-facing proof of why Hebrew needs special handling.