Skip to content

Commit 1375b0b

Browse files
week_4: address CodeRabbit findings on OWASP#957
- run_librarian: fix stale docstring (C.2 rerank is built), guard the semantic retrieve/rerank per-section so one bad section can't abort the dry-run batch - knowledge_source: log ValidationError.errors(include_input=False), never the raw queue row (no content leak) - build_golden_dataset: fail loud on ambiguous ASVS->CRE mappings instead of silently picking one (matches build_explicit/build_update) - candidate_retriever: stable descending sort so tied cosine scores are deterministic - cross_encoder: zip(..., strict=True) (B905) - config_loader_test: annotate OVERRIDES as ClassVar (RUF012) - schemas: centralize the schema_version check across the three envelopes - section_validator: extract _validate_or_raise helper for both call sites - requirements: pin sentence-transformers>=5.0,<6.0 (tested with 5.6) - evaluate_librarian: drop the no-op --dry_run flag (harness never writes)
1 parent 72152e5 commit 1375b0b

10 files changed

Lines changed: 68 additions & 37 deletions

File tree

application/cmd/cre_main.py

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1035,10 +1035,10 @@ def run_librarian(
10351035
10361036
For each knowledge-queue section: try the deterministic explicit-CRE fast
10371037
path (C.0.5); on no/ambiguous reference, run the semantic retriever (C.1)
1038-
and log the top-K candidate shortlist. The cross-encoder rerank (C.2, W4),
1039-
decision/threshold routing (C.3-C.4, W5) and graph writes (W8) are not
1040-
built yet, so this is dry-run only: it never writes a link. ``--run_librarian``
1041-
without writes behaves identically and warns.
1038+
and rerank its top-K shortlist with the cross-encoder (C.2, W4), then log
1039+
the reranked candidates. Decision/threshold routing (C.3-C.4, W5) and graph
1040+
writes (W8) are not built yet, so this is dry-run only: it never writes a
1041+
link. ``--run_librarian`` without writes behaves identically and warns.
10421042
10431043
Ops note: this is opt-in CLI only — it is not on the ``Procfile`` and is
10441044
wired into neither the web app nor the background worker (that lands W8).
@@ -1134,9 +1134,19 @@ def run_librarian(
11341134
logger.info("[explicit] %s -> %s", section.chunk_id, resolution.cre_ids[0])
11351135
continue
11361136

1137+
try:
1138+
audit = retriever.retrieve(section.text)
1139+
audit = reranker.rerank(section.text, audit)
1140+
except Exception as exc: # noqa: BLE001 - one bad section must not abort the batch
1141+
rejected += 1
1142+
logger.warning(
1143+
"[semantic] %s skipped: retrieval/rerank failed: %s",
1144+
section.chunk_id,
1145+
exc,
1146+
)
1147+
continue
1148+
11371149
semantic += 1
1138-
audit = retriever.retrieve(section.text)
1139-
audit = reranker.rerank(section.text, audit)
11401150
top = ", ".join(
11411151
f"{c.cre_id}:{c.score_rerank:.3f}" for c in audit.reranked
11421152
)

application/tests/librarian/config_loader_test.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import os
22
import unittest
33
from dataclasses import FrozenInstanceError
4+
from typing import ClassVar, Dict
45
from unittest import mock
56

67
from application.utils.librarian.config_loader import LibrarianConfig, load_config
@@ -27,7 +28,7 @@ def test_config_is_frozen(self):
2728

2829

2930
class TestConfigLoaderOverrides(unittest.TestCase):
30-
OVERRIDES = {
31+
OVERRIDES: ClassVar[Dict[str, str]] = {
3132
"CRE_LIBRARIAN_RETRIEVER_BACKEND": "pgvector",
3233
"CRE_LIBRARIAN_CROSSENCODER_MODEL": "cross-encoder/other",
3334
"CRE_LIBRARIAN_TOP_K_RETRIEVAL": "50",

application/utils/librarian/candidate_retriever.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,9 @@ def retrieve(self, text: str) -> RetrievalAudit:
157157
# Top-K by descending score. argsort is ascending, so take the tail
158158
# and reverse; cap at pool size when the hub is smaller than K.
159159
k = min(self._top_k, len(self._pool.cre_ids))
160-
top_idx = np.argsort(scores)[-k:][::-1]
160+
# Stable descending sort so tied cosine scores keep hub (index) order —
161+
# deterministic across runs. argsort(-scores) descends; kind="stable".
162+
top_idx = np.argsort(-scores, kind="stable")[:k]
161163

162164
candidates: List[CreCandidate] = [
163165
CreCandidate(

application/utils/librarian/cross_encoder.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ def rerank(self, text: str, audit: RetrievalAudit) -> RetrievalAudit:
100100

101101
reranked = [
102102
c.model_copy(update={"score_rerank": float(s)})
103-
for c, s in zip(candidates, scores)
103+
for c, s in zip(candidates, scores, strict=True)
104104
]
105105
# Highest cross-encoder score first, then keep only the best top_n.
106106
# Python's sort is stable, so ties preserve C.1's cosine order.

application/utils/librarian/knowledge_source.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,11 @@ def items(self) -> Iterator[KnowledgeQueueItem]:
3939
try:
4040
yield KnowledgeQueueItem.model_validate_json(line)
4141
except ValidationError as exc:
42+
# Log only error locations/types, never the raw input —
43+
# queue rows can carry sensitive content.
4244
logger.warning(
4345
"Skipping malformed knowledge_queue row at line %d: %s",
4446
line_no,
45-
exc,
47+
exc.errors(include_input=False),
4648
)
4749
continue

application/utils/librarian/schemas.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,15 @@
2222
_SCHEMA_VERSION_RE = re.compile(r"^0\.\d+\.\d+$")
2323

2424

25+
def _require_schema_version(value: str) -> None:
26+
"""Shared envelope check — every RFC envelope pins schema_version to 0.x.y.
27+
28+
Centralized so the regex and message can't drift across envelopes.
29+
"""
30+
if not _SCHEMA_VERSION_RE.match(value):
31+
raise ValueError(r"schema_version must match ^0\.\d+\.\d+$")
32+
33+
2534
# ---------- Enums (RFC) ----------
2635

2736

@@ -212,8 +221,7 @@ class KnowledgeItem(BaseModel):
212221

213222
@model_validator(mode="after")
214223
def _rfc_rules(self) -> "KnowledgeItem":
215-
if not _SCHEMA_VERSION_RE.match(self.schema_version):
216-
raise ValueError(r"schema_version must match ^0\.\d+\.\d+$")
224+
_require_schema_version(self.schema_version)
217225
if self.status == KnowledgeStatus.accepted and self.content is None:
218226
raise ValueError("status=accepted requires content")
219227
if self.status == KnowledgeStatus.rejected and self.rejection is None:
@@ -238,8 +246,7 @@ class LinkProposal(BaseModel):
238246

239247
@model_validator(mode="after")
240248
def _schema_version_pattern(self) -> "LinkProposal":
241-
if not _SCHEMA_VERSION_RE.match(self.schema_version):
242-
raise ValueError(r"schema_version must match ^0\.\d+\.\d+$")
249+
_require_schema_version(self.schema_version)
243250
return self
244251

245252

@@ -262,8 +269,7 @@ class ReviewItem(BaseModel):
262269

263270
@model_validator(mode="after")
264271
def _schema_version_pattern(self) -> "ReviewItem":
265-
if not _SCHEMA_VERSION_RE.match(self.schema_version):
266-
raise ValueError(r"schema_version must match ^0\.\d+\.\d+$")
272+
_require_schema_version(self.schema_version)
267273
return self
268274

269275

application/utils/librarian/section_validator.py

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,9 @@
2828
"""
2929

3030
from dataclasses import dataclass
31-
from typing import Any, Dict, Optional, Union
31+
from typing import Any, Dict, Optional, Type, TypeVar, Union
3232

33-
from pydantic import ValidationError
33+
from pydantic import BaseModel, ValidationError
3434

3535
from application.utils.librarian.schemas import (
3636
KnowledgeItem,
@@ -42,6 +42,8 @@
4242
SourceType,
4343
)
4444

45+
_ModelT = TypeVar("_ModelT", bound=BaseModel)
46+
4547
KNOWLEDGE_LABEL = "KNOWLEDGE"
4648

4749
# MVP scope: golden dataset and CRE hub vectors are English-only.
@@ -82,6 +84,17 @@ class Section:
8284
locator: Locator
8385

8486

87+
def _validate_or_raise(model_cls: Type[_ModelT], raw: Any) -> _ModelT:
88+
"""Coerce a raw row into ``model_cls``, converting Pydantic's ValidationError
89+
into a typed MalformedKnowledgeItemError so it never escapes this module."""
90+
if isinstance(raw, model_cls):
91+
return raw
92+
try:
93+
return model_cls.model_validate(raw)
94+
except ValidationError as exc:
95+
raise MalformedKnowledgeItemError(str(exc)) from exc
96+
97+
8598
def _require_text(text: str) -> str:
8699
if not text or not text.strip():
87100
raise EmptyTextError("section text is empty or whitespace-only")
@@ -107,11 +120,7 @@ def section_from_queue_row(
107120
108121
Raises a SectionValidationError subclass on any rejection.
109122
"""
110-
if not isinstance(row, KnowledgeQueueItem):
111-
try:
112-
row = KnowledgeQueueItem.model_validate(row)
113-
except ValidationError as exc:
114-
raise MalformedKnowledgeItemError(str(exc)) from exc
123+
row = _validate_or_raise(KnowledgeQueueItem, row)
115124

116125
if row.llm_label != KNOWLEDGE_LABEL:
117126
raise NotKnowledgeError(
@@ -150,11 +159,7 @@ def section_from_knowledge_item(
150159
151160
Raises a SectionValidationError subclass on any rejection.
152161
"""
153-
if not isinstance(item, KnowledgeItem):
154-
try:
155-
item = KnowledgeItem.model_validate(item)
156-
except ValidationError as exc:
157-
raise MalformedKnowledgeItemError(str(exc)) from exc
162+
item = _validate_or_raise(KnowledgeItem, item)
158163

159164
if item.status != KnowledgeStatus.accepted:
160165
raise NotKnowledgeError(

requirements.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,7 @@ redis
9292
ruamel.yaml
9393
ruamel.yaml.clib
9494
scikit-learn
95-
sentence-transformers
95+
sentence-transformers>=5.0.0,<6.0.0
9696
Shapely
9797
six
9898
smmap

scripts/build_golden_dataset.py

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -195,19 +195,27 @@
195195

196196

197197
def _fetch_asvs_cre(conn: sqlite3.Connection, section_id: str) -> Optional[str]:
198-
row = conn.execute(
198+
rows = conn.execute(
199199
"""
200-
SELECT c.external_id
200+
SELECT DISTINCT c.external_id
201201
FROM node n
202202
JOIN cre_node_links l ON l.node = n.id
203203
JOIN cre c ON c.id = l.cre
204204
WHERE n.name LIKE '%ASVS%' AND n.section_id = ?
205205
ORDER BY c.external_id
206-
LIMIT 1
207206
""",
208207
(section_id,),
209-
).fetchone()
210-
return row[0] if row else None
208+
).fetchall()
209+
if not rows:
210+
return None
211+
if len(rows) > 1:
212+
# Fail loud rather than silently pick one — an ambiguous section would
213+
# bake a wrong CRE into the golden set (matches build_explicit/build_update).
214+
raise ValueError(
215+
f"ASVS section {section_id!r} maps to multiple CREs "
216+
f"({', '.join(r[0] for r in rows)}); cannot pick a golden CRE"
217+
)
218+
return rows[0][0]
211219

212220

213221
def build_positive_asvs(conn: sqlite3.Connection) -> List[Dict]:

scripts/evaluate_librarian.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -184,9 +184,6 @@ def main(argv: List[str]) -> int:
184184
parser.add_argument("--threshold", type=float, default=cfg.link_threshold)
185185
parser.add_argument("--top_k_retrieval", type=int, default=cfg.top_k_retrieval)
186186
parser.add_argument("--top_k_rerank", type=int, default=cfg.top_k_rerank)
187-
parser.add_argument(
188-
"--dry_run", action="store_true", help="no writes (always true pre-W8)"
189-
)
190187
parser.add_argument(
191188
"--no_hub_firewall",
192189
action="store_true",

0 commit comments

Comments
 (0)