[Bugfix] Restore truncate_prompt_tokens for Jina rerank/score online (#49963)
Signed-off-by: Umut Polat <52835619+umut-polat@users.noreply.github.com>
Umut Polat committed
Jul 27, 2026 at 17:52 UTC
27d7061ef62b3d853b4d335baf8ea9b3c56d9bf3
2 files changed
+69
-1
tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py
new
+59
@@ -0,0 +1,59 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for JinaRankingIOProcessor online request building."""
4
+
5
+from unittest.mock import MagicMock
6
+
7
+import pytest
8
+
9
+from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
10
+from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessor
11
+from vllm.entrypoints.pooling.scoring.protocol import RerankRequest
12
+from vllm.entrypoints.pooling.scoring.typing import ScoringData
13
+
14
+pytestmark = pytest.mark.skip_global_cleanup
15
+
16
+
17
+def test_online_forwards_truncate_prompt_tokens_to_proxy(monkeypatch):
18
+ """The proxy request handed to the base factory must carry
19
+ truncate_prompt_tokens/truncation_side from the real request.
20
+
21
+ JinaRankingIOProcessor swaps ctx.request for a proxy
22
+ PoolingCompletionRequest before delegating to the base factory, which
23
+ reads truncation off ctx.request. Dropping the fields on the proxy
24
+ silently disables truncate_prompt_tokens for Jina rerank/score.
25
+ """
26
+ proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor)
27
+ proc.valid_inputs_online = MagicMock(
28
+ return_value=ScoringData(data_1=["query"], data_2=["doc"])
29
+ )
30
+ proc._get_token_limits = MagicMock(return_value=(0, 0))
31
+ proc.ensure_str = MagicMock(side_effect=lambda data: list(data))
32
+ proc.format_docs_prompts_func = MagicMock(return_value="formatted prompt")
33
+
34
+ captured: dict[str, object] = {}
35
+
36
+ def _spy_base(self, ctx):
37
+ captured["truncate_prompt_tokens"] = ctx.request.truncate_prompt_tokens
38
+ captured["truncation_side"] = ctx.request.truncation_side
39
+ return []
40
+
41
+ monkeypatch.setattr(PoolingIOProcessor, "get_request_factory_online", _spy_base)
42
+
43
+ request = RerankRequest(
44
+ model="m",
45
+ query="query",
46
+ documents=["doc"],
47
+ truncate_prompt_tokens=512,
48
+ truncation_side="left",
49
+ )
50
+ ctx = MagicMock()
51
+ ctx.request = request
52
+ ctx.prompt_extras = None
53
+
54
+ proc.get_request_factory_online(ctx)
55
+
56
+ assert captured["truncate_prompt_tokens"] == 512
57
+ assert captured["truncation_side"] == "left"
58
+ # The real request is restored after delegating.
59
+ assert ctx.request is request
vllm/entrypoints/pooling/scoring/io_processor.py
+10
-1
@@ -820,7 +820,16 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM
820
for q, d in zip(queries, docs)
821
]
822
823
- ctx.request = PoolingCompletionRequest(task="token_embed", input=prompts)
823
+ # Forward truncation from the real request: the base factory reads
824
+ # these off ctx.request, so omitting them here silently drops
825
+ # truncate_prompt_tokens for Jina rerank/score (unlike the embed and
826
+ # bi/cross-encoder paths, which read them from the real request).
827
+ ctx.request = PoolingCompletionRequest(
828
+ task="token_embed",
829
+ input=prompts,
830
+ truncate_prompt_tokens=request.truncate_prompt_tokens,
831
+ truncation_side=request.truncation_side,
832
+ )
833
requests = PoolingIOProcessor.get_request_factory_online(self, ctx)
834
ctx.request = request
835
return requests