Merge branch 'main' of https://github.com/meta-llama/llama-stack into add_nemo_customizer

tests/unit/cli/test_stack_config.py

@@ -0,0 +1,127 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from datetime import datetime
+
+import pytest
+import yaml
+
+from llama_stack.distribution.configure import (
+    LLAMA_STACK_RUN_CONFIG_VERSION,
+    parse_and_maybe_upgrade_config,
+)
+
+
+@pytest.fixture
+def up_to_date_config():
+    return yaml.safe_load(
+        """
+        version: {version}
+        image_name: foo
+        apis_to_serve: []
+        built_at: {built_at}
+        providers:
+          inference:
+            - provider_id: provider1
+              provider_type: inline::meta-reference
+              config: {{}}
+          safety:
+            - provider_id: provider1
+              provider_type: inline::meta-reference
+              config:
+                llama_guard_shield:
+                  model: Llama-Guard-3-1B
+                  excluded_categories: []
+                  disable_input_check: false
+                  disable_output_check: false
+                enable_prompt_guard: false
+          memory:
+            - provider_id: provider1
+              provider_type: inline::meta-reference
+              config: {{}}
+    """.format(version=LLAMA_STACK_RUN_CONFIG_VERSION, built_at=datetime.now().isoformat())
+    )
+
+
+@pytest.fixture
+def old_config():
+    return yaml.safe_load(
+        """
+        image_name: foo
+        built_at: {built_at}
+        apis_to_serve: []
+        routing_table:
+          inference:
+            - provider_type: remote::ollama
+              config:
+                host: localhost
+                port: 11434
+              routing_key: Llama3.2-1B-Instruct
+            - provider_type: inline::meta-reference
+              config:
+                model: Llama3.1-8B-Instruct
+              routing_key: Llama3.1-8B-Instruct
+          safety:
+            - routing_key: ["shield1", "shield2"]
+              provider_type: inline::meta-reference
+              config:
+                llama_guard_shield:
+                  model: Llama-Guard-3-1B
+                  excluded_categories: []
+                  disable_input_check: false
+                  disable_output_check: false
+                enable_prompt_guard: false
+          memory:
+            - routing_key: vector
+              provider_type: inline::meta-reference
+              config: {{}}
+        api_providers:
+          telemetry:
+            provider_type: noop
+            config: {{}}
+    """.format(built_at=datetime.now().isoformat())
+    )
+
+
+@pytest.fixture
+def invalid_config():
+    return yaml.safe_load(
+        """
+        routing_table: {}
+        api_providers: {}
+    """
+    )
+
+
+def test_parse_and_maybe_upgrade_config_up_to_date(up_to_date_config):
+    result = parse_and_maybe_upgrade_config(up_to_date_config)
+    assert result.version == LLAMA_STACK_RUN_CONFIG_VERSION
+    assert "inference" in result.providers
+
+
+def test_parse_and_maybe_upgrade_config_old_format(old_config):
+    result = parse_and_maybe_upgrade_config(old_config)
+    assert result.version == LLAMA_STACK_RUN_CONFIG_VERSION
+    assert all(api in result.providers for api in ["inference", "safety", "memory", "telemetry"])
+    safety_provider = result.providers["safety"][0]
+    assert safety_provider.provider_type == "inline::meta-reference"
+    assert "llama_guard_shield" in safety_provider.config
+
+    inference_providers = result.providers["inference"]
+    assert len(inference_providers) == 2
+    assert {x.provider_id for x in inference_providers} == {
+        "remote::ollama-00",
+        "inline::meta-reference-01",
+    }
+
+    ollama = inference_providers[0]
+    assert ollama.provider_type == "remote::ollama"
+    assert ollama.config["port"] == 11434
+
+
+def test_parse_and_maybe_upgrade_config_invalid(invalid_config):
+    with pytest.raises(KeyError):
+        parse_and_maybe_upgrade_config(invalid_config)

tests/unit/models/test_prompt_adapter.py

@@ -0,0 +1,288 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import asyncio
+import unittest
+
+from llama_stack.apis.inference import (
+    ChatCompletionRequest,
+    CompletionMessage,
+    StopReason,
+    SystemMessage,
+    ToolCall,
+    ToolConfig,
+    UserMessage,
+)
+from llama_stack.models.llama.datatypes import (
+    BuiltinTool,
+    ToolDefinition,
+    ToolParamDefinition,
+    ToolPromptFormat,
+)
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    chat_completion_request_to_messages,
+    chat_completion_request_to_prompt,
+)
+
+MODEL = "Llama3.1-8B-Instruct"
+MODEL3_2 = "Llama3.2-3B-Instruct"
+
+
+class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        asyncio.get_running_loop().set_debug(False)
+
+    async def test_system_default(self):
+        content = "Hello !"
+        request = ChatCompletionRequest(
+            model=MODEL,
+            messages=[
+                UserMessage(content=content),
+            ],
+        )
+        messages = chat_completion_request_to_messages(request, MODEL)
+        self.assertEqual(len(messages), 2)
+        self.assertEqual(messages[-1].content, content)
+        self.assertTrue("Cutting Knowledge Date: December 2023" in messages[0].content)
+
+    async def test_system_builtin_only(self):
+        content = "Hello !"
+        request = ChatCompletionRequest(
+            model=MODEL,
+            messages=[
+                UserMessage(content=content),
+            ],
+            tools=[
+                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+                ToolDefinition(tool_name=BuiltinTool.brave_search),
+            ],
+        )
+        messages = chat_completion_request_to_messages(request, MODEL)
+        self.assertEqual(len(messages), 2)
+        self.assertEqual(messages[-1].content, content)
+        self.assertTrue("Cutting Knowledge Date: December 2023" in messages[0].content)
+        self.assertTrue("Tools: brave_search" in messages[0].content)
+
+    async def test_system_custom_only(self):
+        content = "Hello !"
+        request = ChatCompletionRequest(
+            model=MODEL,
+            messages=[
+                UserMessage(content=content),
+            ],
+            tools=[
+                ToolDefinition(
+                    tool_name="custom1",
+                    description="custom1 tool",
+                    parameters={
+                        "param1": ToolParamDefinition(
+                            param_type="str",
+                            description="param1 description",
+                            required=True,
+                        ),
+                    },
+                )
+            ],
+            tool_config=ToolConfig(tool_prompt_format=ToolPromptFormat.json),
+        )
+        messages = chat_completion_request_to_messages(request, MODEL)
+        self.assertEqual(len(messages), 3)
+        self.assertTrue("Environment: ipython" in messages[0].content)
+
+        self.assertTrue("Return function calls in JSON format" in messages[1].content)
+        self.assertEqual(messages[-1].content, content)
+
+    async def test_system_custom_and_builtin(self):
+        content = "Hello !"
+        request = ChatCompletionRequest(
+            model=MODEL,
+            messages=[
+                UserMessage(content=content),
+            ],
+            tools=[
+                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+                ToolDefinition(tool_name=BuiltinTool.brave_search),
+                ToolDefinition(
+                    tool_name="custom1",
+                    description="custom1 tool",
+                    parameters={
+                        "param1": ToolParamDefinition(
+                            param_type="str",
+                            description="param1 description",
+                            required=True,
+                        ),
+                    },
+                ),
+            ],
+        )
+        messages = chat_completion_request_to_messages(request, MODEL)
+        self.assertEqual(len(messages), 3)
+
+        self.assertTrue("Environment: ipython" in messages[0].content)
+        self.assertTrue("Tools: brave_search" in messages[0].content)
+
+        self.assertTrue("Return function calls in JSON format" in messages[1].content)
+        self.assertEqual(messages[-1].content, content)
+
+    async def test_completion_message_encoding(self):
+        request = ChatCompletionRequest(
+            model=MODEL3_2,
+            messages=[
+                UserMessage(content="hello"),
+                CompletionMessage(
+                    content="",
+                    stop_reason=StopReason.end_of_turn,
+                    tool_calls=[
+                        ToolCall(
+                            tool_name="custom1",
+                            arguments={"param1": "value1"},
+                            call_id="123",
+                        )
+                    ],
+                ),
+            ],
+            tools=[
+                ToolDefinition(
+                    tool_name="custom1",
+                    description="custom1 tool",
+                    parameters={
+                        "param1": ToolParamDefinition(
+                            param_type="str",
+                            description="param1 description",
+                            required=True,
+                        ),
+                    },
+                ),
+            ],
+            tool_config=ToolConfig(tool_prompt_format=ToolPromptFormat.python_list),
+        )
+        prompt = await chat_completion_request_to_prompt(request, request.model)
+        self.assertIn('[custom1(param1="value1")]', prompt)
+
+        request.model = MODEL
+        request.tool_config.tool_prompt_format = ToolPromptFormat.json
+        prompt = await chat_completion_request_to_prompt(request, request.model)
+        self.assertIn(
+            '{"type": "function", "name": "custom1", "parameters": {"param1": "value1"}}',
+            prompt,
+        )
+
+    async def test_user_provided_system_message(self):
+        content = "Hello !"
+        system_prompt = "You are a pirate"
+        request = ChatCompletionRequest(
+            model=MODEL,
+            messages=[
+                SystemMessage(content=system_prompt),
+                UserMessage(content=content),
+            ],
+            tools=[
+                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+            ],
+        )
+        messages = chat_completion_request_to_messages(request, MODEL)
+        self.assertEqual(len(messages), 2, messages)
+        self.assertTrue(messages[0].content.endswith(system_prompt))
+
+        self.assertEqual(messages[-1].content, content)
+
+    async def test_repalce_system_message_behavior_builtin_tools(self):
+        content = "Hello !"
+        system_prompt = "You are a pirate"
+        request = ChatCompletionRequest(
+            model=MODEL,
+            messages=[
+                SystemMessage(content=system_prompt),
+                UserMessage(content=content),
+            ],
+            tools=[
+                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+            ],
+            tool_config=ToolConfig(
+                tool_choice="auto",
+                tool_prompt_format="python_list",
+                system_message_behavior="replace",
+            ),
+        )
+        messages = chat_completion_request_to_messages(request, MODEL3_2)
+        self.assertEqual(len(messages), 2, messages)
+        self.assertTrue(messages[0].content.endswith(system_prompt))
+        self.assertIn("Environment: ipython", messages[0].content)
+        self.assertEqual(messages[-1].content, content)
+
+    async def test_repalce_system_message_behavior_custom_tools(self):
+        content = "Hello !"
+        system_prompt = "You are a pirate"
+        request = ChatCompletionRequest(
+            model=MODEL,
+            messages=[
+                SystemMessage(content=system_prompt),
+                UserMessage(content=content),
+            ],
+            tools=[
+                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+                ToolDefinition(
+                    tool_name="custom1",
+                    description="custom1 tool",
+                    parameters={
+                        "param1": ToolParamDefinition(
+                            param_type="str",
+                            description="param1 description",
+                            required=True,
+                        ),
+                    },
+                ),
+            ],
+            tool_config=ToolConfig(
+                tool_choice="auto",
+                tool_prompt_format="python_list",
+                system_message_behavior="replace",
+            ),
+        )
+        messages = chat_completion_request_to_messages(request, MODEL3_2)
+
+        self.assertEqual(len(messages), 2, messages)
+        self.assertTrue(messages[0].content.endswith(system_prompt))
+        self.assertIn("Environment: ipython", messages[0].content)
+        self.assertEqual(messages[-1].content, content)
+
+    async def test_replace_system_message_behavior_custom_tools_with_template(self):
+        content = "Hello !"
+        system_prompt = "You are a pirate {{ function_description }}"
+        request = ChatCompletionRequest(
+            model=MODEL,
+            messages=[
+                SystemMessage(content=system_prompt),
+                UserMessage(content=content),
+            ],
+            tools=[
+                ToolDefinition(tool_name=BuiltinTool.code_interpreter),
+                ToolDefinition(
+                    tool_name="custom1",
+                    description="custom1 tool",
+                    parameters={
+                        "param1": ToolParamDefinition(
+                            param_type="str",
+                            description="param1 description",
+                            required=True,
+                        ),
+                    },
+                ),
+            ],
+            tool_config=ToolConfig(
+                tool_choice="auto",
+                tool_prompt_format="python_list",
+                system_message_behavior="replace",
+            ),
+        )
+        messages = chat_completion_request_to_messages(request, MODEL3_2)
+
+        self.assertEqual(len(messages), 2, messages)
+        self.assertIn("Environment: ipython", messages[0].content)
+        self.assertIn("You are a pirate", messages[0].content)
+        # function description is present in the system prompt
+        self.assertIn('"name": "custom1"', messages[0].content)
+        self.assertEqual(messages[-1].content, content)

