Spaces:
Runtime error
Runtime error
| 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"] | |