Upload 5 files

tests/conftest.py

@@ -0,0 +1,47 @@
+"""Tests run with a fake embedder and a temp data dir: no model download, no Groq key needed."""
+import os
+import re
+import tempfile
+import zlib
+
+_TMP = tempfile.mkdtemp(prefix="rag-test-")
+os.environ["DATA_DIR"] = _TMP
+os.environ["USE_RERANKER"] = "false"
+os.environ["USE_QUERY_REWRITE"] = "false"
+os.environ["GROQ_API_KEY"] = "test-key"
+
+import numpy as np  # noqa: E402
+import pytest  # noqa: E402
+
+from app import store  # noqa: E402
+from app.bm25 import bm25  # noqa: E402
+from app.registry import get_registry  # noqa: E402
+
+
+class FakeEmbedder:
+    """Hashed bag-of-words vectors: deterministic, similar words -> similar vectors."""
+
+    dim = 256
+
+    def get_sentence_embedding_dimension(self):
+        return self.dim
+
+    def encode(self, texts, normalize_embeddings=True, **kw):
+        out = np.zeros((len(texts), self.dim), dtype="float32")
+        for i, t in enumerate(texts):
+            for tok in re.findall(r"\w+", t.lower()):
+                out[i, zlib.crc32(tok.encode()) % self.dim] += 1
+        norms = np.linalg.norm(out, axis=1, keepdims=True)
+        norms[norms == 0] = 1
+        return out / norms
+
+
+store.set_embedder(FakeEmbedder())
+
+
+@pytest.fixture(autouse=True)
+def clean_state():
+    store.clear_all()
+    get_registry().clear()
+    bm25.mark_dirty()
+    yield

tests/test_api.py

@@ -0,0 +1,38 @@
+from fastapi.testclient import TestClient
+
+from app import rag
+from app.main import app
+
+client = TestClient(app)  # no `with`: skips model warm-up
+
+
+def upload(name, text):
+    return client.post("/documents", files={"file": (name, text.encode(), "text/plain")})
+
+
+def test_upload_list_ask_delete(monkeypatch):
+    r = upload("notes.txt", "First page about apples.\fSecond page about oranges and juice.")
+    assert r.status_code == 200 and r.json()["duplicate"] is False
+
+    docs = client.get("/documents").json()
+    assert docs[0]["status"] == "ready" and docs[0]["pages"] == 2
+
+    assert upload("notes.txt", "First page about apples.\fSecond page about oranges and juice.").json()["duplicate"]
+
+    monkeypatch.setattr(rag.llm, "chat", lambda *a, **k: "Oranges and juice [1]")
+    r = client.post("/ask", json={"question": "what is on page 2?", "history": []})
+    body = r.json()
+    assert r.status_code == 200 and body["mode"] == "page_lookup" and body["sources"][0]["page"] == 2
+
+    assert client.delete(f"/documents/{docs[0]['doc_id']}").status_code == 200
+    assert client.get("/documents").json() == []
+
+
+def test_ask_without_documents_gives_409():
+    r = client.post("/ask", json={"question": "hello?"})
+    assert r.status_code == 409
+
+
+def test_bad_file_type():
+    r = client.post("/documents", files={"file": ("x.exe", b"abc", "application/octet-stream")})
+    assert r.status_code == 400

tests/test_pipeline.py

@@ -0,0 +1,146 @@
+import pytest
+
+from app import ingest, rag, retrieval, store
+from app.registry import get_registry
+
+
+def make_doc(name: str, pages: list[str]) -> str:
+    data = "\f".join(pages).encode()
+    entry, dup = ingest.submit_document(name, data)
+    assert not dup
+    ingest.process_document(entry["doc_id"], entry["filename"], data)
+    return entry["doc_id"]
+
+
+def big_doc_pages(n=30):
+    return [
+        f"Page {i} of the handbook. The company policy covers general procedures, "
+        f"office conduct and reporting lines. Section {i} describes routine item {i}. "
+        "Employees should follow the standard procedure at all times."
+        for i in range(1, n + 1)
+    ]
+
+
+def test_ingest_registers_pages_and_chunks():
+    doc = make_doc("handbook.txt", big_doc_pages(30))
+    entry = get_registry().get(doc)
+    assert entry["status"] == "ready" and entry["pages"] == 30 and entry["chunks"] >= 30
+
+
+def test_duplicate_upload_is_detected():
+    data = b"hello world, this is a document with enough text to count."
+    e1, d1 = ingest.submit_document("a.txt", data)
+    e2, d2 = ingest.submit_document("copy-of-a.txt", data)
+    assert not d1 and d2 and e1["doc_id"] == e2["doc_id"]
+
+
+def test_page_lookup_returns_exact_page():
+    doc = make_doc("handbook.txt", big_doc_pages(30))
+    chunks = retrieval.fetch_pages([doc], [20])
+    assert chunks and all(c["page"] == 20 for c in chunks)
+    assert "Page 20 of the handbook" in chunks[0]["text"]
+
+
+def test_small_document_not_crowded_out_by_big_one():
+    """The bug: a large PDF used to hog every result slot."""
+    big = make_doc("handbook.txt", big_doc_pages(60))
+    small = make_doc("zoology.txt", ["Zebra migration happens every year across the savanna plains."])
+
+    hits = retrieval.hybrid_search("zebra migration", [big, small], top_k=6)
+    assert hits[0]["doc_id"] == small
+
+    # generic query that matches both: results must come from BOTH documents
+    hits = retrieval.hybrid_search("policy procedure migration", [big, small], top_k=6)
+    assert {h["doc_id"] for h in hits} == {big, small}
+
+
+def test_scope_to_one_document():
+    big = make_doc("handbook.txt", big_doc_pages(10))
+    small = make_doc("zoology.txt", ["Zebra migration happens every year across the plains."])
+    hits = retrieval.hybrid_search("zebra", [big], top_k=6)
+    assert all(h["doc_id"] == big for h in hits)
+
+
+def test_search_with_page_filter():
+    doc = make_doc("handbook.txt", big_doc_pages(30))
+    hits = retrieval.hybrid_search("routine item", [doc], pages=[5, 6], top_k=6)
+    assert hits and {h["page"] for h in hits} <= {5, 6}
+
+
+def test_delete_removes_everything():
+    doc = make_doc("handbook.txt", big_doc_pages(5))
+    assert ingest.delete_document(doc)
+    assert retrieval.fetch_pages([doc], [1]) == []
+    assert get_registry().get(doc) is None
+
+
+def test_empty_pages_are_reported(monkeypatch):
+    pages = big_doc_pages(5)
+    pages[2] = ""  # page 3 has no text
+    doc = make_doc("scan.txt", pages)
+    assert get_registry().get(doc)["empty_pages"] == [3]
+
+    monkeypatch.setattr(rag.llm, "chat", lambda *a, **k: "unused")
+    out = rag.answer_question("what is on page 3?", [doc])
+    assert out["mode"] == "page_lookup"
+    assert "no extractable text" in out["answer"]
+
+
+def test_page_beyond_document_length(monkeypatch):
+    doc = make_doc("short.txt", big_doc_pages(5))
+    monkeypatch.setattr(rag.llm, "chat", lambda *a, **k: "unused")
+    out = rag.answer_question("what is written on page 99", [doc])
+    assert "only 5 pages" in out["answer"]
+
+
+def test_answer_page_lookup_builds_context_with_page_text(monkeypatch):
+    doc = make_doc("handbook.txt", big_doc_pages(30))
+    seen = {}
+
+    def fake_chat(messages, **kw):
+        seen["prompt"] = messages[-1]["content"]
+        return "It says routine item 20 [1]."
+
+    monkeypatch.setattr(rag.llm, "chat", fake_chat)
+    out = rag.answer_question("20 no ke page pe kya likha hai", [doc])
+    assert out["mode"] == "page_lookup" and out["pages"] == [20]
+    assert "Page 20 of the handbook" in seen["prompt"]
+    assert out["sources"][0]["page"] == 20
+
+
+def test_question_naming_a_document_narrows_scope(monkeypatch):
+    a = make_doc("alpha_guide.txt", ["Alpha guide explains the red button procedure."])
+    b = make_doc("beta_guide.txt", ["Beta guide explains the blue button procedure."])
+    monkeypatch.setattr(rag.llm, "chat", lambda *a_, **k: "ok")
+    out = rag.answer_question("what does the beta guide say about the button", None)
+    assert out["documents"] == ["beta_guide.txt"]
+    assert {s["doc_id"] for s in out["sources"]} == {b}
+
+
+def test_failed_ingest_marks_document_failed():
+    entry, _ = ingest.submit_document("empty.txt", b"\f\f\f   ")
+    ingest.process_document(entry["doc_id"], "empty.txt", b"\f\f\f   ")
+    e = get_registry().get(entry["doc_id"])
+    assert e["status"] == "failed" and "No extractable text" in e["error"]
+
+
+def test_pdf_end_to_end(tmp_path):
+    reportlab = pytest.importorskip("reportlab")
+    from reportlab.pdfgen import canvas
+
+    path = tmp_path / "book.pdf"
+    c = canvas.Canvas(str(path))
+    for i in range(1, 26):
+        if i != 21:  # page 21 is blank
+            c.drawString(72, 750, f"Chapter text on page {i}. Unique marker kiwi{i}.")
+            c.drawString(72, 730, "The book explains many interesting subjects in detail.")
+        c.showPage()
+    c.save()
+
+    data = path.read_bytes()
+    entry, _ = ingest.submit_document("book.pdf", data)
+    ingest.process_document(entry["doc_id"], "book.pdf", data)
+    e = get_registry().get(entry["doc_id"])
+    assert e["status"] == "ready" and e["pages"] == 25 and e["empty_pages"] == [21]
+    chunks = retrieval.fetch_pages([entry["doc_id"]], [20])
+    assert "kiwi20" in chunks[0]["text"]