tests/unit/models/test_system_prompts.py

@@ -0,0 +1,82 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# top-level folder for each specific model found within the models/ directory at
+# the top-level of this source tree.
+
+import textwrap
+import unittest
+from datetime import datetime
+
+from llama_stack.models.llama.llama3.prompt_templates import (
+    BuiltinToolGenerator,
+    FunctionTagCustomToolGenerator,
+    JsonCustomToolGenerator,
+    PythonListCustomToolGenerator,
+    SystemDefaultGenerator,
+)
+
+
+class PromptTemplateTests(unittest.TestCase):
+    def check_generator_output(self, generator):
+        for example in generator.data_examples():
+            pt = generator.gen(example)
+            text = pt.render()
+            # print(text)  # debugging
+            if not example:
+                continue
+            for tool in example:
+                assert tool.tool_name in text
+
+    def test_system_default(self):
+        generator = SystemDefaultGenerator()
+        today = datetime.now().strftime("%d %B %Y")
+        expected_text = f"Cutting Knowledge Date: December 2023\nToday Date: {today}"
+        assert expected_text.strip("\n") == generator.gen(generator.data_examples()[0]).render()
+
+    def test_system_builtin_only(self):
+        generator = BuiltinToolGenerator()
+        expected_text = textwrap.dedent(
+            """
+            Environment: ipython
+            Tools: brave_search, wolfram_alpha
+            """
+        )
+        assert expected_text.strip("\n") == generator.gen(generator.data_examples()[0]).render()
+
+    def test_system_custom_only(self):
+        self.maxDiff = None
+        generator = JsonCustomToolGenerator()
+        self.check_generator_output(generator)
+
+    def test_system_custom_function_tag(self):
+        self.maxDiff = None
+        generator = FunctionTagCustomToolGenerator()
+        self.check_generator_output(generator)
+
+    def test_llama_3_2_system_zero_shot(self):
+        generator = PythonListCustomToolGenerator()
+        self.check_generator_output(generator)
+
+    def test_llama_3_2_provided_system_prompt(self):
+        generator = PythonListCustomToolGenerator()
+        user_system_prompt = textwrap.dedent(
+            """
+            Overriding message.
+
+            {{ function_description }}
+            """
+        )
+        example = generator.data_examples()[0]
+
+        pt = generator.gen(example, user_system_prompt)
+        text = pt.render()
+        assert "Overriding message." in text
+        assert '"name": "get_weather"' in text

tests/unit/providers/inference/test_remote_vllm.py

