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

tests/integration/batches/conftest.py

@@ -13,6 +13,7 @@ from contextlib import contextmanager
 from io import BytesIO
 
 import pytest
+
 from llama_stack_api import OpenAIFilePurpose
 
 

tests/integration/files/test_files.py

@@ -9,9 +9,9 @@ from unittest.mock import patch
 
 import pytest
 import requests
-from llama_stack_api import OpenAIFilePurpose
 
 from llama_stack.core.datatypes import User
+from llama_stack_api import OpenAIFilePurpose
 
 purpose = OpenAIFilePurpose.ASSISTANTS
 

tests/integration/inference/test_provider_data_routing.py

@@ -15,6 +15,9 @@ that enables routing based on provider_data alone.
 from unittest.mock import AsyncMock, patch
 
 import pytest
+
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
+from llama_stack.core.telemetry.telemetry import MetricEvent
 from llama_stack_api import (
     Api,
     OpenAIAssistantMessageParam,
@@ -23,9 +26,6 @@ from llama_stack_api import (
     OpenAIChoice,
 )
 
-from llama_stack.core.library_client import LlamaStackAsLibraryClient
-from llama_stack.core.telemetry.telemetry import MetricEvent
-
 
 class OpenAIChatCompletionWithMetrics(OpenAIChatCompletion):
     metrics: list[MetricEvent] | None = None

tests/integration/post_training/test_post_training.py

@@ -9,6 +9,8 @@ import time
 import uuid
 
 import pytest
+
+from llama_stack.log import get_logger
 from llama_stack_api import (
     DataConfig,
     DatasetFormat,
@@ -18,8 +20,6 @@ from llama_stack_api import (
     TrainingConfig,
 )
 
-from llama_stack.log import get_logger
-
 # Configure logging
 logger = get_logger(name=__name__, category="post_training")
 

tests/integration/safety/test_llama_guard.py

@@ -12,9 +12,9 @@ import warnings
 from collections.abc import Generator
 
 import pytest
-from llama_stack_api import ViolationLevel
 
 from llama_stack.models.llama.sku_types import CoreModelId
+from llama_stack_api import ViolationLevel
 
 # Llama Guard models available for text and vision shields
 LLAMA_GUARD_TEXT_MODELS = [CoreModelId.llama_guard_4_12b.value]

tests/integration/safety/test_safety.py

@@ -7,6 +7,7 @@ import base64
 import mimetypes
 
 import pytest
+
 from llama_stack_api import ViolationLevel
 
 CODE_SCANNER_ENABLED_PROVIDERS = {"ollama", "together", "fireworks"}

tests/integration/safety/test_vision_safety.py

@@ -9,6 +9,7 @@ import mimetypes
 import os
 
 import pytest
+
 from llama_stack_api import ViolationLevel
 
 VISION_SHIELD_ENABLED_PROVIDERS = {"together"}

tests/integration/tool_runtime/test_registration.py

@@ -7,9 +7,9 @@
 import re
 
 import pytest
-from llama_stack_api import ToolGroupNotFoundError
 
 from llama_stack.core.library_client import LlamaStackAsLibraryClient
+from llama_stack_api import ToolGroupNotFoundError
 from tests.common.mcp import MCP_TOOLGROUP_ID, make_mcp_server
 
 

tests/integration/vector_io/test_openai_vector_stores.py

@@ -8,12 +8,12 @@ import time
 from io import BytesIO
 
 import pytest
-from llama_stack_api import Chunk, ExpiresAfter
 from llama_stack_client import BadRequestError
 from openai import BadRequestError as OpenAIBadRequestError
 
 from llama_stack.core.library_client import LlamaStackAsLibraryClient
 from llama_stack.log import get_logger
+from llama_stack_api import Chunk, ExpiresAfter
 
 from ..conftest import vector_provider_wrapper
 

tests/integration/vector_io/test_vector_io.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 import pytest
+
 from llama_stack_api import Chunk
 
 from ..conftest import vector_provider_wrapper

tests/unit/conversations/test_conversations.py

@@ -8,7 +8,6 @@ import tempfile
 from pathlib import Path
 
 import pytest
-from llama_stack_api import OpenAIResponseInputMessageContentText, OpenAIResponseMessage
 from openai.types.conversations.conversation import Conversation as OpenAIConversation
 from openai.types.conversations.conversation_item import ConversationItem as OpenAIConversationItem
 from pydantic import TypeAdapter
@@ -25,6 +24,7 @@ from llama_stack.core.storage.datatypes import (
     StorageConfig,
 )
 from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
+from llama_stack_api import OpenAIResponseInputMessageContentText, OpenAIResponseMessage
 
 
 @pytest.fixture

tests/unit/core/routers/test_safety_router.py

@@ -6,10 +6,9 @@
 
 from unittest.mock import AsyncMock
 
-from llama_stack_api import ListShieldsResponse, ModerationObject, ModerationObjectResults, Shield
-
 from llama_stack.core.datatypes import SafetyConfig
 from llama_stack.core.routers.safety import SafetyRouter
+from llama_stack_api import ListShieldsResponse, ModerationObject, ModerationObjectResults, Shield
 
 
 async def test_run_moderation_uses_default_shield_when_model_missing():

tests/unit/core/routers/test_vector_io.py

@@ -7,6 +7,8 @@
 from unittest.mock import AsyncMock, Mock
 
 import pytest
+
+from llama_stack.core.routers.vector_io import VectorIORouter
 from llama_stack_api import (
     ModelNotFoundError,
     ModelType,
@@ -14,8 +16,6 @@ from llama_stack_api import (
     OpenAICreateVectorStoreRequestWithExtraBody,
 )
 
-from llama_stack.core.routers.vector_io import VectorIORouter
-
 
 async def test_single_provider_auto_selection():
     # provider_id automatically selected during vector store create() when only one provider available
@@ -127,7 +127,8 @@ async def test_update_vector_store_same_provider_id_succeeds():
 
 
 async def test_create_vector_store_with_unknown_embedding_model_raises_error():
-    """Test that creating a vector store with an unknown embedding model raises ModelNotFoundError."""
+    """Test that creating a vector store with an unknown embedding model raises
+    FoundError."""
     mock_routing_table = Mock(impls_by_provider_id={"provider": "mock"})
     mock_routing_table.get_object_by_identifier = AsyncMock(return_value=None)
 

tests/unit/core/test_stack_validation.py

@@ -9,10 +9,10 @@
 from unittest.mock import AsyncMock
 
 import pytest
-from llama_stack_api import Api, ListModelsResponse, ListShieldsResponse, Model, ModelType, Shield
 
 from llama_stack.core.datatypes import QualifiedModel, SafetyConfig, StackRunConfig, StorageConfig, VectorStoresConfig
 from llama_stack.core.stack import validate_safety_config, validate_vector_stores_config
+from llama_stack_api import Api, ListModelsResponse, ListShieldsResponse, Model, ModelType, Shield
 
 
 class TestVectorStoresValidation:

tests/unit/distribution/routers/test_routing_tables.py

@@ -9,6 +9,14 @@
 from unittest.mock import AsyncMock
 
 import pytest
+
+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
+from llama_stack.core.routing_tables.models import ModelsRoutingTable
+from llama_stack.core.routing_tables.scoring_functions import ScoringFunctionsRoutingTable
+from llama_stack.core.routing_tables.shields import ShieldsRoutingTable
+from llama_stack.core.routing_tables.toolgroups import ToolGroupsRoutingTable
 from llama_stack_api import (
     URL,
     Api,
@@ -25,14 +33,6 @@ from llama_stack_api import (
     URIDataSource,
 )
 
-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
-from llama_stack.core.routing_tables.models import ModelsRoutingTable
-from llama_stack.core.routing_tables.scoring_functions import ScoringFunctionsRoutingTable
-from llama_stack.core.routing_tables.shields import ShieldsRoutingTable
-from llama_stack.core.routing_tables.toolgroups import ToolGroupsRoutingTable
-
 
 class Impl:
     def __init__(self, api: Api):

tests/unit/distribution/test_api_recordings.py

@@ -9,6 +9,14 @@ from pathlib import Path
 from unittest.mock import patch
 
 import pytest
+from openai import AsyncOpenAI
+
+from llama_stack.testing.api_recorder import (
+    APIRecordingMode,
+    ResponseStorage,
+    api_recording,
+    normalize_inference_request,
+)
 
 # Import the real Pydantic response types instead of using Mocks
 from llama_stack_api import (
@@ -19,14 +27,6 @@ from llama_stack_api import (
     OpenAIEmbeddingsResponse,
     OpenAIEmbeddingUsage,
 )
-from openai import AsyncOpenAI
-
-from llama_stack.testing.api_recorder import (
-    APIRecordingMode,
-    ResponseStorage,
-    api_recording,
-    normalize_inference_request,
-)
 
 
 @pytest.fixture

tests/unit/distribution/test_distribution.py

@@ -9,7 +9,6 @@ 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
@@ -23,6 +22,7 @@ from llama_stack.core.storage.datatypes import (
     SqlStoreReference,
     StorageConfig,
 )
+from llama_stack_api import ProviderSpec
 
 
 class SampleConfig(BaseModel):
@@ -395,9 +395,8 @@ 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_api import Api
 
         # No importlib patch needed, should not import module when type of `config` is BuildConfig or DistributionSpec
         build_config = BuildConfig(
@@ -457,9 +456,8 @@ 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_api import ProviderSpec
 
         fake_spec = ProviderSpec(
             api=Api.inference,
@@ -595,9 +593,8 @@ 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_api import ProviderSpec
 
         spec1 = ProviderSpec(
             api=Api.inference,
@@ -644,9 +641,8 @@ 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_api import ProviderSpec
 
         spec1 = ProviderSpec(
             api=Api.inference,
@@ -693,9 +689,8 @@ 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_api import ProviderSpec
 
         # Module returns both inline and remote variants
         spec1 = ProviderSpec(
@@ -833,9 +828,8 @@ 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_api import ProviderSpec
 
         inference_spec = ProviderSpec(
             api=Api.inference,

tests/unit/files/test_files.py

@@ -6,7 +6,6 @@
 
 
 import pytest
-from llama_stack_api import OpenAIFilePurpose, Order, ResourceNotFoundError
 
 from llama_stack.core.access_control.access_control import default_policy
 from llama_stack.core.storage.datatypes import SqliteSqlStoreConfig, SqlStoreReference
@@ -15,6 +14,7 @@ from llama_stack.providers.inline.files.localfs import (
     LocalfsFilesImplConfig,
 )
 from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
+from llama_stack_api import OpenAIFilePurpose, Order, ResourceNotFoundError
 
 
 class MockUploadFile:

tests/unit/providers/batches/test_reference.py

@@ -58,6 +58,7 @@ import json
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+
 from llama_stack_api import BatchObject, ConflictError, ResourceNotFoundError
 
 

tests/unit/providers/batches/test_reference_idempotency.py

@@ -43,6 +43,7 @@ Key Behaviors Tested:
 import asyncio
 
 import pytest
+
 from llama_stack_api import ConflictError
 
 

tests/unit/providers/files/test_s3_files.py

@@ -8,6 +8,7 @@ from unittest.mock import patch
 
 import pytest
 from botocore.exceptions import ClientError
+
 from llama_stack_api import OpenAIFilePurpose, ResourceNotFoundError
 
 

tests/unit/providers/files/test_s3_files_auth.py

@@ -7,10 +7,10 @@
 from unittest.mock import patch
 
 import pytest
-from llama_stack_api import OpenAIFilePurpose, ResourceNotFoundError
 
 from llama_stack.core.datatypes import User
 from llama_stack.providers.remote.files.s3.files import S3FilesImpl
+from llama_stack_api import OpenAIFilePurpose, ResourceNotFoundError
 
 
 async def test_listing_hides_other_users_file(s3_provider, sample_text_file):

tests/unit/providers/inference/test_bedrock_adapter.py

@@ -8,11 +8,11 @@ from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
-from llama_stack_api import OpenAIChatCompletionRequestWithExtraBody
 from openai import AuthenticationError
 
 from llama_stack.providers.remote.inference.bedrock.bedrock import BedrockInferenceAdapter
 from llama_stack.providers.remote.inference.bedrock.config import BedrockConfig
+from llama_stack_api import OpenAIChatCompletionRequestWithExtraBody
 
 
 def test_adapter_initialization():

tests/unit/providers/inference/test_remote_vllm.py

@@ -9,6 +9,11 @@ import time
 from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch
 
 import pytest
+
+from llama_stack.core.routers.inference import InferenceRouter
+from llama_stack.core.routing_tables.models import ModelsRoutingTable
+from llama_stack.providers.remote.inference.vllm.config import VLLMInferenceAdapterConfig
+from llama_stack.providers.remote.inference.vllm.vllm import VLLMInferenceAdapter
 from llama_stack_api import (
     HealthStatus,
     Model,
@@ -22,11 +27,6 @@ from llama_stack_api import (
     ToolChoice,
 )
 
-from llama_stack.core.routers.inference import InferenceRouter
-from llama_stack.core.routing_tables.models import ModelsRoutingTable
-from llama_stack.providers.remote.inference.vllm.config import VLLMInferenceAdapterConfig
-from llama_stack.providers.remote.inference.vllm.vllm import VLLMInferenceAdapter
-
 # 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

tests/unit/providers/inline/agents/meta_reference/responses/test_streaming.py

@@ -7,12 +7,12 @@
 from unittest.mock import AsyncMock
 
 import pytest
-from llama_stack_api import ToolDef
 
 from llama_stack.providers.inline.agents.meta_reference.responses.streaming import (
     convert_tooldef_to_chat_tool,
 )
 from llama_stack.providers.inline.agents.meta_reference.responses.types import ChatCompletionContext
+from llama_stack_api import ToolDef
 
 
 @pytest.fixture

tests/unit/providers/nvidia/test_datastore.py

@@ -8,10 +8,10 @@ import os
 from unittest.mock import patch
 
 import pytest
-from llama_stack_api import Dataset, DatasetPurpose, ResourceType, URIDataSource
 
 from llama_stack.providers.remote.datasetio.nvidia.config import NvidiaDatasetIOConfig
 from llama_stack.providers.remote.datasetio.nvidia.datasetio import NvidiaDatasetIOAdapter
+from llama_stack_api import Dataset, DatasetPurpose, ResourceType, URIDataSource
 
 
 @pytest.fixture

tests/unit/providers/nvidia/test_eval.py

@@ -8,6 +8,10 @@ import os
 from unittest.mock import MagicMock, patch
 
 import pytest
+
+from llama_stack.models.llama.sku_types import CoreModelId
+from llama_stack.providers.remote.eval.nvidia.config import NVIDIAEvalConfig
+from llama_stack.providers.remote.eval.nvidia.eval import NVIDIAEvalImpl
 from llama_stack_api import (
     Benchmark,
     BenchmarkConfig,
@@ -20,10 +24,6 @@ from llama_stack_api import (
     TopPSamplingStrategy,
 )
 
-from llama_stack.models.llama.sku_types import CoreModelId
-from llama_stack.providers.remote.eval.nvidia.config import NVIDIAEvalConfig
-from llama_stack.providers.remote.eval.nvidia.eval import NVIDIAEvalImpl
-
 MOCK_DATASET_ID = "default/test-dataset"
 MOCK_BENCHMARK_ID = "test-benchmark"
 

tests/unit/providers/nvidia/test_parameters.py

@@ -9,6 +9,12 @@ import warnings
 from unittest.mock import patch
 
 import pytest
+
+from llama_stack.core.library_client import convert_pydantic_to_json_value
+from llama_stack.providers.remote.post_training.nvidia.post_training import (
+    NvidiaPostTrainingAdapter,
+    NvidiaPostTrainingConfig,
+)
 from llama_stack_api import (
     DataConfig,
     DatasetFormat,
@@ -19,12 +25,6 @@ from llama_stack_api import (
     TrainingConfig,
 )
 
-from llama_stack.core.library_client import convert_pydantic_to_json_value
-from llama_stack.providers.remote.post_training.nvidia.post_training import (
-    NvidiaPostTrainingAdapter,
-    NvidiaPostTrainingConfig,
-)
-
 
 class TestNvidiaParameters:
     @pytest.fixture(autouse=True)

tests/unit/providers/nvidia/test_rerank_inference.py

@@ -8,11 +8,11 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import aiohttp
 import pytest
-from llama_stack_api import ModelType
 
 from llama_stack.providers.remote.inference.nvidia.config import NVIDIAConfig
 from llama_stack.providers.remote.inference.nvidia.nvidia import NVIDIAInferenceAdapter
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack_api import ModelType
 
 
 class MockResponse:

tests/unit/providers/nvidia/test_safety.py

@@ -9,6 +9,9 @@ from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+
+from llama_stack.providers.remote.safety.nvidia.config import NVIDIASafetyConfig
+from llama_stack.providers.remote.safety.nvidia.nvidia import NVIDIASafetyAdapter
 from llama_stack_api import (
     OpenAIAssistantMessageParam,
     OpenAIUserMessageParam,
@@ -18,9 +21,6 @@ from llama_stack_api import (
     ViolationLevel,
 )
 
-from llama_stack.providers.remote.safety.nvidia.config import NVIDIASafetyConfig
-from llama_stack.providers.remote.safety.nvidia.nvidia import NVIDIASafetyAdapter
-
 
 class FakeNVIDIASafetyAdapter(NVIDIASafetyAdapter):
     """Test implementation that provides the required shield_store."""

tests/unit/providers/nvidia/test_supervised_fine_tuning.py

@@ -9,15 +9,6 @@ import warnings
 from unittest.mock import patch
 
 import pytest
-from llama_stack_api import (
-    DataConfig,
-    DatasetFormat,
-    LoraFinetuningConfig,
-    OptimizerConfig,
-    OptimizerType,
-    QATFinetuningConfig,
-    TrainingConfig,
-)
 
 from llama_stack.core.library_client import convert_pydantic_to_json_value
 from llama_stack.providers.remote.post_training.nvidia.post_training import (
@@ -27,6 +18,15 @@ from llama_stack.providers.remote.post_training.nvidia.post_training import (
     NvidiaPostTrainingJob,
     NvidiaPostTrainingJobStatusResponse,
 )
+from llama_stack_api import (
+    DataConfig,
+    DatasetFormat,
+    LoraFinetuningConfig,
+    OptimizerConfig,
+    OptimizerType,
+    QATFinetuningConfig,
+    TrainingConfig,
+)
 
 
 @pytest.fixture

tests/unit/providers/test_bedrock.py

@@ -7,10 +7,9 @@
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, PropertyMock, patch
 
-from llama_stack_api import OpenAIChatCompletionRequestWithExtraBody
-
 from llama_stack.providers.remote.inference.bedrock.bedrock import BedrockInferenceAdapter
 from llama_stack.providers.remote.inference.bedrock.config import BedrockConfig
+from llama_stack_api import OpenAIChatCompletionRequestWithExtraBody
 
 
 def test_can_create_adapter():

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

@@ -10,12 +10,12 @@ from typing import Any
 from unittest.mock import AsyncMock, MagicMock, Mock, PropertyMock, patch
 
 import pytest
-from llama_stack_api import Model, ModelType, OpenAIChatCompletionRequestWithExtraBody, OpenAIUserMessageParam
 from pydantic import BaseModel, Field
 
 from llama_stack.core.request_headers import request_provider_data_context
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack_api import Model, ModelType, OpenAIChatCompletionRequestWithExtraBody, OpenAIUserMessageParam
 
 
 class OpenAIMixinImpl(OpenAIMixin):

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

@@ -4,12 +4,11 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_stack_api import OpenAIAssistantMessageParam, OpenAIUserMessageParam
-
 from llama_stack.models.llama.datatypes import RawTextItem
 from llama_stack.providers.utils.inference.prompt_adapter import (
     convert_openai_message_to_raw_message,
 )
+from llama_stack_api import OpenAIAssistantMessageParam, OpenAIUserMessageParam
 
 
 class TestConvertOpenAIMessageToRawMessage:

tests/unit/providers/utils/memory/test_vector_store.py

@@ -7,9 +7,9 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from llama_stack_api import URL, RAGDocument, TextContentItem
 
 from llama_stack.providers.utils.memory.vector_store import content_from_data_and_mime_type, content_from_doc
+from llama_stack_api import URL, RAGDocument, TextContentItem
 
 
 async def test_content_from_doc_with_url():

tests/unit/providers/utils/test_model_registry.py

@@ -34,9 +34,9 @@
 #
 
 import pytest
-from llama_stack_api import Model
 
 from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper, ProviderModelEntry
+from llama_stack_api import Model
 
 
 @pytest.fixture

tests/unit/providers/vector_io/conftest.py

@@ -9,7 +9,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import numpy as np
 import pytest
-from llama_stack_api import Chunk, ChunkMetadata, QueryChunksResponse, VectorStore
 
 from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
 from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig
@@ -19,6 +18,7 @@ from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import SQLiteV
 from llama_stack.providers.remote.vector_io.pgvector.config import PGVectorVectorIOConfig
 from llama_stack.providers.remote.vector_io.pgvector.pgvector import PGVectorIndex, PGVectorVectorIOAdapter
 from llama_stack.providers.utils.kvstore import register_kvstore_backends
+from llama_stack_api import Chunk, ChunkMetadata, QueryChunksResponse, VectorStore
 
 EMBEDDING_DIMENSION = 768
 COLLECTION_PREFIX = "test_collection"

tests/unit/providers/vector_io/test_faiss.py

@@ -9,13 +9,13 @@ from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
-from llama_stack_api import Chunk, Files, HealthStatus, QueryChunksResponse, VectorStore
 
 from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig
 from llama_stack.providers.inline.vector_io.faiss.faiss import (
     FaissIndex,
     FaissVectorIOAdapter,
 )
+from llama_stack_api import Chunk, Files, HealthStatus, QueryChunksResponse, VectorStore
 
 # This test is a unit test for the FaissVectorIOAdapter class. This should only contain
 # tests which are specific to this class. More general (API-level) tests should be placed in

tests/unit/providers/vector_io/test_sqlite_vec.py

@@ -8,13 +8,13 @@ import asyncio
 
 import numpy as np
 import pytest
-from llama_stack_api import Chunk, QueryChunksResponse
 
 from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import (
     SQLiteVecIndex,
     SQLiteVecVectorIOAdapter,
     _create_sqlite_connection,
 )
+from llama_stack_api import Chunk, QueryChunksResponse
 
 # 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/unit/providers/vector_io/test_vector_io_openai_vector_stores.py

@@ -10,6 +10,8 @@ from unittest.mock import AsyncMock, patch
 
 import numpy as np
 import pytest
+
+from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import VECTOR_DBS_PREFIX
 from llama_stack_api import (
     Chunk,
     OpenAICreateVectorStoreFileBatchRequestWithExtraBody,
@@ -21,8 +23,6 @@ from llama_stack_api import (
     VectorStoreNotFoundError,
 )
 
-from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import VECTOR_DBS_PREFIX
-
 # This test is a unit test for the inline VectorIO providers. This should only contain
 # tests which are specific to this class. More general (API-level) tests should be placed in
 # tests/integration/vector_io/
@@ -255,10 +255,9 @@ async def test_insert_chunks_with_missing_document_id(vector_io_adapter):
 
 async def test_document_id_with_invalid_type_raises_error():
     """Ensure TypeError is raised when document_id is not a string."""
-    from llama_stack_api import Chunk
-
     # Integer document_id should raise TypeError
     from llama_stack.providers.utils.vector_io.vector_utils import generate_chunk_id
+    from llama_stack_api import Chunk
 
     chunk = Chunk(content="test", chunk_id=generate_chunk_id("test", "test"), metadata={"document_id": 12345})
     with pytest.raises(TypeError) as exc_info:

tests/unit/providers/vector_io/test_vector_utils.py

@@ -4,9 +4,8 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_stack_api import Chunk, ChunkMetadata
-
 from llama_stack.providers.utils.vector_io.vector_utils import generate_chunk_id
+from llama_stack_api import Chunk, ChunkMetadata
 
 # This test is a unit test for the chunk_utils.py helpers. This should only contain
 # tests which are specific to this file. More general (API-level) tests should be placed in

tests/unit/rag/test_rag_query.py

@@ -7,9 +7,9 @@
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
-from llama_stack_api import Chunk, ChunkMetadata, QueryChunksResponse, RAGQueryConfig
 
 from llama_stack.providers.inline.tool_runtime.rag.memory import MemoryToolRuntimeImpl
+from llama_stack_api import Chunk, ChunkMetadata, QueryChunksResponse, RAGQueryConfig
 
 
 class TestRagQuery:

tests/unit/rag/test_vector_store.py

@@ -12,7 +12,6 @@ from unittest.mock import AsyncMock, MagicMock
 
 import numpy as np
 import pytest
-from llama_stack_api import Chunk, OpenAIEmbeddingData, OpenAIEmbeddingsRequestWithExtraBody, RAGDocument
 
 from llama_stack.providers.utils.memory.vector_store import (
     URL,
@@ -22,6 +21,7 @@ from llama_stack.providers.utils.memory.vector_store import (
     make_overlapped_chunks,
 )
 from llama_stack.providers.utils.vector_io.vector_utils import generate_chunk_id
+from llama_stack_api import Chunk, OpenAIEmbeddingData, OpenAIEmbeddingsRequestWithExtraBody, RAGDocument
 
 DUMMY_PDF_PATH = Path(os.path.abspath(__file__)).parent / "fixtures" / "dummy.pdf"
 # Depending on the machine, this can get parsed a couple of ways

tests/unit/registry/test_registry.py

@@ -6,7 +6,6 @@
 
 
 import pytest
-from llama_stack_api import Model, VectorStore
 
 from llama_stack.core.datatypes import VectorStoreWithOwner
 from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
@@ -16,6 +15,7 @@ from llama_stack.core.store.registry import (
     DiskDistributionRegistry,
 )
 from llama_stack.providers.utils.kvstore import kvstore_impl, register_kvstore_backends
+from llama_stack_api import Model, VectorStore
 
 
 @pytest.fixture
@@ -303,9 +303,8 @@ async def test_double_registration_different_objects(disk_dist_registry):
 
 async def test_double_registration_with_cache(cached_disk_dist_registry):
     """Test double registration behavior with caching enabled."""
-    from llama_stack_api import ModelType
-
     from llama_stack.core.datatypes import ModelWithOwner
+    from llama_stack_api import ModelType
 
     model1 = ModelWithOwner(
         identifier="test_model",

tests/unit/registry/test_registry_acl.py

@@ -5,10 +5,9 @@
 # the root directory of this source tree.
 
 
-from llama_stack_api import ModelType
-
 from llama_stack.core.datatypes import ModelWithOwner, User
 from llama_stack.core.store.registry import CachedDiskDistributionRegistry
+from llama_stack_api import ModelType
 
 
 async def test_registry_cache_with_acl(cached_disk_dist_registry):

tests/unit/server/test_access_control.py

@@ -8,12 +8,12 @@ from unittest.mock import MagicMock, Mock, patch
 
 import pytest
 import yaml
-from llama_stack_api import Api, ModelType
 from pydantic import TypeAdapter, ValidationError
 
 from llama_stack.core.access_control.access_control import AccessDeniedError, is_action_allowed
 from llama_stack.core.datatypes import AccessRule, ModelWithOwner, User
 from llama_stack.core.routing_tables.models import ModelsRoutingTable
+from llama_stack_api import Api, ModelType
 
 
 class AsyncMock(MagicMock):

tests/unit/server/test_resolver.py

@@ -9,7 +9,6 @@ import sys
 from typing import Any, Protocol
 from unittest.mock import AsyncMock, MagicMock
 
-from llama_stack_api import Inference, InlineProviderSpec, ProviderSpec
 from pydantic import BaseModel, Field
 
 from llama_stack.core.datatypes import Api, Provider, StackRunConfig
@@ -27,6 +26,7 @@ from llama_stack.core.storage.datatypes import (
 )
 from llama_stack.providers.utils.kvstore import register_kvstore_backends
 from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
+from llama_stack_api import Inference, InlineProviderSpec, ProviderSpec
 
 
 def add_protocol_methods(cls: type, protocol: type[Protocol]) -> None:

tests/unit/server/test_sse.py

@@ -9,9 +9,9 @@ import logging  # allow-direct-logging
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
-from llama_stack_api import PaginatedResponse
 
 from llama_stack.core.server.server import create_dynamic_typed_route, create_sse_event, sse_generator
+from llama_stack_api import PaginatedResponse
 
 
 @pytest.fixture

tests/unit/tools/test_tools_json_schema.py

@@ -9,10 +9,10 @@ Unit tests for JSON Schema-based tool definitions.
 Tests the new input_schema and output_schema fields.
 """
 
-from llama_stack_api import ToolDef
 from pydantic import ValidationError
 
 from llama_stack.models.llama.datatypes import BuiltinTool, ToolDefinition
+from llama_stack_api import ToolDef
 
 
 class TestToolDefValidation:

tests/unit/utils/inference/test_inference_store.py

@@ -7,6 +7,10 @@
 import time
 
 import pytest
+
+from llama_stack.core.storage.datatypes import InferenceStoreReference, SqliteSqlStoreConfig
+from llama_stack.providers.utils.inference.inference_store import InferenceStore
+from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
 from llama_stack_api import (
     OpenAIAssistantMessageParam,
     OpenAIChatCompletion,
@@ -15,10 +19,6 @@ from llama_stack_api import (
     Order,
 )
 
-from llama_stack.core.storage.datatypes import InferenceStoreReference, SqliteSqlStoreConfig
-from llama_stack.providers.utils.inference.inference_store import InferenceStore
-from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
-
 
 @pytest.fixture(autouse=True)
 def setup_backends(tmp_path):

tests/unit/utils/responses/test_responses_store.py

@@ -9,11 +9,11 @@ from tempfile import TemporaryDirectory
 from uuid import uuid4
 
 import pytest
-from llama_stack_api import OpenAIMessageParam, OpenAIResponseInput, OpenAIResponseObject, OpenAIUserMessageParam, Order
 
 from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqliteSqlStoreConfig
 from llama_stack.providers.utils.responses.responses_store import ResponsesStore
 from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
+from llama_stack_api import OpenAIMessageParam, OpenAIResponseInput, OpenAIResponseObject, OpenAIUserMessageParam, Order
 
 
 def build_store(db_path: str, policy: list | None = None) -> ResponsesStore: