feat: support Space rerank models (#2358)

Co-authored-by: chan <dadachann@users.noreply.github.com>
This commit is contained in:
Hyu
2026-07-26 16:37:34 +08:00
committed by GitHub
parent 38e35d328a
commit dd8d1007a1
7 changed files with 171 additions and 5 deletions
@@ -1147,6 +1147,46 @@ class TestInvokeRerank:
assert results[0]['relevance_score'] == 1.0
assert results[1]['relevance_score'] == 0.0
@pytest.mark.asyncio
@pytest.mark.parametrize(
('model_extra_args', 'expected_url'),
[
({'rerank_path': 'reranks'}, 'https://gateway.example.com/v1/reranks'),
({'rerank_url': 'https://rerank.example.com/api/rerank'}, 'https://rerank.example.com/api/rerank'),
],
)
async def test_invoke_rerank_openai_compatible_endpoint_override(self, model_extra_args, expected_url):
"""Endpoint configuration controls routing and is not sent in the Cohere body."""
requester = litellmchat.LiteLLMRequester(
ap=Mock(),
config={
'base_url': 'https://gateway.example.com/v1/',
'custom_llm_provider': 'openai',
},
)
model = MockRuntimeRerankModel('Qwen3-Reranker-8B', 'test-api-key')
model.model_entity.extra_args = model_extra_args
mock_resp = Mock()
mock_resp.raise_for_status = Mock()
mock_resp.json = Mock(return_value={'results': [{'index': 0, 'relevance_score': 0.8}]})
mock_client = AsyncMock()
mock_client.post = AsyncMock(return_value=mock_resp)
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)
with patch('httpx.AsyncClient', return_value=mock_client):
await requester.invoke_rerank(model=model, query='query', documents=['document'])
assert mock_client.post.call_args.args[0] == expected_url
payload = mock_client.post.call_args.kwargs['json']
assert payload == {
'model': 'Qwen3-Reranker-8B',
'query': 'query',
'documents': ['document'],
'top_n': 1,
}
class TestConvertMessages:
"""Test _convert_messages method"""
@@ -8,7 +8,8 @@ and error handling without calling real LLM APIs.
from __future__ import annotations
import pytest
from unittest.mock import Mock
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.provider.modelmgr import requester
@@ -62,6 +63,48 @@ async def test_model_manager_skips_space_sync_when_disabled(mock_app_for_modelmg
app.space_service.get_models.assert_not_called()
@pytest.mark.asyncio
async def test_sync_new_models_from_space_creates_rerank_models(mock_app_for_modelmgr):
"""Space rerank entries are discovered and persisted under the shared provider."""
app = mock_app_for_modelmgr
provider = persistence_model.ModelProvider(
uuid='space-provider',
name='LangBot Space',
requester='space-chat-completions',
base_url='https://api.langbot.cloud/v1',
api_keys=['space-key'],
)
app.persistence_mgr.execute_async = AsyncMock(return_value=_make_mock_result([provider], first_item=provider))
app.space_service.get_models = AsyncMock(
return_value=[
SimpleNamespace(
uuid='rerank-model-uuid',
model_id='Qwen3-Reranker-8B',
category='rerank',
featured_order=10,
)
]
)
app.llm_model_service.get_llm_models = AsyncMock(return_value=[])
app.embedding_models_service.get_embedding_models = AsyncMock(return_value=[])
app.rerank_models_service = AsyncMock()
app.rerank_models_service.get_rerank_models = AsyncMock(return_value=[])
model_mgr = ModelManager(app)
await model_mgr.sync_new_models_from_space()
app.rerank_models_service.create_rerank_model.assert_awaited_once_with(
{
'uuid': 'rerank-model-uuid',
'name': 'Qwen3-Reranker-8B',
'provider_uuid': 'space-provider',
'extra_args': {},
'prefered_ranking': 10,
},
preserve_uuid=True,
)
# ============================================================================
# Model Loading Tests
# ============================================================================