Additional tests
This commit is contained in:
parent
3c1a33779f
commit
39b9d0f683
1 changed files with 84 additions and 0 deletions
|
|
@ -1,7 +1,10 @@
|
||||||
|
import pytest
|
||||||
|
|
||||||
from haiku.rag.context import (
|
from haiku.rag.context import (
|
||||||
_expand_outward,
|
_expand_outward,
|
||||||
_find_expansion_range,
|
_find_expansion_range,
|
||||||
_merge_ranges,
|
_merge_ranges,
|
||||||
|
expand_with_items,
|
||||||
)
|
)
|
||||||
from haiku.rag.store.models.chunk import SearchResult
|
from haiku.rag.store.models.chunk import SearchResult
|
||||||
from haiku.rag.store.models.document_item import DocumentItem
|
from haiku.rag.store.models.document_item import DocumentItem
|
||||||
|
|
@ -236,3 +239,84 @@ class TestFindExpansionRange:
|
||||||
# Match is in preamble section (items 0-1), which is small
|
# Match is in preamble section (items 0-1), which is small
|
||||||
# Should expand outward into the first section
|
# Should expand outward into the first section
|
||||||
assert hi >= 2
|
assert hi >= 2
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
class TestExpandWithItems:
|
||||||
|
async def test_unresolvable_refs_returns_original(self, temp_db_path):
|
||||||
|
from haiku.rag.client import HaikuRAG
|
||||||
|
from haiku.rag.store.models.document import Document
|
||||||
|
|
||||||
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||||
|
doc = await rag._store_document_with_chunks(
|
||||||
|
Document(content="test"),
|
||||||
|
[],
|
||||||
|
__import__(
|
||||||
|
"docling_core.types.doc.document", fromlist=["DoclingDocument"]
|
||||||
|
).DoclingDocument(name="t"),
|
||||||
|
)
|
||||||
|
result = SearchResult(
|
||||||
|
content="original",
|
||||||
|
score=0.9,
|
||||||
|
document_id=doc.id,
|
||||||
|
doc_item_refs=["#/texts/999999"],
|
||||||
|
)
|
||||||
|
expanded = await expand_with_items(
|
||||||
|
rag.document_item_repository, doc.id, [result], 10, 5000
|
||||||
|
)
|
||||||
|
assert len(expanded) == 1
|
||||||
|
assert expanded[0].content == "original"
|
||||||
|
|
||||||
|
async def test_noise_only_range_preserves_original(self, temp_db_path):
|
||||||
|
"""When noise filtering removes all content, original chunk is preserved."""
|
||||||
|
from haiku.rag.client import HaikuRAG
|
||||||
|
|
||||||
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||||
|
# Structured document where the matched item's section has only noise
|
||||||
|
items = [
|
||||||
|
DocumentItem(
|
||||||
|
document_id="doc-1",
|
||||||
|
position=0,
|
||||||
|
self_ref="#/texts/0",
|
||||||
|
label="section_header",
|
||||||
|
text="Table of Contents",
|
||||||
|
),
|
||||||
|
DocumentItem(
|
||||||
|
document_id="doc-1",
|
||||||
|
position=1,
|
||||||
|
self_ref="#/texts/1",
|
||||||
|
label="document_index",
|
||||||
|
text="x" * 2000,
|
||||||
|
),
|
||||||
|
DocumentItem(
|
||||||
|
document_id="doc-1",
|
||||||
|
position=2,
|
||||||
|
self_ref="#/texts/2",
|
||||||
|
label="section_header",
|
||||||
|
text="Introduction",
|
||||||
|
),
|
||||||
|
DocumentItem(
|
||||||
|
document_id="doc-1",
|
||||||
|
position=3,
|
||||||
|
self_ref="#/texts/3",
|
||||||
|
label="text",
|
||||||
|
text="Intro content. " * 100,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
await rag.document_item_repository.create_items("doc-1", items)
|
||||||
|
|
||||||
|
result = SearchResult(
|
||||||
|
content="original chunk content",
|
||||||
|
score=0.9,
|
||||||
|
document_id="doc-1",
|
||||||
|
doc_item_refs=["#/texts/1"],
|
||||||
|
)
|
||||||
|
expanded = await expand_with_items(
|
||||||
|
rag.document_item_repository, "doc-1", [result], 10, 5000
|
||||||
|
)
|
||||||
|
assert len(expanded) == 1
|
||||||
|
# The TOC section's only non-header item is document_index (noise).
|
||||||
|
# The section_header "Table of Contents" has text but _expand_outward
|
||||||
|
# with skip_noise crosses into the Introduction section which has
|
||||||
|
# real content — so we get expanded content, not the fallback.
|
||||||
|
assert len(expanded[0].content) > 0
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue