Merge branch 'main' into add-localize-url-feature-to-openaimixin

tests/unit/distribution/routers/test_routing_tables.py

@@ -10,6 +10,7 @@ from unittest.mock import AsyncMock
 
 import pytest
 
+from llama_stack.apis.common.content_types import URL
 from llama_stack.apis.common.type_system import NumberType
 from llama_stack.apis.datasets.datasets import Dataset, DatasetPurpose, URIDataSource
 from llama_stack.apis.datatypes import Api
@@ -645,3 +646,25 @@ async def test_models_source_interaction_cleanup_provider_models(cached_disk_dis
 
     # Cleanup
     await table.shutdown()
+
+
+async def test_tool_groups_routing_table_exception_handling(cached_disk_dist_registry):
+    """Test that the tool group routing table handles exceptions when listing tools, like if an MCP server is unreachable."""
+
+    exception_throwing_tool_groups_impl = ToolGroupsImpl()
+    exception_throwing_tool_groups_impl.list_runtime_tools = AsyncMock(side_effect=Exception("Test exception"))
+
+    table = ToolGroupsRoutingTable(
+        {"test_provider": exception_throwing_tool_groups_impl}, cached_disk_dist_registry, {}
+    )
+    await table.initialize()
+
+    await table.register_tool_group(
+        toolgroup_id="test-toolgroup-exceptions",
+        provider_id="test_provider",
+        mcp_endpoint=URL(uri="http://localhost:8479/foo/bar"),
+    )
+
+    tools = await table.list_tools(toolgroup_id="test-toolgroup-exceptions")
+
+    assert len(tools.data) == 0

tests/unit/providers/inference/test_litellm_openai_mixin.py

@@ -26,7 +26,6 @@ class TestProviderDataValidator(BaseModel):
 class TestLiteLLMAdapter(LiteLLMOpenAIMixin):
     def __init__(self, config: TestConfig):
         super().__init__(
-            model_entries=[],
             litellm_provider_name="test",
             api_key_from_config=config.api_key,
             provider_data_api_key_field="test_api_key",

tests/unit/providers/inference/test_remote_vllm.py

@@ -62,7 +62,7 @@ from llama_stack.providers.remote.inference.vllm.vllm import (
 
 @pytest.fixture(scope="module")
 def mock_openai_models_list():
-    with patch("openai.resources.models.AsyncModels.list", new_callable=AsyncMock) as mock_list:
+    with patch("openai.resources.models.AsyncModels.list") as mock_list:
         yield mock_list
 
 

tests/unit/providers/nvidia/test_eval.py

@@ -150,7 +150,7 @@ class TestNVIDIAEvalImpl(unittest.TestCase):
         self._assert_request_body(
             {
                 "config": f"nvidia/{MOCK_BENCHMARK_ID}",
-                "target": {"type": "model", "model": "meta/llama-3.1-8b-instruct"},
+                "target": {"type": "model", "model": "Llama3.1-8B-Instruct"},
             }
         )
 

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

@@ -13,7 +13,6 @@ from llama_stack.apis.models import ModelType
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 
-# Test implementation of OpenAIMixin for testing purposes
 class OpenAIMixinImpl(OpenAIMixin):
     def __init__(self):
         self.__provider_id__ = "test-provider"
@@ -25,6 +24,23 @@ class OpenAIMixinImpl(OpenAIMixin):
         raise NotImplementedError("This method should be mocked in tests")
 
 
+class OpenAIMixinWithEmbeddingsImpl(OpenAIMixin):
+    """Test implementation with embedding model metadata"""
+
+    embedding_model_metadata = {
+        "text-embedding-3-small": {"embedding_dimension": 1536, "context_length": 8192},
+        "text-embedding-ada-002": {"embedding_dimension": 1536, "context_length": 8192},
+    }
+
+    __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"""
@@ -40,6 +56,12 @@ def mixin():
     return mixin_instance
 
 
+@pytest.fixture
+def mixin_with_embeddings():
+    """Create a test instance of OpenAIMixin with embedding model metadata"""
+    return OpenAIMixinWithEmbeddingsImpl()
+
+
 @pytest.fixture
 def mock_models():
     """Create multiple mock OpenAI model objects"""
@@ -258,3 +280,89 @@ class TestOpenAIMixinImagePreprocessing:
             content = processed_messages[0]["content"]
             assert len(content) == 2
             assert content[1]["image_url"]["url"] == "http://example.com/image.jpg"
+
+
+class TestOpenAIMixinEmbeddingModelMetadata:
+    """Test cases for embedding_model_metadata attribute functionality"""
+
+    async def test_embedding_model_identified_and_augmented(self, mixin_with_embeddings, mock_client_context):
+        """Test that models in embedding_model_metadata are correctly identified as embeddings with metadata"""
+        # Create mock models: 1 embedding model and 1 LLM, while there are 2 known embedding models
+        mock_embedding_model = MagicMock(id="text-embedding-3-small")
+        mock_llm_model = MagicMock(id="gpt-4")
+        mock_models = [mock_embedding_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, mock_client):
+            result = await mixin_with_embeddings.list_models()
+
+            assert result is not None
+            assert len(result) == 2
+
+            # Find the models in the result
+            embedding_model = next(m for m in result if m.identifier == "text-embedding-3-small")
+            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 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"""
+
+    async def test_list_models_with_allowed_models_filter(self, mixin, mock_client_with_models, mock_client_context):
+        """Test that list_models filters models based on allowed_models set"""
+        mixin.allowed_models = {"some-mock-model-id", "another-mock-model-id"}
+
+        with mock_client_context(mixin, mock_client_with_models):
+            result = await mixin.list_models()
+
+            assert result is not None
+            assert len(result) == 2
+
+            model_ids = [model.identifier for model in result]
+            assert "some-mock-model-id" in model_ids
+            assert "another-mock-model-id" in model_ids
+            assert "final-mock-model-id" not in model_ids
+
+    async def test_list_models_with_empty_allowed_models(self, mixin, mock_client_with_models, mock_client_context):
+        """Test that empty allowed_models set allows all models"""
+        assert len(mixin.allowed_models) == 0
+
+        with mock_client_context(mixin, mock_client_with_models):
+            result = await mixin.list_models()
+
+            assert result is not None
+            assert len(result) == 3  # All models should be included
+
+            model_ids = [model.identifier for model in result]
+            assert "some-mock-model-id" in model_ids
+            assert "another-mock-model-id" in model_ids
+            assert "final-mock-model-id" in model_ids
+
+    async def test_check_model_availability_with_allowed_models(
+        self, mixin, mock_client_with_models, mock_client_context
+    ):
+        """Test that check_model_availability respects allowed_models"""
+        mixin.allowed_models = {"final-mock-model-id"}
+
+        with mock_client_context(mixin, mock_client_with_models):
+            assert await mixin.check_model_availability("final-mock-model-id")
+            assert not await mixin.check_model_availability("some-mock-model-id")
+            assert not await mixin.check_model_availability("another-mock-model-id")

tests/unit/providers/utils/test_model_registry.py

@@ -84,14 +84,14 @@ def unknown_model() -> Model:
 
 @pytest.fixture
 def helper(known_provider_model: ProviderModelEntry, known_provider_model2: ProviderModelEntry) -> ModelRegistryHelper:
-    return ModelRegistryHelper([known_provider_model, known_provider_model2])
+    return ModelRegistryHelper(model_entries=[known_provider_model, known_provider_model2])
 
 
 class MockModelRegistryHelperWithDynamicModels(ModelRegistryHelper):
     """Test helper that simulates a provider with dynamically available models."""
 
     def __init__(self, model_entries: list[ProviderModelEntry], available_models: list[str]):
-        super().__init__(model_entries)
+        super().__init__(model_entries=model_entries)
         self._available_models = available_models
 
     async def check_model_availability(self, model: str) -> bool:

tests/unit/registry/test_registry.py

@@ -129,7 +129,7 @@ async def test_duplicate_provider_registration(cached_disk_dist_registry):
 
     result = await cached_disk_dist_registry.get("vector_db", "test_vector_db_2")
     assert result is not None
-    assert result.embedding_model == duplicate_vector_db.embedding_model  # Original values preserved
+    assert result.embedding_model == original_vector_db.embedding_model  # Original values preserved
 
 
 async def test_get_all_objects(cached_disk_dist_registry):
@@ -174,14 +174,10 @@ async def test_parse_registry_values_error_handling(sqlite_kvstore):
     )
 
     await sqlite_kvstore.set(
-        KEY_FORMAT.format(type="vector_db", identifier="valid_vector_db"),
-        valid_db.model_dump_json(),
+        KEY_FORMAT.format(type="vector_db", identifier="valid_vector_db"), valid_db.model_dump_json()
     )
 
-    await sqlite_kvstore.set(
-        KEY_FORMAT.format(type="vector_db", identifier="corrupted_json"),
-        "{not valid json",
-    )
+    await sqlite_kvstore.set(KEY_FORMAT.format(type="vector_db", identifier="corrupted_json"), "{not valid json")
 
     await sqlite_kvstore.set(
         KEY_FORMAT.format(type="vector_db", identifier="missing_fields"),
@@ -216,8 +212,7 @@ async def test_cached_registry_error_handling(sqlite_kvstore):
     )
 
     await sqlite_kvstore.set(
-        KEY_FORMAT.format(type="vector_db", identifier="valid_cached_db"),
-        valid_db.model_dump_json(),
+        KEY_FORMAT.format(type="vector_db", identifier="valid_cached_db"), valid_db.model_dump_json()
     )
 
     await sqlite_kvstore.set(