Blame view

tests/test_reranker_dashscope_backend.py 8.41 KB
d31c7f65   tangwang   补充云服务reranker
1
2
  from __future__ import annotations
  
0d3e73ba   tangwang   rerank mini batch
3
4
5
6
  import time
  
  import pytest
  
d31c7f65   tangwang   补充云服务reranker
7
8
9
  from reranker.backends import get_rerank_backend
  from reranker.backends.dashscope_rerank import DashScopeRerankBackend
  
99b72698   tangwang   测试回归钩子梳理
10
11
  pytestmark = [pytest.mark.rerank, pytest.mark.regression]
  
d31c7f65   tangwang   补充云服务reranker
12
  
0d3e73ba   tangwang   rerank mini batch
13
14
15
16
17
18
19
20
  @pytest.fixture(autouse=True)
  def _clear_global_dashscope_key(monkeypatch):
      # Prevent accidental pass-through from unrelated global key.
      monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False)
  
  
  def test_dashscope_backend_factory_loads(monkeypatch):
      monkeypatch.setenv("TEST_RERANK_DASHSCOPE_API_KEY", "test-key")
d31c7f65   tangwang   补充云服务reranker
21
22
23
24
25
      backend = get_rerank_backend(
          "dashscope_rerank",
          {
              "model_name": "qwen3-rerank",
              "endpoint": "https://dashscope.aliyuncs.com/compatible-api/v1/reranks",
0d3e73ba   tangwang   rerank mini batch
26
              "api_key_env": "TEST_RERANK_DASHSCOPE_API_KEY",
d31c7f65   tangwang   补充云服务reranker
27
28
29
30
31
32
          },
      )
      assert isinstance(backend, DashScopeRerankBackend)
  
  
  def test_dashscope_backend_score_with_meta_dedup_and_restore(monkeypatch):
0d3e73ba   tangwang   rerank mini batch
33
      monkeypatch.setenv("TEST_RERANK_DASHSCOPE_API_KEY", "test-key")
d31c7f65   tangwang   补充云服务reranker
34
35
36
37
      backend = DashScopeRerankBackend(
          {
              "model_name": "qwen3-rerank",
              "endpoint": "https://dashscope.aliyuncs.com/compatible-api/v1/reranks",
0d3e73ba   tangwang   rerank mini batch
38
              "api_key_env": "TEST_RERANK_DASHSCOPE_API_KEY",
d31c7f65   tangwang   补充云服务reranker
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
              "top_n_cap": 0,
          }
      )
  
      def _fake_post(query: str, docs: list[str], top_n: int):
          assert query == "wireless mouse"
          # deduplicated docs
          assert docs == ["doc-a", "doc-b"]
          assert top_n == 2
          return {
              "results": [
                  {"index": 1, "relevance_score": 0.9},
                  {"index": 0, "relevance_score": 0.2},
              ]
          }
  
      monkeypatch.setattr(backend, "_post_rerank", _fake_post)
      scores, meta = backend.score_with_meta(
          query="wireless mouse",
          docs=["doc-a", "doc-b", "doc-a", "", "   ", None],
          normalize=True,
      )
  
      assert scores == [0.2, 0.9, 0.2, 0.0, 0.0, 0.0]
      assert meta["input_docs"] == 6
      assert meta["usable_docs"] == 3
      assert meta["unique_docs"] == 2
      assert meta["top_n"] == 2
      assert meta["response_results"] == 2
      assert meta["backend"] == "dashscope_rerank"
  
  
  def test_dashscope_backend_top_n_cap_and_normalize_fallback(monkeypatch):
0d3e73ba   tangwang   rerank mini batch
72
      monkeypatch.setenv("TEST_RERANK_DASHSCOPE_API_KEY", "test-key")
d31c7f65   tangwang   补充云服务reranker
73
74
75
76
      backend = DashScopeRerankBackend(
          {
              "model_name": "qwen3-rerank",
              "endpoint": "https://dashscope.aliyuncs.com/compatible-api/v1/reranks",
0d3e73ba   tangwang   rerank mini batch
77
              "api_key_env": "TEST_RERANK_DASHSCOPE_API_KEY",
d31c7f65   tangwang   补充云服务reranker
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
              "top_n_cap": 1,
          }
      )
  
      def _fake_post(query: str, docs: list[str], top_n: int):
          assert query == "q"
          assert len(docs) == 2
          assert top_n == 1
          # Only top-1 returned, score outside [0,1] to trigger sigmoid fallback
          return {"results": [{"index": 1, "score": 3.0}]}
  
      monkeypatch.setattr(backend, "_post_rerank", _fake_post)
      scores_norm, _ = backend.score_with_meta(query="q", docs=["a", "b"], normalize=True)
      scores_raw, _ = backend.score_with_meta(query="q", docs=["a", "b"], normalize=False)
  
      assert scores_norm[0] == 0.0
      assert 0.95 < scores_norm[1] < 0.96
      assert scores_raw == [0.0, 3.0]
  
  
  def test_dashscope_backend_score_with_meta_topn_request(monkeypatch):
0d3e73ba   tangwang   rerank mini batch
99
      monkeypatch.setenv("TEST_RERANK_DASHSCOPE_API_KEY", "test-key")
