From dd8d1007a112033212cb33b3c769fd221916ec08 Mon Sep 17 00:00:00 2001 From: Hyu Date: Sun, 26 Jul 2026 16:37:34 +0800 Subject: [PATCH] feat: support Space rerank models (#2358) Co-authored-by: chan --- src/langbot/pkg/api/http/service/provider.py | 9 ++++ src/langbot/pkg/provider/modelmgr/modelmgr.py | 30 +++++++++++++ .../modelmgr/requesters/litellmchat.py | 11 +++-- .../api/service/test_provider_service.py | 38 ++++++++++++++++ .../api/service/test_space_service.py | 3 +- tests/unit_tests/provider/test_litellmchat.py | 40 +++++++++++++++++ .../unit_tests/provider/test_model_manager.py | 45 ++++++++++++++++++- 7 files changed, 171 insertions(+), 5 deletions(-) diff --git a/src/langbot/pkg/api/http/service/provider.py b/src/langbot/pkg/api/http/service/provider.py index 598d72e8d..6caa54883 100644 --- a/src/langbot/pkg/api/http/service/provider.py +++ b/src/langbot/pkg/api/http/service/provider.py @@ -232,8 +232,15 @@ class ModelProviderService: llm_models = await self.ap.llm_model_service.get_llm_models_by_provider(provider_uuid) embedding_models = await self.ap.embedding_models_service.get_embedding_models_by_provider(provider_uuid) + rerank_service = getattr(self.ap, 'rerank_models_service', None) + rerank_models = ( + await rerank_service.get_rerank_models_by_provider(provider_uuid) + if rerank_service is not None + else [] + ) existing_llm_names = {model['name'] for model in llm_models} existing_embedding_names = {model['name'] for model in embedding_models} + existing_rerank_names = {model['name'] for model in rerank_models} filtered_models = [] for model in scanned_models: @@ -260,6 +267,8 @@ class ModelProviderService: 'already_added': ( model_name in existing_embedding_names if scanned_type == 'embedding' + else model_name in existing_rerank_names + if scanned_type == 'rerank' else model_name in existing_llm_names ), } diff --git a/src/langbot/pkg/provider/modelmgr/modelmgr.py b/src/langbot/pkg/provider/modelmgr/modelmgr.py index e3e20e026..501090b4c 100644 --- a/src/langbot/pkg/provider/modelmgr/modelmgr.py +++ b/src/langbot/pkg/provider/modelmgr/modelmgr.py @@ -193,6 +193,9 @@ class ModelManager: existing_embedding_models = { m['uuid']: m for m in await self.ap.embedding_models_service.get_embedding_models() } + existing_rerank_models = { + m['uuid']: m for m in await self.ap.rerank_models_service.get_rerank_models() + } created = 0 updated = 0 @@ -258,6 +261,33 @@ class ModelManager: await self.ap.embedding_models_service.update_embedding_model(space_model.uuid, dict(desired)) updated += 1 + elif space_model.category == 'rerank': + existing = existing_rerank_models.get(space_model.uuid) + if existing is None: + await self.ap.rerank_models_service.create_rerank_model( + { + 'uuid': space_model.uuid, + 'name': space_model.model_id, + 'provider_uuid': space_model_provider.uuid, + 'extra_args': {}, + 'prefered_ranking': space_model.featured_order, + }, + preserve_uuid=True, + ) + created += 1 + elif existing.get('provider_uuid') == space_model_provider.uuid: + desired = { + 'name': space_model.model_id, + 'provider_uuid': space_model_provider.uuid, + 'prefered_ranking': space_model.featured_order, + } + if ( + existing.get('name') != desired['name'] + or existing.get('prefered_ranking') != desired['prefered_ranking'] + ): + await self.ap.rerank_models_service.update_rerank_model(space_model.uuid, dict(desired)) + updated += 1 + if created or updated: self.ap.logger.info(f'Synced models from LangBot Space: {created} added, {updated} updated.') diff --git a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py index efad67250..9d8568437 100644 --- a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py +++ b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py @@ -943,16 +943,21 @@ class LiteLLMRequester(requester.ProviderAPIRequester): if api_key: headers['Authorization'] = f'Bearer {api_key}' + request_args = dict(extra_args) + rerank_url = request_args.pop('rerank_url', None) + rerank_path = request_args.pop('rerank_path', 'rerank') + payload: dict[str, typing.Any] = { 'model': model_name, 'query': query, 'documents': documents, 'top_n': top_n, } - if extra_args: - payload.update(extra_args) + if request_args: + payload.update(request_args) - rerank_url = f'{base_url}/rerank' + if not rerank_url: + rerank_url = f'{base_url}/{str(rerank_path).strip("/")}' try: async with httpx.AsyncClient(timeout=timeout) as client: diff --git a/tests/unit_tests/api/service/test_provider_service.py b/tests/unit_tests/api/service/test_provider_service.py index 8b308af8d..acc00b357 100644 --- a/tests/unit_tests/api/service/test_provider_service.py +++ b/tests/unit_tests/api/service/test_provider_service.py @@ -777,6 +777,44 @@ class TestModelProviderServiceScanProviderModels: assert len(result['models']) == 1 assert result['models'][0]['type'] == 'llm' + async def test_scan_provider_marks_existing_rerank_model(self): + """Rerank scan results use the rerank service when computing already_added.""" + ap = SimpleNamespace() + ap.persistence_mgr = SimpleNamespace() + ap.model_mgr = SimpleNamespace() + ap.llm_model_service = SimpleNamespace() + ap.embedding_models_service = SimpleNamespace() + ap.rerank_models_service = SimpleNamespace() + + provider = _create_mock_provider(provider_uuid='rerank-scan-uuid') + ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_mock_result([], first_item=provider)) + ap.persistence_mgr.serialize_model = Mock( + return_value={ + 'uuid': 'rerank-scan-uuid', + 'name': 'New API', + 'requester': 'new-api-chat-completions', + 'base_url': 'https://new-api.example.com/v1', + 'api_keys': ['key'], + } + ) + + runtime_provider = Mock() + runtime_provider.token_mgr.get_token.return_value = 'token' + runtime_provider.requester.scan_models = AsyncMock( + return_value={'models': [{'id': 'Qwen3-Reranker-8B', 'type': 'rerank'}]} + ) + ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider) + ap.llm_model_service.get_llm_models_by_provider = AsyncMock(return_value=[]) + ap.embedding_models_service.get_embedding_models_by_provider = AsyncMock(return_value=[]) + ap.rerank_models_service.get_rerank_models_by_provider = AsyncMock( + return_value=[{'name': 'Qwen3-Reranker-8B'}] + ) + + result = await ModelProviderService(ap).scan_provider_models('rerank-scan-uuid', model_type='rerank') + + assert result['models'][0]['type'] == 'rerank' + assert result['models'][0]['already_added'] is True + async def test_scan_provider_not_implemented_raises_error(self): """Raises ValueError when scan not implemented.""" # Setup diff --git a/tests/unit_tests/api/service/test_space_service.py b/tests/unit_tests/api/service/test_space_service.py index f48b18937..f02a18b5b 100644 --- a/tests/unit_tests/api/service/test_space_service.py +++ b/tests/unit_tests/api/service/test_space_service.py @@ -693,7 +693,7 @@ class TestSpaceServiceGetModels: 'uuid': 'uuid-2', 'model_id': 'model-2', 'provider': 'provider-2', - 'category': 'chat', + 'category': 'rerank', 'status': 'active', }, ] @@ -714,6 +714,7 @@ class TestSpaceServiceGetModels: # Verify assert len(result) == 2 + assert result[1].category == 'rerank' async def test_get_models_api_error(self): """Raises ValueError on API error.""" diff --git a/tests/unit_tests/provider/test_litellmchat.py b/tests/unit_tests/provider/test_litellmchat.py index a22a589a7..d0332f726 100644 --- a/tests/unit_tests/provider/test_litellmchat.py +++ b/tests/unit_tests/provider/test_litellmchat.py @@ -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""" diff --git a/tests/unit_tests/provider/test_model_manager.py b/tests/unit_tests/provider/test_model_manager.py index 015fd5450..9e7a5c09a 100644 --- a/tests/unit_tests/provider/test_model_manager.py +++ b/tests/unit_tests/provider/test_model_manager.py @@ -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 # ============================================================================