Skip to content

feat(workstream-e): LLM re-rank + LangGraph decision flow for cheatsh… - #1014

Open
shreeshtripurwarcomp23-coder wants to merge 3 commits into
OWASP:mainfrom
shreeshtripurwarcomp23-coder:feat/workstream-e-llm-rerank-langgraph
Open

feat(workstream-e): LLM re-rank + LangGraph decision flow for cheatsh…#1014
shreeshtripurwarcomp23-coder wants to merge 3 commits into
OWASP:mainfrom
shreeshtripurwarcomp23-coder:feat/workstream-e-llm-rerank-langgraph

Conversation

@shreeshtripurwarcomp23-coder

Copy link
Copy Markdown
Contributor

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 of CandidateCRE
(the contract Workstream D's retrieve_candidate_cres is expected to
return), rerank_candidates_with_llm() runs a small LangGraph flow that:

  1. Sends the cheat sheet's title/summary/headings and the candidate CREs to
    an LLM, asking it to score (0-1) and justify each candidate.
  2. Validates the response against a strict Pydantic schema and drops any
    cre_id the model invents that wasn't in the original shortlist.
  3. Falls back to retrieval-only scoring — deterministically, per candidate —
    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.
  4. Buckets every result into high / medium / low confidence via
    classify_confidence(), using the RFC's bootstrap thresholds
    (>= 0.85 / >= 0.70), overridable by env var for later recalibration.
  5. Attaches an audit RerankTrace (model, prompt version, UTC timestamp,
    fallback flag/reason) to every result.

Why CandidateCRE is defined here

Workstream D (retrieve_candidate_cres) hasn't landed yet. CandidateCRE
(cre_id, score, text) mirrors the RFC's contract exactly, so wiring in
the real Workstream D output later is a drop-in swap.

Design choices worth flagging for review

  • LLM call is dependency-injected (llm_score_fn), matching the seam
    pattern already used in embed_alignment.py (ai_client) and
    librarian/cross_encoder.py (score_fn). Production defaults to a
    lazily-imported LiteLLM call; the entire test suite runs with a stub,
    no network/API key required.
  • Added langgraph to requirements.txt (>=0.2,<1).
  • Did not touch prompt_client.py — kept this fully decoupled.

Note for maintainers

application/utils/librarian/ (already on main) implements a similar
retrieve → 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.

…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.
@coderabbitai

coderabbitai Bot commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro Plus

Run ID: 2b6112fb-bbfd-422a-ad53-64b4adc29ceb

📥 Commits

Reviewing files that changed from the base of the PR and between f921d11 and be5a9f0.

📒 Files selected for processing (2)
  • application/tests/cheatsheet_rerank_test.py
  • application/utils/external_project_parsers/parsers/cheatsheet_rerank.py
🚧 Files skipped from review as they are similar to previous changes (2)
  • application/tests/cheatsheet_rerank_test.py
  • application/utils/external_project_parsers/parsers/cheatsheet_rerank.py

Summary by CodeRabbit

  • New Features

    • Added intelligent reranking of cheatsheet-to-CRE matches using language model scoring.
    • Added confidence labels, review indicators, explanations, and ranking metadata.
    • Added configurable result limits, scoring models, timeouts, and confidence thresholds.
    • Added reliable fallback to retrieval-based ranking when enhanced scoring is unavailable.
  • Documentation

    • Added guidance covering reranking behavior, validation, fallback handling, and configuration.
  • Tests

    • Added comprehensive coverage for successful ranking, invalid responses, failures, timeouts, and fallback scenarios.

Walkthrough

Changes

Adds 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

Layer / File(s) Summary
Reranking contracts and LLM scoring
application/utils/external_project_parsers/parsers/cheatsheet_rerank.py, docs/rfc-llm-rerank.md, requirements.txt, application/tests/cheatsheet_rerank_test.py
Defines candidate and ranked-result contracts, confidence classification, prompts, strict response schemas, LiteLLM scoring, timeout handling, RFC details, the LangGraph dependency, and confidence tests.
Graph execution and fallback
application/utils/external_project_parsers/parsers/cheatsheet_rerank.py, application/tests/cheatsheet_rerank_test.py, docs/rfc-llm-rerank.md
Adds graph state processing, response validation, hallucinated-ID filtering, retrieval-score fallback, trace creation, sorting, truncation, and success or failure path tests.
Public reranking entrypoint
application/utils/external_project_parsers/parsers/cheatsheet_rerank.py, application/tests/cheatsheet_rerank_test.py
Adds input validation, dependency injection, model and timeout configuration, empty-input handling, timestamp generation, and the ranked-result return path.

Estimated code review effort: 4 (Complex) | ~45 minutes

Mergeability Score: ⚪ Minimal · up to be5a9

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

  • OWASP/OpenCRE#922: Defines related Librarian reranking audit and schema contracts.
  • OWASP/OpenCRE#947: Uses related LiteLLM JSON-schema handling, malformed-response validation, fallback behavior, and mocked LLM tests.
  • OWASP/OpenCRE#957: Implements a related cross-encoder reranking pipeline for cheatsheet and CRE candidates.

Suggested reviewers: northdpole, pa04rth, paoga87

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly identifies the main changes: LLM re-ranking and the LangGraph decision flow for Workstream E.
Description check ✅ Passed The description accurately relates to the implementation, fallback behavior, validation, testing, and dependency changes in the pull request.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 5

🧹 Nitpick comments (3)
application/utils/external_project_parsers/parsers/cheatsheet_rerank.py (2)

58-63: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Read thresholds inside classify_confidence and validate the ordering.

Two consequences follow from parsing the environment at import time:

  • If an operator sets CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD=high, float() raises ValueError during 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, medium

Then 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 win

Use START when upgrading to LangGraph 1.x

StateGraph, add_conditional_edges, and END remain supported. Replace deprecated set_entry_point("rerank") with graph.add_edge(START, "rerank") and import START. requirements.txt currently restricts LangGraph to <1 and 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 | 🔵 Trivial

Add 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

📥 Commits

Reviewing files that changed from the base of the PR and between 71f6c81 and 3d7e706.

📒 Files selected for processing (4)
  • application/tests/cheatsheet_rerank_test.py
  • application/utils/external_project_parsers/parsers/cheatsheet_rerank.py
  • docs/rfc-llm-rerank.md
  • requirements.txt

Comment thread application/tests/cheatsheet_rerank_test.py
Comment thread application/utils/external_project_parsers/parsers/cheatsheet_rerank.py Outdated
Comment thread docs/rfc-llm-rerank.md Outdated
Comment thread requirements.txt Outdated
- 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.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Validate top_n and timeout_seconds before graph execution.

top_n=1.5 passes Line 490 and later fails when _node_classify slices the ranked list. top_n=True silently limits results to one item. Non-positive, non-finite, or boolean timeout_seconds values 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 boolean timeout_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

📥 Commits

Reviewing files that changed from the base of the PR and between 3d7e706 and f921d11.

📒 Files selected for processing (4)
  • application/tests/cheatsheet_rerank_test.py
  • application/utils/external_project_parsers/parsers/cheatsheet_rerank.py
  • docs/rfc-llm-rerank.md
  • requirements.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.
@shreeshtripurwarcomp23-coder

Copy link
Copy Markdown
Contributor Author

@robvanderveer @Pa04rth I'd love a review on this, thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants