diff --git a/application/tests/cheatsheet_rerank_test.py b/application/tests/cheatsheet_rerank_test.py new file mode 100644 index 000000000..f45e65aef --- /dev/null +++ b/application/tests/cheatsheet_rerank_test.py @@ -0,0 +1,278 @@ +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, +) + +# 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( + 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_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) + 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 + ) + 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): + 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": []} + + started = time.monotonic() + results = rerank_candidates_with_llm( + _record(), + _candidates(), + llm_score_fn=slow_stub, + 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)) + + 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..d87171261 --- /dev/null +++ b/application/utils/external_project_parsers/parsers/cheatsheet_rerank.py @@ -0,0 +1,529 @@ +""" +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 math +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 + +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_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). +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 _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 (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}") + + 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: + 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, 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. + """ + 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, + timeout=timeout, + ) + 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``. + + 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) + + +# --------------------------------------------------------------------------- +# 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) + 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(): + """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, 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.add_edge(START, "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 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: + 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( + { + "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..5149cfcbe --- /dev/null +++ b/docs/rfc-llm-rerank.md @@ -0,0 +1,108 @@ +# 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` + +## 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` + +--- + +## 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 + +```text +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..99c6ad88d 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>=1.0.10,<2 numpy scipy scikit-learn