@@ -0,0 +1,234 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import asyncio
+import json
+import logging
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import Any, Dict
+from unittest.mock import AsyncMock, patch
+
+import pytest
+import pytest_asyncio
+from openai.types.chat.chat_completion_chunk import (
+    ChatCompletionChunk as OpenAIChatCompletionChunk,
+)
+from openai.types.chat.chat_completion_chunk import (
+    Choice as OpenAIChoice,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChoiceDelta as OpenAIChoiceDelta,
+)
+from openai.types.model import Model as OpenAIModel
+
+from llama_stack.apis.inference import ToolChoice, ToolConfig
+from llama_stack.apis.models import Model
+from llama_stack.models.llama.datatypes import StopReason
+from llama_stack.providers.remote.inference.vllm.config import VLLMInferenceAdapterConfig
+from llama_stack.providers.remote.inference.vllm.vllm import (
+    VLLMInferenceAdapter,
+    _process_vllm_chat_completion_stream_response,
+)
+
+# These are unit test for the remote vllm provider
+# implementation. This should only contain tests which are specific to
+# the implementation details of those classes. More general
+# (API-level) tests should be placed in tests/integration/inference/
+#
+# How to run this test:
+#
+# pytest tests/unit/providers/inference/test_remote_vllm.py \
+# -v -s --tb=short --disable-warnings
+
+
+class MockInferenceAdapterWithSleep:
+    def __init__(self, sleep_time: int, response: Dict[str, Any]):
+        self.httpd = None
+
+        class DelayedRequestHandler(BaseHTTPRequestHandler):
+            # ruff: noqa: N802
+            def do_POST(self):
+                time.sleep(sleep_time)
+                self.send_response(code=200)
+                self.end_headers()
+                self.wfile.write(json.dumps(response).encode("utf-8"))
+
+        self.request_handler = DelayedRequestHandler
+
+    def __enter__(self):
+        httpd = HTTPServer(("", 0), self.request_handler)
+        self.httpd = httpd
+        host, port = httpd.server_address
+        httpd_thread = threading.Thread(target=httpd.serve_forever)
+        httpd_thread.daemon = True  # stop server if this thread terminates
+        httpd_thread.start()
+
+        config = VLLMInferenceAdapterConfig(url=f"http://{host}:{port}")
+        inference_adapter = VLLMInferenceAdapter(config)
+        return inference_adapter
+
+    def __exit__(self, _exc_type, _exc_value, _traceback):
+        if self.httpd:
+            self.httpd.shutdown()
+            self.httpd.server_close()
+
+
+@pytest.fixture(scope="module")
+def mock_openai_models_list():
+    with patch("openai.resources.models.AsyncModels.list", new_callable=AsyncMock) as mock_list:
+        yield mock_list
+
+
+@pytest_asyncio.fixture(scope="module")
+async def vllm_inference_adapter():
+    config = VLLMInferenceAdapterConfig(url="http://mocked.localhost:12345")
+    inference_adapter = VLLMInferenceAdapter(config)
+    inference_adapter.model_store = AsyncMock()
+    await inference_adapter.initialize()
+    return inference_adapter
+
+
+@pytest.mark.asyncio
+async def test_register_model_checks_vllm(mock_openai_models_list, vllm_inference_adapter):
+    async def mock_openai_models():
+        yield OpenAIModel(id="foo", created=1, object="model", owned_by="test")
+
+    mock_openai_models_list.return_value = mock_openai_models()
+
+    foo_model = Model(identifier="foo", provider_resource_id="foo", provider_id="vllm-inference")
+
+    await vllm_inference_adapter.register_model(foo_model)
+    mock_openai_models_list.assert_called()
+
+
+@pytest.mark.asyncio
+async def test_old_vllm_tool_choice(vllm_inference_adapter):
+    """
+    Test that we set tool_choice to none when no tools are in use
+    to support older versions of vLLM
+    """
+    mock_model = Model(identifier="mock-model", provider_resource_id="mock-model", provider_id="vllm-inference")
+    vllm_inference_adapter.model_store.get_model.return_value = mock_model
+
+    with patch.object(vllm_inference_adapter, "_nonstream_chat_completion") as mock_nonstream_completion:
+        # No tools but auto tool choice
+        await vllm_inference_adapter.chat_completion(
+            "mock-model",
+            [],
+            stream=False,
+            tools=None,
+            tool_config=ToolConfig(tool_choice=ToolChoice.auto),
+        )
+        mock_nonstream_completion.assert_called()
+        request = mock_nonstream_completion.call_args.args[0]
+        # Ensure tool_choice gets converted to none for older vLLM versions
+        assert request.tool_config.tool_choice == ToolChoice.none
+
+
+@pytest.mark.asyncio
+async def test_tool_call_delta_empty_tool_call_buf():
+    """
+    Test that we don't generate extra chunks when processing a
+    tool call response that didn't call any tools. Previously we would
+    emit chunks with spurious ToolCallParseStatus.succeeded or
+    ToolCallParseStatus.failed when processing chunks that didn't
+    actually make any tool calls.
+    """
+
+    async def mock_stream():
+        delta = OpenAIChoiceDelta(content="", tool_calls=None)
+        choices = [OpenAIChoice(delta=delta, finish_reason="stop", index=0)]
+        mock_chunk = OpenAIChatCompletionChunk(
+            id="chunk-1",
+            created=1,
+            model="foo",
+            object="chat.completion.chunk",
+            choices=choices,
+        )
+        for chunk in [mock_chunk]:
+            yield chunk
+
+    chunks = [chunk async for chunk in _process_vllm_chat_completion_stream_response(mock_stream())]
+    assert len(chunks) == 1
+    assert chunks[0].event.stop_reason == StopReason.end_of_turn
+
+
+@pytest.mark.asyncio
+async def test_process_vllm_chat_completion_stream_response_no_choices():
+    """
+    Test that we don't error out when vLLM returns no choices for a
+    completion request. This can happen when there's an error thrown
+    in vLLM for example.
+    """
+
+    async def mock_stream():
+        choices = []
+        mock_chunk = OpenAIChatCompletionChunk(
+            id="chunk-1",
+            created=1,
+            model="foo",
+            object="chat.completion.chunk",
+            choices=choices,
+        )
+        for chunk in [mock_chunk]:
+            yield chunk
+
+    chunks = [chunk async for chunk in _process_vllm_chat_completion_stream_response(mock_stream())]
+    assert len(chunks) == 0
+
+
+def test_chat_completion_doesnt_block_event_loop(caplog):
+    loop = asyncio.new_event_loop()
+    loop.set_debug(True)
+    caplog.set_level(logging.WARNING)
+
+    # Log when event loop is blocked for more than 200ms
+    loop.slow_callback_duration = 0.5
+    # Sleep for 500ms in our delayed http response
+    sleep_time = 0.5
+
+    mock_model = Model(identifier="mock-model", provider_resource_id="mock-model", provider_id="vllm-inference")
+    mock_response = {
+        "id": "chatcmpl-abc123",
+        "object": "chat.completion",
+        "created": 1,
+        "modle": "mock-model",
+        "choices": [
+            {
+                "message": {"content": ""},
+                "logprobs": None,
+                "finish_reason": "stop",
+                "index": 0,
+            }
+        ],
+    }
+
+    async def do_chat_completion():
+        await inference_adapter.chat_completion(
+            "mock-model",
+            [],
+            stream=False,
+            tools=None,
+            tool_config=ToolConfig(tool_choice=ToolChoice.auto),
+        )
+
+    with MockInferenceAdapterWithSleep(sleep_time, mock_response) as inference_adapter:
+        inference_adapter.model_store = AsyncMock()
+        inference_adapter.model_store.get_model.return_value = mock_model
+        loop.run_until_complete(inference_adapter.initialize())
+
+        # Clear the logs so far and run the actual chat completion we care about
+        caplog.clear()
+        loop.run_until_complete(do_chat_completion())
+
+    # Ensure we don't have any asyncio warnings in the captured log
+    # records from our chat completion call. A message gets logged
+    # here any time we exceed the slow_callback_duration configured
+    # above.
+    asyncio_warnings = [record.message for record in caplog.records if record.name == "asyncio"]
+    assert not asyncio_warnings

tests/unit/providers/test_configs.py

@@ -0,0 +1,50 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import pytest
+from pydantic import BaseModel
+
+from llama_stack.distribution.distribution import get_provider_registry, providable_apis
+from llama_stack.distribution.utils.dynamic import instantiate_class_type
+
+
+class TestProviderConfigurations:
+    """Test suite for testing provider configurations across all API types."""
+
+    def test_all_api_providers_exist(self):
+        provider_registry = get_provider_registry()
+        for api in providable_apis():
+            providers = provider_registry.get(api, {})
+            assert providers, f"No providers found for API type: {api}"
+
+    @pytest.mark.parametrize("api", providable_apis())
+    def test_api_providers(self, api):
+        provider_registry = get_provider_registry()
+        providers = provider_registry.get(api, {})
+        assert providers, f"No providers found for API type: {api}"
+
+        failures = []
+        for provider_type, provider_spec in providers.items():
+            try:
+                self._verify_provider_config(provider_type, provider_spec)
+            except Exception as e:
+                failures.append(f"Failed to verify {provider_type} config: {str(e)}")
+
+        if failures:
+            pytest.fail("\n".join(failures))
+
+    def _verify_provider_config(self, provider_type, provider_spec):
+        """Helper method to verify a single provider configuration."""
+        # Get the config class
+        config_class_name = provider_spec.config_class
+        config_type = instantiate_class_type(config_class_name)
+
+        assert issubclass(config_type, BaseModel), f"{config_class_name} is not a subclass of BaseModel"
+
+        assert hasattr(config_type, "sample_run_config"), f"{config_class_name} does not have sample_run_config method"
+
+        sample_config = config_type.sample_run_config(__distro_dir__="foobarbaz")
+        assert isinstance(sample_config, dict), f"{config_class_name}.sample_run_config() did not return a dict"

tests/unit/providers/vector_io/conftest.py

@@ -0,0 +1,42 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import random
+
+import numpy as np
+import pytest
+
+from llama_stack.apis.vector_io import Chunk
+
+EMBEDDING_DIMENSION = 384
+
+
+@pytest.fixture
+def vector_db_id() -> str:
+    return f"test-vector-db-{random.randint(1, 100)}"
+
+
+@pytest.fixture(scope="session")
+def embedding_dimension() -> int:
+    return EMBEDDING_DIMENSION
+
+
+@pytest.fixture(scope="session")
+def sample_chunks():
+    """Generates chunks that force multiple batches for a single document to expose ID conflicts."""
+    n, k = 10, 3
+    sample = [
+        Chunk(content=f"Sentence {i} from document {j}", metadata={"document_id": f"document-{j}"})
+        for j in range(k)
+        for i in range(n)
+    ]
+    return sample
+
+
+@pytest.fixture(scope="session")
+def sample_embeddings(sample_chunks):
+    np.random.seed(42)
+    return np.array([np.random.rand(EMBEDDING_DIMENSION).astype(np.float32) for _ in sample_chunks])

tests/unit/providers/vector_io/test_qdrant.py

@@ -0,0 +1,135 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import asyncio
+import os
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+import pytest_asyncio
+
+from llama_stack.apis.inference import EmbeddingsResponse, Inference
+from llama_stack.apis.vector_io import (
+    QueryChunksResponse,
+    VectorDB,
+    VectorDBStore,
+)
+from llama_stack.providers.inline.vector_io.qdrant.config import (
+    QdrantVectorIOConfig as InlineQdrantVectorIOConfig,
+)
+from llama_stack.providers.remote.vector_io.qdrant.qdrant import (
+    QdrantVectorIOAdapter,
+)
+
+# This test is a unit test for the QdrantVectorIOAdapter class. This should only contain
+# tests which are specific to this class. More general (API-level) tests should be placed in
+# tests/integration/vector_io/
+#
+# How to run this test:
+#
+# pytest tests/unit/providers/vector_io/test_qdrant.py \
+# -v -s --tb=short --disable-warnings --asyncio-mode=auto
+
+
+@pytest.fixture
+def qdrant_config(tmp_path) -> InlineQdrantVectorIOConfig:
+    return InlineQdrantVectorIOConfig(path=os.path.join(tmp_path, "qdrant.db"))
+
+
+@pytest.fixture(scope="session")
+def loop():
+    return asyncio.new_event_loop()
+
+
+@pytest.fixture
+def mock_vector_db(vector_db_id) -> MagicMock:
+    mock_vector_db = MagicMock(spec=VectorDB)
+    mock_vector_db.embedding_model = "embedding_model"
+    mock_vector_db.identifier = vector_db_id
+    return mock_vector_db
+
+
+@pytest.fixture
+def mock_vector_db_store(mock_vector_db) -> MagicMock:
+    mock_store = MagicMock(spec=VectorDBStore)
+    mock_store.get_vector_db = AsyncMock(return_value=mock_vector_db)
+    return mock_store
+
+
+@pytest.fixture
+def mock_api_service(sample_embeddings):
+    mock_api_service = MagicMock(spec=Inference)
+    mock_api_service.embeddings = AsyncMock(return_value=EmbeddingsResponse(embeddings=sample_embeddings))
+    return mock_api_service
+
+
+@pytest_asyncio.fixture
+async def qdrant_adapter(qdrant_config, mock_vector_db_store, mock_api_service, loop) -> QdrantVectorIOAdapter:
+    adapter = QdrantVectorIOAdapter(config=qdrant_config, inference_api=mock_api_service)
+    adapter.vector_db_store = mock_vector_db_store
+    await adapter.initialize()
+    yield adapter
+    await adapter.shutdown()
+
+
+__QUERY = "Sample query"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("max_query_chunks, expected_chunks", [(2, 2), (100, 30)])
+async def test_qdrant_adapter_returns_expected_chunks(
+    qdrant_adapter: QdrantVectorIOAdapter,
+    vector_db_id,
+    sample_chunks,
+    sample_embeddings,
+    max_query_chunks,
+    expected_chunks,
+) -> None:
+    assert qdrant_adapter is not None
+    await qdrant_adapter.insert_chunks(vector_db_id, sample_chunks)
+
+    index = await qdrant_adapter._get_and_cache_vector_db_index(vector_db_id=vector_db_id)
+    assert index is not None
+
+    response = await qdrant_adapter.query_chunks(
+        query=__QUERY,
+        vector_db_id=vector_db_id,
+        params={"max_chunks": max_query_chunks},
+    )
+    assert isinstance(response, QueryChunksResponse)
+    assert len(response.chunks) == expected_chunks
+
+
+# To by-pass attempt to convert a Mock to JSON
+def _prepare_for_json(value: Any) -> str:
+    return str(value)
+
+
+@patch("llama_stack.providers.utils.telemetry.trace_protocol._prepare_for_json", new=_prepare_for_json)
+@pytest.mark.asyncio
+async def test_qdrant_register_and_unregister_vector_db(
+    qdrant_adapter: QdrantVectorIOAdapter,
+    mock_vector_db,
+    sample_chunks,
+) -> None:
+    # Initially, no collections
+    vector_db_id = mock_vector_db.identifier
+    assert len((await qdrant_adapter.client.get_collections()).collections) == 0
+
+    # Register does not create a collection
+    assert not (await qdrant_adapter.client.collection_exists(vector_db_id))
+    await qdrant_adapter.register_vector_db(mock_vector_db)
+    assert not (await qdrant_adapter.client.collection_exists(vector_db_id))
+
+    # First insert creates the collection
+    await qdrant_adapter.insert_chunks(vector_db_id, sample_chunks)
+    assert await qdrant_adapter.client.collection_exists(vector_db_id)
+
+    # Unregister deletes the collection
+    await qdrant_adapter.unregister_vector_db(vector_db_id)
+    assert not (await qdrant_adapter.client.collection_exists(vector_db_id))
+    assert len((await qdrant_adapter.client.get_collections()).collections) == 0

tests/unit/providers/vector_io/test_sqlite_vec.py

@@ -0,0 +1,115 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import asyncio
+import sqlite3
+
+import numpy as np
+import pytest
+import pytest_asyncio
+import sqlite_vec
+
+from llama_stack.apis.vector_io import Chunk, QueryChunksResponse
+from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import (
+    SQLiteVecIndex,
+    SQLiteVecVectorIOAdapter,
+    generate_chunk_id,
+)
+
+# This test is a unit test for the SQLiteVecVectorIOAdapter class. This should only contain
+# tests which are specific to this class. More general (API-level) tests should be placed in
+# tests/integration/vector_io/
+#
+# How to run this test:
+#
+# pytest tests/unit/providers/vector_io/test_sqlite_vec.py \
+# -v -s --tb=short --disable-warnings --asyncio-mode=auto
+
+SQLITE_VEC_PROVIDER = "sqlite_vec"
+
+
+@pytest.fixture(scope="session")
+def loop():
+    return asyncio.new_event_loop()
+
+
+@pytest.fixture(scope="session", autouse=True)
+def sqlite_connection(loop):
+    conn = sqlite3.connect(":memory:")
+    try:
+        conn.enable_load_extension(True)
+        sqlite_vec.load(conn)
+        yield conn
+    finally:
+        conn.close()
+
+
+@pytest_asyncio.fixture(scope="session", autouse=True)
+async def sqlite_vec_index(sqlite_connection, embedding_dimension):
+    return await SQLiteVecIndex.create(dimension=embedding_dimension, connection=sqlite_connection, bank_id="test_bank")
+
+
+@pytest.mark.asyncio
+async def test_add_chunks(sqlite_vec_index, sample_chunks, sample_embeddings):
+    await sqlite_vec_index.add_chunks(sample_chunks, sample_embeddings, batch_size=2)
+    cur = sqlite_vec_index.connection.cursor()
+    cur.execute(f"SELECT COUNT(*) FROM {sqlite_vec_index.metadata_table}")
+    count = cur.fetchone()[0]
+    assert count == len(sample_chunks)
+
+
+@pytest.mark.asyncio
+async def test_query_chunks(sqlite_vec_index, sample_chunks, sample_embeddings, embedding_dimension):
+    await sqlite_vec_index.add_chunks(sample_chunks, sample_embeddings)
+    query_embedding = np.random.rand(embedding_dimension).astype(np.float32)
+    response = await sqlite_vec_index.query(query_embedding, k=2, score_threshold=0.0)
+    assert isinstance(response, QueryChunksResponse)
+    assert len(response.chunks) == 2
+
+
+@pytest.mark.asyncio
+async def test_chunk_id_conflict(sqlite_vec_index, sample_chunks, embedding_dimension):
+    """Test that chunk IDs do not conflict across batches when inserting chunks."""
+    # Reduce batch size to force multiple batches for same document
+    # since there are 10 chunks per document and batch size is 2
+    batch_size = 2
+    sample_embeddings = np.random.rand(len(sample_chunks), embedding_dimension).astype(np.float32)
+
+    await sqlite_vec_index.add_chunks(sample_chunks, sample_embeddings, batch_size=batch_size)
+
+    cur = sqlite_vec_index.connection.cursor()
+
+    # Retrieve all chunk IDs to check for duplicates
+    cur.execute(f"SELECT id FROM {sqlite_vec_index.metadata_table}")
+    chunk_ids = [row[0] for row in cur.fetchall()]
+    cur.close()
+
+    # Ensure all chunk IDs are unique
+    assert len(chunk_ids) == len(set(chunk_ids)), "Duplicate chunk IDs detected across batches!"
+
+
+@pytest.fixture(scope="session")
+async def sqlite_vec_adapter(sqlite_connection):
+    config = type("Config", (object,), {"db_path": ":memory:"})  # Mock config with in-memory database
+    adapter = SQLiteVecVectorIOAdapter(config=config, inference_api=None)
+    await adapter.initialize()
+    yield adapter
+    await adapter.shutdown()
+
+
+def test_generate_chunk_id():
+    chunks = [
+        Chunk(content="test", metadata={"document_id": "doc-1"}),
+        Chunk(content="test ", metadata={"document_id": "doc-1"}),
+        Chunk(content="test 3", metadata={"document_id": "doc-1"}),
+    ]
+
+    chunk_ids = sorted([generate_chunk_id(chunk.metadata["document_id"], chunk.content) for chunk in chunks])
+    assert chunk_ids == [
+        "177a1368-f6a8-0c50-6e92-18677f2c3de3",
+        "bc744db3-1b25-0a9c-cdff-b6ba3df73c36",
+        "f68df25d-d9aa-ab4d-5684-64a233add20d",
+    ]

tests/unit/rag/test_vector_store.py

