Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 99 additions & 0 deletions application/tests/librarian/decision_engine_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
"""Hermetic tests for C.4 — the decision engine (Week 6).

Table-driven over every (confidence, candidates, flag) combination the rule can
see, plus reason-code precedence and the input guards. No key, DB, or model.
"""

import dataclasses
import math
import unittest

from application.utils.librarian.decision_engine import (
ENGINE_NAME,
DecisionError,
DecisionResult,
decide,
)
from application.utils.librarian.schemas import Decision, ReasonCode

TAU = 0.8
CANDS = ("616-305", "764-507", "611-909")


class DecideTest(unittest.TestCase):
def test_links_when_confident_and_unflagged(self):
r = decide(0.95, CANDS, threshold=TAU)
self.assertEqual(r.decision, Decision.linked)
self.assertIsNone(r.reason_code)
self.assertEqual(r.cre_ids, ("616-305",)) # only the top-1 is linked

def test_confidence_exactly_at_threshold_links(self):
# link iff confidence >= threshold — the boundary is inclusive.
r = decide(TAU, CANDS, threshold=TAU)
self.assertEqual(r.decision, Decision.linked)
self.assertIsNone(r.reason_code)

def test_just_below_threshold_reviews(self):
r = decide(TAU - 1e-9, CANDS, threshold=TAU)
self.assertEqual(r.decision, Decision.review)
self.assertEqual(r.reason_code, ReasonCode.below_threshold)
self.assertEqual(r.cre_ids, ("616-305",)) # best-guess suggestion kept

def test_no_candidates_reviews_even_when_confident(self):
r = decide(0.99, (), threshold=TAU)
self.assertEqual(r.decision, Decision.review)
self.assertEqual(r.reason_code, ReasonCode.no_candidates)
self.assertEqual(r.cre_ids, ()) # nothing to suggest

def test_adversarial_flag_reviews_even_when_confident(self):
r = decide(0.99, CANDS, threshold=TAU, adversarial=True)
self.assertEqual(r.decision, Decision.review)
self.assertEqual(r.reason_code, ReasonCode.adversarial_flag)

def test_update_ambiguous_flag_reviews_even_when_confident(self):
r = decide(0.99, CANDS, threshold=TAU, update_ambiguous=True)
self.assertEqual(r.decision, Decision.review)
self.assertEqual(r.reason_code, ReasonCode.update_ambiguous)

def test_precedence_no_candidates_beats_everything(self):
# empty shortlist + a flag + high confidence -> still NO_CANDIDATES.
r = decide(0.99, (), threshold=TAU, adversarial=True, update_ambiguous=True)
self.assertEqual(r.reason_code, ReasonCode.no_candidates)

def test_precedence_adversarial_beats_below_threshold(self):
r = decide(0.10, CANDS, threshold=TAU, adversarial=True)
self.assertEqual(r.reason_code, ReasonCode.adversarial_flag)

def test_precedence_adversarial_beats_update_ambiguous(self):
r = decide(0.99, CANDS, threshold=TAU, adversarial=True, update_ambiguous=True)
self.assertEqual(r.reason_code, ReasonCode.adversarial_flag)

def test_confidence_is_carried_through(self):
for conf in (0.0, 0.42, 0.8, 1.0):
self.assertEqual(decide(conf, CANDS, threshold=TAU).confidence, conf)


class GuardTest(unittest.TestCase):
def test_bad_threshold_rejected(self):
for bad in (-0.1, 1.1, math.nan, math.inf):
with self.assertRaises(DecisionError):
decide(0.5, CANDS, threshold=bad)

def test_bad_confidence_rejected(self):
for bad in (-0.1, 1.1, math.nan, math.inf):
with self.assertRaises(DecisionError):
decide(bad, CANDS, threshold=TAU)


class ResultTest(unittest.TestCase):
def test_engine_name_is_versioned(self):
self.assertRegex(ENGINE_NAME, r"^decision-engine/\d+\.\d+\.\d+$")

def test_result_is_frozen(self):
r = decide(0.95, CANDS, threshold=TAU)
with self.assertRaises(dataclasses.FrozenInstanceError):
r.confidence = 0.1 # type: ignore[misc]


if __name__ == "__main__":
unittest.main()
169 changes: 169 additions & 0 deletions application/tests/librarian/temperature_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
"""Tests for C.3 temperature-scaling calibration (Week 5).

Hermetic and deterministic: synthetic candidate-logit shortlists + 0/1 labels
only — no cross-encoder, DB, or embedding key. Covers the softmax-over-shortlist
confidence, temperature flatten/sharpen behaviour, the NLL fit recovering a known
T and reducing NLL, ECE on perfectly- and mis-calibrated data (hand-checked), and
every guard.
"""

import math
import unittest

import numpy as np
from scipy.special import softmax

from application.utils.librarian.calibration.temperature import (
CALIBRATOR_NAME,
CalibrationError,
DegenerateLabelsError,
TemperatureScaler,
expected_calibration_error,
fit_temperature,
negative_log_likelihood,
)


class TemperatureScalerTest(unittest.TestCase):
def test_confidence_is_softmax_top_mass(self) -> None:
logits = [2.0, 1.0, 0.0]
expected = float(softmax(np.array(logits)).max())
self.assertAlmostEqual(TemperatureScaler(1.0).confidence(logits), expected, 9)

def test_probabilities_sum_to_one_and_peak_at_top(self) -> None:
p = TemperatureScaler(1.0).probabilities([3.0, 1.0, -2.0])
self.assertAlmostEqual(float(p.sum()), 1.0, places=9)
self.assertEqual(int(np.argmax(p)), 0) # highest logit -> highest prob

def test_high_temperature_flattens_toward_uniform(self) -> None:
logits = [3.0, 1.0, 0.0]
hot = TemperatureScaler(1000.0).confidence(logits)
base = TemperatureScaler(1.0).confidence(logits)
self.assertLess(hot, base)
self.assertAlmostEqual(hot, 1.0 / 3.0, places=2) # ~uniform over 3

def test_low_temperature_sharpens_toward_one(self) -> None:
logits = [3.0, 1.0, 0.0]
cold = TemperatureScaler(0.1).confidence(logits)
base = TemperatureScaler(1.0).confidence(logits)
self.assertGreater(cold, base)
self.assertGreater(cold, 0.99)

def test_single_candidate_confidence_is_one(self) -> None:
# softmax over one element is always 1.0, at any temperature.
self.assertAlmostEqual(TemperatureScaler(3.0).confidence([0.42]), 1.0, 9)

def test_empty_shortlist_rejected(self) -> None:
with self.assertRaises(CalibrationError):
TemperatureScaler(1.0).confidence([])

def test_non_positive_or_nonfinite_temperature_rejected(self) -> None:
for bad in (0.0, -1.0, float("inf"), float("nan")):
with self.assertRaises(CalibrationError):
TemperatureScaler(bad)


class NegativeLogLikelihoodTest(unittest.TestCase):
def test_matches_hand_computed_single_pair(self) -> None:
# one shortlist, top-1 correct: loss = -log(softmax([2,0]).max())
conf = float(softmax(np.array([2.0, 0.0])).max())
self.assertAlmostEqual(
negative_log_likelihood([[2.0, 0.0]], [1.0], 1.0), -math.log(conf), 6
)

def test_length_mismatch_rejected(self) -> None:
with self.assertRaises(CalibrationError):
negative_log_likelihood([[1.0, 0.0], [2.0, 0.0]], [1.0], 1.0)

def test_non_positive_temperature_rejected(self) -> None:
with self.assertRaises(CalibrationError):
negative_log_likelihood([[1.0, 0.0]], [1.0], 0.0)