tests/test_units.py

@@ -0,0 +1,55 @@
+from app.bm25 import BM25Index
+from app.chunking import chunk_text
+from app.query import extract_pages, match_docs
+
+
+def test_chunks_respect_size_and_overlap():
+    text = " ".join(f"This is sentence number {i} about topic {i % 5}." for i in range(80))
+    chunks = chunk_text(text, size=300, overlap=80)
+    assert len(chunks) > 5
+    assert all(len(c) <= 300 + 80 for c in chunks)
+    # overlap: the end of one chunk reappears at the start of the next
+    assert chunks[0].split(". ")[-1] in chunks[1]
+
+
+def test_chunk_handles_giant_unbroken_text():
+    chunks = chunk_text("word " * 2000, size=500, overlap=50)
+    assert chunks and all(len(c) <= 560 for c in chunks)
+
+
+def test_extract_pages_variants():
+    assert extract_pages("what is written on page 20?") == [20]
+    assert extract_pages("page no. 7 pe kya hai") == [7]
+    assert extract_pages("20 no ke page pe kya likha hai") == [20]
+    assert extract_pages("20th page summary") == [20]
+    assert extract_pages("explain pages 20-23") == [20, 21, 22, 23]
+    assert extract_pages("page 5 to 7") == [5, 6, 7]
+    assert extract_pages("20 se 22 pages") == [20, 21, 22]
+    assert extract_pages("page 3, 9 and 12") == [3, 9, 12]
+    assert extract_pages("what is the refund policy?") == []
+    assert extract_pages("who was born in 1990?") == []
+
+
+def test_match_docs_by_filename():
+    docs = [
+        {"doc_id": "a", "filename": "HR_Policy.pdf"},
+        {"doc_id": "b", "filename": "annual-report-2024.pdf"},
+    ]
+    assert match_docs("what does the hr policy say about leave", docs) == ["a"]
+    assert match_docs("summary of annual report 2024", docs) == ["b"]
+    assert match_docs("what about leave", docs) == []
+
+
+def test_bm25_scopes_by_doc_and_page():
+    idx = BM25Index()
+    idx.build(
+        [
+            {"id": "1", "doc_id": "A", "page": 1, "text": "zebra migration across the plains"},
+            {"id": "2", "doc_id": "B", "page": 2, "text": "zebra stripes and lions"},
+            {"id": "3", "doc_id": "B", "page": 3, "text": "completely unrelated tax filing"},
+        ]
+    )
+    idx._dirty = False
+    assert {h["id"] for h in idx.search("zebra", limit=5)} == {"1", "2"}
+    assert [h["id"] for h in idx.search("zebra", doc_ids=["B"])] == ["2"]
+    assert idx.search("zebra", pages=[1])[0]["id"] == "1"