@@ -0,0 +1,78 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import base64
+import mimetypes
+import os
+from pathlib import Path
+
+import pytest
+
+from llama_stack.apis.tools import RAGDocument
+from llama_stack.providers.utils.memory.vector_store import URL, content_from_doc
+
+DUMMY_PDF_PATH = Path(os.path.abspath(__file__)).parent / "fixtures" / "dummy.pdf"
+# Depending on the machine, this can get parsed a couple of ways
+DUMMY_PDF_TEXT_CHOICES = ["Dummy PDF file", "Dumm y PDF file"]
+
+
+def read_file(file_path: str) -> bytes:
+    with open(file_path, "rb") as file:
+        return file.read()
+
+
+def data_url_from_file(file_path: str) -> str:
+    with open(file_path, "rb") as file:
+        file_content = file.read()
+
+    base64_content = base64.b64encode(file_content).decode("utf-8")
+    mime_type, _ = mimetypes.guess_type(file_path)
+
+    data_url = f"data:{mime_type};base64,{base64_content}"
+
+    return data_url
+
+
+class TestVectorStore:
+    @pytest.mark.asyncio
+    async def test_returns_content_from_pdf_data_uri(self):
+        data_uri = data_url_from_file(DUMMY_PDF_PATH)
+        doc = RAGDocument(
+            document_id="dummy",
+            content=data_uri,
+            mime_type="application/pdf",
+            metadata={},
+        )
+        content = await content_from_doc(doc)
+        assert content in DUMMY_PDF_TEXT_CHOICES
+
+    @pytest.mark.asyncio
+    async def test_downloads_pdf_and_returns_content(self):
+        # Using GitHub to host the PDF file
+        url = "https://raw.githubusercontent.com/meta-llama/llama-stack/da035d69cfca915318eaf485770a467ca3c2a238/llama_stack/providers/tests/memory/fixtures/dummy.pdf"
+        doc = RAGDocument(
+            document_id="dummy",
+            content=url,
+            mime_type="application/pdf",
+            metadata={},
+        )
+        content = await content_from_doc(doc)
+        assert content in DUMMY_PDF_TEXT_CHOICES
+
+    @pytest.mark.asyncio
+    async def test_downloads_pdf_and_returns_content_with_url_object(self):
+        # Using GitHub to host the PDF file
+        url = "https://raw.githubusercontent.com/meta-llama/llama-stack/da035d69cfca915318eaf485770a467ca3c2a238/llama_stack/providers/tests/memory/fixtures/dummy.pdf"
+        doc = RAGDocument(
+            document_id="dummy",
+            content=URL(
+                uri=url,
+            ),
+            mime_type="application/pdf",
+            metadata={},
+        )
+        content = await content_from_doc(doc)
+        assert content in DUMMY_PDF_TEXT_CHOICES

tests/unit/registry/test_registry.py

@@ -0,0 +1,199 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+
+import pytest
+import pytest_asyncio
+
+from llama_stack.apis.inference import Model
+from llama_stack.apis.vector_dbs import VectorDB
+from llama_stack.distribution.store.registry import (
+    CachedDiskDistributionRegistry,
+    DiskDistributionRegistry,
+)
+from llama_stack.providers.utils.kvstore import kvstore_impl
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+
+
+@pytest.fixture
+def config():
+    config = SqliteKVStoreConfig(db_path="/tmp/test_registry.db")
+    if os.path.exists(config.db_path):
+        os.remove(config.db_path)
+    return config
+
+
+@pytest_asyncio.fixture(scope="function")
+async def registry(config):
+    registry = DiskDistributionRegistry(await kvstore_impl(config))
+    await registry.initialize()
+    return registry
+
+
+@pytest_asyncio.fixture(scope="function")
+async def cached_registry(config):
+    registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
+    await registry.initialize()
+    return registry
+
+
+@pytest.fixture
+def sample_vector_db():
+    return VectorDB(
+        identifier="test_vector_db",
+        embedding_model="all-MiniLM-L6-v2",
+        embedding_dimension=384,
+        provider_resource_id="test_vector_db",
+        provider_id="test-provider",
+    )
+
+
+@pytest.fixture
+def sample_model():
+    return Model(
+        identifier="test_model",
+        provider_resource_id="test_model",
+        provider_id="test-provider",
+    )
+
+
+@pytest.mark.asyncio
+async def test_registry_initialization(registry):
+    # Test empty registry
+    result = await registry.get("nonexistent", "nonexistent")
+    assert result is None
+
+
+@pytest.mark.asyncio
+async def test_basic_registration(registry, sample_vector_db, sample_model):
+    print(f"Registering {sample_vector_db}")
+    await registry.register(sample_vector_db)
+    print(f"Registering {sample_model}")
+    await registry.register(sample_model)
+    print("Getting vector_db")
+    result_vector_db = await registry.get("vector_db", "test_vector_db")
+    assert result_vector_db is not None
+    assert result_vector_db.identifier == sample_vector_db.identifier
+    assert result_vector_db.embedding_model == sample_vector_db.embedding_model
+    assert result_vector_db.provider_id == sample_vector_db.provider_id
+
+    result_model = await registry.get("model", "test_model")
+    assert result_model is not None
+    assert result_model.identifier == sample_model.identifier
+    assert result_model.provider_id == sample_model.provider_id
+
+
+@pytest.mark.asyncio
+async def test_cached_registry_initialization(config, sample_vector_db, sample_model):
+    # First populate the disk registry
+    disk_registry = DiskDistributionRegistry(await kvstore_impl(config))
+    await disk_registry.initialize()
+    await disk_registry.register(sample_vector_db)
+    await disk_registry.register(sample_model)
+
+    # Test cached version loads from disk
+    cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
+    await cached_registry.initialize()
+
+    result_vector_db = await cached_registry.get("vector_db", "test_vector_db")
+    assert result_vector_db is not None
+    assert result_vector_db.identifier == sample_vector_db.identifier
+    assert result_vector_db.embedding_model == sample_vector_db.embedding_model
+    assert result_vector_db.embedding_dimension == sample_vector_db.embedding_dimension
+    assert result_vector_db.provider_id == sample_vector_db.provider_id
+
+
+@pytest.mark.asyncio
+async def test_cached_registry_updates(config):
+    cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
+    await cached_registry.initialize()
+
+    new_vector_db = VectorDB(
+        identifier="test_vector_db_2",
+        embedding_model="all-MiniLM-L6-v2",
+        embedding_dimension=384,
+        provider_resource_id="test_vector_db_2",
+        provider_id="baz",
+    )
+    await cached_registry.register(new_vector_db)
+
+    # Verify in cache
+    result_vector_db = await cached_registry.get("vector_db", "test_vector_db_2")
+    assert result_vector_db is not None
+    assert result_vector_db.identifier == new_vector_db.identifier
+    assert result_vector_db.provider_id == new_vector_db.provider_id
+
+    # Verify persisted to disk
+    new_registry = DiskDistributionRegistry(await kvstore_impl(config))
+    await new_registry.initialize()
+    result_vector_db = await new_registry.get("vector_db", "test_vector_db_2")
+    assert result_vector_db is not None
+    assert result_vector_db.identifier == new_vector_db.identifier
+    assert result_vector_db.provider_id == new_vector_db.provider_id
+
+
+@pytest.mark.asyncio
+async def test_duplicate_provider_registration(config):
+    cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
+    await cached_registry.initialize()
+
+    original_vector_db = VectorDB(
+        identifier="test_vector_db_2",
+        embedding_model="all-MiniLM-L6-v2",
+        embedding_dimension=384,
+        provider_resource_id="test_vector_db_2",
+        provider_id="baz",
+    )
+    await cached_registry.register(original_vector_db)
+
+    duplicate_vector_db = VectorDB(
+        identifier="test_vector_db_2",
+        embedding_model="different-model",
+        embedding_dimension=384,
+        provider_resource_id="test_vector_db_2",
+        provider_id="baz",  # Same provider_id
+    )
+    await cached_registry.register(duplicate_vector_db)
+
+    result = await cached_registry.get("vector_db", "test_vector_db_2")
+    assert result is not None
+    assert result.embedding_model == original_vector_db.embedding_model  # Original values preserved
+
+
+@pytest.mark.asyncio
+async def test_get_all_objects(config):
+    cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
+    await cached_registry.initialize()
+
+    # Create multiple test banks
+    test_vector_dbs = [
+        VectorDB(
+            identifier=f"test_vector_db_{i}",
+            embedding_model="all-MiniLM-L6-v2",
+            embedding_dimension=384,
+            provider_resource_id=f"test_vector_db_{i}",
+            provider_id=f"provider_{i}",
+        )
+        for i in range(3)
+    ]
+
+    # Register all vector_dbs
+    for vector_db in test_vector_dbs:
+        await cached_registry.register(vector_db)
+
+    # Test get_all retrieval
+    all_results = await cached_registry.get_all()
+    assert len(all_results) == 3
+
+    # Verify each vector_db was stored correctly
+    for original_vector_db in test_vector_dbs:
+        matching_vector_dbs = [v for v in all_results if v.identifier == original_vector_db.identifier]
+        assert len(matching_vector_dbs) == 1
+        stored_vector_db = matching_vector_dbs[0]
+        assert stored_vector_db.embedding_model == original_vector_db.embedding_model
+        assert stored_vector_db.provider_id == original_vector_db.provider_id
+        assert stored_vector_db.embedding_dimension == original_vector_db.embedding_dimension

tests/unit/registry/test_registry_acl.py

@@ -0,0 +1,151 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+import shutil
+import tempfile
+
+import pytest
+
+from llama_stack.apis.models import ModelType
+from llama_stack.distribution.datatypes import ModelWithACL
+from llama_stack.distribution.server.auth import AccessAttributes
+from llama_stack.distribution.store.registry import CachedDiskDistributionRegistry
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl
+
+
+@pytest.fixture(scope="function")
+async def kvstore():
+    temp_dir = tempfile.mkdtemp()
+    db_path = os.path.join(temp_dir, "test_registry_acl.db")
+    kvstore_config = SqliteKVStoreConfig(db_path=db_path)
+    kvstore = SqliteKVStoreImpl(kvstore_config)
+    await kvstore.initialize()
+    yield kvstore
+    shutil.rmtree(temp_dir)
+
+
+@pytest.fixture(scope="function")
+async def registry(kvstore):
+    registry = CachedDiskDistributionRegistry(kvstore)
+    await registry.initialize()
+    return registry
+
+
+@pytest.mark.asyncio
+async def test_registry_cache_with_acl(registry):
+    model = ModelWithACL(
+        identifier="model-acl",
+        provider_id="test-provider",
+        provider_resource_id="model-acl-resource",
+        model_type=ModelType.llm,
+        access_attributes=AccessAttributes(roles=["admin"], teams=["ai-team"]),
+    )
+
+    success = await registry.register(model)
+    assert success
+
+    cached_model = registry.get_cached("model", "model-acl")
+    assert cached_model is not None
+    assert cached_model.identifier == "model-acl"
+    assert cached_model.access_attributes.roles == ["admin"]
+    assert cached_model.access_attributes.teams == ["ai-team"]
+
+    fetched_model = await registry.get("model", "model-acl")
+    assert fetched_model is not None
+    assert fetched_model.identifier == "model-acl"
+    assert fetched_model.access_attributes.roles == ["admin"]
+
+    model.access_attributes = AccessAttributes(roles=["admin", "user"], projects=["project-x"])
+    await registry.update(model)
+
+    updated_cached = registry.get_cached("model", "model-acl")
+    assert updated_cached is not None
+    assert updated_cached.access_attributes.roles == ["admin", "user"]
+    assert updated_cached.access_attributes.projects == ["project-x"]
+    assert updated_cached.access_attributes.teams is None
+
+    new_registry = CachedDiskDistributionRegistry(registry.kvstore)
+    await new_registry.initialize()
+
+    new_model = await new_registry.get("model", "model-acl")
+    assert new_model is not None
+    assert new_model.identifier == "model-acl"
+    assert new_model.access_attributes.roles == ["admin", "user"]
+    assert new_model.access_attributes.projects == ["project-x"]
+    assert new_model.access_attributes.teams is None
+
+
+@pytest.mark.asyncio
+async def test_registry_empty_acl(registry):
+    model = ModelWithACL(
+        identifier="model-empty-acl",
+        provider_id="test-provider",
+        provider_resource_id="model-resource",
+        model_type=ModelType.llm,
+        access_attributes=AccessAttributes(),
+    )
+
+    await registry.register(model)
+
+    cached_model = registry.get_cached("model", "model-empty-acl")
+    assert cached_model is not None
+    assert cached_model.access_attributes is not None
+    assert cached_model.access_attributes.roles is None
+    assert cached_model.access_attributes.teams is None
+    assert cached_model.access_attributes.projects is None
+    assert cached_model.access_attributes.namespaces is None
+
+    all_models = await registry.get_all()
+    assert len(all_models) == 1
+
+    model = ModelWithACL(
+        identifier="model-no-acl",
+        provider_id="test-provider",
+        provider_resource_id="model-resource-2",
+        model_type=ModelType.llm,
+    )
+
+    await registry.register(model)
+
+    cached_model = registry.get_cached("model", "model-no-acl")
+    assert cached_model is not None
+    assert cached_model.access_attributes is None
+
+    all_models = await registry.get_all()
+    assert len(all_models) == 2
+
+
+@pytest.mark.asyncio
+async def test_registry_serialization(registry):
+    attributes = AccessAttributes(
+        roles=["admin", "researcher"],
+        teams=["ai-team", "ml-team"],
+        projects=["project-a", "project-b"],
+        namespaces=["prod", "staging"],
+    )
+
+    model = ModelWithACL(
+        identifier="model-serialize",
+        provider_id="test-provider",
+        provider_resource_id="model-resource",
+        model_type=ModelType.llm,
+        access_attributes=attributes,
+    )
+
+    await registry.register(model)
+
+    new_registry = CachedDiskDistributionRegistry(registry.kvstore)
+    await new_registry.initialize()
+
+    loaded_model = await new_registry.get("model", "model-serialize")
+    assert loaded_model is not None
+
+    assert loaded_model.access_attributes.roles == ["admin", "researcher"]
+    assert loaded_model.access_attributes.teams == ["ai-team", "ml-team"]
+    assert loaded_model.access_attributes.projects == ["project-a", "project-b"]
+    assert loaded_model.access_attributes.namespaces == ["prod", "staging"]

tests/unit/server/test_access_control.py