class FitTemperatureTest(unittest.TestCase):
def _synthetic(self, true_t: float, n: int = 4000, k: int = 5):
"""Shortlists whose top-1 correctness is drawn from softmax(logits/true_t).

Data generated with true_t flattens the softmax, so the raw (T=1)
confidence is over-confident and the fit should recover T ~ true_t.
Seeded -> stable.
"""
rng = np.random.default_rng(0)
sets, labels = [], []
for _ in range(n):
z = rng.normal(0.0, 3.0, size=k)
z[::-1].sort() # descending so index 0 is the top-1 (argmax)
p_correct = softmax(z / true_t).max()
sets.append(z.tolist())
labels.append(1.0 if rng.random() < p_correct else 0.0)
return sets, labels

def test_recovers_known_temperature(self) -> None:
sets, labels = self._synthetic(true_t=2.0)
scaler = fit_temperature(sets, labels)
self.assertAlmostEqual(scaler.temperature, 2.0, delta=0.5)

def test_fit_reduces_nll_versus_T1(self) -> None:
sets, labels = self._synthetic(true_t=2.0)
scaler = fit_temperature(sets, labels)
self.assertLess(
negative_log_likelihood(sets, labels, scaler.temperature),
negative_log_likelihood(sets, labels, 1.0),
)

def test_single_class_labels_raise_degenerate(self) -> None:
with self.assertRaises(DegenerateLabelsError):
fit_temperature([[2.0, 0.0], [3.0, 1.0]], [1.0, 1.0])

def test_non_binary_labels_rejected(self) -> None:
with self.assertRaises(CalibrationError):
fit_temperature([[2.0, 0.0], [3.0, 1.0]], [0.0, 2.0])


class ExpectedCalibrationErrorTest(unittest.TestCase):
def test_perfectly_calibrated_is_near_zero(self) -> None:
# bin [0.5,0.6): five correct, five wrong -> acc 0.5 == conf 0.5.
self.assertAlmostEqual(
expected_calibration_error([0.5] * 10, [1.0, 0.0] * 5), 0.0, places=6
)

def test_confidently_wrong_is_large(self) -> None:
self.assertAlmostEqual(
expected_calibration_error([0.9] * 10, [0.0] * 10), 0.9, places=6
)

def test_hand_checked_two_bin_value(self) -> None:
# [0.2,0.3): conf .2 acc 0 gap .2 ; [0.8,0.9): conf .8 acc 1 gap .2
# ECE = 0.5*0.2 + 0.5*0.2 = 0.20
self.assertAlmostEqual(
expected_calibration_error([0.2, 0.2, 0.8, 0.8], [0.0, 0.0, 1.0, 1.0]),
0.20,
places=6,
)

def test_length_mismatch_and_empty_rejected(self) -> None:
with self.assertRaises(CalibrationError):
expected_calibration_error([0.5, 0.6], [1.0])
with self.assertRaises(CalibrationError):
expected_calibration_error([], [])
with self.assertRaises(CalibrationError):
expected_calibration_error([0.5], [1.0], n_bins=0)


class EndToEndTest(unittest.TestCase):
def test_flat_shortlists_are_low_confidence(self) -> None:
# A near-tie shortlist (bad match) -> low top-1 confidence.
self.assertLess(TemperatureScaler(1.0).confidence([0.1, 0.0, -0.1]), 0.45)

def test_peaked_shortlist_is_high_confidence(self) -> None:
# A clear winner -> high top-1 confidence.
self.assertGreater(TemperatureScaler(1.0).confidence([6.0, -1.0, -3.0]), 0.9)


class MetadataTest(unittest.TestCase):
def test_calibrator_name_is_versioned(self) -> None:
self.assertEqual(CALIBRATOR_NAME, "temperature-scaling/0.2.0")


if __name__ == "__main__":
unittest.main()
4 changes: 3 additions & 1 deletion application/utils/librarian/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@
explicit-link resolution).
W3 (C.1): candidate retriever (in-memory + pgvector) + pipeline switch.
W4 (C.2): cross-encoder reranker — re-sorts the C.1 shortlist, fills reranked[].
Calibration + decision routing (C.3-C.4, W5-W6) onward is not built yet.
W5 (C.3): confidence calibration — temperature scaling maps a rerank logit to
an honest probability (fit by NLL on the golden set, gated ECE < 0.10).
Decision routing (C.4, W6) onward is not built yet.

