fix(retriever): return top_k documents from the chunk retriever

[Yigtwxx]

ChucksHybridSearchRetriever.hybrid_search documents top_k as "Number of documents to return", but the reciprocal-rank-fusion query is limited to top_k chunk rows. Those rows are then grouped by document, so a single chunk-dense document can consume the whole window and every other match is silently dropped before the caller sees it.

Description

Two changes in app/retriever/chunks_hybrid_search.py, which have to land together.

1. The fused query is limited to n_results (top_k * 5, the candidate-pool size the function already computes) instead of top_k. The document cap is applied after grouping, where it always was.
2. A per-document chunk cap is enforced on the assembled result. The fusion window is measured in chunks, so widening it lets one document contribute far more passages than before; _MAX_FETCH_CHUNKS_PER_DOC previously bounded only the surrounding context chunks, never the matched ones.

Symptom

Ask for ten documents, get one.

ConnectorService._combined_rrf_search calls this retriever with retriever_top_k = top_k * 2 for every connector surface, so the loss scales with the request.

Root cause

n_results = top_k * 5   # "Fetch extra chunks for better document-level fusion"
...
semantic_search_cte  = ... .limit(n_results)
keyword_search_cte   = ... .limit(n_results)
...
final_query = ... .order_by(text("score DESC")).limit(top_k)   # chunk rows
...
doc_ids = doc_order[:top_k]   # "Keep only top_k documents"

Both CTE legs collect top_k * 5 candidates, and then the fusion query throws that pool away by taking top_k rows. Those rows are chunks. Grouping N chunks yields at most N documents, so doc_order[:top_k] on the last line can never truncate anything: it runs, but the slice is always a no-op. The two comments describe an intent the code does not implement.

Three things in the tree say this is a mistake rather than a design choice:

• The sibling module app/retriever/documents_hybrid_search.py:323 ends with the same .limit(top_k), but there the rows are documents, so it is correct. The chunk retriever inherited a limit whose unit changed underneath it.
• The newer retrieval path does it the other way round: app/agents/chat/multi_agent_chat/shared/retrieval/hybrid_search.py:194 limits the fused query to candidate_pool (top_k * _CANDIDATE_MULTIPLIER, also × 5) and truncates to top_k documents afterwards, at :233.
• That newer module has test_top_k_caps_the_number_of_documents. The legacy retriever has no equivalent — test_optimized_chunk_retriever.py only asserted len(results) >= 1.

This code is live: ConnectorService builds the retriever at app/services/connector_service.py:28 and :210, _combined_rrf_search is the shared path behind roughly fifteen connector searches, and ConnectorService is constructed on the streaming chat path (app/tasks/chat/streaming/flows/shared/pre_stream_setup.py:17).

Why the second change is required, not scope creep

With only the limit fixed, the existing test test_per_doc_chunk_limit_respected fails:

E   AssertionError: assert 35 <= 20

The per-document fetch filter is rn <= _MAX_FETCH_CHUNKS_PER_DOC OR Chunk.id.in_(matched_chunk_ids), so matched chunks were always exempt from the cap. That was invisible while the fusion window held at most top_k chunks in total — the matched set was small enough to fit under rn <= 20 on its own. Widen the window and a 35-chunk document returns all 35, which is a context-size regression rather than a fix.

_cap_chunks_per_document bounds each document's contribution at _MAX_FETCH_CHUNKS_PER_DOC, keeping matched (citable) chunks ahead of surrounding context and preserving reading order. matched_chunk_ids is recomputed from what survives, so it never advertises a chunk that is not in the payload. This mirrors _reading_order in the newer module, which keeps _MAX_PASSAGES_PER_DOC chunks per document for the same reason.

Motivation and Context

No linked issue — found while comparing the legacy retriever against the newer shared/retrieval path.

Screenshots

Not applicable — no UI change.

API Changes

• This PR includes API changes

Change Type

• Bug fix
• New feature
• Performance improvement
• Refactoring
• Documentation
• Dependency/Build system
• Breaking change
• Other (specify):

Testing Performed

• Tested locally
• Manual/QA verification

Two tests were added to tests/integration/retriever/test_optimized_chunk_retriever.py, using the existing seed_large_doc fixture (one document with 35 matching chunks, one with a single matching chunk):

• test_top_k_counts_documents_not_chunks — with top_k=10, both documents must come back.
• test_top_k_still_caps_the_number_of_documents — with top_k=1, exactly one comes back, so the cap itself is still exercised.

Against dev, the first one fails:

    assert seed_large_doc["large_doc"].id in returned
>   assert seed_large_doc["small_doc"].id in returned
E   assert 12 in {11}