@@ -0,0 +1,240 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+import shutil
+import tempfile
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+
+from llama_stack.apis.datatypes import Api
+from llama_stack.apis.models import ModelType
+from llama_stack.distribution.datatypes import AccessAttributes, ModelWithACL
+from llama_stack.distribution.routers.routing_tables import ModelsRoutingTable
+from llama_stack.distribution.store.registry import CachedDiskDistributionRegistry
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl
+
+
+class AsyncMock(MagicMock):
+    async def __call__(self, *args, **kwargs):
+        return super(AsyncMock, self).__call__(*args, **kwargs)
+
+
+def _return_model(model):
+    return model
+
+
+@pytest.fixture
+async def test_setup():
+    temp_dir = tempfile.mkdtemp()
+    db_path = os.path.join(temp_dir, "test_access_control.db")
+    kvstore_config = SqliteKVStoreConfig(db_path=db_path)
+    kvstore = SqliteKVStoreImpl(kvstore_config)
+    await kvstore.initialize()
+    registry = CachedDiskDistributionRegistry(kvstore)
+    await registry.initialize()
+
+    mock_inference = Mock()
+    mock_inference.__provider_spec__ = MagicMock()
+    mock_inference.__provider_spec__.api = Api.inference
+    mock_inference.register_model = AsyncMock(side_effect=_return_model)
+    routing_table = ModelsRoutingTable(
+        impls_by_provider_id={"test_provider": mock_inference},
+        dist_registry=registry,
+    )
+    yield registry, routing_table
+    shutil.rmtree(temp_dir)
+
+
+@pytest.mark.asyncio
+@patch("llama_stack.distribution.routers.routing_tables.get_auth_attributes")
+async def test_access_control_with_cache(mock_get_auth_attributes, test_setup):
+    registry, routing_table = test_setup
+    model_public = ModelWithACL(
+        identifier="model-public",
+        provider_id="test_provider",
+        provider_resource_id="model-public",
+        model_type=ModelType.llm,
+    )
+    model_admin_only = ModelWithACL(
+        identifier="model-admin",
+        provider_id="test_provider",
+        provider_resource_id="model-admin",
+        model_type=ModelType.llm,
+        access_attributes=AccessAttributes(roles=["admin"]),
+    )
+    model_data_scientist = ModelWithACL(
+        identifier="model-data-scientist",
+        provider_id="test_provider",
+        provider_resource_id="model-data-scientist",
+        model_type=ModelType.llm,
+        access_attributes=AccessAttributes(roles=["data-scientist", "researcher"], teams=["ml-team"]),
+    )
+    await registry.register(model_public)
+    await registry.register(model_admin_only)
+    await registry.register(model_data_scientist)
+
+    mock_get_auth_attributes.return_value = {"roles": ["admin"], "teams": ["management"]}
+    all_models = await routing_table.list_models()
+    assert len(all_models.data) == 2
+
+    model = await routing_table.get_model("model-public")
+    assert model.identifier == "model-public"
+    model = await routing_table.get_model("model-admin")
+    assert model.identifier == "model-admin"
+    with pytest.raises(ValueError):
+        await routing_table.get_model("model-data-scientist")
+
+    mock_get_auth_attributes.return_value = {"roles": ["data-scientist"], "teams": ["other-team"]}
+    all_models = await routing_table.list_models()
+    assert len(all_models.data) == 1
+    assert all_models.data[0].identifier == "model-public"
+    model = await routing_table.get_model("model-public")
+    assert model.identifier == "model-public"
+    with pytest.raises(ValueError):
+        await routing_table.get_model("model-admin")
+    with pytest.raises(ValueError):
+        await routing_table.get_model("model-data-scientist")
+
+    mock_get_auth_attributes.return_value = {"roles": ["data-scientist"], "teams": ["ml-team"]}
+    all_models = await routing_table.list_models()
+    assert len(all_models.data) == 2
+    model_ids = [m.identifier for m in all_models.data]
+    assert "model-public" in model_ids
+    assert "model-data-scientist" in model_ids
+    assert "model-admin" not in model_ids
+    model = await routing_table.get_model("model-public")
+    assert model.identifier == "model-public"
+    model = await routing_table.get_model("model-data-scientist")
+    assert model.identifier == "model-data-scientist"
+    with pytest.raises(ValueError):
+        await routing_table.get_model("model-admin")
+
+
+@pytest.mark.asyncio
+@patch("llama_stack.distribution.routers.routing_tables.get_auth_attributes")
+async def test_access_control_and_updates(mock_get_auth_attributes, test_setup):
+    registry, routing_table = test_setup
+    model_public = ModelWithACL(
+        identifier="model-updates",
+        provider_id="test_provider",
+        provider_resource_id="model-updates",
+        model_type=ModelType.llm,
+    )
+    await registry.register(model_public)
+    mock_get_auth_attributes.return_value = {
+        "roles": ["user"],
+    }
+    model = await routing_table.get_model("model-updates")
+    assert model.identifier == "model-updates"
+    model_public.access_attributes = AccessAttributes(roles=["admin"])
+    await registry.update(model_public)
+    mock_get_auth_attributes.return_value = {
+        "roles": ["user"],
+    }
+    with pytest.raises(ValueError):
+        await routing_table.get_model("model-updates")
+    mock_get_auth_attributes.return_value = {
+        "roles": ["admin"],
+    }
+    model = await routing_table.get_model("model-updates")
+    assert model.identifier == "model-updates"
+
+
+@pytest.mark.asyncio
+@patch("llama_stack.distribution.routers.routing_tables.get_auth_attributes")
+async def test_access_control_empty_attributes(mock_get_auth_attributes, test_setup):
+    registry, routing_table = test_setup
+    model = ModelWithACL(
+        identifier="model-empty-attrs",
+        provider_id="test_provider",
+        provider_resource_id="model-empty-attrs",
+        model_type=ModelType.llm,
+        access_attributes=AccessAttributes(),
+    )
+    await registry.register(model)
+    mock_get_auth_attributes.return_value = {
+        "roles": [],
+    }
+    result = await routing_table.get_model("model-empty-attrs")
+    assert result.identifier == "model-empty-attrs"
+    all_models = await routing_table.list_models()
+    model_ids = [m.identifier for m in all_models.data]
+    assert "model-empty-attrs" in model_ids
+
+
+@pytest.mark.asyncio
+@patch("llama_stack.distribution.routers.routing_tables.get_auth_attributes")
+async def test_no_user_attributes(mock_get_auth_attributes, test_setup):
+    registry, routing_table = test_setup
+    model_public = ModelWithACL(
+        identifier="model-public-2",
+        provider_id="test_provider",
+        provider_resource_id="model-public-2",
+        model_type=ModelType.llm,
+    )
+    model_restricted = ModelWithACL(
+        identifier="model-restricted",
+        provider_id="test_provider",
+        provider_resource_id="model-restricted",
+        model_type=ModelType.llm,
+        access_attributes=AccessAttributes(roles=["admin"]),
+    )
+    await registry.register(model_public)
+    await registry.register(model_restricted)
+    mock_get_auth_attributes.return_value = None
+    model = await routing_table.get_model("model-public-2")
+    assert model.identifier == "model-public-2"
+
+    with pytest.raises(ValueError):
+        await routing_table.get_model("model-restricted")
+
+    all_models = await routing_table.list_models()
+    assert len(all_models.data) == 1
+    assert all_models.data[0].identifier == "model-public-2"
+
+
+@pytest.mark.asyncio
+@patch("llama_stack.distribution.routers.routing_tables.get_auth_attributes")
+async def test_automatic_access_attributes(mock_get_auth_attributes, test_setup):
+    """Test that newly created resources inherit access attributes from their creator."""
+    registry, routing_table = test_setup
+
+    # Set creator's attributes
+    creator_attributes = {"roles": ["data-scientist"], "teams": ["ml-team"], "projects": ["llama-3"]}
+    mock_get_auth_attributes.return_value = creator_attributes
+
+    # Create model without explicit access attributes
+    model = ModelWithACL(
+        identifier="auto-access-model",
+        provider_id="test_provider",
+        provider_resource_id="auto-access-model",
+        model_type=ModelType.llm,
+    )
+    await routing_table.register_object(model)
+
+    # Verify the model got creator's attributes
+    registered_model = await routing_table.get_model("auto-access-model")
+    assert registered_model.access_attributes is not None
+    assert registered_model.access_attributes.roles == ["data-scientist"]
+    assert registered_model.access_attributes.teams == ["ml-team"]
+    assert registered_model.access_attributes.projects == ["llama-3"]
+
+    # Verify another user without matching attributes can't access it
+    mock_get_auth_attributes.return_value = {"roles": ["engineer"], "teams": ["infra-team"]}
+    with pytest.raises(ValueError):
+        await routing_table.get_model("auto-access-model")
+
+    # But a user with matching attributes can
+    mock_get_auth_attributes.return_value = {
+        "roles": ["data-scientist", "engineer"],
+        "teams": ["ml-team", "platform-team"],
+        "projects": ["llama-3"],
+    }
+    model = await routing_table.get_model("auto-access-model")
+    assert model.identifier == "auto-access-model"

tests/unit/server/test_auth.py

@@ -0,0 +1,206 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from llama_stack.distribution.server.auth import AuthenticationMiddleware
+
+
+class MockResponse:
+    def __init__(self, status_code, json_data):
+        self.status_code = status_code
+        self._json_data = json_data
+
+    def json(self):
+        return self._json_data
+
+
+@pytest.fixture
+def mock_auth_endpoint():
+    return "http://mock-auth-service/validate"
+
+
+@pytest.fixture
+def valid_api_key():
+    return "valid_api_key_12345"
+
+
+@pytest.fixture
+def invalid_api_key():
+    return "invalid_api_key_67890"
+
+
+@pytest.fixture
+def app(mock_auth_endpoint):
+    app = FastAPI()
+    app.add_middleware(AuthenticationMiddleware, auth_endpoint=mock_auth_endpoint)
+
+    @app.get("/test")
+    def test_endpoint():
+        return {"message": "Authentication successful"}
+
+    return app
+
+
+@pytest.fixture
+def client(app):
+    return TestClient(app)
+
+
+@pytest.fixture
+def mock_scope():
+    return {
+        "type": "http",
+        "path": "/models/list",
+        "headers": [
+            (b"content-type", b"application/json"),
+            (b"authorization", b"Bearer test-api-key"),
+            (b"user-agent", b"test-user-agent"),
+        ],
+        "query_string": b"limit=100&offset=0",
+    }
+
+
+@pytest.fixture
+def mock_middleware(mock_auth_endpoint):
+    mock_app = AsyncMock()
+    return AuthenticationMiddleware(mock_app, mock_auth_endpoint), mock_app
+
+
+async def mock_post_success(*args, **kwargs):
+    return MockResponse(200, {"message": "Authentication successful"})
+
+
+async def mock_post_failure(*args, **kwargs):
+    return MockResponse(401, {"message": "Authentication failed"})
+
+
+async def mock_post_exception(*args, **kwargs):
+    raise Exception("Connection error")
+
+
+def test_missing_auth_header(client):
+    response = client.get("/test")
+    assert response.status_code == 401
+    assert "Missing or invalid Authorization header" in response.json()["error"]["message"]
+
+
+def test_invalid_auth_header_format(client):
+    response = client.get("/test", headers={"Authorization": "InvalidFormat token123"})
+    assert response.status_code == 401
+    assert "Missing or invalid Authorization header" in response.json()["error"]["message"]
+
+
+@patch("httpx.AsyncClient.post", new=mock_post_success)
+def test_valid_authentication(client, valid_api_key):
+    response = client.get("/test", headers={"Authorization": f"Bearer {valid_api_key}"})
+    assert response.status_code == 200
+    assert response.json() == {"message": "Authentication successful"}
+
+
+@patch("httpx.AsyncClient.post", new=mock_post_failure)
+def test_invalid_authentication(client, invalid_api_key):
+    response = client.get("/test", headers={"Authorization": f"Bearer {invalid_api_key}"})
+    assert response.status_code == 401
+    assert "Authentication failed" in response.json()["error"]["message"]
+
+
+@patch("httpx.AsyncClient.post", new=mock_post_exception)
+def test_auth_service_error(client, valid_api_key):
+    response = client.get("/test", headers={"Authorization": f"Bearer {valid_api_key}"})
+    assert response.status_code == 401
+    assert "Authentication service error" in response.json()["error"]["message"]
+
+
+def test_auth_request_payload(client, valid_api_key, mock_auth_endpoint):
+    with patch("httpx.AsyncClient.post") as mock_post:
+        mock_response = MockResponse(200, {"message": "Authentication successful"})
+        mock_post.return_value = mock_response
+
+        client.get(
+            "/test?param1=value1&param2=value2",
+            headers={
+                "Authorization": f"Bearer {valid_api_key}",
+                "User-Agent": "TestClient",
+                "Content-Type": "application/json",
+            },
+        )
+
+        # Check that the auth endpoint was called with the correct payload
+        call_args = mock_post.call_args
+        assert call_args is not None
+
+        url, kwargs = call_args[0][0], call_args[1]
+        assert url == mock_auth_endpoint
+
+        payload = kwargs["json"]
+        assert payload["api_key"] == valid_api_key
+        assert payload["request"]["path"] == "/test"
+        assert "authorization" not in payload["request"]["headers"]
+        assert "param1" in payload["request"]["params"]
+        assert "param2" in payload["request"]["params"]
+
+
+@pytest.mark.asyncio
+async def test_auth_middleware_with_access_attributes(mock_middleware, mock_scope):
+    middleware, mock_app = mock_middleware
+    mock_receive = AsyncMock()
+    mock_send = AsyncMock()
+
+    with patch("httpx.AsyncClient") as mock_client:
+        mock_client_instance = AsyncMock()
+        mock_client.return_value.__aenter__.return_value = mock_client_instance
+
+        mock_client_instance.post.return_value = MockResponse(
+            200,
+            {
+                "access_attributes": {
+                    "roles": ["admin", "user"],
+                    "teams": ["ml-team"],
+                    "projects": ["project-x", "project-y"],
+                }
+            },
+        )
+
+        await middleware(mock_scope, mock_receive, mock_send)
+
+        assert "user_attributes" in mock_scope
+        assert mock_scope["user_attributes"]["roles"] == ["admin", "user"]
+        assert mock_scope["user_attributes"]["teams"] == ["ml-team"]
+        assert mock_scope["user_attributes"]["projects"] == ["project-x", "project-y"]
+
+        mock_app.assert_called_once_with(mock_scope, mock_receive, mock_send)
+
+
+@pytest.mark.asyncio
+async def test_auth_middleware_no_attributes(mock_middleware, mock_scope):
+    """Test middleware behavior with no access attributes"""
+    middleware, mock_app = mock_middleware
+    mock_receive = AsyncMock()
+    mock_send = AsyncMock()
+
+    with patch("httpx.AsyncClient") as mock_client:
+        mock_client_instance = AsyncMock()
+        mock_client.return_value.__aenter__.return_value = mock_client_instance
+
+        mock_client_instance.post.return_value = MockResponse(
+            200,
+            {
+                "message": "Authentication successful"
+                # No access_attributes
+            },
+        )
+
+        await middleware(mock_scope, mock_receive, mock_send)
+
+        assert "user_attributes" in mock_scope
+        attributes = mock_scope["user_attributes"]
+        assert "namespaces" in attributes
+        assert attributes["namespaces"] == ["test-api-key"]

