feat(bench): parametric retrieval flags for experimental matrix - #37
Conversation
Adds --retrieval-top-k / --context-format / --adjacent-turns / --llm-query-expansion / --reranker / --multihop-decompose to locomo_runner.py so the forthcoming C1-C6 experiment matrix can toggle each architectural lever independently. Defaults preserve baseline behaviour — the flag additions alone don't change any existing run's output. Each config gets recorded in the output JSON's meta block so the scorecard doc can attribute per-category gains correctly.
📝 WalkthroughWalkthroughEnhanced the benchmarking runner with multihop query decomposition via LLM, optional reranking, adjacent-turn context injection, and global turn-index tracking to improve retrieval-augmented question answering capability and evaluation fidelity. Changes
Sequence DiagramsequenceDiagram
participant User as Input Query
participant LLM as LLM Decomposer
participant Ret as Retriever
participant Rer as Reranker
participant Ctx as Context Builder
participant QA as QA Processor
User->>LLM: Optional: decompose into<br/>sub-queries
LLM-->>Ret: Sub-queries (or original query)
Ret->>Ret: Retrieve hits for each query<br/>(union/dedupe)
Ret-->>Rer: Retrieved candidates
Rer->>Rer: Optional: rerank by relevance
Rer-->>Ctx: Ranked hits with turn_idx
Ctx->>Ctx: Build adjacent_turns_map<br/>from neighboring turns
Ctx-->>QA: Context with formatted dates<br/>+ adjacent turn text
QA->>QA: Optional: LLM query expansion
QA-->>User: Final QA response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| if context_format == "session_date": | ||
| prefix = f"[Session date: {dt}] " if dt else "" | ||
| elif context_format == "both": | ||
| prefix = f"[Session date: {dt}] [{dt}] " if dt else "" |
There was a problem hiding this comment.
WARNING: Duplicate date prefix in "both" context format. This line outputs the same date string twice, which adds unnecessary noise to the LLM context.
| raw = await vmem.search(query, limit=top_k) | ||
| hits = [{"text": r["text"], "metadata": r.get("metadata", {}), | ||
| "score": r.get("similarity", 0.0)} for r in raw] | ||
| if reranker is not None and getattr(reranker, "available", False): |
There was a problem hiding this comment.
SUGGESTION: Safe attribute check on reranker. If reranker object doesn't have an available attribute, this will silently skip reranking. Consider checking for existence explicitly before using getattr.
| seen_texts.add(t) | ||
| all_hits.append(h) | ||
| # Cap at retrieval_top_k, preserving order | ||
| hits = all_hits[:retrieval_top_k] |
There was a problem hiding this comment.
WARNING: Multihop hits are unsorted when truncated. Results are preserved in first seen order across sub-queries, not by relevance score. Higher scoring hits from later sub-queries can be dropped when truncating.
| continue | ||
| idx = int(idx) | ||
| neighbours = [] | ||
| for offset in range(-adjacent_turns, adjacent_turns + 1): |
There was a problem hiding this comment.
SUGGESTION: Range includes 2adjacent_turns neighbours. The help text says ±N, but this loop will return exactly 2N neighbours. This is technically correct but documentation should clarify that it's full window size.
| lines = [l.strip() for l in raw.splitlines() if l.strip()] | ||
| if len(lines) >= 2: | ||
| return lines | ||
| except Exception: |
There was a problem hiding this comment.
SUGGESTION: Exception swallowed silently in query decomposition. At minimum log a warning when decomposition fails for debugging purposes.
Code Review SummaryStatus: 5 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
SUGGESTION
Files Reviewed (1 file)
Fix these issues in Kilo Cloud Reviewed by seed-2-0-pro-260328 · 173,765 tokens |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
benchmarks/locomo_runner.py (1)
337-348: Surface fallback paths for the LLM-assisted features.Both blocks swallow every exception and silently revert to the baseline path. That keeps the benchmark running, but it also means a run can be recorded with
llm_query_expansion=trueormultihop_decompose=trueeven when the feature never actually executed. A warning or counter here would make the experiment matrix auditable.Also applies to: 399-407
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@benchmarks/locomo_runner.py`:
- Around line 202-207: The adjacent-turn expansion is crossing session
boundaries because global_idx is compared numerically in _build_adjacent_map();
restrict adjacency to the same session by checking dia_id (or session id) before
treating numeric neighbors as adjacent. Update _build_adjacent_map() (and any
code that uses turn_index) to only consider neighbor indices when
turn_index[str(neighbor_idx)]["dia_id"] ==
turn_index[str(global_idx)]["dia_id"], or else compute adjacency using
session-local indices rather than the global_idx, so last-turn/first-turn of
adjacent sessions are never linked.
- Around line 413-421: The current dedupe uses the hit's "text" which can
collapse distinct dialogue turns; change the dedupe key to turn identity instead
(preferably a tuple of h.get("dia_id") and h.get("turn_idx") ) when building
seen_texts/seen_ids before appending to all_hits; fall back to h.get("text")
only if both dia_id and turn_idx are missing. Update the loop that iterates
sq_hits (and the seen_texts variable name if desired) to compute this identity
per hit and use it for membership checks and adding to the seen set so distinct
turns with identical text are not discarded.
- Around line 411-425: The benchmark is incorrectly using retrieval_top_k as the
effective evaluation K instead of trimming candidates to the intended top_k:
when multihop_decompose is true you build all_hits from subquery results using
retrieval_top_k and then slice to retrieval_top_k, but nowhere do you trim to
the final top_k before building the answer context/evidence_hits; update the
logic in the block that computes hits (affecting multihop_decompose path and the
else path that calls _retrieve) so that after collecting candidate hits you
always trim hits = hits[:top_k] (or compute evidence_hits from hits[:top_k])
before any downstream context construction or evidence selection, ensuring both
branches use top_k for evaluation while still using retrieval_top_k only as the
candidate pool size for _retrieve/_decompose_query.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: cad2bb2f-d2b5-4a9a-8fb8-af4a9a493200
📒 Files selected for processing (1)
benchmarks/locomo_runner.py
| turn_index[str(global_idx)] = { | ||
| "datetime": dt, | ||
| "text": f"[{speaker}] {text}", | ||
| "speaker": speaker, | ||
| "dia_id": dia_id, | ||
| } |
There was a problem hiding this comment.
Keep adjacent-turn expansion inside the same session.
global_idx spans every session, and _build_adjacent_map() only checks numeric adjacency. A hit on the last turn of one session can therefore pull the first turn of the next session as a “neighbor”, which injects unrelated context into the prompt.
Suggested fix
turn_index[str(global_idx)] = {
+ "session": session_key,
"datetime": dt,
"text": f"[{speaker}] {text}",
"speaker": speaker,
"dia_id": dia_id,
} for hit in hits:
meta = hit.get("metadata", {}) or {}
idx = meta.get("turn_idx")
if idx is None:
continue
idx = int(idx)
+ current_session = turn_index.get(str(idx), {}).get("session")
neighbours = []
for offset in range(-adjacent_turns, adjacent_turns + 1):
if offset == 0:
continue
ni = idx + offset
@@
if ni in seen_indices or ni in primary_indices:
continue
if str(ni) not in turn_index:
continue
+ if turn_index[str(ni)].get("session") != current_session:
+ continue
neighbours.append(turn_index[str(ni)]["text"])
seen_indices.add(ni)Also applies to: 224-255
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@benchmarks/locomo_runner.py` around lines 202 - 207, The adjacent-turn
expansion is crossing session boundaries because global_idx is compared
numerically in _build_adjacent_map(); restrict adjacency to the same session by
checking dia_id (or session id) before treating numeric neighbors as adjacent.
Update _build_adjacent_map() (and any code that uses turn_index) to only
consider neighbor indices when turn_index[str(neighbor_idx)]["dia_id"] ==
turn_index[str(global_idx)]["dia_id"], or else compute adjacency using
session-local indices rather than the global_idx, so last-turn/first-turn of
adjacent sessions are never linked.
| if multihop_decompose: | ||
| sub_queries = await _decompose_query(client, ollama_url, retrieval_query) | ||
| seen_texts: set[str] = set() | ||
| all_hits: list[dict] = [] | ||
| for sq in sub_queries: | ||
| sq_hits = await _retrieve(strategy, sq, vmem, retrieval_top_k, reranker) | ||
| for h in sq_hits: | ||
| t = h.get("text", "") | ||
| if t not in seen_texts: | ||
| seen_texts.add(t) | ||
| all_hits.append(h) | ||
| # Cap at retrieval_top_k, preserving order | ||
| hits = all_hits[:retrieval_top_k] | ||
| else: | ||
| hits = await _retrieve(strategy, retrieval_query, vmem, retrieval_top_k, reranker) |
There was a problem hiding this comment.
--retrieval-top-k is changing the effective evaluation K too.
Line 425 retrieves retrieval_top_k hits, but nothing trims back to top_k before building the answer context or computing evidence_hits. That means the new flag widens the benchmarked K instead of just the retrieval candidate pool, which changes results and breaks the CLI contract.
Suggested fix
if multihop_decompose:
sub_queries = await _decompose_query(client, ollama_url, retrieval_query)
seen_texts: set[str] = set()
all_hits: list[dict] = []
for sq in sub_queries:
sq_hits = await _retrieve(strategy, sq, vmem, retrieval_top_k, reranker)
for h in sq_hits:
t = h.get("text", "")
if t not in seen_texts:
seen_texts.add(t)
all_hits.append(h)
# Cap at retrieval_top_k, preserving order
- hits = all_hits[:retrieval_top_k]
+ candidate_hits = all_hits[:retrieval_top_k]
else:
- hits = await _retrieve(strategy, retrieval_query, vmem, retrieval_top_k, reranker)
+ candidate_hits = await _retrieve(strategy, retrieval_query, vmem, retrieval_top_k, reranker)
+
+ hits = candidate_hits[:top_k]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@benchmarks/locomo_runner.py` around lines 411 - 425, The benchmark is
incorrectly using retrieval_top_k as the effective evaluation K instead of
trimming candidates to the intended top_k: when multihop_decompose is true you
build all_hits from subquery results using retrieval_top_k and then slice to
retrieval_top_k, but nowhere do you trim to the final top_k before building the
answer context/evidence_hits; update the logic in the block that computes hits
(affecting multihop_decompose path and the else path that calls _retrieve) so
that after collecting candidate hits you always trim hits = hits[:top_k] (or
compute evidence_hits from hits[:top_k]) before any downstream context
construction or evidence selection, ensuring both branches use top_k for
evaluation while still using retrieval_top_k only as the candidate pool size for
_retrieve/_decompose_query.
| seen_texts: set[str] = set() | ||
| all_hits: list[dict] = [] | ||
| for sq in sub_queries: | ||
| sq_hits = await _retrieve(strategy, sq, vmem, retrieval_top_k, reranker) | ||
| for h in sq_hits: | ||
| t = h.get("text", "") | ||
| if t not in seen_texts: | ||
| seen_texts.add(t) | ||
| all_hits.append(h) |
There was a problem hiding this comment.
Deduplicate multihop hits by turn identity, not by text.
Using text as the dedupe key will collapse distinct turns that happen to contain the same utterance, which is common in dialogue. That can drop the real evidence hit and attach adjacent context from the wrong turn. Prefer dia_id or turn_idx, with text only as a last-resort fallback.
Suggested fix
- seen_texts: set[str] = set()
+ seen_hit_ids: set[str] = set()
all_hits: list[dict] = []
for sq in sub_queries:
sq_hits = await _retrieve(strategy, sq, vmem, retrieval_top_k, reranker)
for h in sq_hits:
- t = h.get("text", "")
- if t not in seen_texts:
- seen_texts.add(t)
+ meta = h.get("metadata", {}) or {}
+ hit_id = str(
+ meta.get("dia_id")
+ or meta.get("turn_idx")
+ or h.get("text", "")
+ )
+ if hit_id not in seen_hit_ids:
+ seen_hit_ids.add(hit_id)
all_hits.append(h)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| seen_texts: set[str] = set() | |
| all_hits: list[dict] = [] | |
| for sq in sub_queries: | |
| sq_hits = await _retrieve(strategy, sq, vmem, retrieval_top_k, reranker) | |
| for h in sq_hits: | |
| t = h.get("text", "") | |
| if t not in seen_texts: | |
| seen_texts.add(t) | |
| all_hits.append(h) | |
| seen_hit_ids: set[str] = set() | |
| all_hits: list[dict] = [] | |
| for sq in sub_queries: | |
| sq_hits = await _retrieve(strategy, sq, vmem, retrieval_top_k, reranker) | |
| for h in sq_hits: | |
| meta = h.get("metadata", {}) or {} | |
| hit_id = str( | |
| meta.get("dia_id") | |
| or meta.get("turn_idx") | |
| or h.get("text", "") | |
| ) | |
| if hit_id not in seen_hit_ids: | |
| seen_hit_ids.add(hit_id) | |
| all_hits.append(h) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@benchmarks/locomo_runner.py` around lines 413 - 421, The current dedupe uses
the hit's "text" which can collapse distinct dialogue turns; change the dedupe
key to turn identity instead (preferably a tuple of h.get("dia_id") and
h.get("turn_idx") ) when building seen_texts/seen_ids before appending to
all_hits; fall back to h.get("text") only if both dia_id and turn_idx are
missing. Update the loop that iterates sq_hits (and the seen_texts variable name
if desired) to compute this identity per hit and use it for membership checks
and adding to the seen set so distinct turns with identical text are not
discarded.
1 participant
Summary
benchmarks/locomo_runner.pyfor the C1–C6 experiment matrix; all defaults preserve baseline behaviour exactly.metablock so scorecard runs are fully attributable.Flags added
--retrieval-top-k INT--top-k--context-format {plain,session_date,both}plain--adjacent-turns INT0--llm-query-expansionexpand_query_llmbefore retrieval--reranker {ms-marco,bge-v2-m3,off}ms-marco--multihop-decomposeIntegration decisions
bge-v2-m3 reranker: raises
NotImplementedErrorif selected (with a clear message pointing the user tohuggingface-cli download). TheCrossEncoderRerankerclass is tightly coupled to the ms-marco tokenizer path; wiring a second model would require a deeper refactor.ms-marcoandoffare fully functional.--rerankerdefault isms-marco: matches current behaviour — the cross-encoder was already being passed toretrieve()whenstrategy=full. Forvector-onlystrategy, reranking is now applied inline in_retrieve()when a reranker is loaded.--adjacent-turnsplumbing:_ingest_conversationnow returns aturn_indexdict alongside the count/elapsed values; this is threaded through_guarded→_process_qaas an explicit kwarg.Test plan
python3 benchmarks/locomo_runner.py --help— all 6 flags visible with correct defaultspython3 -c "import ast; ast.parse(open('benchmarks/locomo_runner.py').read())"— clean parseSummary by CodeRabbit