d31c7f65   tangwang   补充云服务reranker
100
101
102
103
      backend = DashScopeRerankBackend(
          {
              "model_name": "qwen3-rerank",
              "endpoint": "https://dashscope.aliyuncs.com/compatible-api/v1/reranks",
0d3e73ba   tangwang   rerank mini batch
104
              "api_key_env": "TEST_RERANK_DASHSCOPE_API_KEY",
d31c7f65   tangwang   补充云服务reranker
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
              "top_n_cap": 0,
          }
      )
  
      def _fake_post(query: str, docs: list[str], top_n: int):
          assert query == "q"
          assert docs == ["d1", "d2", "d3"]
          assert top_n == 2
          return {"results": [{"index": 2, "relevance_score": 0.8}, {"index": 0, "relevance_score": 0.3}]}
  
      monkeypatch.setattr(backend, "_post_rerank", _fake_post)
      scores, meta = backend.score_with_meta_topn(query="q", docs=["d1", "d2", "d3"], top_n=2)
      assert scores == [0.3, 0.0, 0.8]
      assert meta["top_n"] == 2
      assert meta["requested_top_n"] == 2
0d3e73ba   tangwang   rerank mini batch
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
  
  
  def test_dashscope_backend_batchsize_concurrent_full_topn(monkeypatch):
      monkeypatch.setenv("TEST_RERANK_DASHSCOPE_API_KEY", "test-key")
      backend = DashScopeRerankBackend(
          {
              "model_name": "qwen3-rerank",
              "endpoint": "https://dashscope.aliyuncs.com/compatible-api/v1/reranks",
              "api_key_env": "TEST_RERANK_DASHSCOPE_API_KEY",
              "top_n_cap": 0,
              "batchsize": 2,
          }
      )
  
      def _fake_post(query: str, docs: list[str], top_n: int):
          assert query == "q"
          # batching path asks every batch for full local list
          assert top_n == len(docs)
          time.sleep(0.05)
          return {
              "results": [
                  {"index": i, "relevance_score": float(i + 1) / 10.0}
                  for i, _ in enumerate(docs)
              ]
          }
  
      monkeypatch.setattr(backend, "_post_rerank", _fake_post)
      start = time.perf_counter()
      scores, meta = backend.score_with_meta(query="q", docs=["d1", "d2", "d3", "d4", "d5", "d6"])
      elapsed = time.perf_counter() - start
  
      # 3 batches * 50ms serial ~=150ms; concurrent should be significantly lower.
      assert elapsed < 0.14
      assert len(scores) == 6
      assert meta["batches"] == 3
      assert meta["batch_concurrency"] == 3
      assert meta["response_results"] == 6
  
  
  def test_dashscope_backend_batchsize_still_effective_when_topn_limited(monkeypatch):
      monkeypatch.setenv("TEST_RERANK_DASHSCOPE_API_KEY", "test-key")
      backend = DashScopeRerankBackend(
          {
              "model_name": "qwen3-rerank",
              "endpoint": "https://dashscope.aliyuncs.com/compatible-api/v1/reranks",
              "api_key_env": "TEST_RERANK_DASHSCOPE_API_KEY",
              "top_n_cap": 0,
              "batchsize": 2,
          }
      )
  
      called = {"count": 0}
  
      def _fake_post(query: str, docs: list[str], top_n: int):
          called["count"] += 1
          # batching remains enabled; each batch asks for full local scores
          assert top_n == len(docs)
          score_map = {"d1": 0.9, "d2": 0.1, "d3": 0.8, "d4": 0.2}
          return {
              "results": [
                  {"index": i, "relevance_score": score_map[doc]}
                  for i, doc in enumerate(docs)
              ]
          }
  
      monkeypatch.setattr(backend, "_post_rerank", _fake_post)
      scores, meta = backend.score_with_meta_topn(query="q", docs=["d1", "d2", "d3", "d4"], top_n=2)
  
      assert called["count"] == 2
      assert scores == [0.9, 0.0, 0.8, 0.0]
      assert meta["batches"] == 2
      assert meta["top_n"] == 2
  
  
  def test_dashscope_backend_batchsize_raises_when_one_batch_fails(monkeypatch):
      monkeypatch.setenv("TEST_RERANK_DASHSCOPE_API_KEY", "test-key")
      backend = DashScopeRerankBackend(
          {
              "model_name": "qwen3-rerank",
              "endpoint": "https://dashscope.aliyuncs.com/compatible-api/v1/reranks",
              "api_key_env": "TEST_RERANK_DASHSCOPE_API_KEY",
              "top_n_cap": 0,
              "batchsize": 2,
          }
      )
  
      def _fake_post(query: str, docs: list[str], top_n: int):
          if docs == ["d3", "d4"]:
              raise RuntimeError("provider temporary error")
          return {
              "results": [
                  {"index": i, "relevance_score": 0.1}
                  for i, _ in enumerate(docs)
              ]
          }
  
      monkeypatch.setattr(backend, "_post_rerank", _fake_post)
  
      with pytest.raises(RuntimeError, match="DashScope rerank batch failed"):
          backend.score_with_meta(query="q", docs=["d1", "d2", "d3", "d4"])
  
  
  def test_dashscope_backend_requires_api_key_env():
      with pytest.raises(ValueError, match="api_key_env is required"):
          DashScopeRerankBackend(
              {
                  "model_name": "qwen3-rerank",
                  "endpoint": "https://dashscope.aliyuncs.com/compatible-api/v1/reranks",
                  "top_n_cap": 0,
              }
          )
  
  
  def test_dashscope_backend_requires_api_key_env_value(monkeypatch):
      monkeypatch.delenv("TEST_RERANK_DASHSCOPE_API_KEY", raising=False)
      with pytest.raises(ValueError, match="set env TEST_RERANK_DASHSCOPE_API_KEY"):
          DashScopeRerankBackend(
              {
                  "model_name": "qwen3-rerank",
                  "endpoint": "https://dashscope.aliyuncs.com/compatible-api/v1/reranks",
                  "api_key_env": "TEST_RERANK_DASHSCOPE_API_KEY",
                  "top_n_cap": 0,
              }
          )