From 3d7e7064a0aa9e2da907fc07ecc7988066c3575a Mon Sep 17 00:00:00 2001 From: Shreesh Date: Wed, 12 Aug 2026 18:39:09 +0000 Subject: [PATCH 1/6] feat(workstream-e): LLM re-rank + LangGraph decision flow for cheatsheet 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 #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. --- application/tests/cheatsheet_rerank_test.py | 229 +++++++++ .../parsers/cheatsheet_rerank.py | 461 ++++++++++++++++++ docs/rfc-llm-rerank.md | 87 ++++ requirements.txt | 1 + 4 files changed, 778 insertions(+) create mode 100644 application/tests/cheatsheet_rerank_test.py create mode 100644 application/utils/external_project_parsers/parsers/cheatsheet_rerank.py create mode 100644 docs/rfc-llm-rerank.md diff --git a/application/tests/cheatsheet_rerank_test.py b/application/tests/cheatsheet_rerank_test.py new file mode 100644 index 000000000..ed75621db --- /dev/null +++ b/application/tests/cheatsheet_rerank_test.py @@ -0,0 +1,229 @@ +import time +import unittest + +from application.defs.cheatsheet_defs import CheatsheetRecord +from application.utils.external_project_parsers.parsers.cheatsheet_rerank import ( + CandidateCRE, + RerankError, + build_rerank_graph, + classify_confidence, + rerank_candidates_with_llm, +) + + +def _record(**overrides) -> CheatsheetRecord: + defaults = dict( + source_id="Secrets_Management_Cheat_Sheet", + title="Secrets Management Cheat Sheet", + hyperlink="https://cheatsheetseries.owasp.org/cheatsheets/Secrets_Management_Cheat_Sheet.html", + summary="Guidance on secure storage, rotation, and operational handling of secrets.", + headings=["Introduction", "Architectural Patterns", "Secret Rotation"], + raw_markdown_path="cheatsheets/Secrets_Management_Cheat_Sheet.md", + ) + defaults.update(overrides) + return CheatsheetRecord(**defaults) + + +def _candidates(): + return [ + CandidateCRE( + cre_id="623-550", score=0.62, text="Operational secret rotation controls." + ), + CandidateCRE(cre_id="123-456", score=0.40, text="Unrelated logging guidance."), + ] + + +class ClassifyConfidenceTest(unittest.TestCase): + def test_high(self): + self.assertEqual(classify_confidence(0.9), "high") + self.assertEqual(classify_confidence(0.85), "high") + + def test_medium(self): + self.assertEqual(classify_confidence(0.7), "medium") + self.assertEqual(classify_confidence(0.84), "medium") + + def test_low(self): + self.assertEqual(classify_confidence(0.0), "low") + self.assertEqual(classify_confidence(0.69), "low") + + def test_out_of_range_raises(self): + with self.assertRaises(RerankError): + classify_confidence(1.5) + with self.assertRaises(RerankError): + classify_confidence(-0.1) + + def test_non_numeric_raises(self): + with self.assertRaises(RerankError): + classify_confidence("high") # type: ignore[arg-type] + + +class RerankCandidatesWithLlmTest(unittest.TestCase): + def test_empty_candidates_returns_empty(self): + self.assertEqual(rerank_candidates_with_llm(_record(), []), []) + + def test_invalid_top_n_raises(self): + with self.assertRaises(RerankError): + rerank_candidates_with_llm(_record(), _candidates(), top_n=0) + + def test_successful_rerank_produces_reason_and_confidence(self): + def stub(system, user, *, model): + self.assertIn("CHEATSHEET_TITLE", user) + self.assertIn("623-550", user) + return { + "ranked": [ + { + "cre_id": "623-550", + "score": 0.91, + "reason": "Directly covers rotation.", + }, + {"cre_id": "123-456", "score": 0.2, "reason": "Off-topic."}, + ] + } + + results = rerank_candidates_with_llm( + _record(), _candidates(), llm_score_fn=stub, top_n=5 + ) + self.assertEqual(len(results), 2) + top = results[0] + self.assertEqual(top.cre_id, "623-550") + self.assertEqual(top.confidence, "high") + self.assertFalse(top.needs_review) + self.assertFalse(top.trace.fallback_used) + self.assertEqual(top.trace.prompt_version, "v1") + self.assertIn("rotation", top.reason.lower()) + self.assertEqual(results[1].confidence, "low") + self.assertTrue(results[1].needs_review) + + def test_top_n_truncates_and_sorts_descending(self): + def stub(system, user, *, model): + return { + "ranked": [ + {"cre_id": "623-550", "score": 0.3, "reason": "r1"}, + {"cre_id": "123-456", "score": 0.95, "reason": "r2"}, + ] + } + + results = rerank_candidates_with_llm( + _record(), _candidates(), llm_score_fn=stub, top_n=1 + ) + self.assertEqual(len(results), 1) + self.assertEqual(results[0].cre_id, "123-456") + + def test_hallucinated_cre_id_is_dropped(self): + def stub(system, user, *, model): + return { + "ranked": [ + {"cre_id": "623-550", "score": 0.9, "reason": "ok"}, + {"cre_id": "999-999", "score": 0.99, "reason": "invented"}, + ] + } + + results = rerank_candidates_with_llm( + _record(), _candidates(), llm_score_fn=stub, top_n=5 + ) + ids = {r.cre_id for r in results} + self.assertNotIn("999-999", ids) + # the un-scored real candidate still gets a retrieval-only entry + self.assertIn("123-456", ids) + + def test_llm_exception_falls_back_to_retrieval_score(self): + def stub(system, user, *, model): + raise RuntimeError("provider unavailable") + + results = rerank_candidates_with_llm( + _record(), _candidates(), llm_score_fn=stub, top_n=5 + ) + self.assertEqual(len(results), 2) + for r in results: + self.assertTrue(r.trace.fallback_used) + self.assertIsNotNone(r.trace.fallback_reason) + self.assertTrue(r.needs_review) + # retrieval ordering preserved (0.62 > 0.40) + self.assertEqual(results[0].cre_id, "623-550") + + def test_llm_timeout_falls_back(self): + def slow_stub(system, user, *, model): + time.sleep(0.2) + return {"ranked": []} + + results = rerank_candidates_with_llm( + _record(), + _candidates(), + llm_score_fn=slow_stub, + top_n=5, + timeout_seconds=0.01, + ) + self.assertEqual(len(results), 2) + self.assertTrue(all(r.trace.fallback_used for r in results)) + + def test_malformed_json_falls_back(self): + def bad_stub(system, user, *, model): + return {"not_ranked_key": []} + + results = rerank_candidates_with_llm( + _record(), _candidates(), llm_score_fn=bad_stub, top_n=5 + ) + self.assertTrue(all(r.trace.fallback_used for r in results)) + + def test_llm_returns_no_valid_candidates_falls_back(self): + def empty_stub(system, user, *, model): + return { + "ranked": [{"cre_id": "not-a-real-id", "score": 0.5, "reason": "x"}] + } + + results = rerank_candidates_with_llm( + _record(), _candidates(), llm_score_fn=empty_stub, top_n=5 + ) + self.assertTrue(all(r.trace.fallback_used for r in results)) + + +class RerankGraphIntegrationTest(unittest.TestCase): + """End-to-end execution of the compiled LangGraph flow (RFC Issue E, Checkpoint E5).""" + + def test_graph_runs_success_path(self): + app = build_rerank_graph() + + def stub(system, user, *, model): + return {"ranked": [{"cre_id": "623-550", "score": 0.88, "reason": "match"}]} + + state = app.invoke( + { + "record": _record(), + "candidates": [_candidates()[0]], + "top_n": 5, + "llm_score_fn": stub, + "model_name": "test-model", + "timeout_seconds": 5.0, + "generated_at": "2026-08-13T00:00:00+00:00", + "fallback_used": False, + "fallback_reason": None, + } + ) + self.assertEqual(len(state["ranked"]), 1) + self.assertEqual(state["ranked"][0].confidence, "high") + + def test_graph_runs_fallback_path(self): + app = build_rerank_graph() + + def failing_stub(system, user, *, model): + raise RuntimeError("boom") + + state = app.invoke( + { + "record": _record(), + "candidates": _candidates(), + "top_n": 5, + "llm_score_fn": failing_stub, + "model_name": "test-model", + "timeout_seconds": 5.0, + "generated_at": "2026-08-13T00:00:00+00:00", + "fallback_used": False, + "fallback_reason": None, + } + ) + self.assertEqual(len(state["ranked"]), 2) + self.assertTrue(all(r.trace.fallback_used for r in state["ranked"])) + + +if __name__ == "__main__": + unittest.main() diff --git a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py new file mode 100644 index 000000000..b44b83a1a --- /dev/null +++ b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py @@ -0,0 +1,461 @@ +""" +RFC Workstream E — LLM Re-Rank and Decision Graph (LangGraph). + +See: docs/rfc/cheatsheets-llm-autonomous-mapping-rfc.md, section 5 +("Workstream E: LLM Re-Rank and Decision Graph (LangGraph)") and the +Issue E checklist in section 12. + +This module owns the "ReRank/Explain -> Threshold" stage of the overall +pipeline (docs/rfc section 7): given a ``CheatsheetRecord`` (Workstream B, +``application/defs/cheatsheet_defs.py``) and the top-k ``CandidateCRE`` +shortlist for it (Workstream D, ``retrieve_candidate_cres``), it asks an LLM +to re-rank and justify the shortlist, assigns a confidence band to each +result, and always returns a usable ``RankedCRE`` list — even when the LLM +call fails, times out, or returns malformed output — by falling back to the +retrieval-only ordering. + +Design notes +------------ +* ``CandidateCRE`` is defined *here* rather than imported from Workstream D + because that workstream's ``retrieve_candidate_cres`` has not landed yet. + The field set (``cre_id``, ``score``, ``text``) mirrors the RFC's + ``CandidateCRE`` contract exactly, so swapping in the real Workstream D + output only requires constructing this same dataclass. +* The LLM call is dependency-injected as ``llm_score_fn`` — a plain + ``(system, user) -> dict`` callable — exactly like the ``ai_client`` seam + in ``application/prompt_client/embed_alignment.py`` and the ``score_fn`` + seam in ``application/utils/librarian/cross_encoder.py``. Production code + never has to inject anything (a LiteLLM-backed default is wired lazily so + this module stays import-light for tests); the test suite and any + harness inject a deterministic stub instead, which keeps the LangGraph + flow hermetically testable. +* Confidence bands and thresholds follow the RFC's bootstrap defaults + (section 11 "Open Questions"): high >= 0.85, medium >= 0.70, else low. + Both are overridable via environment variables so they can be + recalibrated later against PR #865-derived precision/recall data without + a code change. +""" + +from __future__ import annotations + +import json +import logging +import os +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any, Callable, Dict, List, Optional, TypedDict + +from pydantic import BaseModel, Field, ValidationError + +from application.defs.cheatsheet_defs import CheatsheetRecord + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Confidence thresholds (RFC section 11 bootstrap defaults; recalibrate via env) +# --------------------------------------------------------------------------- +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") +) + +# Identifiers persisted into RerankTrace for the RFC audit trail (mirrors +# RETRIEVER_NAME / RERANKER_NAME conventions used elsewhere in the codebase). +RERANKER_NAME = "llm-cheatsheet-reranker" +PROMPT_VERSION = "v1" + +DEFAULT_TOP_N = 5 +DEFAULT_TIMEOUT_SECONDS = 30.0 +DEFAULT_MODEL_ENV_VAR = "CRE_CHEATSHEET_RERANK_MODEL" +DEFAULT_MODEL_FALLBACK = "gemini/gemini-2.5-flash" + +REASON_MAX_LENGTH = 400 + + +class RerankError(ValueError): + """Base class for reranker construction/usage failures.""" + + +def classify_confidence(score: float) -> str: + """ + Map a 0-1 re-rank score to a confidence band ("high" | "medium" | "low"). + + Thresholds are the RFC's bootstrap defaults and are recalibratable via + ``CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD`` / ``_MEDIUM_THRESHOLD``. + """ + if not isinstance(score, (int, float)) or isinstance(score, bool): + raise RerankError(f"score must be a number, got {score!r}") + if not (0.0 <= float(score) <= 1.0): + raise RerankError(f"score must be in [0, 1], got {score!r}") + + if score >= HIGH_CONFIDENCE_THRESHOLD: + return "high" + if score >= MEDIUM_CONFIDENCE_THRESHOLD: + return "medium" + return "low" + + +# --------------------------------------------------------------------------- +# Data contracts +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class CandidateCRE: + """ + One retrieval-stage candidate for a CheatsheetRecord. + + Mirrors the RFC's Workstream D output contract. ``text`` is optional + context (e.g. the CRE's embeddings_content) given to the LLM so it can + judge fit; when absent the LLM is told only the cre_id, which degrades + rationale quality but never breaks the flow. + """ + + cre_id: str + score: float + text: str = "" + + +@dataclass(frozen=True) +class RerankTrace: + """Audit metadata captured for every rerank run (RFC Issue E, criterion 3).""" + + model: str + prompt_version: str + generated_at: str + fallback_used: bool + fallback_reason: Optional[str] = None + + +@dataclass(frozen=True) +class RankedCRE: + """One re-ranked, explained candidate — Workstream E's output contract.""" + + cre_id: str + score: float + retrieval_score: float + confidence: str + reason: str + needs_review: bool + trace: RerankTrace + + +# --------------------------------------------------------------------------- +# LLM structured-output schema (strict; mirrors embed_alignment.AlignmentPayload) +# --------------------------------------------------------------------------- + + +class _RerankItem(BaseModel): + cre_id: str + score: float = Field(ge=0.0, le=1.0) + reason: str = "" + + +class _RerankPayload(BaseModel): + ranked: List[_RerankItem] + + +def rerank_response_json_schema() -> Dict[str, Any]: + """Provider-friendly JSON schema for strict structured LLM outputs.""" + return _RerankPayload.model_json_schema() + + +# --------------------------------------------------------------------------- +# Prompting +# --------------------------------------------------------------------------- + + +def _system_prompt() -> str: + return ( + "You map an OWASP cheat sheet to the Common Requirement (CRE) entries " + "it best satisfies. You will be given the cheat sheet's title, summary, " + "and headings, plus a shortlist of candidate CREs with their ids. " + "Score how well each candidate CRE matches the cheat sheet's content on " + "a 0.0-1.0 scale (1.0 = the cheat sheet is clearly authoritative " + "guidance for that CRE), and give a short one-sentence reason for each " + "score, grounded in the cheat sheet's actual headings/summary. " + "Only score cre_ids given to you; never invent new ones. " + "Return ONLY valid JSON of the form " + '{"ranked": [{"cre_id": "...", "score": 0.0, "reason": "..."}]}, ' + "one entry per candidate given." + ) + + +def _user_payload(record: CheatsheetRecord, candidates: List[CandidateCRE]) -> str: + lines = [ + f"CHEATSHEET_TITLE: {record.title}", + f"CHEATSHEET_SUMMARY: {record.summary}", + "CHEATSHEET_HEADINGS: " + "; ".join(record.headings), + "", + "CANDIDATE_CRES (cre_id | text):", + ] + for c in candidates: + text_preview = (c.text or "")[:800] + lines.append(f"{c.cre_id} | {text_preview}") + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Default (production) LLM call — lazy litellm import so this module stays +# import-light and hermetically testable without a real LLM dependency. +# --------------------------------------------------------------------------- + + +def _default_model_name() -> str: + return os.environ.get( + DEFAULT_MODEL_ENV_VAR, + os.environ.get("CRE_LLM_CHAT_MODEL", DEFAULT_MODEL_FALLBACK), + ) + + +def default_llm_score_fn(system: str, user: str, *, model: str) -> Dict[str, Any]: + """Production LLM call via LiteLLM. Raises on any failure; callers must + handle fallback (this function intentionally does not swallow errors).""" + try: + import litellm # type: ignore + except ImportError as exc: # pragma: no cover - exercised only without litellm + raise RerankError("litellm package is required for LLM re-rank calls") from exc + + resp = litellm.completion( + model=model, + messages=[ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ], + response_format={"type": "json_object"}, + temperature=0.2, + ) + choices = getattr(resp, "choices", None) + if not choices: + raise RerankError("LLM response contained no choices") + content = choices[0].message.content + if isinstance(content, list): # some providers return content blocks + content = "".join( + b.get("text", "") if isinstance(b, dict) else str(b) for b in content + ) + return json.loads(content) + + +def _call_with_timeout( + fn: Callable[[], Dict[str, Any]], timeout_seconds: float +) -> Dict[str, Any]: + """Run ``fn`` with a hard wall-clock timeout so a hung LLM call can never + block the pipeline; raises on timeout or on any exception from ``fn``.""" + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(fn) + try: + return future.result(timeout=timeout_seconds) + except FutureTimeoutError as exc: + raise RerankError( + f"LLM re-rank call exceeded {timeout_seconds}s timeout" + ) from exc + + +# --------------------------------------------------------------------------- +# LangGraph flow: rerank -> (success: classify) | (failure: fallback -> classify) +# --------------------------------------------------------------------------- + + +class _RerankState(TypedDict, total=False): + record: CheatsheetRecord + candidates: List[CandidateCRE] + top_n: int + llm_score_fn: Callable[..., Dict[str, Any]] + model_name: str + timeout_seconds: float + generated_at: str + scored: Dict[str, Dict[str, Any]] # cre_id -> {"score": float, "reason": str} + fallback_used: bool + fallback_reason: Optional[str] + ranked: List[RankedCRE] + + +def _node_llm_rerank(state: _RerankState) -> _RerankState: + """Call the LLM, validate its output, and record per-candidate scores. + + On any failure (LLM error, timeout, malformed JSON, schema violation) + this node records the reason and leaves ``scored`` empty; the + conditional edge below routes to the fallback node instead of raising. + """ + record = state["record"] + candidates = state["candidates"] + llm_score_fn = state["llm_score_fn"] + model_name = state["model_name"] + timeout_seconds = state["timeout_seconds"] + + system = _system_prompt() + user = _user_payload(record, candidates) + + try: + raw = _call_with_timeout( + lambda: llm_score_fn(system, user, model=model_name), timeout_seconds + ) + payload = _RerankPayload.model_validate(raw) + except (RerankError, ValidationError, json.JSONDecodeError, TypeError) as exc: + logger.warning("LLM re-rank failed for %s: %s", record.source_id, exc) + state["fallback_reason"] = f"{type(exc).__name__}: {exc}"[:REASON_MAX_LENGTH] + state["scored"] = {} + return state + except Exception as exc: # defensive: never let an unexpected error crash the run + logger.warning( + "LLM re-rank failed unexpectedly for %s: %s", record.source_id, exc + ) + state["fallback_reason"] = f"unexpected:{type(exc).__name__}: {exc}"[ + :REASON_MAX_LENGTH + ] + state["scored"] = {} + return state + + known_ids = {c.cre_id for c in candidates} + scored: Dict[str, Dict[str, Any]] = {} + for item in payload.ranked: + if item.cre_id not in known_ids: + logger.info( + "Dropping hallucinated cre_id %r not in candidate shortlist for %s", + item.cre_id, + record.source_id, + ) + continue + scored[item.cre_id] = { + "score": item.score, + "reason": item.reason[:REASON_MAX_LENGTH], + } + + if not scored: + state["fallback_reason"] = "LLM returned no valid scored candidates" + + state["scored"] = scored + return state + + +def _route_after_rerank(state: _RerankState) -> str: + return "classify" if state.get("scored") else "fallback" + + +def _node_fallback(state: _RerankState) -> _RerankState: + """Retrieval-only scoring: use each candidate's raw similarity as-is.""" + state["fallback_used"] = True + state["scored"] = { + c.cre_id: { + "score": max(0.0, min(1.0, c.score)), + "reason": "Retrieval-only score (LLM re-rank unavailable).", + } + for c in state["candidates"] + } + return state + + +def _node_classify(state: _RerankState) -> _RerankState: + candidates = state["candidates"] + scored = state["scored"] + fallback_used = state.get("fallback_used", False) + fallback_reason = state.get("fallback_reason") + trace = RerankTrace( + model=state["model_name"], + prompt_version=PROMPT_VERSION, + generated_at=state["generated_at"], + fallback_used=fallback_used, + fallback_reason=fallback_reason if fallback_used else None, + ) + + ranked: List[RankedCRE] = [] + for c in candidates: + entry = scored.get(c.cre_id) + if entry is None: + # LLM succeeded overall but skipped this one candidate: fall back + # to its retrieval score individually rather than dropping it. + entry = { + "score": max(0.0, min(1.0, c.score)), + "reason": "Not scored by reranker; using retrieval score.", + } + confidence = classify_confidence(entry["score"]) + ranked.append( + RankedCRE( + cre_id=c.cre_id, + score=entry["score"], + retrieval_score=c.score, + confidence=confidence, + reason=entry["reason"], + needs_review=(confidence == "low") or fallback_used, + trace=trace, + ) + ) + + ranked.sort(key=lambda r: r.score, reverse=True) + state["ranked"] = ranked[: state["top_n"]] + return state + + +def build_rerank_graph(): + """Compile and return the Workstream E LangGraph flow. + + Nodes: ``rerank`` -> (``classify`` | ``fallback`` -> ``classify``) -> END. + Exposed standalone so it can be inspected, visualized, or exercised + directly in integration tests without going through the convenience + wrapper below. + """ + from langgraph.graph import StateGraph, END + + graph = StateGraph(_RerankState) + graph.add_node("rerank", _node_llm_rerank) + graph.add_node("fallback", _node_fallback) + graph.add_node("classify", _node_classify) + + graph.set_entry_point("rerank") + graph.add_conditional_edges( + "rerank", _route_after_rerank, {"classify": "classify", "fallback": "fallback"} + ) + graph.add_edge("fallback", "classify") + graph.add_edge("classify", END) + + return graph.compile() + + +# --------------------------------------------------------------------------- +# Public entrypoint (RFC function-level API, section 6) +# --------------------------------------------------------------------------- + + +def rerank_candidates_with_llm( + record: CheatsheetRecord, + candidates: List[CandidateCRE], + *, + llm_score_fn: Optional[Callable[..., Dict[str, Any]]] = None, + top_n: int = DEFAULT_TOP_N, + model_name: Optional[str] = None, + timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS, +) -> List[RankedCRE]: + """ + Re-rank ``candidates`` for ``record`` via the LangGraph flow above. + + ``llm_score_fn`` defaults to a LiteLLM-backed call + (:func:`default_llm_score_fn`); tests and harnesses should inject a + deterministic stub instead. Never raises on LLM failure — falls back to + retrieval-only ordering and marks the trace accordingly. + """ + if not candidates: + return [] + if top_n <= 0: + raise RerankError(f"top_n must be > 0, got {top_n}") + + resolved_model = model_name or _default_model_name() + score_fn = llm_score_fn or default_llm_score_fn + + app = build_rerank_graph() + result = app.invoke( + { + "record": record, + "candidates": candidates, + "top_n": top_n, + "llm_score_fn": score_fn, + "model_name": resolved_model, + "timeout_seconds": timeout_seconds, + "generated_at": datetime.now(timezone.utc).isoformat(), + "fallback_used": False, + "fallback_reason": None, + } + ) + return result["ranked"] diff --git a/docs/rfc-llm-rerank.md b/docs/rfc-llm-rerank.md new file mode 100644 index 000000000..d02c1e822 --- /dev/null +++ b/docs/rfc-llm-rerank.md @@ -0,0 +1,87 @@ +# RFC Workstream E — LLM Re-Rank and Decision Graph (LangGraph) + +This document explains the implementation and behavior of RFC Workstream E +(LLM Re-Rank and Decision Graph) from the Cheatsheet to CRE Mapping RFC. + +The goal of this module is to take the top-k CRE candidates retrieved for a +cheat sheet (Workstream D) and turn them into an explained, confidence-scored +shortlist that Workstream F can persist to `suggestions.json` for human +review. + +The implementation is located in: + +* `application/utils/external_project_parsers/parsers/cheatsheet_rerank.py` + +--- + +## Sources for more context + +* RFC: `docs/rfc/cheatsheets-llm-autonomous-mapping-rfc.md` +* Workstream B (structured extraction) doc: `docs/rfc-structured-extraction.md` + +--- + +## What Workstream E implements + +Given a `CheatsheetRecord` (Workstream B's contract, +`application/defs/cheatsheet_defs.py`) and a list of `CandidateCRE` (the +contract Workstream D's `retrieve_candidate_cres` is expected to return — +defined locally here since Workstream D has not landed yet), the module +exposes: + +* `rerank_candidates_with_llm(record, candidates, ...) -> list[RankedCRE]` — + the public entrypoint. Runs the LangGraph flow described below and always + returns a usable, sorted, confidence-scored shortlist. +* `classify_confidence(score: float) -> str` — maps a 0-1 score to + `"high"` / `"medium"` / `"low"` using the RFC's bootstrap thresholds + (`>= 0.85` high, `>= 0.70` medium, else low), overridable via + `CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD` / `CRE_CHEATSHEET_RERANK_MEDIUM_THRESHOLD`. +* `build_rerank_graph()` — compiles and returns the raw LangGraph app, for + direct inspection or integration testing. + +### The LangGraph flow + +``` +START -> rerank --(success)--> classify -> END + \--(failure)--> fallback -> classify -> END +``` + +* **`rerank`** — builds a prompt from the cheat sheet's title/summary/headings + and the candidate CREs, calls the injected `llm_score_fn`, and validates the + response against a strict Pydantic schema (`_RerankPayload`, mirroring + `application/prompt_client/embed_alignment.py`'s `AlignmentPayload` + pattern). Any candidate `cre_id` the LLM invents that isn't in the original + shortlist is dropped and logged, never trusted. +* **`fallback`** — runs whenever the LLM call raises, times out + (`timeout_seconds`, default 30s, enforced with a hard wall-clock cutoff), + returns malformed JSON, or scores zero valid candidates. It scores every + candidate using its raw retrieval similarity instead, so the pipeline never + crashes and never silently drops a cheat sheet. +* **`classify`** — assigns a confidence band and a `needs_review` flag + (`true` when confidence is `"low"` or the run used the fallback path) to + every candidate, attaches an audit `RerankTrace` (model name, prompt + version, UTC timestamp, whether fallback was used and why), sorts + descending by score, and truncates to `top_n` (default 5). + +### Dependency injection / testability + +The LLM call is injected as `llm_score_fn: (system, user, *, model) -> dict`, +the same seam pattern used elsewhere in this codebase (`ai_client` in +`embed_alignment.py`, `score_fn` in `application/utils/librarian/cross_encoder.py`). +Production code defaults to `default_llm_score_fn`, a thin LiteLLM wrapper +lazily imported so this module has no hard LLM dependency; tests inject a +deterministic stub, which keeps the graph — including both the success and +fallback paths — hermetically testable without any network or API key. See +`application/tests/cheatsheet_rerank_test.py`. + +### What this module deliberately does not do + +* It does not call Workstream D's retrieval — callers supply `CandidateCRE`s. +* It does not write `suggestions.json` — that's Workstream F + (`build_suggestions` / `write_suggestions_json`), which is expected to + consume `RankedCRE.reason` as the suggestion's `reason` field and + `RankedCRE.confidence` as its `confidence` field. +* It does not decide auto-link vs. review on its own beyond the + `needs_review` hint — Phase 1 is review-first for every suggestion + regardless (RFC section 11), so `needs_review` is informational, not a + gate. diff --git a/requirements.txt b/requirements.txt index c55834701..41bf7c9da 100644 --- a/requirements.txt +++ b/requirements.txt @@ -44,6 +44,7 @@ python-markdown-maker # chat (/rest/v1/completion) — embed prompt via LiteLLM, match with sklearn litellm +langgraph>=0.2,<1 numpy scipy scikit-learn From f921d110b35facf85ae3f48d5fbcf02b01adfcd6 Mon Sep 17 00:00:00 2001 From: Shreesh Date: Wed, 12 Aug 2026 19:15:28 +0000 Subject: [PATCH 2/6] fix(workstream-e): address review findings on rerank/LangGraph module - 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. --- application/tests/cheatsheet_rerank_test.py | 21 +++- .../parsers/cheatsheet_rerank.py | 104 +++++++++++++----- docs/rfc-llm-rerank.md | 23 +++- requirements.txt | 2 +- 4 files changed, 119 insertions(+), 31 deletions(-) diff --git a/application/tests/cheatsheet_rerank_test.py b/application/tests/cheatsheet_rerank_test.py index ed75621db..5e8594b83 100644 --- a/application/tests/cheatsheet_rerank_test.py +++ b/application/tests/cheatsheet_rerank_test.py @@ -10,6 +10,13 @@ rerank_candidates_with_llm, ) +# LangGraph's first StateGraph().compile() in a process pays a one-time +# lazy-import/compile cost (observed ~0.5s), unrelated to anything under +# test. Pay it here, at module load, so timing-sensitive assertions (e.g. +# test_llm_timeout_falls_back) measure only our own timeout mechanism, both +# in isolation and as part of the full suite. +build_rerank_graph() + def _record(**overrides) -> CheatsheetRecord: defaults = dict( @@ -121,10 +128,13 @@ def stub(system, user, *, model): results = rerank_candidates_with_llm( _record(), _candidates(), llm_score_fn=stub, top_n=5 ) - ids = {r.cre_id for r in results} - self.assertNotIn("999-999", ids) - # the un-scored real candidate still gets a retrieval-only entry - self.assertIn("123-456", ids) + by_id = {r.cre_id: r for r in results} + self.assertNotIn("999-999", by_id) + # the un-scored real candidate still gets a retrieval-only entry, + # and must always be flagged for review since it was never actually + # judged by the reranker (regardless of its confidence band). + self.assertIn("123-456", by_id) + self.assertTrue(by_id["123-456"].needs_review) def test_llm_exception_falls_back_to_retrieval_score(self): def stub(system, user, *, model): @@ -146,6 +156,7 @@ def slow_stub(system, user, *, model): time.sleep(0.2) return {"ranked": []} + started = time.monotonic() results = rerank_candidates_with_llm( _record(), _candidates(), @@ -153,6 +164,8 @@ def slow_stub(system, user, *, model): top_n=5, timeout_seconds=0.01, ) + elapsed = time.monotonic() - started + self.assertLess(elapsed, 0.15) # well under the 0.2s stub delay self.assertEqual(len(results), 2) self.assertTrue(all(r.trace.fallback_used for r in results)) diff --git a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py index b44b83a1a..75908e663 100644 --- a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py +++ b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py @@ -43,6 +43,7 @@ import os from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError from dataclasses import dataclass +from functools import partial from datetime import datetime, timezone from typing import Any, Callable, Dict, List, Optional, TypedDict @@ -55,12 +56,10 @@ # --------------------------------------------------------------------------- # Confidence thresholds (RFC section 11 bootstrap defaults; recalibrate via env) # --------------------------------------------------------------------------- -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") -) +_HIGH_THRESHOLD_ENV_VAR = "CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD" +_MEDIUM_THRESHOLD_ENV_VAR = "CRE_CHEATSHEET_RERANK_MEDIUM_THRESHOLD" +_DEFAULT_HIGH_THRESHOLD = "0.85" +_DEFAULT_MEDIUM_THRESHOLD = "0.70" # Identifiers persisted into RerankTrace for the RFC audit trail (mirrors # RETRIEVER_NAME / RERANKER_NAME conventions used elsewhere in the codebase). @@ -79,21 +78,44 @@ class RerankError(ValueError): """Base class for reranker construction/usage failures.""" +def _resolve_threshold(env_var: str, default: str) -> float: + raw = os.environ.get(env_var, default) + try: + value = float(raw) + except (TypeError, ValueError) as exc: + raise RerankError(f"{env_var}={raw!r} is not a valid float") from exc + if not (0.0 <= value <= 1.0): + raise RerankError(f"{env_var}={value!r} must be in [0, 1]") + return value + + def classify_confidence(score: float) -> str: """ Map a 0-1 re-rank score to a confidence band ("high" | "medium" | "low"). - Thresholds are the RFC's bootstrap defaults and are recalibratable via - ``CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD`` / ``_MEDIUM_THRESHOLD``. + Thresholds are the RFC's bootstrap defaults (high >= 0.85, medium >= 0.70) + and are recalibratable via ``CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD`` / + ``CRE_CHEATSHEET_RERANK_MEDIUM_THRESHOLD``. Thresholds are re-read from + the environment on every call (rather than cached at import time) so + overrides — including ones set after this module is imported, as in + tests — always take effect. """ if not isinstance(score, (int, float)) or isinstance(score, bool): raise RerankError(f"score must be a number, got {score!r}") if not (0.0 <= float(score) <= 1.0): raise RerankError(f"score must be in [0, 1], got {score!r}") - if score >= HIGH_CONFIDENCE_THRESHOLD: + high = _resolve_threshold(_HIGH_THRESHOLD_ENV_VAR, _DEFAULT_HIGH_THRESHOLD) + medium = _resolve_threshold(_MEDIUM_THRESHOLD_ENV_VAR, _DEFAULT_MEDIUM_THRESHOLD) + if high < medium: + raise RerankError( + f"{_HIGH_THRESHOLD_ENV_VAR}={high!r} must be >= " + f"{_MEDIUM_THRESHOLD_ENV_VAR}={medium!r}" + ) + + if score >= high: return "high" - if score >= MEDIUM_CONFIDENCE_THRESHOLD: + if score >= medium: return "medium" return "low" @@ -211,9 +233,18 @@ def _default_model_name() -> str: ) -def default_llm_score_fn(system: str, user: str, *, model: str) -> Dict[str, Any]: +def default_llm_score_fn( + system: str, user: str, *, model: str, timeout: Optional[float] = None +) -> Dict[str, Any]: """Production LLM call via LiteLLM. Raises on any failure; callers must - handle fallback (this function intentionally does not swallow errors).""" + handle fallback (this function intentionally does not swallow errors). + + ``timeout``, when given, is passed straight through to LiteLLM so the + underlying HTTP request itself is bounded — the wall-clock cutoff in + ``_call_with_timeout`` protects the pipeline either way, but a + request-level timeout lets the worker thread actually terminate instead + of continuing to block on the socket after we've stopped waiting on it. + """ try: import litellm # type: ignore except ImportError as exc: # pragma: no cover - exercised only without litellm @@ -227,6 +258,7 @@ def default_llm_score_fn(system: str, user: str, *, model: str) -> Dict[str, Any ], response_format={"type": "json_object"}, temperature=0.2, + timeout=timeout, ) choices = getattr(resp, "choices", None) if not choices: @@ -243,15 +275,26 @@ def _call_with_timeout( fn: Callable[[], Dict[str, Any]], timeout_seconds: float ) -> Dict[str, Any]: """Run ``fn`` with a hard wall-clock timeout so a hung LLM call can never - block the pipeline; raises on timeout or on any exception from ``fn``.""" - with ThreadPoolExecutor(max_workers=1) as pool: - future = pool.submit(fn) - try: - return future.result(timeout=timeout_seconds) - except FutureTimeoutError as exc: - raise RerankError( - f"LLM re-rank call exceeded {timeout_seconds}s timeout" - ) from exc + block the pipeline; raises on timeout or on any exception from ``fn``. + + Uses an explicit (non-context-manager) executor so a timeout returns to + the caller immediately instead of blocking on ``shutdown(wait=True)`` + for a thread that is still running the (now-abandoned) call. + """ + pool = ThreadPoolExecutor(max_workers=1) + future = pool.submit(fn) + try: + return future.result(timeout=timeout_seconds) + except FutureTimeoutError as exc: + pool.shutdown(wait=False) + raise RerankError( + f"LLM re-rank call exceeded {timeout_seconds}s timeout" + ) from exc + except Exception: + pool.shutdown(wait=False) + raise + else: + pool.shutdown(wait=False) # --------------------------------------------------------------------------- @@ -364,9 +407,13 @@ def _node_classify(state: _RerankState) -> _RerankState: ranked: List[RankedCRE] = [] for c in candidates: entry = scored.get(c.cre_id) + per_candidate_fallback = entry is None if entry is None: # LLM succeeded overall but skipped this one candidate: fall back # to its retrieval score individually rather than dropping it. + # This is unscored, unexplained data — always flag it for review + # even if the raw retrieval score happens to land in a + # medium/high band. entry = { "score": max(0.0, min(1.0, c.score)), "reason": "Not scored by reranker; using retrieval score.", @@ -379,7 +426,9 @@ def _node_classify(state: _RerankState) -> _RerankState: retrieval_score=c.score, confidence=confidence, reason=entry["reason"], - needs_review=(confidence == "low") or fallback_used, + needs_review=(confidence == "low") + or fallback_used + or per_candidate_fallback, trace=trace, ) ) @@ -397,14 +446,14 @@ def build_rerank_graph(): directly in integration tests without going through the convenience wrapper below. """ - from langgraph.graph import StateGraph, END + from langgraph.graph import StateGraph, START, END graph = StateGraph(_RerankState) graph.add_node("rerank", _node_llm_rerank) graph.add_node("fallback", _node_fallback) graph.add_node("classify", _node_classify) - graph.set_entry_point("rerank") + graph.add_edge(START, "rerank") graph.add_conditional_edges( "rerank", _route_after_rerank, {"classify": "classify", "fallback": "fallback"} ) @@ -442,7 +491,12 @@ def rerank_candidates_with_llm( raise RerankError(f"top_n must be > 0, got {top_n}") resolved_model = model_name or _default_model_name() - score_fn = llm_score_fn or default_llm_score_fn + if llm_score_fn is not None: + score_fn = llm_score_fn + else: + # Bind the request-level timeout only for the built-in LiteLLM path; + # injected stubs are not required to accept a ``timeout`` kwarg. + score_fn = partial(default_llm_score_fn, timeout=timeout_seconds) app = build_rerank_graph() result = app.invoke( diff --git a/docs/rfc-llm-rerank.md b/docs/rfc-llm-rerank.md index d02c1e822..5149cfcbe 100644 --- a/docs/rfc-llm-rerank.md +++ b/docs/rfc-llm-rerank.md @@ -12,6 +12,27 @@ The implementation is located in: * `application/utils/external_project_parsers/parsers/cheatsheet_rerank.py` +## Acceptance criteria + +- [ ] **Valid LLM result structure**: every `RankedCRE` returned has a + non-empty `reason`, a `score` in `[0, 1]`, and a `confidence` of + `"high"` / `"medium"` / `"low"` — verified by + `test_successful_rerank_produces_reason_and_confidence`. +- [ ] **Deterministic fallback**: an LLM exception, timeout, malformed JSON, + or an all-hallucinated response never raises out of + `rerank_candidates_with_llm` — it always returns one `RankedCRE` per + input candidate, each with `trace.fallback_used == True` — verified by + `test_llm_exception_falls_back_to_retrieval_score`, + `test_llm_timeout_falls_back`, `test_malformed_json_falls_back`, and + `test_llm_returns_no_valid_candidates_falls_back`. +- [ ] **Auditable trace**: every result's `trace` carries `model`, + `prompt_version`, an ISO-8601 UTC `generated_at`, and + `fallback_used`/`fallback_reason` — verified by the same tests above. +- [ ] **Workstream F compatibility**: `RankedCRE.cre_id`, `.score`, + `.confidence`, and `.reason` map 1:1 onto the RFC's + `candidate_cres[]` entries in `suggestions.json` (section 4), so + Workstream F can serialize a `RankedCRE` list directly. + --- ## Sources for more context @@ -41,7 +62,7 @@ exposes: ### The LangGraph flow -``` +```text START -> rerank --(success)--> classify -> END \--(failure)--> fallback -> classify -> END ``` diff --git a/requirements.txt b/requirements.txt index 41bf7c9da..99c6ad88d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -44,7 +44,7 @@ python-markdown-maker # chat (/rest/v1/completion) — embed prompt via LiteLLM, match with sklearn litellm -langgraph>=0.2,<1 +langgraph>=1.0.10,<2 numpy scipy scikit-learn From be5a9f0dbc49d1e8238efef81d88df58b4f5a888 Mon Sep 17 00:00:00 2001 From: Shreesh Date: Thu, 13 Aug 2026 02:02:44 +0000 Subject: [PATCH 3/6] fix(workstream-e): validate top_n/timeout_seconds before graph execution 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. --- application/tests/cheatsheet_rerank_test.py | 36 +++++++++++++++++++ .../parsers/cheatsheet_rerank.py | 18 ++++++++-- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/application/tests/cheatsheet_rerank_test.py b/application/tests/cheatsheet_rerank_test.py index 5e8594b83..f45e65aef 100644 --- a/application/tests/cheatsheet_rerank_test.py +++ b/application/tests/cheatsheet_rerank_test.py @@ -72,6 +72,42 @@ def test_invalid_top_n_raises(self): with self.assertRaises(RerankError): rerank_candidates_with_llm(_record(), _candidates(), top_n=0) + def test_float_top_n_raises(self): + # a float would otherwise pass the "> 0" check and crash later with + # an opaque TypeError from list slicing deep inside the graph. + with self.assertRaises(RerankError): + rerank_candidates_with_llm(_record(), _candidates(), top_n=2.5) + + def test_boolean_top_n_raises(self): + # bool is an int subclass in Python; reject it explicitly rather + # than silently treating True/False as 1/0. + with self.assertRaises(RerankError): + rerank_candidates_with_llm(_record(), _candidates(), top_n=True) + + def test_zero_timeout_seconds_raises(self): + with self.assertRaises(RerankError): + rerank_candidates_with_llm(_record(), _candidates(), timeout_seconds=0) + + def test_infinite_timeout_seconds_raises(self): + # an infinite timeout would defeat the whole point of the timeout + # guard and could hang the pipeline forever on a stuck LLM call. + with self.assertRaises(RerankError): + rerank_candidates_with_llm( + _record(), _candidates(), timeout_seconds=float("inf") + ) + + def test_boolean_timeout_seconds_raises(self): + with self.assertRaises(RerankError): + rerank_candidates_with_llm(_record(), _candidates(), timeout_seconds=True) + + def test_invalid_params_raise_even_with_empty_candidates(self): + # validation must happen before the empty-candidates early return, + # not be silently skipped by it. + with self.assertRaises(RerankError): + rerank_candidates_with_llm(_record(), [], top_n=0) + with self.assertRaises(RerankError): + rerank_candidates_with_llm(_record(), [], timeout_seconds=-1) + def test_successful_rerank_produces_reason_and_confidence(self): def stub(system, user, *, model): self.assertIn("CHEATSHEET_TITLE", user) diff --git a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py index 75908e663..d87171261 100644 --- a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py +++ b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py @@ -40,6 +40,7 @@ import json import logging +import math import os from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError from dataclasses import dataclass @@ -485,10 +486,23 @@ def rerank_candidates_with_llm( deterministic stub instead. Never raises on LLM failure — falls back to retrieval-only ordering and marks the trace accordingly. """ - if not candidates: - return [] + if not isinstance(top_n, int) or isinstance(top_n, bool): + raise RerankError(f"top_n must be a non-boolean int, got {top_n!r}") if top_n <= 0: raise RerankError(f"top_n must be > 0, got {top_n}") + if not isinstance(timeout_seconds, (int, float)) or isinstance( + timeout_seconds, bool + ): + raise RerankError( + f"timeout_seconds must be a non-boolean number, got {timeout_seconds!r}" + ) + if not math.isfinite(timeout_seconds) or timeout_seconds <= 0: + raise RerankError( + f"timeout_seconds must be a finite number > 0, got {timeout_seconds!r}" + ) + + if not candidates: + return [] resolved_model = model_name or _default_model_name() if llm_score_fn is not None: From 8bc30bb4db82b06e53e9b0493769834584cdc21f Mon Sep 17 00:00:00 2001 From: Shreesh Date: Thu, 13 Aug 2026 18:54:47 +0000 Subject: [PATCH 4/6] refactor(workstream-e): pivot from LLM reranker to LLM rationale generator Per @Abhijeet2409's closing comment on #944 and the merged Module C (#991), candidate retrieval and reranking for the mapping pipeline now live in application/utils/librarian/ (C.1 retrieve, C.2 cross-encoder rerank, C.3 calibration, C.4 decide+emit). This PR's original LLM-rerank approach duplicated C.2 rather than adding to it. What Module C's merged code does not do, and cannot do with a cross-encoder, is explain its pick in prose: - schemas.ProposedLink.rationale is a real field in the RFC wire contract. - emitter.py's _proposed_links() sets rationale=None unconditionally -- no code path anywhere in Module C populates it. - decision_engine.decide() only ever surfaces a top-1 cre_id per chunk (candidate_cre_ids[:1]), so 'explain the pick' is one candidate, not a shortlist. This commit retargets the module at that gap instead: - generate_link_rationale(section_text, cre_id, cre_text, score, ...) runs a small LangGraph flow (generate -> format | generate -> fallback -> format) that asks an LLM for one short, grounded sentence explaining why a CRE matches a cheat sheet, given Module C's calibrated score. - Same fallback discipline as before: any LLM error, timeout (hard wall-clock cutoff), or malformed/empty output produces a short, honest, score-only rationale instead of raising -- a link is never blocked by an LLM hiccup. - classify_confidence is kept as an independent, human-facing utility (Module C's decide() thresholds numerically and has no notion of bands); it is no longer on the path that decides linked vs. review -- that's Module C's call alone. - No dependency on CheatsheetRecord/Workstream D's contract anymore -- operates on plain (section_text, cre_id, cre_text, score), matching what Module C's own Section/CreCandidate objects would supply directly. - Does not import or modify application/utils/librarian/ -- that's an actively developed, separately owned module. Wiring this in as the source of ProposedLink.rationale is proposed for discussion with the Module C maintainers, not applied unilaterally. 18 tests (was 22): dropped tests tied to the old multi-candidate rerank API (top_n sorting/truncation, hallucinated-id handling across a shortlist -- no longer applicable to a single-candidate rationale), kept and adapted everything else (confidence boundaries, validation-before-LLM- call, LLM exception/timeout/malformed/empty-output fallback, two end-to-end graph runs). docs/rfc-llm-rerank.md rewritten to match, with links to #944 and #991. --- application/tests/cheatsheet_rerank_test.py | 264 ++++------ .../parsers/cheatsheet_rerank.py | 449 ++++++++---------- docs/rfc-llm-rerank.md | 146 +++--- 3 files changed, 357 insertions(+), 502 deletions(-) diff --git a/application/tests/cheatsheet_rerank_test.py b/application/tests/cheatsheet_rerank_test.py index f45e65aef..ba2aa345a 100644 --- a/application/tests/cheatsheet_rerank_test.py +++ b/application/tests/cheatsheet_rerank_test.py @@ -1,13 +1,11 @@ import time import unittest -from application.defs.cheatsheet_defs import CheatsheetRecord from application.utils.external_project_parsers.parsers.cheatsheet_rerank import ( - CandidateCRE, RerankError, - build_rerank_graph, + build_rationale_graph, classify_confidence, - rerank_candidates_with_llm, + generate_link_rationale, ) # LangGraph's first StateGraph().compile() in a process pays a one-time @@ -15,29 +13,15 @@ # test. Pay it here, at module load, so timing-sensitive assertions (e.g. # test_llm_timeout_falls_back) measure only our own timeout mechanism, both # in isolation and as part of the full suite. -build_rerank_graph() +build_rationale_graph() -def _record(**overrides) -> CheatsheetRecord: - defaults = dict( - source_id="Secrets_Management_Cheat_Sheet", - title="Secrets Management Cheat Sheet", - hyperlink="https://cheatsheetseries.owasp.org/cheatsheets/Secrets_Management_Cheat_Sheet.html", - summary="Guidance on secure storage, rotation, and operational handling of secrets.", - headings=["Introduction", "Architectural Patterns", "Secret Rotation"], - raw_markdown_path="cheatsheets/Secrets_Management_Cheat_Sheet.md", - ) - defaults.update(overrides) - return CheatsheetRecord(**defaults) - - -def _candidates(): - return [ - CandidateCRE( - cre_id="623-550", score=0.62, text="Operational secret rotation controls." - ), - CandidateCRE(cre_id="123-456", score=0.40, text="Unrelated logging guidance."), - ] +SECTION_TEXT = ( + "Secrets Management Cheat Sheet: guidance on secure storage, rotation, " + "and operational handling of secrets." +) +CRE_ID = "623-550" +CRE_TEXT = "Operational secret rotation controls." class ClassifyConfidenceTest(unittest.TestCase): @@ -64,183 +48,141 @@ def test_non_numeric_raises(self): classify_confidence("high") # type: ignore[arg-type] -class RerankCandidatesWithLlmTest(unittest.TestCase): - def test_empty_candidates_returns_empty(self): - self.assertEqual(rerank_candidates_with_llm(_record(), []), []) - - def test_invalid_top_n_raises(self): +class GenerateLinkRationaleTest(unittest.TestCase): + def test_empty_cre_id_raises(self): with self.assertRaises(RerankError): - rerank_candidates_with_llm(_record(), _candidates(), top_n=0) + generate_link_rationale(SECTION_TEXT, "", CRE_TEXT, 0.9) - def test_float_top_n_raises(self): - # a float would otherwise pass the "> 0" check and crash later with - # an opaque TypeError from list slicing deep inside the graph. - with self.assertRaises(RerankError): - rerank_candidates_with_llm(_record(), _candidates(), top_n=2.5) + def test_invalid_score_raises_before_llm_call(self): + called = {"n": 0} + + def stub(system, user, *, model): + called["n"] += 1 + return {"rationale": "x"} - def test_boolean_top_n_raises(self): - # bool is an int subclass in Python; reject it explicitly rather - # than silently treating True/False as 1/0. with self.assertRaises(RerankError): - rerank_candidates_with_llm(_record(), _candidates(), top_n=True) + generate_link_rationale( + SECTION_TEXT, CRE_ID, CRE_TEXT, 1.5, llm_rationale_fn=stub + ) + self.assertEqual(called["n"], 0) def test_zero_timeout_seconds_raises(self): with self.assertRaises(RerankError): - rerank_candidates_with_llm(_record(), _candidates(), timeout_seconds=0) + generate_link_rationale( + SECTION_TEXT, CRE_ID, CRE_TEXT, 0.9, timeout_seconds=0 + ) def test_infinite_timeout_seconds_raises(self): - # an infinite timeout would defeat the whole point of the timeout - # guard and could hang the pipeline forever on a stuck LLM call. with self.assertRaises(RerankError): - rerank_candidates_with_llm( - _record(), _candidates(), timeout_seconds=float("inf") + generate_link_rationale( + SECTION_TEXT, + CRE_ID, + CRE_TEXT, + 0.9, + timeout_seconds=float("inf"), ) def test_boolean_timeout_seconds_raises(self): with self.assertRaises(RerankError): - rerank_candidates_with_llm(_record(), _candidates(), timeout_seconds=True) - - def test_invalid_params_raise_even_with_empty_candidates(self): - # validation must happen before the empty-candidates early return, - # not be silently skipped by it. - with self.assertRaises(RerankError): - rerank_candidates_with_llm(_record(), [], top_n=0) - with self.assertRaises(RerankError): - rerank_candidates_with_llm(_record(), [], timeout_seconds=-1) - - def test_successful_rerank_produces_reason_and_confidence(self): - def stub(system, user, *, model): - self.assertIn("CHEATSHEET_TITLE", user) - self.assertIn("623-550", user) - return { - "ranked": [ - { - "cre_id": "623-550", - "score": 0.91, - "reason": "Directly covers rotation.", - }, - {"cre_id": "123-456", "score": 0.2, "reason": "Off-topic."}, - ] - } - - results = rerank_candidates_with_llm( - _record(), _candidates(), llm_score_fn=stub, top_n=5 - ) - self.assertEqual(len(results), 2) - top = results[0] - self.assertEqual(top.cre_id, "623-550") - self.assertEqual(top.confidence, "high") - self.assertFalse(top.needs_review) - self.assertFalse(top.trace.fallback_used) - self.assertEqual(top.trace.prompt_version, "v1") - self.assertIn("rotation", top.reason.lower()) - self.assertEqual(results[1].confidence, "low") - self.assertTrue(results[1].needs_review) - - def test_top_n_truncates_and_sorts_descending(self): - def stub(system, user, *, model): - return { - "ranked": [ - {"cre_id": "623-550", "score": 0.3, "reason": "r1"}, - {"cre_id": "123-456", "score": 0.95, "reason": "r2"}, - ] - } - - results = rerank_candidates_with_llm( - _record(), _candidates(), llm_score_fn=stub, top_n=1 - ) - self.assertEqual(len(results), 1) - self.assertEqual(results[0].cre_id, "123-456") + generate_link_rationale( + SECTION_TEXT, CRE_ID, CRE_TEXT, 0.9, timeout_seconds=True + ) - def test_hallucinated_cre_id_is_dropped(self): + def test_successful_generation_produces_rationale_and_trace(self): def stub(system, user, *, model): - return { - "ranked": [ - {"cre_id": "623-550", "score": 0.9, "reason": "ok"}, - {"cre_id": "999-999", "score": 0.99, "reason": "invented"}, - ] - } + self.assertIn("CHEATSHEET_TEXT", user) + self.assertIn(CRE_ID, user) + return {"rationale": "Both cover secret rotation controls directly."} - results = rerank_candidates_with_llm( - _record(), _candidates(), llm_score_fn=stub, top_n=5 + result = generate_link_rationale( + SECTION_TEXT, CRE_ID, CRE_TEXT, 0.91, llm_rationale_fn=stub ) - by_id = {r.cre_id: r for r in results} - self.assertNotIn("999-999", by_id) - # the un-scored real candidate still gets a retrieval-only entry, - # and must always be flagged for review since it was never actually - # judged by the reranker (regardless of its confidence band). - self.assertIn("123-456", by_id) - self.assertTrue(by_id["123-456"].needs_review) - - def test_llm_exception_falls_back_to_retrieval_score(self): + self.assertEqual(result.cre_id, CRE_ID) + self.assertIn("rotation", result.rationale.lower()) + self.assertEqual(result.confidence, "high") + self.assertFalse(result.fallback_used) + self.assertFalse(result.trace.fallback_used) + self.assertEqual(result.trace.prompt_version, "v2") + self.assertIsNone(result.trace.fallback_reason) + + def test_llm_exception_falls_back_to_score_only_rationale(self): def stub(system, user, *, model): raise RuntimeError("provider unavailable") - results = rerank_candidates_with_llm( - _record(), _candidates(), llm_score_fn=stub, top_n=5 + result = generate_link_rationale( + SECTION_TEXT, CRE_ID, CRE_TEXT, 0.42, llm_rationale_fn=stub ) - self.assertEqual(len(results), 2) - for r in results: - self.assertTrue(r.trace.fallback_used) - self.assertIsNotNone(r.trace.fallback_reason) - self.assertTrue(r.needs_review) - # retrieval ordering preserved (0.62 > 0.40) - self.assertEqual(results[0].cre_id, "623-550") + self.assertTrue(result.fallback_used) + self.assertTrue(result.trace.fallback_used) + self.assertIsNotNone(result.trace.fallback_reason) + self.assertIn("0.42", result.rationale) def test_llm_timeout_falls_back(self): def slow_stub(system, user, *, model): time.sleep(0.2) - return {"ranked": []} + return {"rationale": "too slow"} started = time.monotonic() - results = rerank_candidates_with_llm( - _record(), - _candidates(), - llm_score_fn=slow_stub, - top_n=5, + result = generate_link_rationale( + SECTION_TEXT, + CRE_ID, + CRE_TEXT, + 0.6, + llm_rationale_fn=slow_stub, timeout_seconds=0.01, ) elapsed = time.monotonic() - started self.assertLess(elapsed, 0.15) # well under the 0.2s stub delay - self.assertEqual(len(results), 2) - self.assertTrue(all(r.trace.fallback_used for r in results)) + self.assertTrue(result.fallback_used) def test_malformed_json_falls_back(self): def bad_stub(system, user, *, model): - return {"not_ranked_key": []} + return {"not_rationale_key": "x"} - results = rerank_candidates_with_llm( - _record(), _candidates(), llm_score_fn=bad_stub, top_n=5 + result = generate_link_rationale( + SECTION_TEXT, CRE_ID, CRE_TEXT, 0.6, llm_rationale_fn=bad_stub ) - self.assertTrue(all(r.trace.fallback_used for r in results)) + self.assertTrue(result.fallback_used) - def test_llm_returns_no_valid_candidates_falls_back(self): + def test_empty_rationale_falls_back(self): def empty_stub(system, user, *, model): - return { - "ranked": [{"cre_id": "not-a-real-id", "score": 0.5, "reason": "x"}] - } + return {"rationale": ""} + + result = generate_link_rationale( + SECTION_TEXT, CRE_ID, CRE_TEXT, 0.6, llm_rationale_fn=empty_stub + ) + self.assertTrue(result.fallback_used) + + def test_missing_cre_text_still_produces_a_rationale(self): + # cre_text may be empty (RFC's original CandidateCRE allowed this); + # the prompt degrades gracefully rather than crashing. + def stub(system, user, *, model): + self.assertIn("", user) + return {"rationale": "Plausible match based on cheat sheet alone."} - results = rerank_candidates_with_llm( - _record(), _candidates(), llm_score_fn=empty_stub, top_n=5 + result = generate_link_rationale( + SECTION_TEXT, CRE_ID, "", 0.75, llm_rationale_fn=stub ) - self.assertTrue(all(r.trace.fallback_used for r in results)) + self.assertFalse(result.fallback_used) + self.assertEqual(result.confidence, "medium") -class RerankGraphIntegrationTest(unittest.TestCase): - """End-to-end execution of the compiled LangGraph flow (RFC Issue E, Checkpoint E5).""" +class RationaleGraphIntegrationTest(unittest.TestCase): + """End-to-end execution of the compiled LangGraph flow.""" def test_graph_runs_success_path(self): - app = build_rerank_graph() + app = build_rationale_graph() def stub(system, user, *, model): - return {"ranked": [{"cre_id": "623-550", "score": 0.88, "reason": "match"}]} + return {"rationale": "match"} state = app.invoke( { - "record": _record(), - "candidates": [_candidates()[0]], - "top_n": 5, - "llm_score_fn": stub, + "section_text": SECTION_TEXT, + "cre_id": CRE_ID, + "cre_text": CRE_TEXT, + "score": 0.88, + "llm_rationale_fn": stub, "model_name": "test-model", "timeout_seconds": 5.0, "generated_at": "2026-08-13T00:00:00+00:00", @@ -248,21 +190,22 @@ def stub(system, user, *, model): "fallback_reason": None, } ) - self.assertEqual(len(state["ranked"]), 1) - self.assertEqual(state["ranked"][0].confidence, "high") + self.assertEqual(state["result"].confidence, "high") + self.assertFalse(state["result"].fallback_used) def test_graph_runs_fallback_path(self): - app = build_rerank_graph() + app = build_rationale_graph() def failing_stub(system, user, *, model): raise RuntimeError("boom") state = app.invoke( { - "record": _record(), - "candidates": _candidates(), - "top_n": 5, - "llm_score_fn": failing_stub, + "section_text": SECTION_TEXT, + "cre_id": CRE_ID, + "cre_text": CRE_TEXT, + "score": 0.3, + "llm_rationale_fn": failing_stub, "model_name": "test-model", "timeout_seconds": 5.0, "generated_at": "2026-08-13T00:00:00+00:00", @@ -270,8 +213,7 @@ def failing_stub(system, user, *, model): "fallback_reason": None, } ) - self.assertEqual(len(state["ranked"]), 2) - self.assertTrue(all(r.trace.fallback_used for r in state["ranked"])) + self.assertTrue(state["result"].fallback_used) if __name__ == "__main__": diff --git a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py index d87171261..a526f91b1 100644 --- a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py +++ b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py @@ -1,39 +1,66 @@ """ -RFC Workstream E — LLM Re-Rank and Decision Graph (LangGraph). +RFC Workstream E — LLM rationale generation for cheat-sheet CRE links. See: docs/rfc/cheatsheets-llm-autonomous-mapping-rfc.md, section 5 -("Workstream E: LLM Re-Rank and Decision Graph (LangGraph)") and the -Issue E checklist in section 12. - -This module owns the "ReRank/Explain -> Threshold" stage of the overall -pipeline (docs/rfc section 7): given a ``CheatsheetRecord`` (Workstream B, -``application/defs/cheatsheet_defs.py``) and the top-k ``CandidateCRE`` -shortlist for it (Workstream D, ``retrieve_candidate_cres``), it asks an LLM -to re-rank and justify the shortlist, assigns a confidence band to each -result, and always returns a usable ``RankedCRE`` list — even when the LLM -call fails, times out, or returns malformed output — by falling back to the -retrieval-only ordering. - -Design notes ------------- -* ``CandidateCRE`` is defined *here* rather than imported from Workstream D - because that workstream's ``retrieve_candidate_cres`` has not landed yet. - The field set (``cre_id``, ``score``, ``text``) mirrors the RFC's - ``CandidateCRE`` contract exactly, so swapping in the real Workstream D - output only requires constructing this same dataclass. -* The LLM call is dependency-injected as ``llm_score_fn`` — a plain - ``(system, user) -> dict`` callable — exactly like the ``ai_client`` seam - in ``application/prompt_client/embed_alignment.py`` and the ``score_fn`` - seam in ``application/utils/librarian/cross_encoder.py``. Production code - never has to inject anything (a LiteLLM-backed default is wired lazily so - this module stays import-light for tests); the test suite and any - harness inject a deterministic stub instead, which keeps the LangGraph - flow hermetically testable. -* Confidence bands and thresholds follow the RFC's bootstrap defaults - (section 11 "Open Questions"): high >= 0.85, medium >= 0.70, else low. - Both are overridable via environment variables so they can be - recalibrated later against PR #865-derived precision/recall data without - a code change. +("Workstream E: LLM Re-Rank and Decision Graph (LangGraph)"). + +## Why this module looks different from the original Workstream E scope + +The RFC originally scoped Workstream E as "LLM re-rank the top-k candidates +from Workstream D." Since then, retrieval *and* reranking for the mapping +pipeline have consolidated onto Module C ("The Librarian", +``application/utils/librarian/``): C.1 retrieves candidates, C.2 reranks them +with a cross-encoder, C.3 calibrates confidence, and C.4 decides + emits an +RFC ``LinkProposal`` or ``ReviewItem`` (see PR #944's closing comment and +PR #991). Building a second, LLM-based reranker on top of that would compete +with C.2 rather than add anything. + +What Module C's merged code does *not* do, and structurally cannot do with a +cross-encoder, is explain its pick in prose. Concretely: + +* ``schemas.ProposedLink.rationale: Optional[str]`` is a real field in the + RFC wire contract. +* ``emitter.py``'s ``_proposed_links()`` sets ``rationale=None`` on every + single link, unconditionally -- there is no code path anywhere in Module C + that populates it. +* ``decision_engine.decide()`` only ever surfaces a single top-1 + ``cre_id`` per chunk (``candidate_cre_ids[:1]``), so the scope of "explain + the pick" is one candidate, not a shortlist. + +This module fills exactly that gap: given the section text and the one CRE +Module C already chose (id, text, and its calibrated score), it asks an LLM +for a short, grounded, one-sentence rationale -- the same LLM-reasoning +capability Workstream E was always meant to contribute, retargeted at the +one place in the pipeline that's actually missing it, instead of duplicating +C.2's scoring job. + +This is a discussion-first proposal, not a fait accompli: wiring +``generate_link_rationale`` into ``application/utils/librarian/emitter.py`` +is left to the Module C maintainers to decide on, since that module is an +actively developed, separately owned GSoC deliverable. This file stays +self-contained and does not import or modify anything under +``application/utils/librarian/``. + +## Design notes (mostly carried over from the original implementation) + +* The LLM call is dependency-injected as ``llm_rationale_fn`` -- the same + seam pattern used throughout this codebase (``ai_client`` in + ``embed_alignment.py``, ``score_fn`` in ``librarian/cross_encoder.py``). + Production defaults to a lazily-imported LiteLLM call; tests inject a + deterministic stub, keeping the whole flow hermetically testable. +* A small LangGraph flow (generate -> format | generate -> fallback -> + format) still backs the public entrypoint, per the RFC's "Decision Graph + (LangGraph)" framing -- now scoped to one node's worth of real work + (generate a rationale) plus its fallback, rather than a multi-stage + rerank pipeline that would have duplicated C.2/C.3/C.4. +* Never raises on LLM failure -- falls back to a short, honest, templated + rationale ("Retrieval/rerank score S; LLM explanation unavailable.") so a + cheat sheet's link is never blocked or degraded by an LLM hiccup. +* ``classify_confidence`` is kept as a small, independently useful utility + for human-facing review UIs (Module C's own ``decide()`` thresholds + numerically and has no notion of confidence *bands*), but is no longer on + the path that decides whether something links or gets reviewed -- that + call is Module C's alone. """ from __future__ import annotations @@ -46,12 +73,10 @@ from dataclasses import dataclass from functools import partial from datetime import datetime, timezone -from typing import Any, Callable, Dict, List, Optional, TypedDict +from typing import Any, Callable, Dict, Optional, TypedDict from pydantic import BaseModel, Field, ValidationError -from application.defs.cheatsheet_defs import CheatsheetRecord - logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- @@ -62,12 +87,11 @@ _DEFAULT_HIGH_THRESHOLD = "0.85" _DEFAULT_MEDIUM_THRESHOLD = "0.70" -# Identifiers persisted into RerankTrace for the RFC audit trail (mirrors -# RETRIEVER_NAME / RERANKER_NAME conventions used elsewhere in the codebase). -RERANKER_NAME = "llm-cheatsheet-reranker" -PROMPT_VERSION = "v1" +# Identifiers persisted into RationaleTrace for the RFC audit trail (mirrors +# RETRIEVER_NAME / RERANKER_NAME conventions in application/utils/librarian/). +RATIONALE_GENERATOR_NAME = "llm-cheatsheet-rationale-generator" +PROMPT_VERSION = "v2" -DEFAULT_TOP_N = 5 DEFAULT_TIMEOUT_SECONDS = 30.0 DEFAULT_MODEL_ENV_VAR = "CRE_CHEATSHEET_RERANK_MODEL" DEFAULT_MODEL_FALLBACK = "gemini/gemini-2.5-flash" @@ -76,7 +100,7 @@ class RerankError(ValueError): - """Base class for reranker construction/usage failures.""" + """Base class for this module's construction/usage failures.""" def _resolve_threshold(env_var: str, default: str) -> float: @@ -92,14 +116,14 @@ def _resolve_threshold(env_var: str, default: str) -> float: def classify_confidence(score: float) -> str: """ - Map a 0-1 re-rank score to a confidence band ("high" | "medium" | "low"). - - Thresholds are the RFC's bootstrap defaults (high >= 0.85, medium >= 0.70) - and are recalibratable via ``CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD`` / - ``CRE_CHEATSHEET_RERANK_MEDIUM_THRESHOLD``. Thresholds are re-read from - the environment on every call (rather than cached at import time) so - overrides — including ones set after this module is imported, as in - tests — always take effect. + Map a 0-1 score to a confidence band ("high" | "medium" | "low"). + + Independent utility for human-facing review UIs -- Module C's own + ``decide()`` thresholds numerically and has no notion of bands; this + does not feed back into that decision. Thresholds are the RFC's + bootstrap defaults (high >= 0.85, medium >= 0.70), recalibratable via + ``CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD`` / + ``CRE_CHEATSHEET_RERANK_MEDIUM_THRESHOLD``, re-read on every call. """ if not isinstance(score, (int, float)) or isinstance(score, bool): raise RerankError(f"score must be a number, got {score!r}") @@ -127,24 +151,8 @@ def classify_confidence(score: float) -> str: @dataclass(frozen=True) -class CandidateCRE: - """ - One retrieval-stage candidate for a CheatsheetRecord. - - Mirrors the RFC's Workstream D output contract. ``text`` is optional - context (e.g. the CRE's embeddings_content) given to the LLM so it can - judge fit; when absent the LLM is told only the cre_id, which degrades - rationale quality but never breaks the flow. - """ - - cre_id: str - score: float - text: str = "" - - -@dataclass(frozen=True) -class RerankTrace: - """Audit metadata captured for every rerank run (RFC Issue E, criterion 3).""" +class RationaleTrace: + """Audit metadata captured for every rationale-generation run.""" model: str prompt_version: str @@ -154,36 +162,32 @@ class RerankTrace: @dataclass(frozen=True) -class RankedCRE: - """One re-ranked, explained candidate — Workstream E's output contract.""" +class LinkRationale: + """ + One CRE link's generated rationale -- meant to fill + ``librarian.schemas.ProposedLink.rationale``, which every current + Module C code path leaves ``None``. + """ cre_id: str - score: float - retrieval_score: float + rationale: str confidence: str - reason: str - needs_review: bool - trace: RerankTrace + fallback_used: bool + trace: RationaleTrace # --------------------------------------------------------------------------- -# LLM structured-output schema (strict; mirrors embed_alignment.AlignmentPayload) +# LLM structured-output schema # --------------------------------------------------------------------------- -class _RerankItem(BaseModel): - cre_id: str - score: float = Field(ge=0.0, le=1.0) - reason: str = "" - - -class _RerankPayload(BaseModel): - ranked: List[_RerankItem] +class _RationalePayload(BaseModel): + rationale: str = Field(min_length=1, max_length=REASON_MAX_LENGTH) -def rerank_response_json_schema() -> Dict[str, Any]: +def rationale_response_json_schema() -> Dict[str, Any]: """Provider-friendly JSON schema for strict structured LLM outputs.""" - return _RerankPayload.model_json_schema() + return _RationalePayload.model_json_schema() # --------------------------------------------------------------------------- @@ -193,36 +197,30 @@ def rerank_response_json_schema() -> Dict[str, Any]: def _system_prompt() -> str: return ( - "You map an OWASP cheat sheet to the Common Requirement (CRE) entries " - "it best satisfies. You will be given the cheat sheet's title, summary, " - "and headings, plus a shortlist of candidate CREs with their ids. " - "Score how well each candidate CRE matches the cheat sheet's content on " - "a 0.0-1.0 scale (1.0 = the cheat sheet is clearly authoritative " - "guidance for that CRE), and give a short one-sentence reason for each " - "score, grounded in the cheat sheet's actual headings/summary. " - "Only score cre_ids given to you; never invent new ones. " - "Return ONLY valid JSON of the form " - '{"ranked": [{"cre_id": "...", "score": 0.0, "reason": "..."}]}, ' - "one entry per candidate given." + "You explain why an OWASP cheat sheet is a good match for a specific " + "Common Requirement (CRE) entry. You will be given the cheat sheet's " + "text, the CRE's id and text, and Module C's calibrated match score. " + "Write ONE short, concrete sentence (max 60 words) grounded in the " + "actual cheat sheet content and CRE text -- do not restate the score, " + "do not invent facts not present in either text. " + 'Return ONLY valid JSON of the form {"rationale": "..."}.' ) -def _user_payload(record: CheatsheetRecord, candidates: List[CandidateCRE]) -> str: +def _user_payload(section_text: str, cre_id: str, cre_text: str, score: float) -> str: lines = [ - f"CHEATSHEET_TITLE: {record.title}", - f"CHEATSHEET_SUMMARY: {record.summary}", - "CHEATSHEET_HEADINGS: " + "; ".join(record.headings), + f"CHEATSHEET_TEXT: {section_text[:2000]}", + "", + f"CRE_ID: {cre_id}", + f"CRE_TEXT: {(cre_text or '')[:800]}", "", - "CANDIDATE_CRES (cre_id | text):", + f"CALIBRATED_SCORE: {score:.3f}", ] - for c in candidates: - text_preview = (c.text or "")[:800] - lines.append(f"{c.cre_id} | {text_preview}") return "\n".join(lines) # --------------------------------------------------------------------------- -# Default (production) LLM call — lazy litellm import so this module stays +# Default (production) LLM call -- lazy litellm import so this module stays # import-light and hermetically testable without a real LLM dependency. # --------------------------------------------------------------------------- @@ -234,22 +232,23 @@ def _default_model_name() -> str: ) -def default_llm_score_fn( +def default_llm_rationale_fn( system: str, user: str, *, model: str, timeout: Optional[float] = None ) -> Dict[str, Any]: """Production LLM call via LiteLLM. Raises on any failure; callers must handle fallback (this function intentionally does not swallow errors). ``timeout``, when given, is passed straight through to LiteLLM so the - underlying HTTP request itself is bounded — the wall-clock cutoff in - ``_call_with_timeout`` protects the pipeline either way, but a - request-level timeout lets the worker thread actually terminate instead - of continuing to block on the socket after we've stopped waiting on it. + underlying HTTP request itself is bounded, letting the worker thread in + ``_call_with_timeout`` actually terminate rather than just being + abandoned. """ try: import litellm # type: ignore except ImportError as exc: # pragma: no cover - exercised only without litellm - raise RerankError("litellm package is required for LLM re-rank calls") from exc + raise RerankError( + "litellm package is required for LLM rationale calls" + ) from exc resp = litellm.completion( model=model, @@ -289,7 +288,7 @@ def _call_with_timeout( except FutureTimeoutError as exc: pool.shutdown(wait=False) raise RerankError( - f"LLM re-rank call exceeded {timeout_seconds}s timeout" + f"LLM rationale call exceeded {timeout_seconds}s timeout" ) from exc except Exception: pool.shutdown(wait=False) @@ -299,197 +298,146 @@ def _call_with_timeout( # --------------------------------------------------------------------------- -# LangGraph flow: rerank -> (success: classify) | (failure: fallback -> classify) +# LangGraph flow: generate -> (success: format) | (failure: fallback -> format) # --------------------------------------------------------------------------- -class _RerankState(TypedDict, total=False): - record: CheatsheetRecord - candidates: List[CandidateCRE] - top_n: int - llm_score_fn: Callable[..., Dict[str, Any]] +class _RationaleState(TypedDict, total=False): + section_text: str + cre_id: str + cre_text: str + score: float + llm_rationale_fn: Callable[..., Dict[str, Any]] model_name: str timeout_seconds: float generated_at: str - scored: Dict[str, Dict[str, Any]] # cre_id -> {"score": float, "reason": str} + rationale_text: Optional[str] fallback_used: bool fallback_reason: Optional[str] - ranked: List[RankedCRE] - + result: LinkRationale -def _node_llm_rerank(state: _RerankState) -> _RerankState: - """Call the LLM, validate its output, and record per-candidate scores. - On any failure (LLM error, timeout, malformed JSON, schema violation) - this node records the reason and leaves ``scored`` empty; the - conditional edge below routes to the fallback node instead of raising. - """ - record = state["record"] - candidates = state["candidates"] - llm_score_fn = state["llm_score_fn"] +def _node_generate(state: _RationaleState) -> _RationaleState: + """Call the LLM and validate its output. On any failure, record the + reason and leave ``rationale_text`` unset; the conditional edge below + routes to the fallback node instead of raising.""" + system = _system_prompt() + user = _user_payload( + state["section_text"], state["cre_id"], state["cre_text"], state["score"] + ) + llm_rationale_fn = state["llm_rationale_fn"] model_name = state["model_name"] timeout_seconds = state["timeout_seconds"] - system = _system_prompt() - user = _user_payload(record, candidates) - try: raw = _call_with_timeout( - lambda: llm_score_fn(system, user, model=model_name), timeout_seconds + lambda: llm_rationale_fn(system, user, model=model_name), timeout_seconds ) - payload = _RerankPayload.model_validate(raw) + payload = _RationalePayload.model_validate(raw) except (RerankError, ValidationError, json.JSONDecodeError, TypeError) as exc: - logger.warning("LLM re-rank failed for %s: %s", record.source_id, exc) + logger.warning("LLM rationale failed for %s: %s", state["cre_id"], exc) state["fallback_reason"] = f"{type(exc).__name__}: {exc}"[:REASON_MAX_LENGTH] - state["scored"] = {} + state["rationale_text"] = None return state except Exception as exc: # defensive: never let an unexpected error crash the run logger.warning( - "LLM re-rank failed unexpectedly for %s: %s", record.source_id, exc + "LLM rationale failed unexpectedly for %s: %s", state["cre_id"], exc ) state["fallback_reason"] = f"unexpected:{type(exc).__name__}: {exc}"[ :REASON_MAX_LENGTH ] - state["scored"] = {} + state["rationale_text"] = None return state - known_ids = {c.cre_id for c in candidates} - scored: Dict[str, Dict[str, Any]] = {} - for item in payload.ranked: - if item.cre_id not in known_ids: - logger.info( - "Dropping hallucinated cre_id %r not in candidate shortlist for %s", - item.cre_id, - record.source_id, - ) - continue - scored[item.cre_id] = { - "score": item.score, - "reason": item.reason[:REASON_MAX_LENGTH], - } - - if not scored: - state["fallback_reason"] = "LLM returned no valid scored candidates" - - state["scored"] = scored + state["rationale_text"] = payload.rationale[:REASON_MAX_LENGTH] return state -def _route_after_rerank(state: _RerankState) -> str: - return "classify" if state.get("scored") else "fallback" +def _route_after_generate(state: _RationaleState) -> str: + return "format" if state.get("rationale_text") else "fallback" -def _node_fallback(state: _RerankState) -> _RerankState: - """Retrieval-only scoring: use each candidate's raw similarity as-is.""" +def _node_fallback(state: _RationaleState) -> _RationaleState: + """Deterministic, honest fallback: no LLM prose, just the score.""" state["fallback_used"] = True - state["scored"] = { - c.cre_id: { - "score": max(0.0, min(1.0, c.score)), - "reason": "Retrieval-only score (LLM re-rank unavailable).", - } - for c in state["candidates"] - } + state["rationale_text"] = ( + f"Retrieval/rerank score {state['score']:.2f}; LLM explanation unavailable." + ) return state -def _node_classify(state: _RerankState) -> _RerankState: - candidates = state["candidates"] - scored = state["scored"] - fallback_used = state.get("fallback_used", False) - fallback_reason = state.get("fallback_reason") - trace = RerankTrace( +def _node_format(state: _RationaleState) -> _RationaleState: + trace = RationaleTrace( model=state["model_name"], prompt_version=PROMPT_VERSION, generated_at=state["generated_at"], - fallback_used=fallback_used, - fallback_reason=fallback_reason if fallback_used else None, + fallback_used=state.get("fallback_used", False), + fallback_reason=( + state.get("fallback_reason") if state.get("fallback_used") else None + ), + ) + state["result"] = LinkRationale( + cre_id=state["cre_id"], + rationale=state["rationale_text"], + confidence=classify_confidence(state["score"]), + fallback_used=state.get("fallback_used", False), + trace=trace, ) - - ranked: List[RankedCRE] = [] - for c in candidates: - entry = scored.get(c.cre_id) - per_candidate_fallback = entry is None - if entry is None: - # LLM succeeded overall but skipped this one candidate: fall back - # to its retrieval score individually rather than dropping it. - # This is unscored, unexplained data — always flag it for review - # even if the raw retrieval score happens to land in a - # medium/high band. - entry = { - "score": max(0.0, min(1.0, c.score)), - "reason": "Not scored by reranker; using retrieval score.", - } - confidence = classify_confidence(entry["score"]) - ranked.append( - RankedCRE( - cre_id=c.cre_id, - score=entry["score"], - retrieval_score=c.score, - confidence=confidence, - reason=entry["reason"], - needs_review=(confidence == "low") - or fallback_used - or per_candidate_fallback, - trace=trace, - ) - ) - - ranked.sort(key=lambda r: r.score, reverse=True) - state["ranked"] = ranked[: state["top_n"]] return state -def build_rerank_graph(): +def build_rationale_graph(): """Compile and return the Workstream E LangGraph flow. - Nodes: ``rerank`` -> (``classify`` | ``fallback`` -> ``classify``) -> END. - Exposed standalone so it can be inspected, visualized, or exercised - directly in integration tests without going through the convenience - wrapper below. + Nodes: ``generate`` -> (``format`` | ``fallback`` -> ``format``) -> END. """ from langgraph.graph import StateGraph, START, END - graph = StateGraph(_RerankState) - graph.add_node("rerank", _node_llm_rerank) + graph = StateGraph(_RationaleState) + graph.add_node("generate", _node_generate) graph.add_node("fallback", _node_fallback) - graph.add_node("classify", _node_classify) + graph.add_node("format", _node_format) - graph.add_edge(START, "rerank") + graph.add_edge(START, "generate") graph.add_conditional_edges( - "rerank", _route_after_rerank, {"classify": "classify", "fallback": "fallback"} + "generate", _route_after_generate, {"format": "format", "fallback": "fallback"} ) - graph.add_edge("fallback", "classify") - graph.add_edge("classify", END) + graph.add_edge("fallback", "format") + graph.add_edge("format", END) return graph.compile() # --------------------------------------------------------------------------- -# Public entrypoint (RFC function-level API, section 6) +# Public entrypoint # --------------------------------------------------------------------------- -def rerank_candidates_with_llm( - record: CheatsheetRecord, - candidates: List[CandidateCRE], +def generate_link_rationale( + section_text: str, + cre_id: str, + cre_text: str, + score: float, *, - llm_score_fn: Optional[Callable[..., Dict[str, Any]]] = None, - top_n: int = DEFAULT_TOP_N, + llm_rationale_fn: Optional[Callable[..., Dict[str, Any]]] = None, model_name: Optional[str] = None, timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS, -) -> List[RankedCRE]: +) -> LinkRationale: """ - Re-rank ``candidates`` for ``record`` via the LangGraph flow above. + Generate a natural-language rationale for one CRE link, via the + LangGraph flow above. - ``llm_score_fn`` defaults to a LiteLLM-backed call - (:func:`default_llm_score_fn`); tests and harnesses should inject a - deterministic stub instead. Never raises on LLM failure — falls back to - retrieval-only ordering and marks the trace accordingly. + Intended to fill ``librarian.schemas.ProposedLink.rationale`` for the + single top-1 candidate Module C's ``decide()`` already chose -- this + does not rerank or re-decide anything. + + ``llm_rationale_fn`` defaults to a LiteLLM-backed call + (:func:`default_llm_rationale_fn`); tests and harnesses should inject a + deterministic stub instead. Never raises on LLM failure -- falls back to + a short, honest, score-only rationale and marks the trace accordingly. """ - if not isinstance(top_n, int) or isinstance(top_n, bool): - raise RerankError(f"top_n must be a non-boolean int, got {top_n!r}") - if top_n <= 0: - raise RerankError(f"top_n must be > 0, got {top_n}") + if not isinstance(cre_id, str) or not cre_id.strip(): + raise RerankError(f"cre_id must be a non-empty string, got {cre_id!r}") if not isinstance(timeout_seconds, (int, float)) or isinstance( timeout_seconds, bool ): @@ -500,25 +448,26 @@ def rerank_candidates_with_llm( raise RerankError( f"timeout_seconds must be a finite number > 0, got {timeout_seconds!r}" ) - - if not candidates: - return [] + # classify_confidence performs the score range/type validation; reuse it + # up front so a bad score fails fast, before any LLM call is attempted. + classify_confidence(score) resolved_model = model_name or _default_model_name() - if llm_score_fn is not None: - score_fn = llm_score_fn + if llm_rationale_fn is not None: + score_fn = llm_rationale_fn else: # Bind the request-level timeout only for the built-in LiteLLM path; # injected stubs are not required to accept a ``timeout`` kwarg. - score_fn = partial(default_llm_score_fn, timeout=timeout_seconds) + score_fn = partial(default_llm_rationale_fn, timeout=timeout_seconds) - app = build_rerank_graph() + app = build_rationale_graph() result = app.invoke( { - "record": record, - "candidates": candidates, - "top_n": top_n, - "llm_score_fn": score_fn, + "section_text": section_text, + "cre_id": cre_id, + "cre_text": cre_text, + "score": float(score), + "llm_rationale_fn": score_fn, "model_name": resolved_model, "timeout_seconds": timeout_seconds, "generated_at": datetime.now(timezone.utc).isoformat(), @@ -526,4 +475,4 @@ def rerank_candidates_with_llm( "fallback_reason": None, } ) - return result["ranked"] + return result["result"] diff --git a/docs/rfc-llm-rerank.md b/docs/rfc-llm-rerank.md index 5149cfcbe..58602971d 100644 --- a/docs/rfc-llm-rerank.md +++ b/docs/rfc-llm-rerank.md @@ -1,108 +1,72 @@ -# RFC Workstream E — LLM Re-Rank and Decision Graph (LangGraph) +# RFC Workstream E — LLM rationale generation for cheat-sheet CRE links -This document explains the implementation and behavior of RFC Workstream E -(LLM Re-Rank and Decision Graph) from the Cheatsheet to CRE Mapping RFC. - -The goal of this module is to take the top-k CRE candidates retrieved for a -cheat sheet (Workstream D) and turn them into an explained, confidence-scored -shortlist that Workstream F can persist to `suggestions.json` for human -review. +**Status update:** this module was originally built as an LLM re-rank step +(Workstream D candidates in, ranked/explained shortlist out). Since then, +candidate retrieval *and* reranking for the whole mapping pipeline have +consolidated onto Module C ("The Librarian") — see PR #944's closing comment +and the merged PR #991. This doc describes the module as it stands now, +retargeted at a gap in Module C rather than duplicating its C.2 reranker. +Wiring it into `application/utils/librarian/emitter.py` is left to the +Module C maintainers to decide on; see the discussion on PR #1014. The implementation is located in: * `application/utils/external_project_parsers/parsers/cheatsheet_rerank.py` -## Acceptance criteria - -- [ ] **Valid LLM result structure**: every `RankedCRE` returned has a - non-empty `reason`, a `score` in `[0, 1]`, and a `confidence` of - `"high"` / `"medium"` / `"low"` — verified by - `test_successful_rerank_produces_reason_and_confidence`. -- [ ] **Deterministic fallback**: an LLM exception, timeout, malformed JSON, - or an all-hallucinated response never raises out of - `rerank_candidates_with_llm` — it always returns one `RankedCRE` per - input candidate, each with `trace.fallback_used == True` — verified by - `test_llm_exception_falls_back_to_retrieval_score`, - `test_llm_timeout_falls_back`, `test_malformed_json_falls_back`, and - `test_llm_returns_no_valid_candidates_falls_back`. -- [ ] **Auditable trace**: every result's `trace` carries `model`, - `prompt_version`, an ISO-8601 UTC `generated_at`, and - `fallback_used`/`fallback_reason` — verified by the same tests above. -- [ ] **Workstream F compatibility**: `RankedCRE.cre_id`, `.score`, - `.confidence`, and `.reason` map 1:1 onto the RFC's - `candidate_cres[]` entries in `suggestions.json` (section 4), so - Workstream F can serialize a `RankedCRE` list directly. - ---- - ## Sources for more context * RFC: `docs/rfc/cheatsheets-llm-autonomous-mapping-rfc.md` -* Workstream B (structured extraction) doc: `docs/rfc-structured-extraction.md` - ---- +* Module C: `application/utils/librarian/` (`schemas.py`, `decision_engine.py`, + `cross_encoder.py`, `emitter.py`) +* PR #944 (Workstream D, closed in favor of Module C): + https://github.com/OWASP/OpenCRE/pull/944 +* PR #991 (Module C week 6b — emitter + pipeline glue): + https://github.com/OWASP/OpenCRE/pull/991 -## What Workstream E implements +## The gap this fills -Given a `CheatsheetRecord` (Workstream B's contract, -`application/defs/cheatsheet_defs.py`) and a list of `CandidateCRE` (the -contract Workstream D's `retrieve_candidate_cres` is expected to return — -defined locally here since Workstream D has not landed yet), the module -exposes: +Module C's merged code has no code path that populates +`schemas.ProposedLink.rationale` — `emitter.py`'s `_proposed_links()` sets +`rationale=None` on every link, always. A cross-encoder (C.2) produces a +similarity score, not prose, so nothing downstream of it can fill that field +without an LLM. `decision_engine.decide()` also only ever surfaces a single +top-1 `cre_id` per chunk, so the scope of "explain the pick" is one +candidate, not a shortlist — this module is scoped accordingly. -* `rerank_candidates_with_llm(record, candidates, ...) -> list[RankedCRE]` — - the public entrypoint. Runs the LangGraph flow described below and always - returns a usable, sorted, confidence-scored shortlist. -* `classify_confidence(score: float) -> str` — maps a 0-1 score to - `"high"` / `"medium"` / `"low"` using the RFC's bootstrap thresholds - (`>= 0.85` high, `>= 0.70` medium, else low), overridable via - `CRE_CHEATSHEET_RERANK_HIGH_THRESHOLD` / `CRE_CHEATSHEET_RERANK_MEDIUM_THRESHOLD`. -* `build_rerank_graph()` — compiles and returns the raw LangGraph app, for - direct inspection or integration testing. +## What it implements -### The LangGraph flow +`generate_link_rationale(section_text, cre_id, cre_text, score, ...) -> LinkRationale` +runs a small LangGraph flow: ```text -START -> rerank --(success)--> classify -> END - \--(failure)--> fallback -> classify -> END +START -> generate --(success)--> format -> END + \--(failure)--> fallback -> format -> END ``` -* **`rerank`** — builds a prompt from the cheat sheet's title/summary/headings - and the candidate CREs, calls the injected `llm_score_fn`, and validates the - response against a strict Pydantic schema (`_RerankPayload`, mirroring - `application/prompt_client/embed_alignment.py`'s `AlignmentPayload` - pattern). Any candidate `cre_id` the LLM invents that isn't in the original - shortlist is dropped and logged, never trusted. -* **`fallback`** — runs whenever the LLM call raises, times out - (`timeout_seconds`, default 30s, enforced with a hard wall-clock cutoff), - returns malformed JSON, or scores zero valid candidates. It scores every - candidate using its raw retrieval similarity instead, so the pipeline never - crashes and never silently drops a cheat sheet. -* **`classify`** — assigns a confidence band and a `needs_review` flag - (`true` when confidence is `"low"` or the run used the fallback path) to - every candidate, attaches an audit `RerankTrace` (model name, prompt - version, UTC timestamp, whether fallback was used and why), sorts - descending by score, and truncates to `top_n` (default 5). - -### Dependency injection / testability - -The LLM call is injected as `llm_score_fn: (system, user, *, model) -> dict`, -the same seam pattern used elsewhere in this codebase (`ai_client` in -`embed_alignment.py`, `score_fn` in `application/utils/librarian/cross_encoder.py`). -Production code defaults to `default_llm_score_fn`, a thin LiteLLM wrapper -lazily imported so this module has no hard LLM dependency; tests inject a -deterministic stub, which keeps the graph — including both the success and -fallback paths — hermetically testable without any network or API key. See -`application/tests/cheatsheet_rerank_test.py`. - -### What this module deliberately does not do - -* It does not call Workstream D's retrieval — callers supply `CandidateCRE`s. -* It does not write `suggestions.json` — that's Workstream F - (`build_suggestions` / `write_suggestions_json`), which is expected to - consume `RankedCRE.reason` as the suggestion's `reason` field and - `RankedCRE.confidence` as its `confidence` field. -* It does not decide auto-link vs. review on its own beyond the - `needs_review` hint — Phase 1 is review-first for every suggestion - regardless (RFC section 11), so `needs_review` is informational, not a - gate. +* **`generate`** — asks an LLM for one short, grounded sentence explaining + why `cre_text` matches `section_text`, given Module C's calibrated + `score`. Validated against a strict Pydantic schema (non-empty, capped + length). +* **`fallback`** — runs on any LLM error, timeout (hard wall-clock cutoff, + default 30s), or malformed/empty output. Produces a short, honest, + score-only rationale ("Retrieval/rerank score 0.42; LLM explanation + unavailable.") so a link is never blocked by an LLM hiccup. +* **`format`** — attaches confidence band (`classify_confidence`, RFC + bootstrap thresholds 0.85/0.70, env-overridable) and an audit + `RationaleTrace` (model, prompt version, UTC timestamp, fallback + flag/reason). + +The LLM call is dependency-injected (`llm_rationale_fn`), same seam pattern +as `embed_alignment.py`'s `ai_client` and `librarian/cross_encoder.py`'s +`score_fn`. Production defaults to a lazily-imported LiteLLM call; the full +test suite runs with a stub, no network/API key required. + +## What this module deliberately does not do + +* It does not retrieve or rerank candidates — Module C's C.1/C.2 own that. +* It does not decide auto-link vs. review — Module C's C.4 `decide()` owns + that; `classify_confidence` here is a separate, human-facing utility only. +* It does not modify `application/utils/librarian/` — that's an actively + developed, separately owned module. Wiring this in as the source of + `ProposedLink.rationale` is proposed, not applied, pending discussion with + the Module C maintainers. From d9cc139f87b45a4f72e9224c9a2dcbddc2f8f2e8 Mon Sep 17 00:00:00 2001 From: Shreesh Date: Wed, 19 Aug 2026 01:37:29 +0000 Subject: [PATCH 5/6] fix(workstream-e): move langgraph to requirements-dev.txt (not prod slug) Per review on #1014: requirements.txt adding langgraph>=1.0.10,<2 is a Heroku-slug/supply-chain decision, not a drive-by import -- shouldn't land on the production web slug without that being agreed on separately. - requirements.txt: langgraph removed entirely. - requirements-dev.txt: langgraph>=1.0.10,<2 added, matching the existing sentence-transformers entry ('ML / Librarian ... never install on Heroku') -- same category of dependency, same treatment. - No code changes needed: langgraph was already lazily imported inside build_rationale_graph(), not at module level, so importing cheatsheet_rerank.py (or calling classify_confidence()) never touched langgraph to begin with. Verified by simulating langgraph's absence (patching builtins.__import__) -- module imports and classify_confidence both work fine; only an actual generate_link_rationale() call raises a clean ImportError, as expected for an optional dev dependency. - Confirmed the module's second question too: grepped the whole repo for 'cheatsheet_rerank' -- the only importer is its own test file. No CLI/cre.py wiring, matching the reviewer's read exactly. - CI's Test workflow runs 'make test' -> install-deps-python -> 'pip install -r requirements-dev.txt', so this doesn't change what CI installs or break any test. - docs/rfc-llm-rerank.md and the module docstring both updated to state the dependency footprint explicitly, so this doesn't need rediscovering next review. --- .../parsers/cheatsheet_rerank.py | 11 +++++++++++ docs/rfc-llm-rerank.md | 11 +++++++++++ requirements-dev.txt | 7 +++++++ requirements.txt | 1 - 4 files changed, 29 insertions(+), 1 deletion(-) diff --git a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py index a526f91b1..212af653b 100644 --- a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py +++ b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py @@ -53,6 +53,17 @@ (LangGraph)" framing -- now scoped to one node's worth of real work (generate a rationale) plus its fallback, rather than a multi-stage rerank pipeline that would have duplicated C.2/C.3/C.4. +* ``langgraph`` is a **dev/CI-only dependency**: it lives in + ``requirements-dev.txt``, not ``requirements.txt``, and is imported + lazily inside :func:`build_rationale_graph` -- importing this module, or + calling :func:`classify_confidence`, never touches ``langgraph`` at all. + This mirrors the existing ``sentence-transformers`` entry in + ``requirements-dev.txt`` ("never install on Heroku") for exactly the + same reason: this module has no CLI/``cre.py`` wiring yet (grep the repo + for ``cheatsheet_rerank`` -- the only importer is its own test file), so + it must not grow the production web slug. If/when this is wired into a + live import path, that's the point to revisit whether ``langgraph`` + belongs in ``requirements.txt`` proper. * Never raises on LLM failure -- falls back to a short, honest, templated rationale ("Retrieval/rerank score S; LLM explanation unavailable.") so a cheat sheet's link is never blocked or degraded by an LLM hiccup. diff --git a/docs/rfc-llm-rerank.md b/docs/rfc-llm-rerank.md index 58602971d..d0a8bdcbb 100644 --- a/docs/rfc-llm-rerank.md +++ b/docs/rfc-llm-rerank.md @@ -13,6 +13,17 @@ The implementation is located in: * `application/utils/external_project_parsers/parsers/cheatsheet_rerank.py` +## Dependency footprint + +`langgraph` lives in `requirements-dev.txt`, not `requirements.txt` — it is +a dev/CI-only dependency, lazily imported inside `build_rationale_graph()`. +Importing this module, or calling `classify_confidence()`, never touches +`langgraph`. This matches the existing `sentence-transformers` entry in +`requirements-dev.txt` (Module C's own ML dependency, annotated "never +install on Heroku"), for the same reason: this module has no CLI/`cre.py` +wiring yet, so it must not grow the production web slug. Revisit this once +it's actually wired into a live import path. + ## Sources for more context * RFC: `docs/rfc/cheatsheets-llm-autonomous-mapping-rfc.md` diff --git a/requirements-dev.txt b/requirements-dev.txt index 4b8784bdc..992664d08 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -6,6 +6,13 @@ sentence-transformers>=5.0.0,<6.0.0 setuptools +# Workstream E cheat-sheet LLM rationale generation (application/utils/ +# external_project_parsers/parsers/cheatsheet_rerank.py). Not on the +# production import path (no CLI/cre.py wiring yet) and lazily imported +# inside the function that needs it, so it's dev/CI-only, same as +# sentence-transformers above -- never install on Heroku. +langgraph>=1.0.10,<2 + # importer / embeddings scrape tooling (not needed for prod chat) playwright nltk diff --git a/requirements.txt b/requirements.txt index 99c6ad88d..c55834701 100644 --- a/requirements.txt +++ b/requirements.txt @@ -44,7 +44,6 @@ python-markdown-maker # chat (/rest/v1/completion) — embed prompt via LiteLLM, match with sklearn litellm -langgraph>=1.0.10,<2 numpy scipy scikit-learn From 922f247f913d1031900e863e62b535b1726299eb Mon Sep 17 00:00:00 2001 From: Shreesh Tripurwar Date: Wed, 19 Aug 2026 08:50:56 +0530 Subject: [PATCH 6/6] Update application/utils/external_project_parsers/parsers/cheatsheet_rerank.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Shreesh Tripurwar --- .../parsers/cheatsheet_rerank.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py index 212af653b..2c70a88c2 100644 --- a/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py +++ b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py @@ -378,6 +378,15 @@ def _node_fallback(state: _RationaleState) -> _RationaleState: def _node_format(state: _RationaleState) -> _RationaleState: + rationale_text = state.get("rationale_text") + if not rationale_text: + # Defensive: the router only reaches `format` after `generate` or + # `fallback` has set this, but keep the contract total. + rationale_text = ( + f"Retrieval/rerank score {state['score']:.2f}; " + "LLM explanation unavailable." + ) + state["fallback_used"] = True trace = RationaleTrace( model=state["model_name"], prompt_version=PROMPT_VERSION, @@ -389,7 +398,7 @@ def _node_format(state: _RationaleState) -> _RationaleState: ) state["result"] = LinkRationale( cre_id=state["cre_id"], - rationale=state["rationale_text"], + rationale=rationale_text, confidence=classify_confidence(state["score"]), fallback_used=state.get("fallback_used", False), trace=trace,