Skip to content

fix(retrieval): skip map-unit count equality on asset filters #413

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
suguanYang merged 3 commits into from
Sep 13, 2026
+268 −4
Merged

fix(retrieval): skip map-unit count equality on asset filters #413

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b921639
fix(retrieval): skip map-unit count equality on asset filters
suguanYang Sep 13, 2026
6ff0a09
Potential fix for pull request finding 'CodeQL / Potentially uninitia…
suguanYang Sep 13, 2026
2118ac1
Potential fix for pull request finding 'CodeQL / Potentially uninitia…
suguanYang Sep 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
246 changes: 246 additions & 0 deletions apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -781,6 +781,252 @@ async def test_classic_route_image_filter_scores_only_units_with_images(
assert results[0]["chunk_type"] == "image"


async def test_classic_route_mixed_section_asset_filters_do_not_mark_index_unusable(
developer_api_client_factory: Callable[
[], AbstractAsyncContextManager[AsyncClient]
],
) -> None:
identifier = uuid4().hex[:8]
namespace = f"classic-mixed-{identifier}"
async with developer_api_client_factory() as api_client:
mixed = await _publish_document(
namespace=namespace,
source_file_name="mixed-assets.pdf",
chunks=[
{
"chunk_id": f"plain-{identifier}",
"type": "text",
"content": "mixedsection plaintext filler with no assets",
"path": "mixed-assets.pdf/Root/Plain/body",
"order": 1,
"metadata": {},
},
{
"chunk_id": f"chart-body-{identifier}",
"type": "text",
"content": "mixedsection chartmarker beside a plot",
"path": "mixed-assets.pdf/Root/Chart/body",
"order": 2,
"metadata": {"connect_to": [{"target": f"chart-{identifier}"}]},
},
{
"chunk_id": f"chart-{identifier}",
"type": "image",
"content": "mixedsection chartmarker plot",
"path": "images/mixed-chart.png",
"order": 3,
"file_path": "images/mixed-chart.png",
"metadata": {},
},
{
"chunk_id": f"grid-body-{identifier}",
"type": "text",
"content": "mixedsection tablemarker beside a grid",
"path": "mixed-assets.pdf/Root/Grid/body",
"order": 4,
"metadata": {"connect_to": [{"target": f"grid-{identifier}"}]},
},
{
"chunk_id": f"grid-{identifier}",
"type": "table",
"content": "<table><tr><td>mixedsection tablemarker grid</td></tr></table>",
"path": "tables/mixed-grid.html",
"order": 5,
"file_path": "tables/mixed-grid.html",
"metadata": {},
},
],
)
for extra_kind, extra_type, extra_path, extra_content in (
("photo", "image", "images/extra-photo.png", "unrelated landscape photo"),
("diagram", "image", "images/extra-diagram.png", "unrelated diagram"),
("sheet", "table", "tables/extra-sheet.html", "<table><tr><td>unrelated sheet</td></tr></table>"),
("grid", "table", "tables/extra-grid.html", "<table><tr><td>unrelated grid</td></tr></table>"),
):
extra_chunk_id = f"{extra_kind}-{identifier}"
await _publish_document(
namespace=namespace,
source_file_name=f"extra-{extra_kind}.pdf",
chunks=[
{
"chunk_id": f"{extra_kind}-body-{identifier}",
"type": "text",
"content": f"unrelated {extra_kind} caption",
"path": f"extra-{extra_kind}.pdf/Root/Section/body",
"order": 1,
"metadata": {"connect_to": [{"target": extra_chunk_id}]},
},
{
"chunk_id": extra_chunk_id,
"type": extra_type,
"content": extra_content,
"path": extra_path,
"order": 2,
"file_path": extra_path,
"metadata": {},
},
],
)
async with contract_db_session() as db:
mixed_units = list(
(
await db.execute(
select(DocumentMapUnit).where(
DocumentMapUnit.document_id == mixed["document_id"]
)
)
).scalars()
)
assert len(mixed_units) >= 3
assert any(unit.has_image and not unit.has_table for unit in mixed_units)
assert any(unit.has_table and not unit.has_image for unit in mixed_units)
assert any(not unit.has_image and not unit.has_table for unit in mixed_units)

image_response = await api_client.post(
"/api/v1/retrieval/query",
json={
"namespace": namespace,
"query": "mixedsection chartmarker",
"top_k": 1,
"use_agentic": False,
"chunk_types": ["image"],
},
)
table_response = await api_client.post(
"/api/v1/retrieval/query",
json={
"namespace": namespace,
"query": "mixedsection tablemarker",
"top_k": 1,
"use_agentic": False,
"chunk_types": ["table"],
},
)
unfiltered_response = await api_client.post(
"/api/v1/retrieval/query",
json={
"namespace": namespace,
"query": "mixedsection plaintext",
"top_k": 1,
"use_agentic": False,
},
)

assert image_response.status_code == 200
image_body = cast(dict[str, object], image_response.json())
image_results = cast(list[dict[str, object]], image_body["results"])
assert image_body["router_used"] == "classic_topk"
assert [row["chunk_id"] for row in image_results] == [f"chart-{identifier}"]
assert image_results[0]["chunk_type"] == "image"

assert table_response.status_code == 200
table_body = cast(dict[str, object], table_response.json())
table_results = cast(list[dict[str, object]], table_body["results"])
assert table_body["router_used"] == "classic_topk"
assert [row["chunk_id"] for row in table_results] == [f"grid-{identifier}"]
assert table_results[0]["chunk_type"] == "table"

assert unfiltered_response.status_code == 200
unfiltered_body = cast(dict[str, object], unfiltered_response.json())
assert unfiltered_body["router_used"] == "classic_topk"


@pytest.mark.parametrize(
"incomplete_index_kind", ["legacy_format", "missing_index"]
)
async def test_classic_route_image_filter_raises_when_index_is_unusable(
developer_api_client_factory: Callable[
[], AbstractAsyncContextManager[AsyncClient]
],
incomplete_index_kind: str,
) -> None:
identifier = uuid4().hex[:8]
namespace = f"classic-image-unusable-{incomplete_index_kind}-{identifier}"
async with developer_api_client_factory() as api_client:
document = await _publish_document(
namespace=namespace,
source_file_name="unusable-image.pdf",
chunks=[
{
"chunk_id": f"plain-{identifier}",
"type": "text",
"content": "unusable image plaintext filler",
"path": "unusable-image.pdf/Root/Plain/body",
"order": 1,
"metadata": {},
},
{
"chunk_id": f"body-{identifier}",
"type": "text",
"content": "unusable image marker next to a chart",
"path": "unusable-image.pdf/Root/Chart/body",
"order": 2,
"metadata": {"connect_to": [{"target": f"chart-{identifier}"}]},
},
{
"chunk_id": f"chart-{identifier}",
"type": "image",
"content": "unusable image marker chart",
"path": "images/unusable-chart.png",
"order": 3,
"file_path": "images/unusable-chart.png",
"metadata": {},
},
],
)
await _publish_document(
namespace=namespace,
source_file_name="extra-image.pdf",
chunks=[
{
"chunk_id": f"extra-body-{identifier}",
"type": "text",
"content": "unrelated extra caption",
"path": "extra-image.pdf/Root/Section/body",
"order": 1,
"metadata": {"connect_to": [{"target": f"extra-{identifier}"}]},
},
{
"chunk_id": f"extra-{identifier}",
"type": "image",
"content": "unrelated extra photo",
"path": "images/extra.png",
"order": 2,
"file_path": "images/extra.png",
"metadata": {},
},
],
)
if incomplete_index_kind == "legacy_format":
await ContractDatabase.execute(
"""
UPDATE document_map_unit_indexes
SET format_version = 1
WHERE document_id = :document_id
""",
{"document_id": document["document_id"]},
)
else:
await ContractDatabase.execute(
"""
DELETE FROM document_map_unit_indexes
WHERE document_id = :document_id
""",
{"document_id": document["document_id"]},
)
with pytest.raises(RuntimeError, match="map-unit index is incomplete"):
await api_client.post(
"/api/v1/retrieval/query",
json={
"namespace": namespace,
"query": "unusable image marker",
"top_k": 1,
"use_agentic": False,
"chunk_types": ["image"],
},
)


async def test_connected_hydration_does_not_load_legacy_job_chunks(
developer_api_client_factory: Callable[
[], AbstractAsyncContextManager[AsyncClient]
Expand Down
26 changes: 22 additions & 4 deletions packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -451,9 +451,13 @@ async def map_unit_discovery(
has_index_storage_mismatch = indexed_unit_count != int(
actual_unit_count or 0
) or indexed_token_count != int(actual_token_count or 0)
# Equality against indexes.unit_count is only valid when unit_rows is the
# full inventory of those revisions. Image/table type filters, signal
# paths, and exclude_sections load a subset, so require the looser
# undercount check used for unfiltered token projection.
has_index_unit_count_mismatch = (
indexed_unit_count < len(unit_rows)
if is_unfiltered_scope
if is_unfiltered_scope or type_clause or signal_paths or exclude_sections
else indexed_unit_count != len(unit_rows)
)
is_index_format_incompatible = any(
Expand Down Expand Up @@ -488,8 +492,9 @@ async def map_unit_discovery(
_token_count,
) in index_parts
)
has_revision_coverage_mismatch = len(index_parts) != len(expected_revisions)
has_unusable_index = (
len(index_parts) != len(expected_revisions)
has_revision_coverage_mismatch
or has_index_unit_count_mismatch
or has_index_storage_mismatch
or is_index_format_incompatible
Expand All @@ -505,11 +510,24 @@ async def map_unit_discovery(
)
except Exception as exc:
logger.warning("retrieval index readiness publish failed: %s", exc)
if has_revision_coverage_mismatch:
unusable_reason = "revision_coverage"
elif is_index_format_incompatible:
unusable_reason = "format"
elif has_index_storage_mismatch:
unusable_reason = "storage"
else:
unusable_reason = "unit_count_mismatch"
chunk_types_label = (
",".join(sorted(chunk_types)) if chunk_types else "none"
)
raise RuntimeError(
"retrieval map-unit index is incomplete or incompatible "
f"(user_id={user_id} namespace={namespace} "
f"(reason={unusable_reason} user_id={user_id} namespace={namespace} "
f"expected_revisions={len(expected_revisions)} "
f"indexed_revisions={len(index_parts)})"
f"indexed_revisions={len(index_parts)} "
f"indexed_unit_count={indexed_unit_count} unit_rows={len(unit_rows)} "
f"unfiltered={is_unfiltered_scope} chunk_types={chunk_types_label})"
)
if has_incomplete_index_statistics:
logger.warning(
Expand Down
Loading