melbinjp/DocQA

1"""A vector must never name the wrong chunk.2 3Once a chunk owns several vectors, `self.chunks[vector_id]` is wrong and quietly4so: it would return real text, from the wrong place, with a page number5attached. That is worse than returning nothing, which is why the mapping is6tested with fake embeddings rather than trusted to review.7"""8import asyncio9 10import numpy as np11import pytest12 13from rag_session import RAGSession14 15 16class FakeModel:17    """Embeds by a marker in the text, so retrieval order is predictable."""18 19    def get_sentence_embedding_dimension(self):20        return 321 22    def encode(self, texts, convert_to_numpy=True):23        out = []24        for t in texts:25            if "APPLE" in t:26                out.append([1.0, 0.0, 0.0])27            elif "BANANA" in t:28                out.append([0.0, 1.0, 0.0])29            else:30                out.append([0.0, 0.0, 1.0])31        return np.array(out, dtype="float32")32 33 34def make_session():35    session = RAGSession(source="doc.pdf", embedding_model=FakeModel())36    chunks = ["chunk zero, mentions APPLE somewhere inside",37              "chunk one, mentions BANANA somewhere inside"]38    pages = [4, 9]39    # Chunk 0 owns three vectors, chunk 1 owns two.40    vector_parents = [0, 0, 0, 1, 1]41    texts = ["whole zero APPLE", "window zero a", "window zero b APPLE",42             "whole one BANANA", "window one BANANA"]43    session.ingest(chunks, FakeModel().encode(texts), pages, vector_parents)44    return session45 46 47def test_a_vector_resolves_to_its_own_chunk_and_page():48    s = make_session()49    r = asyncio.run(s.query("APPLE", k=2))50    assert r[0]["text"].startswith("chunk zero")51    assert r[0]["page"] == 452 53 54def test_the_other_chunk_is_not_confused_with_it():55    s = make_session()56    r = asyncio.run(s.query("BANANA", k=2))57    assert r[0]["text"].startswith("chunk one")58    assert r[0]["page"] == 959 60 61def test_several_matching_windows_produce_one_result_not_several():62    """Chunk zero owns two APPLE vectors. It is still one answer."""63    s = make_session()64    r = asyncio.run(s.query("APPLE", k=5))65    texts = [x["text"] for x in r]66    assert len(texts) == len(set(texts)), f"duplicated chunk in results: {texts}"67 68 69def test_a_wrong_length_mapping_is_refused_rather_than_guessed():70    s = RAGSession(source="d", embedding_model=FakeModel())71    with pytest.raises(ValueError):72        s.ingest(["a", "b"], FakeModel().encode(["x", "y", "z"]), [1, 2], [0, 1])73 74 75def test_the_default_mapping_is_one_to_one():76    """A caller that knows nothing about windows still gets sane behaviour."""77    s = RAGSession(source="d", embedding_model=FakeModel())78    s.ingest(["only APPLE chunk"], FakeModel().encode(["only APPLE chunk"]), [7])79    r = asyncio.run(s.query("APPLE", k=1))80    assert r[0]["text"] == "only APPLE chunk" and r[0]["page"] == 781 82 83def test_a_second_ingest_does_not_point_at_the_first_documents_chunks():84    """Parent indices arrive relative to the call; they must be offset."""85    s = make_session()86    s.ingest(["chunk two, mentions CHERRY"], FakeModel().encode(["CHERRY here"]),87             [11], [0])88    r = asyncio.run(s.query("CHERRY", k=1))89    assert r[0]["text"].startswith("chunk two")90    assert r[0]["page"] == 1191