tests/unit/server/test_replace_env_vars.py

@@ -0,0 +1,66 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+import unittest
+
+from llama_stack.distribution.stack import replace_env_vars
+
+
+class TestReplaceEnvVars(unittest.TestCase):
+    def setUp(self):
+        # Clear any existing environment variables we'll use in tests
+        for var in ["TEST_VAR", "EMPTY_VAR", "ZERO_VAR"]:
+            if var in os.environ:
+                del os.environ[var]
+
+        # Set up test environment variables
+        os.environ["TEST_VAR"] = "test_value"
+        os.environ["EMPTY_VAR"] = ""
+        os.environ["ZERO_VAR"] = "0"
+
+    def test_simple_replacement(self):
+        self.assertEqual(replace_env_vars("${env.TEST_VAR}"), "test_value")
+
+    def test_default_value_when_not_set(self):
+        self.assertEqual(replace_env_vars("${env.NOT_SET:default}"), "default")
+
+    def test_default_value_when_set(self):
+        self.assertEqual(replace_env_vars("${env.TEST_VAR:default}"), "test_value")
+
+    def test_default_value_when_empty(self):
+        self.assertEqual(replace_env_vars("${env.EMPTY_VAR:default}"), "default")
+
+    def test_conditional_value_when_set(self):
+        self.assertEqual(replace_env_vars("${env.TEST_VAR+conditional}"), "conditional")
+
+    def test_conditional_value_when_not_set(self):
+        self.assertEqual(replace_env_vars("${env.NOT_SET+conditional}"), "")
+
+    def test_conditional_value_when_empty(self):
+        self.assertEqual(replace_env_vars("${env.EMPTY_VAR+conditional}"), "")
+
+    def test_conditional_value_with_zero(self):
+        self.assertEqual(replace_env_vars("${env.ZERO_VAR+conditional}"), "conditional")
+
+    def test_mixed_syntax(self):
+        self.assertEqual(replace_env_vars("${env.TEST_VAR:default} and ${env.NOT_SET+conditional}"), "test_value and ")
+        self.assertEqual(
+            replace_env_vars("${env.NOT_SET:default} and ${env.TEST_VAR+conditional}"), "default and conditional"
+        )
+
+    def test_nested_structures(self):
+        data = {
+            "key1": "${env.TEST_VAR:default}",
+            "key2": ["${env.NOT_SET:default}", "${env.TEST_VAR+conditional}"],
+            "key3": {"nested": "${env.NOT_SET+conditional}"},
+        }
+        expected = {"key1": "test_value", "key2": ["default", "conditional"], "key3": {"nested": ""}}
+        self.assertEqual(replace_env_vars(data), expected)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/server/test_resolver.py

@@ -0,0 +1,117 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import inspect
+import sys
+from typing import Any, Dict, Protocol
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from pydantic import BaseModel, Field
+
+from llama_stack.apis.inference import Inference
+from llama_stack.distribution.datatypes import (
+    Api,
+    Provider,
+    StackRunConfig,
+)
+from llama_stack.distribution.resolver import resolve_impls
+from llama_stack.distribution.routers.routers import InferenceRouter
+from llama_stack.distribution.routers.routing_tables import ModelsRoutingTable
+from llama_stack.providers.datatypes import InlineProviderSpec, ProviderSpec
+
+
+def add_protocol_methods(cls: type, protocol: type[Protocol]) -> None:
+    """Dynamically add protocol methods to a class by inspecting the protocol."""
+    for name, value in inspect.getmembers(protocol):
+        if inspect.isfunction(value) and hasattr(value, "__webmethod__"):
+            # Get the signature
+            sig = inspect.signature(value)
+
+            # Create an async function with the same signature that returns a MagicMock
+            async def mock_impl(*args, **kwargs):
+                return MagicMock()
+
+            # Set the signature on our mock implementation
+            mock_impl.__signature__ = sig
+            # Add it to the class
+            setattr(cls, name, mock_impl)
+
+
+class SampleConfig(BaseModel):
+    foo: str = Field(
+        default="bar",
+        description="foo",
+    )
+
+    @classmethod
+    def sample_run_config(cls, **kwargs: Any) -> Dict[str, Any]:
+        return {
+            "foo": "baz",
+        }
+
+
+class SampleImpl:
+    def __init__(self, config: SampleConfig, deps: Dict[Api, Any], provider_spec: ProviderSpec = None):
+        self.__provider_id__ = "test_provider"
+        self.__provider_spec__ = provider_spec
+        self.__provider_config__ = config
+        self.__deps__ = deps
+        self.foo = config.foo
+
+    async def initialize(self):
+        pass
+
+
+@pytest.mark.asyncio
+async def test_resolve_impls_basic():
+    # Create a real provider spec
+    provider_spec = InlineProviderSpec(
+        api=Api.inference,
+        provider_type="sample",
+        module="test_module",
+        config_class="test_resolver.SampleConfig",
+        api_dependencies=[],
+    )
+
+    # Create provider registry with our provider
+    provider_registry = {Api.inference: {provider_spec.provider_type: provider_spec}}
+
+    run_config = StackRunConfig(
+        image_name="test_image",
+        providers={
+            "inference": [
+                Provider(
+                    provider_id="sample_provider",
+                    provider_type="sample",
+                    config=SampleConfig.sample_run_config(),
+                )
+            ]
+        },
+    )
+
+    dist_registry = MagicMock()
+
+    mock_module = MagicMock()
+    impl = SampleImpl(SampleConfig(foo="baz"), {}, provider_spec)
+    add_protocol_methods(SampleImpl, Inference)
+
+    mock_module.get_provider_impl = AsyncMock(return_value=impl)
+    sys.modules["test_module"] = mock_module
+
+    impls = await resolve_impls(run_config, provider_registry, dist_registry)
+
+    assert Api.inference in impls
+    assert isinstance(impls[Api.inference], InferenceRouter)
+
+    table = impls[Api.inference].routing_table
+    assert isinstance(table, ModelsRoutingTable)
+
+    impl = table.impls_by_provider_id["sample_provider"]
+    assert isinstance(impl, SampleImpl)
+    assert impl.foo == "baz"
+    assert impl.__provider_id__ == "sample_provider"
+    assert impl.__provider_spec__ == provider_spec