b8bf846bb7
Collision detection replaced the lookup's dict comprehension with `setdefault`, which also flipped a chunk found by several searches from its last occurrence to its first. The copies differ in everything the expansion window decides, figures included, so that silently changed what a citation renders. The rules are separate now: a repeated (source, chunk_id) takes the later result, a chunk_id under two sources is still refused.
102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
from haiku.rag.store.models.chunk import SearchResult
|
|
from haiku.rag.store.models.citation import resolve_citations
|
|
|
|
|
|
def _result(
|
|
chunk_id: str,
|
|
chunk_ids: list[str] | None = None,
|
|
document_meta: dict | None = None,
|
|
chunk_meta: dict | None = None,
|
|
) -> SearchResult:
|
|
return SearchResult(
|
|
content="content",
|
|
score=0.9,
|
|
chunk_id=chunk_id,
|
|
chunk_ids=chunk_ids or [],
|
|
document_id="doc-1",
|
|
document_uri="test://doc",
|
|
document_meta=document_meta or {},
|
|
chunk_meta=chunk_meta or {},
|
|
)
|
|
|
|
|
|
def test_resolve_citations_copies_merged_chunk_ids():
|
|
result = _result("c1", chunk_ids=["c1", "c2"])
|
|
citations = resolve_citations(["c1"], [result])
|
|
assert len(citations) == 1
|
|
assert citations[0].chunk_id == "c1"
|
|
assert citations[0].chunk_ids == ["c1", "c2"]
|
|
|
|
|
|
def test_resolve_citations_falls_back_to_chunk_id():
|
|
result = _result("c1")
|
|
citations = resolve_citations(["c1"], [result])
|
|
assert len(citations) == 1
|
|
assert citations[0].chunk_ids == ["c1"]
|
|
|
|
|
|
def test_resolve_citations_strips_brackets():
|
|
result = _result("c1")
|
|
citations = resolve_citations(["[c1]"], [result])
|
|
assert len(citations) == 1
|
|
assert citations[0].chunk_id == "c1"
|
|
|
|
|
|
def test_resolve_citations_skips_unknown_ids():
|
|
result = _result("c1")
|
|
citations = resolve_citations(["c1", "missing"], [result])
|
|
assert len(citations) == 1
|
|
|
|
|
|
def test_resolve_citations_copies_document_meta():
|
|
result = _result(
|
|
"c1", document_meta={"source_url": "https://example.org/report/view"}
|
|
)
|
|
citations = resolve_citations(["c1"], [result])
|
|
assert citations[0].document_meta == {
|
|
"source_url": "https://example.org/report/view"
|
|
}
|
|
|
|
|
|
def test_resolve_citations_copies_chunk_meta():
|
|
result = _result("c1", chunk_meta={"para_no": "12", "speaker": "MR SMITH"})
|
|
citations = resolve_citations(["c1"], [result])
|
|
assert citations[0].chunk_meta == {"para_no": "12", "speaker": "MR SMITH"}
|
|
|
|
|
|
def test_a_repeated_chunk_is_cited_from_its_last_occurrence():
|
|
"""One chunk is returned by several searches, each expanded against what
|
|
that search found in the same document, so the copies differ in everything
|
|
the window decides. The later entry supplies them."""
|
|
earlier = SearchResult(
|
|
content="narrow window",
|
|
score=0.9,
|
|
chunk_id="c1",
|
|
chunk_ids=["c1"],
|
|
document_id="doc-1",
|
|
document_uri="test://doc",
|
|
doc_item_refs=["#/texts/4"],
|
|
page_numbers=[2],
|
|
headings=["Maintenance"],
|
|
)
|
|
later = SearchResult(
|
|
content="wider window",
|
|
score=0.4,
|
|
chunk_id="c1",
|
|
chunk_ids=["c1", "c2"],
|
|
document_id="doc-1",
|
|
document_uri="test://doc",
|
|
doc_item_refs=["#/texts/4", "#/texts/5", "#/pictures/0"],
|
|
page_numbers=[2, 3],
|
|
headings=["Maintenance", "Calibration"],
|
|
)
|
|
|
|
[citation] = resolve_citations(["c1"], [earlier, later])
|
|
|
|
assert citation.content == "wider window"
|
|
assert citation.chunk_ids == ["c1", "c2"]
|
|
assert citation.doc_item_refs == ["#/texts/4", "#/texts/5", "#/pictures/0"]
|
|
assert citation.page_numbers == [2, 3]
|
|
assert citation.headings == ["Maintenance", "Calibration"]
|
|
# The window decides which figures travel with the citation.
|
|
assert citation.picture_refs == ["#/pictures/0"]
|