Vendored RFC JSON schemas live under ``_rfc_schemas/``. They are pinned to
upstream/owasp-graph @ 2b1437987768d5ed20fe9ee721ab9a898c4b84af (PR #734).
Expand Down
13 changes: 13 additions & 0 deletions application/utils/librarian/calibration/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
"""Module C.3 — confidence calibration (Week 5).

C.2 (the cross-encoder, W4) emits a raw ranking logit per candidate — great for
ordering, meaningless as confidence (a +1.5 is not "82% sure"). C.3 turns that
logit into an honest probability via **temperature scaling**: ``p = sigmoid(z/T)``
with a single scalar ``T`` fit by negative-log-likelihood on the golden set, and
proves the result honest with **ECE < 0.10**.

The W6 decision engine thresholds that probability (auto-link vs. human review),
so calibration is what makes the threshold trustworthy. Kept dependency-light
(numpy + scipy) and model-free so it stays hermetically testable — mirrors the
C.1/C.2 seams.
"""
Comment on lines +1 to +13

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Docstring describes the rejected sigmoid approach, not the actual softmax-over-shortlist implementation.

This package docstring says calibration is p = sigmoid(z/T) on a single logit. But temperature.py's own module docstring explicitly explains why that single-logit sigmoid approach doesn't work and why it instead calibrates softmax(logits / T) over the whole shortlist. This file should describe the actual algorithm to avoid misleading readers.

📝 Suggested fix
-logit into an honest probability via **temperature scaling**: ``p = sigmoid(z/T)``
-with a single scalar ``T`` fit by negative-log-likelihood on the golden set, and
+logit into an honest probability via **temperature scaling**: calibrating the
+*softmax over the whole shortlist*, ``p = softmax(logits / T)``, with confidence
+being the top-1 mass — a single scalar ``T`` fit by negative-log-likelihood on
+the golden set, and
 proves the result honest with **ECE < 0.10**.
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
"""Module C.3confidence calibration (Week 5).
C.2 (the cross-encoder, W4) emits a raw ranking logit per candidategreat for
ordering, meaningless as confidence (a +1.5 is not "82% sure"). C.3 turns that
logit into an honest probability via **temperature scaling**: ``p = sigmoid(z/T)``
with a single scalar ``T`` fit by negative-log-likelihood on the golden set, and
proves the result honest with **ECE < 0.10**.
The W6 decision engine thresholds that probability (auto-link vs. human review),
so calibration is what makes the threshold trustworthy. Kept dependency-light
(numpy + scipy) and model-free so it stays hermetically testablemirrors the
C.1/C.2 seams.
"""
"""Module C.3confidence calibration (Week 5).
C.2 (the cross-encoder, W4) emits a raw ranking logit per candidategreat for
ordering, meaningless as confidence (a +1.5 is not "82% sure"). C.3 turns that
logit into an honest probability via **temperature scaling**: calibrating the
*softmax over the whole shortlist*, ``p = softmax(logits / T)``, with confidence
being the top-1 massa single scalar ``T`` fit by negative-log-likelihood on
the golden set, and proves the result honest with **ECE < 0.10**.
The W6 decision engine thresholds that probability (auto-link vs. human review),
so calibration is what makes the threshold trustworthy. Kept dependency-light
(numpy + scipy) and model-free so it stays hermetically testablemirrors the
C.1/C.2 seams.
"""
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@application/utils/librarian/calibration/__init__.py` around lines 1 - 13,
Update the package docstring in the module-level documentation to describe the
implemented shortlist-wide softmax calibration used by temperature.py, replacing
the single-logit sigmoid formula and related claims. Explain that logits are
scaled by a fitted scalar temperature and normalized across each candidate
shortlist, while preserving the existing purpose, NLL fitting, and ECE context.

Loading
Loading