Add rerank models to the dynamic model list; Fix integration tests

tests/unit/providers/nvidia/test_rerank_inference.py

@@ -4,11 +4,12 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import aiohttp
 import pytest
 
+from llama_stack.apis.models import ModelType
 from llama_stack.providers.remote.inference.nvidia.config import NVIDIAConfig
 from llama_stack.providers.remote.inference.nvidia.nvidia import NVIDIAInferenceAdapter
 
@@ -170,3 +171,35 @@ async def test_client_error():
     with patch("aiohttp.ClientSession", return_value=mock_session):
         with pytest.raises(ConnectionError, match="Failed to connect.*Network error"):
             await adapter.rerank(model="test-model", query="q", items=["a"])
+
+
+async def test_list_models_adds_rerank_models():
+    """Test that list_models adds rerank models to the dynamic model list."""
+    adapter = create_adapter()
+    adapter.__provider_id__ = "nvidia"
+
+    # Mock the list_models from the superclass to return some dynamic models
+    base_models = [
+        MagicMock(identifier="llm-1", model_type=ModelType.llm),
+        MagicMock(identifier="embedding-1", model_type=ModelType.embedding),
+    ]
+
+    with patch.object(NVIDIAInferenceAdapter.__bases__[0], "list_models", return_value=base_models):
+        result = await adapter.list_models()
+
+        assert result is not None
+
+        # Check that the rerank models are added
+        model_ids = [m.identifier for m in result]
+        assert "nv-rerank-qa-mistral-4b:1" in model_ids
+        assert "nvidia/nv-rerankqa-mistral-4b-v3" in model_ids
+        assert "nvidia/llama-3.2-nv-rerankqa-1b-v2" in model_ids
+
+        rerank_models = [m for m in result if m.model_type == ModelType.rerank]
+
+        assert len(rerank_models) == 3
+
+        for rerank_model in rerank_models:
+            assert rerank_model.provider_id == "nvidia"
+            assert rerank_model.metadata == {}
+            assert rerank_model.identifier in adapter._model_cache

tests/unit/providers/utils/inference/test_openai_mixin.py

@@ -35,6 +35,40 @@ class OpenAIMixinWithEmbeddingsImpl(OpenAIMixinImpl):
     }
 
 
+class OpenAIMixinWithRerankImpl(OpenAIMixin):
+    """Test implementation with rerank model list"""
+
+    rerank_model_list = ["rerank-model-1", "rerank-model-2"]
+
+    def __init__(self):
+        self.__provider_id__ = "test-provider"
+
+    def get_api_key(self) -> str:
+        raise NotImplementedError("This method should be mocked in tests")
+
+    def get_base_url(self) -> str:
+        raise NotImplementedError("This method should be mocked in tests")
+
+
+class OpenAIMixinWithEmbeddingsAndRerankImpl(OpenAIMixin):
+    """Test implementation with both embedding model metadata and rerank model list"""
+
+    embedding_model_metadata = {
+        "text-embedding-3-small": {"embedding_dimension": 1536, "context_length": 8192},
+        "text-embedding-ada-002": {"embedding_dimension": 1536, "context_length": 8192},
+    }
+
+    rerank_model_list = ["rerank-model-1", "rerank-model-2"]
+
+    __provider_id__ = "test-provider"
+
+    def get_api_key(self) -> str:
+        raise NotImplementedError("This method should be mocked in tests")
+
+    def get_base_url(self) -> str:
+        raise NotImplementedError("This method should be mocked in tests")
+
+
 @pytest.fixture
 def mixin():
     """Create a test instance of OpenAIMixin with mocked model_store"""
@@ -56,6 +90,18 @@ def mixin_with_embeddings():
     return OpenAIMixinWithEmbeddingsImpl()
 
 
+@pytest.fixture
+def mixin_with_rerank():
+    """Create a test instance of OpenAIMixin with rerank model list"""
+    return OpenAIMixinWithRerankImpl()
+
+
+@pytest.fixture
+def mixin_with_embeddings_and_rerank():
+    """Create a test instance of OpenAIMixin with both embedding model metadata and rerank model list"""
+    return OpenAIMixinWithEmbeddingsAndRerankImpl()
+
+
 @pytest.fixture
 def mock_models():
     """Create multiple mock OpenAI model objects"""
@@ -317,6 +363,96 @@ class TestOpenAIMixinEmbeddingModelMetadata:
             assert llm_model.provider_resource_id == "gpt-4"
 
 
+class TestOpenAIMixinRerankModelList:
+    """Test cases for rerank_model_list attribute functionality"""
+
+    async def test_rerank_model_identified(self, mixin_with_rerank, mock_client_context):
+        """Test that models in rerank_model_list are correctly identified as rerank models"""
+        # Create mock models: 1 rerank model and 1 LLM
+        mock_rerank_model = MagicMock(id="rerank-model-1")
+        mock_llm_model = MagicMock(id="gpt-4")
+        mock_models = [mock_rerank_model, mock_llm_model]
+
+        mock_client = MagicMock()
+
+        async def mock_models_list():
+            for model in mock_models:
+                yield model
+
+        mock_client.models.list.return_value = mock_models_list()
+
+        with mock_client_context(mixin_with_rerank, mock_client):
+            result = await mixin_with_rerank.list_models()
+
+            assert result is not None
+            assert len(result) == 2
+
+            # Find the models in the result
+            rerank_model = next(m for m in result if m.identifier == "rerank-model-1")
+            llm_model = next(m for m in result if m.identifier == "gpt-4")
+
+            # Check rerank model
+            assert rerank_model.model_type == ModelType.rerank
+            assert rerank_model.metadata == {}  # No metadata for rerank models
+            assert rerank_model.provider_id == "test-provider"
+            assert rerank_model.provider_resource_id == "rerank-model-1"
+
+            # Check LLM model
+            assert llm_model.model_type == ModelType.llm
+            assert llm_model.metadata == {}  # No metadata for LLMs
+            assert llm_model.provider_id == "test-provider"
+            assert llm_model.provider_resource_id == "gpt-4"
+
+
+class TestOpenAIMixinMixedModelTypes:
+    """Test cases for mixed model types (LLM, embedding, rerank)"""
+
+    async def test_mixed_model_types_identification(self, mixin_with_embeddings_and_rerank, mock_client_context):
+        """Test that LLM, embedding, and rerank models are correctly identified with proper types and metadata"""
+        # Create mock models: 1 embedding, 1 rerank, 1 LLM
+        mock_embedding_model = MagicMock(id="text-embedding-3-small")
+        mock_rerank_model = MagicMock(id="rerank-model-1")
+        mock_llm_model = MagicMock(id="gpt-4")
+        mock_models = [mock_embedding_model, mock_rerank_model, mock_llm_model]
+
+        mock_client = MagicMock()
+
+        async def mock_models_list():
+            for model in mock_models:
+                yield model
+
+        mock_client.models.list.return_value = mock_models_list()
+
+        with mock_client_context(mixin_with_embeddings_and_rerank, mock_client):
+            result = await mixin_with_embeddings_and_rerank.list_models()
+
+            assert result is not None
+            assert len(result) == 3
+
+            # Find the models in the result
+            embedding_model = next(m for m in result if m.identifier == "text-embedding-3-small")
+            rerank_model = next(m for m in result if m.identifier == "rerank-model-1")
+            llm_model = next(m for m in result if m.identifier == "gpt-4")
+
+            # Check embedding model
+            assert embedding_model.model_type == ModelType.embedding
+            assert embedding_model.metadata == {"embedding_dimension": 1536, "context_length": 8192}
+            assert embedding_model.provider_id == "test-provider"
+            assert embedding_model.provider_resource_id == "text-embedding-3-small"
+
+            # Check rerank model
+            assert rerank_model.model_type == ModelType.rerank
+            assert rerank_model.metadata == {}  # No metadata for rerank models
+            assert rerank_model.provider_id == "test-provider"
+            assert rerank_model.provider_resource_id == "rerank-model-1"
+
+            # Check LLM model
+            assert llm_model.model_type == ModelType.llm
+            assert llm_model.metadata == {}  # No metadata for LLMs
+            assert llm_model.provider_id == "test-provider"
+            assert llm_model.provider_resource_id == "gpt-4"
+
+
 class TestOpenAIMixinAllowedModels:
     """Test cases for allowed_models filtering functionality"""