64 lines
2 KiB
Python
64 lines
2 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"}
|