Merge remote-tracking branch 'upstream/main' into add-mcp-authentication-param

tests/unit/distribution/routers/test_routing_tables.py

@@ -9,15 +9,22 @@
 from unittest.mock import AsyncMock
 
 import pytest
+from llama_stack_api import (
+    URL,
+    Api,
+    Dataset,
+    DatasetPurpose,
+    ListToolDefsResponse,
+    Model,
+    ModelNotFoundError,
+    ModelType,
+    NumberType,
+    Shield,
+    ToolDef,
+    ToolGroup,
+    URIDataSource,
+)
 
-from llama_stack.apis.common.content_types import URL
-from llama_stack.apis.common.errors import ModelNotFoundError
-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
-from llama_stack.apis.models import Model, ModelType
-from llama_stack.apis.shields.shields import Shield
-from llama_stack.apis.tools import ListToolDefsResponse, ToolDef, ToolGroup
 from llama_stack.core.datatypes import RegistryEntrySource
 from llama_stack.core.routing_tables.benchmarks import BenchmarksRoutingTable
 from llama_stack.core.routing_tables.datasets import DatasetsRoutingTable

tests/unit/distribution/test_api_recordings.py

@@ -9,10 +9,9 @@ from pathlib import Path
 from unittest.mock import patch
 
 import pytest
-from openai import AsyncOpenAI
 
 # Import the real Pydantic response types instead of using Mocks
-from llama_stack.apis.inference import (
+from llama_stack_api import (
     OpenAIAssistantMessageParam,
     OpenAIChatCompletion,
     OpenAIChoice,
@@ -20,6 +19,8 @@ from llama_stack.apis.inference import (
     OpenAIEmbeddingsResponse,
     OpenAIEmbeddingUsage,
 )
+from openai import AsyncOpenAI
+
 from llama_stack.testing.api_recorder import (
     APIRecordingMode,
     ResponseStorage,

tests/unit/distribution/test_distribution.py

@@ -9,6 +9,7 @@ from unittest.mock import patch
 
 import pytest
 import yaml
+from llama_stack_api import ProviderSpec
 from pydantic import BaseModel, Field, ValidationError
 
 from llama_stack.core.datatypes import Api, Provider, StackRunConfig
@@ -22,7 +23,6 @@ from llama_stack.core.storage.datatypes import (
     SqlStoreReference,
     StorageConfig,
 )
-from llama_stack.providers.datatypes import ProviderSpec
 
 
 class SampleConfig(BaseModel):
@@ -312,7 +312,7 @@ pip_packages:
         """Test loading an external provider from a module (success path)."""
         from types import SimpleNamespace
 
-        from llama_stack.providers.datatypes import Api, ProviderSpec
+        from llama_stack_api import Api, ProviderSpec
 
         # Simulate a provider module with get_provider_spec
         fake_spec = ProviderSpec(
@@ -395,8 +395,9 @@ pip_packages:
 
     def test_external_provider_from_module_building(self, mock_providers):
         """Test loading an external provider from a module during build (building=True, partial spec)."""
+        from llama_stack_api import Api
+
         from llama_stack.core.datatypes import BuildConfig, BuildProvider, DistributionSpec
-        from llama_stack.providers.datatypes import Api
 
         # No importlib patch needed, should not import module when type of `config` is BuildConfig or DistributionSpec
         build_config = BuildConfig(
@@ -456,8 +457,9 @@ class TestGetExternalProvidersFromModule:
         """Test provider with module containing version spec (e.g., package==1.0.0)."""
         from types import SimpleNamespace
 
+        from llama_stack_api import ProviderSpec
+
         from llama_stack.core.distribution import get_external_providers_from_module
-        from llama_stack.providers.datatypes import ProviderSpec
 
         fake_spec = ProviderSpec(
             api=Api.inference,
@@ -593,8 +595,9 @@ class TestGetExternalProvidersFromModule:
         """Test when get_provider_spec returns a list of specs."""
         from types import SimpleNamespace
 
+        from llama_stack_api import ProviderSpec
+
         from llama_stack.core.distribution import get_external_providers_from_module
-        from llama_stack.providers.datatypes import ProviderSpec
 
         spec1 = ProviderSpec(
             api=Api.inference,
@@ -641,8 +644,9 @@ class TestGetExternalProvidersFromModule:
         """Test that list return filters specs by provider_type."""
         from types import SimpleNamespace
 
+        from llama_stack_api import ProviderSpec
+
         from llama_stack.core.distribution import get_external_providers_from_module
-        from llama_stack.providers.datatypes import ProviderSpec
 
         spec1 = ProviderSpec(
             api=Api.inference,
@@ -689,8 +693,9 @@ class TestGetExternalProvidersFromModule:
         """Test that list return adds multiple different provider_types when config requests them."""
         from types import SimpleNamespace
 
+        from llama_stack_api import ProviderSpec
+
         from llama_stack.core.distribution import get_external_providers_from_module
-        from llama_stack.providers.datatypes import ProviderSpec
 
         # Module returns both inline and remote variants
         spec1 = ProviderSpec(
@@ -828,8 +833,9 @@ class TestGetExternalProvidersFromModule:
         """Test multiple APIs with providers."""
         from types import SimpleNamespace
 
+        from llama_stack_api import ProviderSpec
+
         from llama_stack.core.distribution import get_external_providers_from_module
-        from llama_stack.providers.datatypes import ProviderSpec
 
         inference_spec = ProviderSpec(
             api=Api.inference,