From a0418c539b5acfc806b5e28b89baca9542a8fbb4 Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Sat, 24 Oct 2026 00:46:42 +0530 Subject: [PATCH 1/5] Add POST /rest/v1/link/suggest: the Librarian's shortlist as an API Fixes #275. Given a CRE Standard document, return the CREs most likely to be its link targets, ranked. This is C.1 of the Librarian exposed directly, so a standard author can ask "where does this section belong?" without running the OIE pipeline, and without a second similarity stack growing beside the one that is evaluated and gated in CI (the maintainer decision on #965). Retrieval only, deliberately: C.2 and C.3 need a torch cross-encoder, which the web process must not load (the same constraint #1161 respects). The score is the retriever's cosine similarity, named as such. Automated mapping with a confidence threshold stays with the worker-side mapping job. factory.build_retriever_for is split out of build_components with no change in behaviour, so the endpoint and the pipeline construct C.1 from one place. Results always carry the public CRE id. The embeddings table keys vectors by the CRE row's internal UUID, so both retriever backends return UUIDs; each candidate is resolved through the graph (internal id first, external second) and a stale embedding for a CRE the graph no longer has is returned with the retriever's label rather than silently dropped. top_k is 1 to 20, default 5, whole numbers only: int() would truncate 2.5 to 2 and read True as 1, and a mistyped limit should be refused, not reinterpreted. Same auth and the same unavailable/rate-limit mapping as /rest/v1/completion, since both call the paid embedding API. 26 hermetic tests: the service with a stub retriever against an in-memory graph (including a UUID-keyed candidate resolving to "287-251"), and the real Flask route with only the builder seam patched. Verified against the real CRE pool in a local cache with a stored vector as the query: the issue's own ASVS V2.9.2 example returns 287-251 first, then 585-408 Challenge nonce cryptography. The paid-API call path is PromptHandler.get_text_embeddings, unchanged and shared with chat. --- .../tests/librarian/link_suggest_test.py | 320 ++++++++++++++++++ application/utils/librarian/factory.py | 62 +++- application/utils/librarian/link_suggest.py | 190 +++++++++++ application/web/web_main.py | 67 ++++ 4 files changed, 621 insertions(+), 18 deletions(-) create mode 100644 application/tests/librarian/link_suggest_test.py create mode 100644 application/utils/librarian/link_suggest.py diff --git a/application/tests/librarian/link_suggest_test.py b/application/tests/librarian/link_suggest_test.py new file mode 100644 index 000000000..e153d368e --- /dev/null +++ b/application/tests/librarian/link_suggest_test.py @@ -0,0 +1,320 @@ +"""Tests for ``POST /rest/v1/link/suggest`` and the service behind it (#275). + +Hermetic: the retriever is a stub returning a fixed ``RetrievalAudit``, so no +embedding call is made. The web tests go through the real Flask route with the +builder seam patched, which is the only way to prove the route's validation, +auth bypass and error mapping without a model or a key. +""" + +from cre_logging import get_logger + +logger = get_logger(__name__) + +import os +import unittest +from unittest import mock + +from application import create_app, sqla +from application.database import db +from application.defs import cre_defs as defs +from application.utils.librarian import link_suggest +from application.utils.librarian.link_suggest import ( + DEFAULT_TOP_K, + MAX_TOP_K, + InvalidStandardError, + clamp_top_k, + standard_query_text, + suggest_links, +) +from application.utils.librarian.schemas import CreCandidate, RetrievalAudit + + +def _audit() -> RetrievalAudit: + return RetrievalAudit( + retriever="stub/1.0.0", + candidates=[ + CreCandidate( + cre_id="287-251", cre_name="from retriever", score_vector=0.91234 + ), + CreCandidate(cre_id="000-001", score_vector=0.5), + ], + reranked=[], + threshold=0.8, + ) + + +class _StubRetriever: + def __init__(self) -> None: + self.calls = [] + + def retrieve(self, text: str) -> RetrievalAudit: + self.calls.append(text) + return _audit() + + +def _standard(**overrides): + doc = { + "doctype": "Standard", + "name": "ASVS", + "sectionID": "V2.9.2", + "section": ( + "Verify that the challenge nonce is at least 64 bits in length, and " + "statistically unique or unique over the lifetime of the cryptographic device." + ), + "hyperlink": "https://github.com/OWASP/ASVS/blob/v4.0.2/4.0/en/0x11-V2-Authentication.md", + "links": [], + } + doc.update(overrides) + return doc + + +class StandardQueryTextTest(unittest.TestCase): + def test_joins_the_text_fields_and_ignores_provenance(self) -> None: + text = standard_query_text( + _standard(subsection="Nonce size", description="Use 64-bit nonces") + ) + self.assertTrue(text.startswith("Verify that the challenge nonce")) + self.assertIn("Nonce size. Use 64-bit nonces", text) + self.assertNotIn("ASVS", text) + self.assertNotIn("V2.9.2", text) + + def test_rejects_non_standard_doctype(self) -> None: + with self.assertRaises(InvalidStandardError): + standard_query_text(_standard(doctype="CRE")) + + def test_rejects_missing_name(self) -> None: + with self.assertRaises(InvalidStandardError): + standard_query_text(_standard(name="")) + + def test_rejects_a_document_with_no_text_to_embed(self) -> None: + with self.assertRaises(InvalidStandardError): + standard_query_text(_standard(section="", subsection="", description="")) + + def test_rejects_non_string_text_fields(self) -> None: + with self.assertRaises(InvalidStandardError): + standard_query_text(_standard(section=["not", "a", "string"])) + + def test_rejects_non_object_bodies(self) -> None: + with self.assertRaises(InvalidStandardError): + standard_query_text(["a", "list"]) # type: ignore[arg-type] + + +class ClampTopKTest(unittest.TestCase): + def test_absent_means_default(self) -> None: + self.assertEqual(clamp_top_k(None), DEFAULT_TOP_K) + self.assertEqual(clamp_top_k(""), DEFAULT_TOP_K) + + def test_parses_and_caps(self) -> None: + self.assertEqual(clamp_top_k("7"), 7) + self.assertEqual(clamp_top_k(99), MAX_TOP_K) + + def test_rejects_non_positive_and_non_integer(self) -> None: + for bad in ("0", "-3", "x", 2.5): + with self.assertRaises(InvalidStandardError): + clamp_top_k(bad) + + +class SuggestLinksTest(unittest.TestCase): + def setUp(self) -> None: + self.app = create_app(mode="test") + self.ctx = self.app.app_context() + self.ctx.push() + sqla.create_all() + self.collection = db.Node_collection() + self.added = self.collection.add_cre( + defs.CRE( + id="287-251", + name="Use a unique challenge nonce of sufficient size", + tags=["Cryptography"], + ) + ) + + def tearDown(self) -> None: + sqla.session.remove() + sqla.drop_all() + self.ctx.pop() + + def test_internal_uuid_candidates_resolve_to_the_public_cre_id(self) -> None: + """The embeddings table keys vectors by the CRE row's internal UUID, so + that is what the retriever returns. A caller must only ever see + "287-251"; this is the bug a stub keyed by external ids would hide.""" + internal_id = self.added.id + + class _UuidRetriever: + def retrieve(self, text): + return RetrievalAudit( + retriever="stub/1.0.0", + candidates=[CreCandidate(cre_id=internal_id, score_vector=0.9)], + reranked=[], + threshold=0.8, + ) + + (result,) = suggest_links( + self.collection, _standard(), retriever=_UuidRetriever() + ) + + self.assertEqual(result["id"], "287-251") + self.assertEqual( + result["name"], "Use a unique challenge nonce of sufficient size" + ) + self.assertEqual(result["tags"], ["Cryptography"]) + + def test_returns_the_issue_275_shape_best_first(self) -> None: + results = suggest_links( + self.collection, _standard(), retriever=_StubRetriever() + ) + + self.assertEqual( + results[0], + { + "doctype": "CRE", + "id": "287-251", + "name": "Use a unique challenge nonce of sufficient size", + "tags": ["Cryptography"], + "score": 0.9123, + }, + ) + + def test_a_candidate_missing_from_the_graph_is_returned_not_dropped(self) -> None: + """A stale embedding for a CRE the graph no longer has is still a + result the caller should see; silently dropping it would hide the + drift that caused it.""" + results = suggest_links( + self.collection, _standard(), retriever=_StubRetriever() + ) + + self.assertEqual(results[1]["id"], "000-001") + self.assertEqual(results[1]["name"], "") + self.assertEqual(results[1]["tags"], []) + + def test_graph_name_wins_over_the_retrievers_label(self) -> None: + results = suggest_links( + self.collection, _standard(), retriever=_StubRetriever() + ) + self.assertNotEqual(results[0]["name"], "from retriever") + + def test_top_k_slices_the_shortlist(self) -> None: + results = suggest_links( + self.collection, _standard(), top_k=1, retriever=_StubRetriever() + ) + self.assertEqual(len(results), 1) + + def test_top_k_outside_bounds_is_refused(self) -> None: + for bad in (0, MAX_TOP_K + 1): + with self.assertRaises(InvalidStandardError): + suggest_links( + self.collection, _standard(), top_k=bad, retriever=_StubRetriever() + ) + + def test_embeds_the_section_text_not_the_whole_document(self) -> None: + stub = _StubRetriever() + suggest_links(self.collection, _standard(), retriever=stub) + self.assertEqual(len(stub.calls), 1) + self.assertTrue(stub.calls[0].startswith("Verify that the challenge nonce")) + + def test_a_supplied_retriever_means_no_builder_is_invoked(self) -> None: + def _must_not_run(database): + raise AssertionError("builder must not be called") + + suggest_links( + self.collection, + _standard(), + retriever=_StubRetriever(), + retriever_factory=_must_not_run, + ) + + def test_invalid_document_fails_before_any_retrieval(self) -> None: + stub = _StubRetriever() + with self.assertRaises(InvalidStandardError): + suggest_links(self.collection, _standard(section=""), retriever=stub) + self.assertEqual(stub.calls, []) + + +class LinkSuggestRouteTest(unittest.TestCase): + """The real Flask route, with only the retriever builder seam patched.""" + + def setUp(self) -> None: + self.app = create_app(mode="test") + self.ctx = self.app.app_context() + self.ctx.push() + sqla.create_all() + os.environ["INSECURE_REQUESTS"] = "True" + os.environ["NO_LOGIN"] = "1" + self.collection = db.Node_collection() + self.collection.add_cre( + defs.CRE( + id="287-251", name="Use a unique challenge nonce", tags=["Cryptography"] + ) + ) + + def tearDown(self) -> None: + os.environ.pop("NO_LOGIN", None) + sqla.session.remove() + sqla.drop_all() + self.ctx.pop() + + def _post(self, body, query=""): + with mock.patch.object( + link_suggest, "_default_retriever", return_value=_StubRetriever() + ): + with self.app.test_client() as c: + return c.post(f"/rest/v1/link/suggest{query}", json=body) + + def test_returns_ranked_cres(self) -> None: + resp = self._post(_standard()) + self.assertEqual(resp.status_code, 200, resp.get_data(as_text=True)) + body = resp.get_json() + self.assertEqual([r["id"] for r in body], ["287-251", "000-001"]) + self.assertEqual(body[0]["doctype"], "CRE") + self.assertEqual(body[0]["tags"], ["Cryptography"]) + + def test_top_k_from_the_query_string(self) -> None: + resp = self._post(_standard(), query="?top_k=1") + self.assertEqual(len(resp.get_json()), 1) + + def test_non_object_body_is_400(self) -> None: + resp = self._post(["not", "an", "object"]) + self.assertEqual(resp.status_code, 400) + + def test_document_without_text_is_400_with_the_reason(self) -> None: + resp = self._post(_standard(section="")) + self.assertEqual(resp.status_code, 400) + self.assertIn("section", resp.get_data(as_text=True)) + + def test_bad_top_k_is_400(self) -> None: + resp = self._post(_standard(), query="?top_k=zero") + self.assertEqual(resp.status_code, 400) + + def test_missing_ml_dependencies_is_503(self) -> None: + with mock.patch.object( + link_suggest, "_default_retriever", side_effect=ImportError("no torch") + ): + with self.app.test_client() as c: + resp = c.post("/rest/v1/link/suggest", json=_standard()) + self.assertEqual(resp.status_code, 503) + + def test_missing_litellm_is_503(self) -> None: + with mock.patch.object( + link_suggest, + "_default_retriever", + side_effect=RuntimeError("litellm is not installed"), + ): + with self.app.test_client() as c: + resp = c.post("/rest/v1/link/suggest", json=_standard()) + self.assertEqual(resp.status_code, 503) + + def test_provider_failure_is_reported_not_swallowed(self) -> None: + class _Boom(Exception): + status_code = 502 + + with mock.patch.object( + link_suggest, "_default_retriever", side_effect=_Boom("upstream down") + ): + with self.app.test_client() as c: + resp = c.post("/rest/v1/link/suggest", json=_standard()) + self.assertEqual(resp.status_code, 502) + self.assertIn("error", resp.get_json()) + + +if __name__ == "__main__": + unittest.main() diff --git a/application/utils/librarian/factory.py b/application/utils/librarian/factory.py index d9f4ba118..2de90e53a 100644 --- a/application/utils/librarian/factory.py +++ b/application/utils/librarian/factory.py @@ -62,26 +62,22 @@ def build_scaler(config: Optional[LibrarianConfig] = None) -> Scaler: return TemperatureScaler(config.temperature) -def build_components( +def build_retriever_for( database: Any, *, config: Optional[LibrarianConfig] = None, embed_fn: Optional[Callable[[str], Sequence[float]]] = None, -) -> LibrarianComponents: - """Construct C.1 + C.2 + C.3 against a connected OpenCRE database. +) -> Any: + """Construct C.1 alone against a connected OpenCRE database. - Args: - database: a connected ``db.Node_collection`` (the caller owns it, the - same way Module B's ``run_noise_filter`` takes a session it did not - open). - config: Module C settings; defaults to ``load_config()``. - embed_fn: text -> embedding. Defaults to the prompt handler's embedder, - which calls the paid embedding API; injectable so a caller can - supply a local or fake embedder. + Returns ``(retriever, cre_embeddings)``: the retriever behind the shared + ``retrieve()`` seam, and the hub embeddings it was built from (the in_memory + pool source, and the id set ``build_components`` records as known CREs). - Raises whatever the pgvector guard raises when that backend is configured - but unavailable — deliberately, rather than falling back to a different - retrieval path and producing silently different rankings. + Split out of ``build_components`` so a caller that only needs the shortlist + (the ``/rest/v1/link/suggest`` endpoint, which runs in the web process and + must not load a torch cross-encoder) can share C.1 exactly, instead of + re-implementing retrieval and drifting from the pipeline. """ from application.defs import cre_defs as defs from application.utils.librarian.candidate_retriever import ( @@ -89,10 +85,6 @@ def build_components( RetrieverBackend, build_retriever, ) - from application.utils.librarian.cross_encoder import ( - CrossEncoderReranker, - build_cross_encoder_score_fn, - ) config = config or load_config() backend = RetrieverBackend(config.retriever_backend) @@ -130,6 +122,40 @@ def build_components( else None ), ) + return retriever, cre_embeddings + + +def build_components( + database: Any, + *, + config: Optional[LibrarianConfig] = None, + embed_fn: Optional[Callable[[str], Sequence[float]]] = None, +) -> LibrarianComponents: + """Construct C.1 + C.2 + C.3 against a connected OpenCRE database. + + Args: + database: a connected ``db.Node_collection`` (the caller owns it, the + same way Module B's ``run_noise_filter`` takes a session it did not + open). + config: Module C settings; defaults to ``load_config()``. + embed_fn: text -> embedding. Defaults to the prompt handler's embedder, + which calls the paid embedding API; injectable so a caller can + supply a local or fake embedder. + + Raises whatever the pgvector guard raises when that backend is configured + but unavailable — deliberately, rather than falling back to a different + retrieval path and producing silently different rankings. + """ + from application.defs import cre_defs as defs + from application.utils.librarian.cross_encoder import ( + CrossEncoderReranker, + build_cross_encoder_score_fn, + ) + + config = config or load_config() + retriever, cre_embeddings = build_retriever_for( + database, config=config, embed_fn=embed_fn + ) # C.2 scores each (section, candidate) pair against the CRE's # embeddings_content — the same text the hub vectors were built from. diff --git a/application/utils/librarian/link_suggest.py b/application/utils/librarian/link_suggest.py new file mode 100644 index 000000000..3bffbcd5b --- /dev/null +++ b/application/utils/librarian/link_suggest.py @@ -0,0 +1,190 @@ +"""``POST /rest/v1/link/suggest``: the Librarian's shortlist as a public API. + +Given a CRE ``Standard`` document (the same JSON shape the importers and the +REST API already speak), return the CREs most likely to be its link targets, +ranked. This is C.1 of the Librarian exposed directly, so a standard author or +maintainer can ask "where does this section belong?" without running the OIE +pipeline, and without a second similarity stack growing beside the one that is +evaluated and gated in CI (#275, and the maintainer decision on #965: do not +build a parallel mapping path). + +What this deliberately is not: the Librarian's *decision*. The pipeline's C.2 +rerank and C.3 calibration run in a worker with a torch cross-encoder loaded; +the web process must not pay that (the same constraint the chat retrieval in +#1161 respects). The score returned here is the retriever's cosine similarity, +named as such, ranked, and suitable for a human to choose from. Automated +mapping with a confidence threshold is ``mapping_job`` territory (#965), which +calls the same retriever and then the rest of the chain. +""" + +from cre_logging import get_logger + +logger = get_logger(__name__) + +from typing import Any, Callable, Dict, List, Mapping, Optional + +from application.defs import cre_defs as defs + +#: Hard ceiling on one request's shortlist. The retriever's own top_k is the +#: pipeline setting (20 by default); a caller may ask for fewer, never more. +MAX_TOP_K = 20 +DEFAULT_TOP_K = 5 + +#: Fields of a Standard that carry text worth embedding, in the order they are +#: joined. ``name`` is excluded on purpose: "ASVS" adds nothing to a section's +#: meaning and would pull every ASVS section toward the same neighbourhood. +_TEXT_FIELDS = ("section", "subsection", "description") + + +class InvalidStandardError(ValueError): + """The request body is not a Standard document this endpoint can embed.""" + + +def standard_query_text(doc: Mapping[str, Any]) -> str: + """Build the text C.1 embeds from a Standard document, or raise. + + Mirrors how the pipeline treats a section: the section heading, the + subsection and the description are what carry the requirement; ids and + hyperlinks are provenance. At least one text field must be non-empty. + """ + if not isinstance(doc, Mapping): + raise InvalidStandardError("body must be a JSON object") + doctype = doc.get("doctype") + if doctype != defs.Credoctypes.Standard.value: + raise InvalidStandardError( + f"doctype must be {defs.Credoctypes.Standard.value!r}, got {doctype!r}" + ) + name = doc.get("name") + if not isinstance(name, str) or not name.strip(): + raise InvalidStandardError("name is required") + parts: List[str] = [] + for field in _TEXT_FIELDS: + value = doc.get(field) + if value is None: + continue + if not isinstance(value, str): + raise InvalidStandardError(f"{field} must be a string") + if value.strip(): + parts.append(value.strip()) + if not parts: + raise InvalidStandardError( + "at least one of section, subsection or description must be non-empty" + ) + return ". ".join(parts) + + +def clamp_top_k(raw: Any) -> int: + """Parse a caller's top_k; absent means the default, bad input raises.""" + if raw is None or raw == "": + return DEFAULT_TOP_K + # Whole numbers only. int() would quietly truncate 2.5 to 2 and read True + # as 1; a limit the caller mistyped should be refused, not reinterpreted. + if isinstance(raw, bool): + raise InvalidStandardError("top_k must be an integer") + if isinstance(raw, int): + value = raw + elif isinstance(raw, float) and raw.is_integer(): + value = int(raw) + elif isinstance(raw, str) and raw.strip().lstrip("-").isdigit(): + value = int(raw.strip()) + else: + raise InvalidStandardError("top_k must be an integer") + if value < 1: + raise InvalidStandardError("top_k must be at least 1") + return min(value, MAX_TOP_K) + + +def _default_retriever(database: Any) -> Any: + """C.1 exactly as the pipeline builds it, minus C.2 and C.3. + + A seam: tests and callers with their own retriever pass one in and this is + never called. Lazy import keeps the web process free of ML imports until a + request actually needs them. + """ + from application.utils.librarian.factory import build_retriever_for + + retriever, _ = build_retriever_for(database) + return retriever + + +def _resolve_cre(database: Any, candidate_id: str) -> Optional[defs.CRE]: + """Find the graph's CRE for a retriever candidate id, whichever space it is in. + + The embeddings table keys CRE vectors by the CRE row's *internal* id, so + both retriever backends hand back UUIDs, while everything a human or an + importer sees is the external id ("287-251"). The eval harness translates + the pool for exactly this reason; here the translation happens per result, + so the endpoint is correct against a pool keyed either way. + """ + for kwargs in ({"internal_id": candidate_id}, {"external_id": candidate_id}): + try: + matches = database.get_CREs(**kwargs) or [] + except Exception: # a malformed id must not take the whole shortlist down + matches = [] + if matches: + found: defs.CRE = matches[0] + return found + return None + + +def suggest_links( + database: Any, + doc: Mapping[str, Any], + *, + top_k: int = DEFAULT_TOP_K, + retriever: Optional[Any] = None, + retriever_factory: Optional[Callable[[Any], Any]] = None, +) -> List[Dict[str, Any]]: + """Rank the CREs most likely to link to ``doc``. + + Returns the response shape #275 asked for, one entry per candidate, best + first: ``{"doctype": "CRE", "id", "name", "tags", "score"}``. ``score`` is + the retriever's cosine similarity in [0, 1]. ``name`` and ``tags`` come from + the graph when the CRE is present; a candidate the retriever knows but the + graph does not (a stale embedding) is still returned, with the retriever's + name and no tags, rather than silently dropped. + """ + text = standard_query_text(doc) + if top_k < 1 or top_k > MAX_TOP_K: + raise InvalidStandardError(f"top_k must be between 1 and {MAX_TOP_K}") + if retriever is None: + # Resolved at call time, not bound as a default, so the module-level + # builder stays the single seam tests and callers can replace. + factory = retriever_factory or _default_retriever + retriever = factory(database) + audit = retriever.retrieve(text) + + results: List[Dict[str, Any]] = [] + for cand in audit.candidates[:top_k]: + cre = _resolve_cre(database, cand.cre_id) + name = ( + (cre.name if cre is not None and cre.name else None) or cand.cre_name or "" + ) + tags: List[str] = ( + sorted(t for t in (cre.tags or []) if t) if cre is not None else [] + ) + results.append( + { + "doctype": defs.Credoctypes.CRE.value, + # Always the public id ("287-251"), whatever the pool was keyed by. + "id": cre.id if cre is not None and cre.id else cand.cre_id, + "name": name, + "tags": tags, + "score": ( + round(cand.score_vector, 4) + if cand.score_vector is not None + else None + ), + } + ) + return results + + +__all__ = [ + "DEFAULT_TOP_K", + "MAX_TOP_K", + "InvalidStandardError", + "clamp_top_k", + "standard_query_text", + "suggest_links", +] diff --git a/application/web/web_main.py b/application/web/web_main.py index 094f818e8..e94494823 100644 --- a/application/web/web_main.py +++ b/application/web/web_main.py @@ -1285,6 +1285,73 @@ def chat_cre() -> Any: return jsonify(response) +@app.route("/rest/v1/link/suggest", methods=["POST"]) +@login_required +def link_suggest() -> Any: + """Rank the CREs a Standard section most likely links to (#275). + + Body: a CRE ``Standard`` JSON document. Optional ``top_k`` (query string or + body, 1 to 20, default 5). Response: a JSON array of CREs, best first, each + with ``doctype``, ``id``, ``name``, ``tags`` and the retriever's cosine + ``score``. Same auth and same unavailable/rate-limit handling as + ``/rest/v1/completion``: both call the paid embedding API. + """ + payload = request.get_json(force=True, silent=True) + if not isinstance(payload, dict): + abort(400, description="body must be a JSON object (a CRE Standard document)") + if posthog: + posthog.capture("link_suggest", "") + + database = db.Node_collection() + # Lazy import to avoid loading heavy prompt/ML dependencies at web boot. + try: + from application.prompt_client import llm_error_utils + from application.utils.librarian import link_suggest as suggest + except ImportError: + abort( + 503, + "Link suggestion is unavailable on this deployment " + "(ML/LLM dependencies not installed).", + ) + + try: + top_k = suggest.clamp_top_k(request.args.get("top_k", payload.get("top_k"))) + results = suggest.suggest_links(database, payload, top_k=top_k) + except suggest.InvalidStandardError as exc: + abort(400, description=str(exc)) + except ImportError: + abort( + 503, + "Link suggestion is unavailable on this deployment " + "(ML/LLM dependencies not installed).", + ) + except RuntimeError as exc: + # PromptHandler raises RuntimeError when litellm is missing. + if "litellm" not in str(exc).lower(): + raise + abort( + 503, + "Link suggestion is unavailable on this deployment " + "(ML/LLM dependencies not installed).", + ) + except Exception as e: + if llm_error_utils.is_rate_limit_error(e): + return ( + jsonify( + { + "error": ( + "The AI service is temporarily rate-limited or out of quota. " + "Please try again in a minute." + ) + } + ), + 503, + ) + status_code = _llm_error_status_code(e) or 500 + return jsonify({"error": f"AI Service Error: {str(e)}"}), status_code + return jsonify(results) + + class CREFlow: """ "This class handles authentication with google's oauth""" From 569a91287f5d469404d80a7a564144e0d101b7ea Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Sat, 24 Oct 2026 01:58:17 +0530 Subject: [PATCH 2/5] link/suggest: generic error body, strict top_k parsing, one lookup per candidate Addresses the review on #1171. The catch-all returned str(e) to the client. suggest_links runs database lookups, config loading and pgvector guards, whose exceptions can carry SQL text, paths and config values; CodeQL flagged it as information exposure. The full exception now goes to the server log and the client gets the status code with a generic message. The completion route has the same pattern; it is not copied here. top_k string parsing used an isdigit-style pre-check that accepted inputs int() rejects ("--3", a superscript digit), so those surfaced as a 500 from the catch-all instead of a 400. int() is now the parser and its failure is converted; two tests cover the two inputs at both the service and the route. _resolve_cre did an internal-id lookup and then an external-id lookup per candidate, and get_CREs warns on each miss, so one stale embedding warned twice on every request. The lookup is now chosen once by the id's shape (public ids are "ddd-ddd", the pattern cre_defs.CRE enforces). --- .../tests/librarian/link_suggest_test.py | 9 ++++- application/utils/librarian/link_suggest.py | 38 +++++++++++++------ application/web/web_main.py | 9 ++++- 3 files changed, 43 insertions(+), 13 deletions(-) diff --git a/application/tests/librarian/link_suggest_test.py b/application/tests/librarian/link_suggest_test.py index e153d368e..834ae40db 100644 --- a/application/tests/librarian/link_suggest_test.py +++ b/application/tests/librarian/link_suggest_test.py @@ -109,7 +109,7 @@ def test_parses_and_caps(self) -> None: self.assertEqual(clamp_top_k(99), MAX_TOP_K) def test_rejects_non_positive_and_non_integer(self) -> None: - for bad in ("0", "-3", "x", 2.5): + for bad in ("0", "-3", "x", 2.5, "--3", "\u00b2", True): with self.assertRaises(InvalidStandardError): clamp_top_k(bad) @@ -303,6 +303,11 @@ def test_missing_litellm_is_503(self) -> None: resp = c.post("/rest/v1/link/suggest", json=_standard()) self.assertEqual(resp.status_code, 503) + def test_malformed_top_k_strings_are_400_not_500(self) -> None: + for bad in ("--3", "\u00b2"): + resp = self._post(_standard(), query=f"?top_k={bad}") + self.assertEqual(resp.status_code, 400, bad) + def test_provider_failure_is_reported_not_swallowed(self) -> None: class _Boom(Exception): status_code = 502 @@ -314,6 +319,8 @@ class _Boom(Exception): resp = c.post("/rest/v1/link/suggest", json=_standard()) self.assertEqual(resp.status_code, 502) self.assertIn("error", resp.get_json()) + # the client gets a status and a generic message, never the raw text + self.assertNotIn("upstream down", resp.get_data(as_text=True)) if __name__ == "__main__": diff --git a/application/utils/librarian/link_suggest.py b/application/utils/librarian/link_suggest.py index 3bffbcd5b..0ea834b0f 100644 --- a/application/utils/librarian/link_suggest.py +++ b/application/utils/librarian/link_suggest.py @@ -21,10 +21,14 @@ logger = get_logger(__name__) +import re from typing import Any, Callable, Dict, List, Mapping, Optional from application.defs import cre_defs as defs +#: The public CRE id shape ("287-251"); everything else is treated as a row id. +_PUBLIC_ID = re.compile(r"\d{3}-\d{3}") + #: Hard ceiling on one request's shortlist. The retriever's own top_k is the #: pipeline setting (20 by default); a caller may ask for fewer, never more. MAX_TOP_K = 20 @@ -85,8 +89,13 @@ def clamp_top_k(raw: Any) -> int: value = raw elif isinstance(raw, float) and raw.is_integer(): value = int(raw) - elif isinstance(raw, str) and raw.strip().lstrip("-").isdigit(): - value = int(raw.strip()) + elif isinstance(raw, str): + # int() is the parser; "--3" and "\u00b2" pass isdigit-style checks and + # then blow up inside int(), so its failure is the signal, not a pre-check. + try: + value = int(raw.strip()) + except ValueError: + raise InvalidStandardError("top_k must be an integer") from None else: raise InvalidStandardError("top_k must be an integer") if value < 1: @@ -116,15 +125,22 @@ def _resolve_cre(database: Any, candidate_id: str) -> Optional[defs.CRE]: the pool for exactly this reason; here the translation happens per result, so the endpoint is correct against a pool keyed either way. """ - for kwargs in ({"internal_id": candidate_id}, {"external_id": candidate_id}): - try: - matches = database.get_CREs(**kwargs) or [] - except Exception: # a malformed id must not take the whole shortlist down - matches = [] - if matches: - found: defs.CRE = matches[0] - return found - return None + # One lookup, chosen by the id's shape: public ids are "ddd-ddd" (the + # pattern cre_defs.CRE enforces), anything else is a row id. A single miss + # then logs once, instead of twice per stale candidate on every request. + kwargs = ( + {"external_id": candidate_id} + if _PUBLIC_ID.fullmatch(candidate_id or "") + else {"internal_id": candidate_id} + ) + try: + matches = database.get_CREs(**kwargs) or [] + except Exception: # a malformed id must not take the whole shortlist down + return None + if not matches: + return None + found: defs.CRE = matches[0] + return found def suggest_links( diff --git a/application/web/web_main.py b/application/web/web_main.py index e94494823..97d5b8625 100644 --- a/application/web/web_main.py +++ b/application/web/web_main.py @@ -1347,8 +1347,15 @@ def link_suggest() -> Any: ), 503, ) + # Never echo the exception: SQLAlchemy and provider errors can carry SQL + # text, paths and config values. Full detail goes to the server log; the + # client gets the status code and a generic message. status_code = _llm_error_status_code(e) or 500 - return jsonify({"error": f"AI Service Error: {str(e)}"}), status_code + logger.error("link_suggest failed with status %s", status_code, exc_info=True) + return ( + jsonify({"error": "Link suggestion failed upstream; see server logs."}), + status_code, + ) return jsonify(results) From 431091b6e57b39ad0d0e115fa2035048fcdfa91a Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Sat, 24 Oct 2026 01:59:04 +0530 Subject: [PATCH 3/5] link/suggest: top_k strings are an optional sign and ASCII digits only int() also accepts "2_0"; a top_k is not a place for underscore literals. The string form is now matched before parsing, so "--3", a superscript digit and "2_0" are all a 400. --- application/tests/librarian/link_suggest_test.py | 4 ++-- application/utils/librarian/link_suggest.py | 16 +++++++++------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/application/tests/librarian/link_suggest_test.py b/application/tests/librarian/link_suggest_test.py index 834ae40db..8d891175c 100644 --- a/application/tests/librarian/link_suggest_test.py +++ b/application/tests/librarian/link_suggest_test.py @@ -109,7 +109,7 @@ def test_parses_and_caps(self) -> None: self.assertEqual(clamp_top_k(99), MAX_TOP_K) def test_rejects_non_positive_and_non_integer(self) -> None: - for bad in ("0", "-3", "x", 2.5, "--3", "\u00b2", True): + for bad in ("0", "-3", "x", 2.5, "--3", "\u00b2", "2_0", " 7 x", True): with self.assertRaises(InvalidStandardError): clamp_top_k(bad) @@ -304,7 +304,7 @@ def test_missing_litellm_is_503(self) -> None: self.assertEqual(resp.status_code, 503) def test_malformed_top_k_strings_are_400_not_500(self) -> None: - for bad in ("--3", "\u00b2"): + for bad in ("--3", "\u00b2", "2_0"): resp = self._post(_standard(), query=f"?top_k={bad}") self.assertEqual(resp.status_code, 400, bad) diff --git a/application/utils/librarian/link_suggest.py b/application/utils/librarian/link_suggest.py index 0ea834b0f..7cfb2a7da 100644 --- a/application/utils/librarian/link_suggest.py +++ b/application/utils/librarian/link_suggest.py @@ -27,7 +27,9 @@ from application.defs import cre_defs as defs #: The public CRE id shape ("287-251"); everything else is treated as a row id. -_PUBLIC_ID = re.compile(r"\d{3}-\d{3}") +_PUBLIC_ID = re.compile(r"[0-9]{3}-[0-9]{3}") +#: What a top_k string may be: an optional sign and ASCII digits, nothing else. +_ASCII_INT = re.compile(r"[+-]?[0-9]+") #: Hard ceiling on one request's shortlist. The retriever's own top_k is the #: pipeline setting (20 by default); a caller may ask for fewer, never more. @@ -90,12 +92,12 @@ def clamp_top_k(raw: Any) -> int: elif isinstance(raw, float) and raw.is_integer(): value = int(raw) elif isinstance(raw, str): - # int() is the parser; "--3" and "\u00b2" pass isdigit-style checks and - # then blow up inside int(), so its failure is the signal, not a pre-check. - try: - value = int(raw.strip()) - except ValueError: - raise InvalidStandardError("top_k must be an integer") from None + # Strict ASCII integer: isdigit() admits "\u00b2" and lstrip("-") admits + # "--3", while int() itself admits "2_0". Only an optional sign and + # ASCII digits are a top_k; everything else is a 400. + if not _ASCII_INT.fullmatch(raw.strip()): + raise InvalidStandardError("top_k must be an integer") + value = int(raw.strip()) else: raise InvalidStandardError("top_k must be an integer") if value < 1: From c5296b85d24e5ea4c94f766881305b569636e40e Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Sun, 25 Oct 2026 20:03:26 +0530 Subject: [PATCH 4/5] link/suggest: one failure path for every exception, log failed CRE lookups Review fixes on the suggest endpoint. - The route no longer has a dedicated RuntimeError clause. It re-raised any RuntimeError whose message did not mention litellm, and a sibling except cannot catch a re-raise from the same try, so a retriever or database RuntimeError skipped the rate-limit check, the status mapping, the generic client body and the server-side log, and the client got Flask's default 500 page. The litellm check now sits at the top of the generic handler; everything else follows the same path as any other failure. - _resolve_cre logs a CRE lookup that raises, with the candidate id and the traceback, before returning the candidate unresolved. A lookup that raises is not a stale embedding, and the log now says which. - Two tests: a non-litellm RuntimeError yields the generic 500 body with no exception text; a raising lookup is logged with the candidate id. --- .../tests/librarian/link_suggest_test.py | 29 +++++++++++++++++++ application/utils/librarian/link_suggest.py | 11 ++++++- application/web/web_main.py | 18 ++++++------ 3 files changed, 48 insertions(+), 10 deletions(-) diff --git a/application/tests/librarian/link_suggest_test.py b/application/tests/librarian/link_suggest_test.py index 8d891175c..681d7eb41 100644 --- a/application/tests/librarian/link_suggest_test.py +++ b/application/tests/librarian/link_suggest_test.py @@ -134,6 +134,16 @@ def tearDown(self) -> None: sqla.drop_all() self.ctx.pop() + def test_a_failing_cre_lookup_is_logged_with_the_candidate(self) -> None: + """A lookup that raises is not a miss: it is logged, named, and the + candidate still comes back unresolved so the shortlist survives.""" + broken = mock.Mock() + broken.get_CREs.side_effect = RuntimeError("database gone") + with self.assertLogs(link_suggest.logger, level="WARNING") as logs: + self.assertIsNone(link_suggest._resolve_cre(broken, "287-251")) + self.assertIn("287-251", logs.output[0]) + self.assertIn("database gone", logs.output[0]) + def test_internal_uuid_candidates_resolve_to_the_public_cre_id(self) -> None: """The embeddings table keys vectors by the CRE row's internal UUID, so that is what the retriever returns. A caller must only ever see @@ -303,6 +313,25 @@ def test_missing_litellm_is_503(self) -> None: resp = c.post("/rest/v1/link/suggest", json=_standard()) self.assertEqual(resp.status_code, 503) + def test_other_runtime_errors_get_the_generic_body_not_flasks_500_page( + self, + ) -> None: + """Only a missing litellm is a 503; any other RuntimeError must go + through the same generic handling as every other failure.""" + with mock.patch.object( + link_suggest, + "_default_retriever", + side_effect=RuntimeError("index file /srv/cre/secret.idx is corrupt"), + ): + with self.app.test_client() as c: + resp = c.post("/rest/v1/link/suggest", json=_standard()) + self.assertEqual(resp.status_code, 500) + self.assertEqual( + resp.get_json(), + {"error": "Link suggestion failed upstream; see server logs."}, + ) + self.assertNotIn("secret.idx", resp.get_data(as_text=True)) + def test_malformed_top_k_strings_are_400_not_500(self) -> None: for bad in ("--3", "\u00b2", "2_0"): resp = self._post(_standard(), query=f"?top_k={bad}") diff --git a/application/utils/librarian/link_suggest.py b/application/utils/librarian/link_suggest.py index 7cfb2a7da..d72545dba 100644 --- a/application/utils/librarian/link_suggest.py +++ b/application/utils/librarian/link_suggest.py @@ -137,7 +137,16 @@ def _resolve_cre(database: Any, candidate_id: str) -> Optional[defs.CRE]: ) try: matches = database.get_CREs(**kwargs) or [] - except Exception: # a malformed id must not take the whole shortlist down + except Exception: + # A malformed id must not take the whole shortlist down, but a lookup + # that raises is not the same as a lookup that misses: say which + # candidate failed so a broken database path does not read as "stale". + logger.warning( + "link_suggest: CRE lookup failed for candidate %r; " + "returning it unresolved", + candidate_id, + exc_info=True, + ) return None if not matches: return None diff --git a/application/web/web_main.py b/application/web/web_main.py index 97d5b8625..4d0da4f97 100644 --- a/application/web/web_main.py +++ b/application/web/web_main.py @@ -1325,16 +1325,16 @@ def link_suggest() -> Any: "Link suggestion is unavailable on this deployment " "(ML/LLM dependencies not installed).", ) - except RuntimeError as exc: - # PromptHandler raises RuntimeError when litellm is missing. - if "litellm" not in str(exc).lower(): - raise - abort( - 503, - "Link suggestion is unavailable on this deployment " - "(ML/LLM dependencies not installed).", - ) except Exception as e: + # PromptHandler raises RuntimeError when litellm is missing. Any other + # RuntimeError (retriever, database) falls through to the generic + # handling below; a dedicated except clause would re-raise past it. + if isinstance(e, RuntimeError) and "litellm" in str(e).lower(): + abort( + 503, + "Link suggestion is unavailable on this deployment " + "(ML/LLM dependencies not installed).", + ) if llm_error_utils.is_rate_limit_error(e): return ( jsonify( From 4cad5cf756d092f4a5b2aa065022e01202ad52ca Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Mon, 26 Oct 2026 14:46:21 +0530 Subject: [PATCH 5/5] link/suggest: a failed CRE lookup is an error, not an unresolved candidate _resolve_cre no longer catches exceptions from get_CREs. The lookup is chosen by id shape and is a plain equality filter either way, so a malformed id simply misses and returns None; it never raises. An exception therefore means the database failed, and swallowing it returned a 200 whose candidates carried raw internal UUIDs. It now reaches the route's single error path: logged server-side with the traceback, generic body to the client. Tests: a raising lookup propagates; malformed and empty ids are misses; a database failure during resolution is a 500 with the generic body and no exception text. --- .../tests/librarian/link_suggest_test.py | 32 +++++++++++++++---- application/utils/librarian/link_suggest.py | 18 +++-------- 2 files changed, 30 insertions(+), 20 deletions(-) diff --git a/application/tests/librarian/link_suggest_test.py b/application/tests/librarian/link_suggest_test.py index 681d7eb41..6a2c1fc15 100644 --- a/application/tests/librarian/link_suggest_test.py +++ b/application/tests/librarian/link_suggest_test.py @@ -134,15 +134,19 @@ def tearDown(self) -> None: sqla.drop_all() self.ctx.pop() - def test_a_failing_cre_lookup_is_logged_with_the_candidate(self) -> None: - """A lookup that raises is not a miss: it is logged, named, and the - candidate still comes back unresolved so the shortlist survives.""" + def test_a_failing_cre_lookup_propagates_instead_of_reading_as_a_miss( + self, + ) -> None: + """A lookup that raises is not a stale embedding: it must not come + back as an unresolved candidate carrying a raw UUID.""" broken = mock.Mock() broken.get_CREs.side_effect = RuntimeError("database gone") - with self.assertLogs(link_suggest.logger, level="WARNING") as logs: - self.assertIsNone(link_suggest._resolve_cre(broken, "287-251")) - self.assertIn("287-251", logs.output[0]) - self.assertIn("database gone", logs.output[0]) + with self.assertRaises(RuntimeError): + link_suggest._resolve_cre(broken, "287-251") + + def test_a_malformed_candidate_id_is_a_miss_not_an_error(self) -> None: + self.assertIsNone(link_suggest._resolve_cre(self.collection, "not-an-id")) + self.assertIsNone(link_suggest._resolve_cre(self.collection, "")) def test_internal_uuid_candidates_resolve_to_the_public_cre_id(self) -> None: """The embeddings table keys vectors by the CRE row's internal UUID, so @@ -332,6 +336,20 @@ def test_other_runtime_errors_get_the_generic_body_not_flasks_500_page( ) self.assertNotIn("secret.idx", resp.get_data(as_text=True)) + def test_a_database_failure_during_resolution_is_a_500_not_a_200(self) -> None: + with mock.patch.object( + db.Node_collection, + "get_CREs", + side_effect=RuntimeError("connection reset by /var/run/pg.sock"), + ): + resp = self._post(_standard()) + self.assertEqual(resp.status_code, 500) + self.assertEqual( + resp.get_json(), + {"error": "Link suggestion failed upstream; see server logs."}, + ) + self.assertNotIn("pg.sock", resp.get_data(as_text=True)) + def test_malformed_top_k_strings_are_400_not_500(self) -> None: for bad in ("--3", "\u00b2", "2_0"): resp = self._post(_standard(), query=f"?top_k={bad}") diff --git a/application/utils/librarian/link_suggest.py b/application/utils/librarian/link_suggest.py index d72545dba..ab9096cc3 100644 --- a/application/utils/librarian/link_suggest.py +++ b/application/utils/librarian/link_suggest.py @@ -135,19 +135,11 @@ def _resolve_cre(database: Any, candidate_id: str) -> Optional[defs.CRE]: if _PUBLIC_ID.fullmatch(candidate_id or "") else {"internal_id": candidate_id} ) - try: - matches = database.get_CREs(**kwargs) or [] - except Exception: - # A malformed id must not take the whole shortlist down, but a lookup - # that raises is not the same as a lookup that misses: say which - # candidate failed so a broken database path does not read as "stale". - logger.warning( - "link_suggest: CRE lookup failed for candidate %r; " - "returning it unresolved", - candidate_id, - exc_info=True, - ) - return None + # No try/except: an id of any shape is a plain equality filter, so a + # malformed id misses (None below) and never raises. A lookup that raises + # means the database failed, and that must reach the route's error + # handling rather than come back as a 200 full of unresolved UUIDs. + matches = database.get_CREs(**kwargs) or [] if not matches: return None found: defs.CRE = matches[0]