The retriever's own perf log makes the change visible: docs=1 before, docs=2 after, from the same query and the same seed data.

Results on this branch:

• tests/integration/retriever + tests/unit/retriever: 12 passed. No existing test was modified.
• Full unit suite: 11 failed, 2994 passed, 1 skipped, identical to a clean dev checkout on this machine (git-tree and knowledge-store tests plus test_pat_fail_closed_static), so none of them come from this diff.
• ruff check and ruff format --check clean on both changed files.

What does not change

• Scoring. The RRF constant, the two search legs, their n_results candidate pools and the score DESC ordering are untouched.
• Document ordering. Documents still come back in first-seen fusion-rank order.
• The result dict shape: chunk_id, content, chunks, document, score, matched_chunk_ids, source.
• Filters: document_type, the date range, workspace scoping, and the exclusion of documents in the deleting state.
• documents_hybrid_search.py, which is correct as it stands.
• _MAX_FETCH_CHUNKS_PER_DOC keeps its value of 20.

Remaining risk

The fused query now returns up to top_k * 5 rows instead of top_k, so it does more work — that is the cost of actually filling the window, and it is the same trade-off the newer shared/retrieval module already makes with an identical multiplier. Per-document payload size is bounded either way by the cap in change 2.

One behaviour genuinely changes: a document that matches more than 20 chunks now returns its 20 best-ranked matches instead of all of them. Before this PR that case could not arise in practice, because the fusion window was never wide enough to produce it.

Checklist

• Follows project coding standards and conventions
• Documentation updated as needed
• Dependencies updated as needed
• No lint/build errors or new warnings
• All relevant tests are passing

High-level PR Summary

This PR fixes a critical bug in ChucksHybridSearchRetriever where requesting N documents would often return fewer documents than expected. The issue occurred because the reciprocal-rank-fusion query was limited to top_k chunk rows instead of documents, allowing a single chunk-dense document to consume the entire result window and silently drop other matches. The fix widens the fusion window to top_k * 5 (the existing candidate pool size) and applies the top_k limit after grouping chunks by document. Additionally, a per-document chunk cap (_MAX_FETCH_CHUNKS_PER_DOC) is enforced to prevent context-size regression when matched chunks exceed 20 per document.

⏱️ Estimated Review Time: 15-30 minutes

💡 Review Order Suggestion

Order	File Path
1	surfsense_backend/tests/integration/retriever/test_optimized_chunk_retriever.py
2	surfsense_backend/app/retriever/chunks_hybrid_search.py

[vercel]

@Yigtwxx is attempting to deploy a commit to the Rohan Verma's projects Team on Vercel.

A member of the Team first needs to authorize it.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: a5ec2a38-3d03-4bde-9ed7-5c45da214dae

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[Yigtwxx]

Reading the red checks, since none of them come from this diff.

• Frontend Quality — FAILURE, and therefore Quality Gate. This PR touches no TypeScript at all. The failure is a single pre-existing format error in surfsense_web/app/(home)/free/[model_slug]/page.tsx, which is already on dev. The biome-check-web hook in .pre-commit-config.yaml sets always_run: true with pass_filenames: false, so it checks the whole surfsense_web tree no matter what a PR changed — the workflow's --from-ref/--to-ref narrowing does not reach it. Measured on a clean LF checkout of dev at a89216059: Checked 1097 files. Found 1 error. Every open PR inherits it.
• Journey — FAILURE. Fails in Build & start backend stack, before any test runs: container surfsense-e2e-celery_worker-1 is unhealthy. db, redis and backend all report healthy; only the worker times out. Infrastructure, not the diff.
• Vercel — FAILURE is Authorization required to deploy, the usual result for a fork PR, and recurseml/analysis — ERROR is the bot erroring on itself.

Green: Unit Tests, Integration Tests, Test Gate, Backend Quality, File Quality, Security Scan, CodeRabbit.

---

Update — both reds now have fixes rather than just explanations.

• The biome error is fixed in fix(web): use ES2022-safe array copies in client code #1707, whose Quality Gate came back green.
• The Journey bring-up is fixed in fix(e2e): probe the celery worker over the broker, not through the app #1708. Root cause: the worker healthcheck runs celery -A app.celery_app inspect ping, and -A makes the CLI import the entire application (embeddings, LiteLLM, docling) before it can send anything — about 40s in CI, against a timeout: 5s. Every probe was killed mid-import, so the worker was reported unhealthy while its own log showed celery@... ready. and no errors. That job has not passed in the last 60 runs, dev included.

Neither is a dependency of this PR; merging them first is just what turns this board green.