feat(workstream-e): LLM re-rank + LangGraph decision flow for cheatsh… - #1014
Conversation
…eet CRE mapping Implements RFC Workstream E (docs/rfc/cheatsheets-llm-autonomous-mapping-rfc.md, Issue E) on top of Workstream B's CheatsheetRecord contract. - cheatsheet_rerank.py: rerank_candidates_with_llm(record, candidates) runs a small LangGraph flow (rerank -> classify | rerank -> fallback -> classify) that scores/explains each candidate CRE via an injected LLM call, validates the response against a strict Pydantic schema, drops any hallucinated cre_id not in the original shortlist, and always falls back to retrieval-only scoring on LLM error, timeout (hard wall-clock cutoff), or malformed output so a cheat sheet is never dropped from the pipeline. - classify_confidence(score) buckets into high/medium/low using the RFC's bootstrap thresholds (0.85 / 0.70), overridable via env vars for later recalibration against PR OWASP#865 precision/recall data. - Every result carries a RerankTrace (model, prompt version, UTC timestamp, fallback flag/reason) for the RFC audit trail. - CandidateCRE is defined locally (mirrors the RFC's Workstream D contract) since retrieve_candidate_cres hasn't landed yet; swapping in the real Workstream D output only requires constructing this same dataclass. - LLM call is dependency-injected (mirrors the ai_client seam in embed_alignment.py and the score_fn seam in librarian/cross_encoder.py), defaulting to a lazily-imported LiteLLM call so the module has no hard LLM dependency and stays hermetically testable. - 16 unit/integration tests: confidence boundaries, successful rerank, top_n sort/truncate, hallucinated-id handling, per-candidate fallback, LLM exception / timeout / malformed-JSON / empty-result fallback paths, and two end-to-end compiled-graph runs (success + fallback). - docs/rfc-llm-rerank.md documents the flow, mirroring the Workstream B doc. - requirements.txt: add langgraph.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
Summary by CodeRabbit
WalkthroughChangesAdds a LangGraph-based LLM reranking pipeline for cheatsheet-to-CRE matching. It defines strict contracts, confidence classification, timeout handling, fallback scoring, audit traces, public configuration, RFC documentation, and unit and integration tests. Cheatsheet candidate reranking
Estimated code review effort: 4 (Complex) | ~45 minutes Mergeability Score: ⚪ Minimal · up to The change adds the LLM re-ranking flow and its tests without any supplied merge-blocking correctness, security, availability, or deployment risk; no actionable merge-blocking risk remains beyond normal checks and review. Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (3)
application/utils/external_project_parsers/parsers/cheatsheet_rerank.py (2)
58-63: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRead thresholds inside
classify_confidenceand validate the ordering.Two consequences follow from parsing the environment at import time:
- If an operator sets
CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD=high,float()raisesValueErrorduring module import. The failure surfaces as an import error far from its cause.- Tests and callers cannot recalibrate thresholds after import, which weakens the "recalibratable via env" claim in the docstring.
The code also never checks that the high threshold is greater than or equal to the medium threshold. If the two values are inverted, the
"medium"band becomes unreachable and no error is reported.♻️ Proposed refactor to resolve thresholds lazily and validate them
-HIGH_CONFIDENCE_THRESHOLD = float( - os.environ.get("CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD", "0.85") -) -MEDIUM_CONFIDENCE_THRESHOLD = float( - os.environ.get("CRE_CHEATSHEET_RERANK_MEDIUM_THRESHOLD", "0.70") -) +DEFAULT_HIGH_CONFIDENCE_THRESHOLD = 0.85 +DEFAULT_MEDIUM_CONFIDENCE_THRESHOLD = 0.70 + + +def _thresholds() -> "tuple[float, float]": + """Resolve (high, medium) thresholds from env with validation.""" + + def _read(env_var: str, default: float) -> float: + raw = os.environ.get(env_var) + if raw is None: + return default + try: + return float(raw) + except ValueError as exc: + raise RerankError(f"{env_var} must be a float, got {raw!r}") from exc + + high = _read( + "CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD", DEFAULT_HIGH_CONFIDENCE_THRESHOLD + ) + medium = _read( + "CRE_CHEATSHEET_RERANK_MEDIUM_THRESHOLD", DEFAULT_MEDIUM_CONFIDENCE_THRESHOLD + ) + if medium > high: + raise RerankError( + f"medium threshold {medium} must not exceed high threshold {high}" + ) + return high, mediumThen use the resolved values in
classify_confidence:- if score >= HIGH_CONFIDENCE_THRESHOLD: + high, medium = _thresholds() + if score >= high: return "high" - if score >= MEDIUM_CONFIDENCE_THRESHOLD: + if score >= medium: return "medium" return "low"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@application/utils/external_project_parsers/parsers/cheatsheet_rerank.py` around lines 58 - 63, Remove the module-level threshold parsing and resolve CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD and CRE_CHEATSHEET_RERANK_MEDIUM_THRESHOLD inside classify_confidence on each call, converting invalid values into a clear validation error there; also validate that the high threshold is greater than or equal to the medium threshold before applying the confidence bands.
400-414: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winUse
STARTwhen upgrading to LangGraph 1.x
StateGraph,add_conditional_edges, andENDremain supported. Replace deprecatedset_entry_point("rerank")withgraph.add_edge(START, "rerank")and importSTART.requirements.txtcurrently restricts LangGraph to<1and does not document the claimed 1.0.10 security fix.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@application/utils/external_project_parsers/parsers/cheatsheet_rerank.py` around lines 400 - 414, Update the graph construction around StateGraph to import START and replace set_entry_point("rerank") with an edge from START to "rerank", preserving the existing conditional and terminal edges. Also update the LangGraph requirement constraint and its accompanying documentation to reflect the intended 1.0.10 security fix instead of restricting the package below version 1.Source: Linters/SAST tools
docs/rfc-llm-rerank.md (1)
6-9: 📐 Maintainability & Code Quality | 🔵 TrivialAdd measurable success criteria.
This RFC defines the goal, context, and constraints. It does not define acceptance criteria that show the reranking flow is complete. Specify observable criteria for valid LLM results, fallback behavior, audit traces, and Workstream F compatibility.
What criteria must pass before this workstream can be accepted? I can provide the Requirements template from
requirements-gate.mdc.As per coding guidelines,
**/*.{md,txt}requires: “Check tickets for completeness (goal, success criteria, context, constraints) before coding. Ask clarifying questions if requirements are missing and offer the Requirements template from requirements-gate.mdc.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/rfc-llm-rerank.md` around lines 6 - 9, Add a measurable acceptance-criteria section to the RFC covering valid LLM result structure, deterministic fallback behavior, auditable reranking traces, and compatibility with Workstream F’s suggestions.json contract. Define observable pass/fail conditions for each criterion, and use the Requirements template from requirements-gate.mdc if available.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@application/tests/cheatsheet_rerank_test.py`:
- Around line 144-158: Update test_llm_timeout_falls_back to record
time.monotonic() immediately before calling rerank_candidates_with_llm and
assert the elapsed duration is well below the slow_stub delay, while preserving
the existing fallback assertions.
In `@application/utils/external_project_parsers/parsers/cheatsheet_rerank.py`:
- Around line 365-385: Update the per-candidate fallback handling in the loop
building RankedCRE results so candidates missing from scored are always marked
needs_review=True, regardless of classify_confidence(entry["score"]). Preserve
the existing confidence classification and fallback_used behavior for normally
scored candidates and whole-run fallbacks.
- Around line 242-254: Update _call_with_timeout to manage ThreadPoolExecutor
explicitly and call shutdown(wait=False) when the timeout or function error
exits, avoiding the context manager’s blocking shutdown; preserve propagation of
successful results and exceptions. Also add an explicit timeout value to the
litellm.completion call in the rerank provider flow so the underlying HTTP
request is bounded and the worker can terminate.
In `@docs/rfc-llm-rerank.md`:
- Line 44: Update the fenced code block in the RFC documentation to specify the
text language by changing its opening fence to use the text language identifier,
while preserving the block’s contents.
In `@requirements.txt`:
- Line 47: Update the langgraph dependency constraint in requirements.txt from
the 0.x range to langgraph>=1.0.10,<2, then run the graph integration tests to
verify compatibility.
---
Nitpick comments:
In `@application/utils/external_project_parsers/parsers/cheatsheet_rerank.py`:
- Around line 58-63: Remove the module-level threshold parsing and resolve
CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD and CRE_CHEATSHEET_RERANK_MEDIUM_THRESHOLD
inside classify_confidence on each call, converting invalid values into a clear
validation error there; also validate that the high threshold is greater than or
equal to the medium threshold before applying the confidence bands.
- Around line 400-414: Update the graph construction around StateGraph to import
START and replace set_entry_point("rerank") with an edge from START to "rerank",
preserving the existing conditional and terminal edges. Also update the
LangGraph requirement constraint and its accompanying documentation to reflect
the intended 1.0.10 security fix instead of restricting the package below
version 1.
In `@docs/rfc-llm-rerank.md`:
- Around line 6-9: Add a measurable acceptance-criteria section to the RFC
covering valid LLM result structure, deterministic fallback behavior, auditable
reranking traces, and compatibility with Workstream F’s suggestions.json
contract. Define observable pass/fail conditions for each criterion, and use the
Requirements template from requirements-gate.mdc if available.
🪄 Autofix
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: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro Plus
Run ID: 7b109444-2dfb-45cc-9ad2-3ad30e5dc35a
📒 Files selected for processing (4)
application/tests/cheatsheet_rerank_test.pyapplication/utils/external_project_parsers/parsers/cheatsheet_rerank.pydocs/rfc-llm-rerank.mdrequirements.txt
- needs_review is now always True for a candidate the LLM never scored (previously it could read False if the retrieval-only fallback score happened to land in a medium/high confidence band). - _call_with_timeout no longer blocks on ThreadPoolExecutor's default shutdown(wait=True) after a timeout; it shuts down with wait=False so a timed-out call returns to the caller immediately instead of waiting for the abandoned thread. Verified: test_llm_timeout_falls_back now asserts the call returns well under the stub's artificial delay. - default_llm_score_fn now accepts an optional timeout and passes it to litellm.completion, bounding the underlying HTTP request itself so the worker thread can actually terminate (not just be abandoned). - classify_confidence now resolves CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD / _MEDIUM_THRESHOLD on every call instead of once at import time (so overrides set after import take effect), and validates high >= medium. - build_rerank_graph uses an explicit START edge instead of set_entry_point (current LangGraph idiom). - requirements.txt: langgraph>=1.0.10,<2 (>=1.0.10 for a security fix; previous >=0.2,<1 pin was already stale against what's actually installed/tested). - docs/rfc-llm-rerank.md: label the flow-diagram fence as text, add a measurable acceptance-criteria checklist tied to specific tests. - Test warms up build_rerank_graph() once at module load so the new timing assertion isn't skewed by LangGraph's one-time first-compile cost when this test runs in isolation.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
application/utils/external_project_parsers/parsers/cheatsheet_rerank.py (1)
478-491: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winValidate
top_nandtimeout_secondsbefore graph execution.
top_n=1.5passes Line 490 and later fails when_node_classifyslices the ranked list.top_n=Truesilently limits results to one item. Non-positive, non-finite, or booleantimeout_secondsvalues can become a provider error or retrieval fallback instead of a caller error. Reject these values before the empty-candidate return.Proposed fix
- if not candidates: - return [] - if top_n <= 0: + if isinstance(top_n, bool) or not isinstance(top_n, int) or top_n <= 0: raise RerankError(f"top_n must be > 0, got {top_n}") + if ( + isinstance(timeout_seconds, bool) + or not isinstance(timeout_seconds, (int, float)) + or not (0 < timeout_seconds < float("inf")) + ): + raise RerankError( + f"timeout_seconds must be a finite value > 0, got {timeout_seconds!r}" + ) + if not candidates: + return []Add contract tests for float and boolean
top_n, and zero, infinite, and booleantimeout_seconds. As per coding guidelines, “Use test-first development for new behavior and importers.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@application/utils/external_project_parsers/parsers/cheatsheet_rerank.py` around lines 478 - 491, Update the reranking entry point containing the candidates early return to validate top_n and timeout_seconds before returning for empty candidates or executing the graph. Require top_n to be a non-boolean integer greater than zero, and timeout_seconds to be a non-boolean finite value greater than zero; raise the established RerankError for invalid inputs. Add contract tests covering float and boolean top_n plus zero, infinite, and boolean timeout_seconds.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@application/utils/external_project_parsers/parsers/cheatsheet_rerank.py`:
- Around line 478-491: Update the reranking entry point containing the
candidates early return to validate top_n and timeout_seconds before returning
for empty candidates or executing the graph. Require top_n to be a non-boolean
integer greater than zero, and timeout_seconds to be a non-boolean finite value
greater than zero; raise the established RerankError for invalid inputs. Add
contract tests covering float and boolean top_n plus zero, infinite, and boolean
timeout_seconds.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro Plus
Run ID: 245dc5ab-d031-409e-9b87-2fef1bea9d85
📒 Files selected for processing (4)
application/tests/cheatsheet_rerank_test.pyapplication/utils/external_project_parsers/parsers/cheatsheet_rerank.pydocs/rfc-llm-rerank.mdrequirements.txt
🚧 Files skipped from review as they are similar to previous changes (1)
- application/tests/cheatsheet_rerank_test.py
rerank_candidates_with_llm previously validated top_n after the empty-candidates early return, and never validated timeout_seconds at all: - A float top_n (e.g. 2.5) passed the '> 0' check and only failed later, deep inside the graph, with an opaque 'TypeError: slice indices must be integers' from list slicing in _node_classify. - A boolean top_n (bool is an int subclass) was silently accepted as 0/1. - timeout_seconds accepted zero, negative, infinite, or boolean values with no validation; an infinite timeout in particular would defeat the timeout guard entirely and could hang the pipeline forever on a stuck LLM call. - Calling with empty candidates bypassed all of the above, since the early return ran before validation. Now both are validated up front (non-boolean int > 0 for top_n; non-boolean, finite number > 0 for timeout_seconds), before the empty-candidates check, raising the existing RerankError. Added 7 contract tests: float/boolean top_n, zero/infinite/boolean timeout_seconds, and that invalid params still raise even with empty candidates.
|
@robvanderveer @Pa04rth I'd love a review on this, thanks! |
Summary
Implements RFC Workstream E — LLM Re-Rank and Decision Graph (LangGraph) —
from
docs/rfc/cheatsheets-llm-autonomous-mapping-rfc.md(Issue E, section 12).Given a
CheatsheetRecord(Workstream B) and a shortlist ofCandidateCRE(the contract Workstream D's
retrieve_candidate_cresis expected toreturn),
rerank_candidates_with_llm()runs a small LangGraph flow that:an LLM, asking it to score (0-1) and justify each candidate.
cre_idthe model invents that wasn't in the original shortlist.on LLM error, timeout (hard wall-clock cutoff, default 30s), or malformed
output, so a cheat sheet can never be dropped from the pipeline by an LLM
hiccup.
high/medium/lowconfidence viaclassify_confidence(), using the RFC's bootstrap thresholds(
>= 0.85/>= 0.70), overridable by env var for later recalibration.RerankTrace(model, prompt version, UTC timestamp,fallback flag/reason) to every result.
Why
CandidateCREis defined hereWorkstream D (
retrieve_candidate_cres) hasn't landed yet.CandidateCRE(
cre_id,score,text) mirrors the RFC's contract exactly, so wiring inthe real Workstream D output later is a drop-in swap.
Design choices worth flagging for review
llm_score_fn), matching the seampattern already used in
embed_alignment.py(ai_client) andlibrarian/cross_encoder.py(score_fn). Production defaults to alazily-imported LiteLLM call; the entire test suite runs with a stub,
no network/API key required.
langgraphtorequirements.txt(>=0.2,<1).prompt_client.py— kept this fully decoupled.Note for maintainers
application/utils/librarian/(already on main) implements a similarretrieve → rerank → decide → emit flow for a generic
knowledge_queue,using a cross-encoder rather than an LLM. This PR is scoped strictly to
the RFC's Workstream E contract for cheat sheets, as assigned. Happy to
align the two if that's useful.
Testing
16 new tests in
application/tests/cheatsheet_rerank_test.py, all passing,alongside the existing cheatsheet test suite. black/flake8 clean.