From 06235c23624c3ca01083ad97b72a13d0963baa13 Mon Sep 17 00:00:00 2001 From: chottokun <29515187+chottokun@users.noreply.github.com> Date: Thu, 15 Oct 2026 14:25:40 +0000 Subject: [PATCH 1/2] refactor: separate concerns in RerankService with characterization tests Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> --- src/app/services/rerank.py | 175 +++++++++++----------- src/tests/test_rerank_characterization.py | 154 +++++++++++++++++++ 2 files changed, 245 insertions(+), 84 deletions(-) create mode 100644 src/tests/test_rerank_characterization.py diff --git a/src/app/services/rerank.py b/src/app/services/rerank.py index 2fefbaf..26e3cea 100644 --- a/src/app/services/rerank.py +++ b/src/app/services/rerank.py @@ -17,6 +17,7 @@ def _calculate_rerank_tokens(model: Any, query: str, documents: List[str]) -> Usage: + """Calculates token usage for reranking query and documents batch.""" with model.tokenizer_lock: tokenizer = model.tokenizer q_tokens = len(tokenizer.encode(query, add_special_tokens=False)) @@ -48,6 +49,7 @@ def _calculate_rerank_tokens(model: Any, query: str, documents: List[str]) -> Us def _sort_and_format_rerank_results( results: List[dict], top_n: Optional[int] ) -> List[RerankData]: + """Sorts rerank results by score descending and returns RerankData list.""" if top_n is not None: sorted_results = heapq.nlargest( top_n, results, key=lambda x: (x["score"], -x["document"]) @@ -59,19 +61,18 @@ def _sort_and_format_rerank_results( class RerankService(BaseRerankService): - """ - Default production implementation of BaseRerankService. - """ + """Default production implementation of BaseRerankService.""" def __init__( self, proxy_to_tei_func: Optional[Any] = None, model_loader: Optional[Any] = None, - ): + ) -> None: self.proxy_to_tei_func = proxy_to_tei_func self.model_loader = model_loader async def create_rerank(self, request: RerankRequest) -> RerankResponse: + """Processes a rerank request by delegating to TEI proxy, Logit Gate, or standard cross-encoder.""" import app.main as main_mod if ( @@ -87,92 +88,98 @@ async def create_rerank(self, request: RerankRequest) -> RerankResponse: proxy_func = getattr(main_mod, "_proxy_to_tei", self.proxy_to_tei_func) if tei_url and proxy_func: - tei_results = await proxy_func( - tei_url, - "/rerank", - {"query": request.query, "texts": request.documents}, - ) - - results = [] - for item in tei_results: - idx = item["index"] - score = item["score"] - result_item = {"document": idx, "score": float(score)} - if request.return_documents: - result_item["text"] = request.documents[idx] - results.append(result_item) - - response_data = _sort_and_format_rerank_results(results, request.top_n) - usage = Usage(prompt_tokens=0, total_tokens=0) - return RerankResponse( - query=request.query, - data=response_data, - model=request.model, - usage=usage, - ) + return await self._process_tei_proxy(request, tei_url, proxy_func) if request.model in LOGIT_GATE_MODELS: - model = get_validated_model( - request.model, - LOGIT_GATE_MODELS, - "logit gate", - loader=self.model_loader, - ) + return self._process_logit_gate(request) + + return self._process_standard_rerank(request) + + async def _process_tei_proxy( + self, request: RerankRequest, tei_url: str, proxy_func: Any + ) -> RerankResponse: + """Proxies rerank evaluation to Text Embeddings Inference (TEI) server.""" + tei_results = await proxy_func( + tei_url, + "/rerank", + {"query": request.query, "texts": request.documents}, + ) - gate_service = LogitGateService(model) - matcher = AsciiMatcher(min_token_len=LOGIT_GATE_ASCII_MIN_TOKEN_LEN) - gate_results = gate_service.predict_margins( - request.query, request.documents - ) + results = [] + for item in tei_results: + idx = item["index"] + score = item["score"] + result_item = {"document": idx, "score": float(score)} + if request.return_documents: + result_item["text"] = request.documents[idx] + results.append(result_item) - beta = ( - 0.0 - if request.use_ascii_boost is False - else LOGIT_GATE_ASCII_BOOST_WEIGHT - ) - containment_scores = matcher.score_documents( - request.query, request.documents - ) + response_data = _sort_and_format_rerank_results(results, request.top_n) + usage = Usage(prompt_tokens=0, total_tokens=0) + return RerankResponse( + query=request.query, + data=response_data, + model=request.model, + usage=usage, + ) - th = ( - request.threshold - if request.threshold is not None - else LOGIT_GATE_THRESHOLD - ) - results = [] - for i in range(len(request.documents)): - delta_z_base = gate_results[i]["logit_margin"] - containment = containment_scores[i] - delta_z_final = delta_z_base + (beta * containment) - score = _sigmoid(delta_z_final) - passed = score >= th - - if request.drop_failed and not passed: - continue - - entropy = _binary_entropy(score) - - result_item = { - "document": i, - "score": float(score), - "passed": bool(passed), - "logit_margin": float(delta_z_final), - "containment_score": float(containment), - "entropy": float(entropy), - } - if request.return_documents: - result_item["text"] = request.documents[i] - results.append(result_item) - - response_data = _sort_and_format_rerank_results(results, request.top_n) - usage = _calculate_rerank_tokens(model, request.query, request.documents) - return RerankResponse( - query=request.query, - data=response_data, - model=request.model, - usage=usage, - ) + def _process_logit_gate(self, request: RerankRequest) -> RerankResponse: + """Processes rerank request using hybrid Logit Gate model and ASCII boost.""" + model = get_validated_model( + request.model, + LOGIT_GATE_MODELS, + "logit gate", + loader=self.model_loader, + ) + + gate_service = LogitGateService(model) + matcher = AsciiMatcher(min_token_len=LOGIT_GATE_ASCII_MIN_TOKEN_LEN) + gate_results = gate_service.predict_margins(request.query, request.documents) + + beta = ( + 0.0 if request.use_ascii_boost is False else LOGIT_GATE_ASCII_BOOST_WEIGHT + ) + containment_scores = matcher.score_documents(request.query, request.documents) + + th = ( + request.threshold if request.threshold is not None else LOGIT_GATE_THRESHOLD + ) + results = [] + for i in range(len(request.documents)): + delta_z_base = gate_results[i]["logit_margin"] + containment = containment_scores[i] + delta_z_final = delta_z_base + (beta * containment) + score = _sigmoid(delta_z_final) + passed = score >= th + + if request.drop_failed and not passed: + continue + + entropy = _binary_entropy(score) + + result_item = { + "document": i, + "score": float(score), + "passed": bool(passed), + "logit_margin": float(delta_z_final), + "containment_score": float(containment), + "entropy": float(entropy), + } + if request.return_documents: + result_item["text"] = request.documents[i] + results.append(result_item) + + response_data = _sort_and_format_rerank_results(results, request.top_n) + usage = _calculate_rerank_tokens(model, request.query, request.documents) + return RerankResponse( + query=request.query, + data=response_data, + model=request.model, + usage=usage, + ) + def _process_standard_rerank(self, request: RerankRequest) -> RerankResponse: + """Processes rerank request using standard cross-encoder model.""" model = get_validated_model( request.model, RERANK_MODELS, diff --git a/src/tests/test_rerank_characterization.py b/src/tests/test_rerank_characterization.py new file mode 100644 index 0000000..5f5f25a --- /dev/null +++ b/src/tests/test_rerank_characterization.py @@ -0,0 +1,154 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock, patch +from fastapi import HTTPException + +from app.schemas import RerankRequest +from app.services.rerank import RerankService, _calculate_rerank_tokens +from app.config import RERANK_MODELS, LOGIT_GATE_MODELS + + +@pytest.mark.anyio +async def test_rerank_tei_proxy_return_documents_toggle(): + """Verify TEI proxy handling when return_documents is True vs False.""" + mock_proxy = AsyncMock( + return_value=[ + {"index": 1, "score": 0.85}, + {"index": 0, "score": 0.42}, + ] + ) + service = RerankService(proxy_to_tei_func=mock_proxy) + + req_with_docs = RerankRequest( + query="python rerank", + documents=["doc A", "doc B"], + model=RERANK_MODELS[0], + return_documents=True, + ) + + with ( + patch("app.services.rerank.RERANK_TEI_URL", "http://tei-rerank"), + patch("app.main._proxy_to_tei", mock_proxy), + ): + res_with_docs = await service.create_rerank(req_with_docs) + + assert len(res_with_docs.data) == 2 + assert res_with_docs.data[0].document == 1 + assert res_with_docs.data[0].text == "doc B" + assert res_with_docs.data[1].document == 0 + assert res_with_docs.data[1].text == "doc A" + + req_no_docs = RerankRequest( + query="python rerank", + documents=["doc A", "doc B"], + model=RERANK_MODELS[0], + return_documents=False, + ) + + with ( + patch("app.services.rerank.RERANK_TEI_URL", "http://tei-rerank"), + patch("app.main._proxy_to_tei", mock_proxy), + ): + res_no_docs = await service.create_rerank(req_no_docs) + + assert len(res_no_docs.data) == 2 + assert res_no_docs.data[0].text is None + assert res_no_docs.data[1].text is None + + +@pytest.mark.anyio +async def test_rerank_logit_gate_ascii_boost_disabled(): + """Verify Logit Gate behavior when use_ascii_boost=False.""" + mock_model = MagicMock() + mock_model.tokenizer_lock = MagicMock() + mock_model.lock = MagicMock() + mock_model.tokenizer.encode.side_effect = lambda text, **kw: [1, 2, 3] + + mock_loader = MagicMock(return_value=mock_model) + service = RerankService(model_loader=mock_loader) + + with ( + patch( + "app.services.rerank.LogitGateService.predict_margins", + return_value=[{"logit_margin": 0.0}, {"logit_margin": 1.0}], + ), + patch("app.services.rerank.RERANK_TEI_URL", None), + ): + req = RerankRequest( + query="test query with ABC-123", + documents=["doc with ABC-123", "another doc"], + model=LOGIT_GATE_MODELS[0], + use_ascii_boost=False, + ) + res = await service.create_rerank(req) + + assert len(res.data) == 2 + for d in res.data: + assert d.logit_margin in (0.0, 1.0) + + +@pytest.mark.anyio +async def test_rerank_logit_gate_custom_threshold_and_drop_failed(): + """Verify Logit Gate threshold filtering when drop_failed is True.""" + mock_model = MagicMock() + mock_model.tokenizer_lock = MagicMock() + mock_model.lock = MagicMock() + mock_model.tokenizer.encode.side_effect = lambda text, **kw: [1, 2, 3] + + mock_loader = MagicMock(return_value=mock_model) + service = RerankService(model_loader=mock_loader) + + with ( + patch( + "app.services.rerank.LogitGateService.predict_margins", + return_value=[{"logit_margin": 2.0}, {"logit_margin": -2.0}], + ), + patch( + "app.services.rerank.AsciiMatcher.score_documents", + return_value=[0.0, 0.0], + ), + patch("app.services.rerank.RERANK_TEI_URL", None), + ): + req = RerankRequest( + query="query", + documents=["high score doc", "low score doc"], + model=LOGIT_GATE_MODELS[0], + threshold=0.5, + drop_failed=True, + ) + res = await service.create_rerank(req) + + assert len(res.data) == 1 + assert res.data[0].document == 0 + assert res.data[0].passed is True + + +def test_calculate_rerank_tokens_fallback_on_exception(): + """Verify _calculate_rerank_tokens falls back to per-document tokenization when batch encoding raises Exception.""" + mock_model = MagicMock() + mock_tokenizer = MagicMock() + mock_model.tokenizer = mock_tokenizer + mock_model.tokenizer_lock = MagicMock() + + mock_tokenizer.encode.side_effect = lambda text, add_special_tokens=False: [1, 2] + mock_tokenizer.num_special_tokens_to_add.return_value = 2 + mock_tokenizer.side_effect = RuntimeError("Batch encoding failed") + + usage = _calculate_rerank_tokens(mock_model, "query string", ["doc one", "doc two"]) + assert usage.prompt_tokens == 12 + assert usage.total_tokens == 12 + + +@pytest.mark.anyio +async def test_rerank_unsupported_model_exception(): + """Verify HTTPException is raised for invalid model.""" + service = RerankService() + req = RerankRequest( + query="query", + documents=["doc"], + model="non-existent-model", + ) + with pytest.raises(HTTPException) as exc_info: + await service.create_rerank(req) + + assert exc_info.value.status_code == 400 + assert "not found for reranks" in exc_info.value.detail From 6d4336da27885718dbef7ed1f8b52bb511b81863 Mon Sep 17 00:00:00 2001 From: chottokun Date: Sun, 18 Oct 2026 15:04:14 +0900 Subject: [PATCH 2/2] test(rerank): add adversarial index bounds validation and edge case tests --- src/app/services/rerank.py | 11 +++ src/tests/test_rerank_characterization.py | 87 +++++++++++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/src/app/services/rerank.py b/src/app/services/rerank.py index 26e3cea..1649b20 100644 --- a/src/app/services/rerank.py +++ b/src/app/services/rerank.py @@ -106,8 +106,19 @@ async def _process_tei_proxy( ) results = [] + num_docs = len(request.documents) for item in tei_results: idx = item["index"] + if ( + isinstance(idx, bool) + or not isinstance(idx, int) + or idx < 0 + or idx >= num_docs + ): + raise HTTPException( + status_code=500, + detail=f"TEI Proxy returned invalid document index: {idx}", + ) score = item["score"] result_item = {"document": idx, "score": float(score)} if request.return_documents: diff --git a/src/tests/test_rerank_characterization.py b/src/tests/test_rerank_characterization.py index 5f5f25a..b79230e 100644 --- a/src/tests/test_rerank_characterization.py +++ b/src/tests/test_rerank_characterization.py @@ -152,3 +152,90 @@ async def test_rerank_unsupported_model_exception(): assert exc_info.value.status_code == 400 assert "not found for reranks" in exc_info.value.detail + + +@pytest.mark.anyio +async def test_rerank_tei_proxy_out_of_bounds_index_raises_500(): + """Adversarial test: TEI proxy returning an out-of-bounds or negative document index must raise HTTPException(500) rather than silently returning wrong document or raw IndexError.""" + mock_proxy = AsyncMock( + return_value=[ + {"index": 99, "score": 0.95}, + ] + ) + service = RerankService(proxy_to_tei_func=mock_proxy) + req = RerankRequest( + query="query", + documents=["doc A", "doc B"], + model=RERANK_MODELS[0], + return_documents=True, + ) + + with ( + patch("app.services.rerank.RERANK_TEI_URL", "http://tei-rerank"), + patch("app.main._proxy_to_tei", mock_proxy), + ): + with pytest.raises(HTTPException) as exc_info: + await service.create_rerank(req) + + assert exc_info.value.status_code == 500 + + +@pytest.mark.anyio +async def test_rerank_tei_proxy_negative_index_raises_500(): + """Adversarial test: TEI proxy returning a negative document index (e.g., -1) must not silently index the last document in Python.""" + mock_proxy = AsyncMock( + return_value=[ + {"index": -1, "score": 0.95}, + ] + ) + service = RerankService(proxy_to_tei_func=mock_proxy) + req = RerankRequest( + query="query", + documents=["doc A", "doc B"], + model=RERANK_MODELS[0], + return_documents=True, + ) + + with ( + patch("app.services.rerank.RERANK_TEI_URL", "http://tei-rerank"), + patch("app.main._proxy_to_tei", mock_proxy), + ): + with pytest.raises(HTTPException) as exc_info: + await service.create_rerank(req) + + assert exc_info.value.status_code == 500 + + +@pytest.mark.anyio +async def test_rerank_logit_gate_all_documents_dropped(): + """Boundary test: When all documents fail threshold and drop_failed=True, returns empty data list without error.""" + mock_model = MagicMock() + mock_model.tokenizer_lock = MagicMock() + mock_model.lock = MagicMock() + mock_model.tokenizer.encode.side_effect = lambda text, **kw: [1, 2] + + mock_loader = MagicMock(return_value=mock_model) + service = RerankService(model_loader=mock_loader) + + with ( + patch( + "app.services.rerank.LogitGateService.predict_margins", + return_value=[{"logit_margin": -5.0}, {"logit_margin": -3.0}], + ), + patch( + "app.services.rerank.AsciiMatcher.score_documents", + return_value=[0.0, 0.0], + ), + patch("app.services.rerank.RERANK_TEI_URL", None), + ): + req = RerankRequest( + query="unanswerable query", + documents=["irrelevant 1", "irrelevant 2"], + model=LOGIT_GATE_MODELS[0], + threshold=0.9, + drop_failed=True, + ) + res = await service.create_rerank(req) + + assert res.data == [] + assert res.usage.total_tokens > 0