Merge branch 'main' into vllm

llama_stack/providers/tests/agents/fixtures.py

@@ -81,14 +81,28 @@ async def agents_stack(request, inference_model, safety_shield):
     inference_models = (
         inference_model if isinstance(inference_model, list) else [inference_model]
     )
-    models = [
-        ModelInput(
-            model_id=model,
-            model_type=ModelType.llm,
-            provider_id=providers["inference"][0].provider_id,
+
+    # NOTE: meta-reference provider needs 1 provider per model, lookup provider_id from provider config
+    model_to_provider_id = {}
+    for provider in providers["inference"]:
+        if "model" in provider.config:
+            model_to_provider_id[provider.config["model"]] = provider.provider_id
+
+    models = []
+    for model in inference_models:
+        if model in model_to_provider_id:
+            provider_id = model_to_provider_id[model]
+        else:
+            provider_id = providers["inference"][0].provider_id
+
+        models.append(
+            ModelInput(
+                model_id=model,
+                model_type=ModelType.llm,
+                provider_id=provider_id,
+            )
         )
-        for model in inference_models
-    ]
+
     models.append(
         ModelInput(
             model_id="all-MiniLM-L6-v2",

llama_stack/providers/tests/agents/test_agents.py

@@ -5,11 +5,31 @@
 # the root directory of this source tree.
 
 import os
+from typing import Dict, List
 
 import pytest
+from llama_models.llama3.api.datatypes import BuiltinTool
 
-from llama_stack.apis.agents import *  # noqa: F403
-from llama_stack.providers.datatypes import *  # noqa: F403
+from llama_stack.apis.agents import (
+    AgentConfig,
+    AgentTool,
+    AgentTurnResponseEventType,
+    AgentTurnResponseStepCompletePayload,
+    AgentTurnResponseStreamChunk,
+    AgentTurnResponseTurnCompletePayload,
+    Attachment,
+    MemoryToolDefinition,
+    SearchEngineType,
+    SearchToolDefinition,
+    ShieldCallStep,
+    StepType,
+    ToolChoice,
+    ToolExecutionStep,
+    Turn,
+)
+from llama_stack.apis.inference import CompletionMessage, SamplingParams, UserMessage
+from llama_stack.apis.safety import ViolationLevel
+from llama_stack.providers.datatypes import Api
 
 # How to run this test:
 #

llama_stack/providers/tests/agents/test_persistence.py

@@ -6,9 +6,9 @@
 
 import pytest
 
-from llama_stack.apis.agents import *  # noqa: F403
-from llama_stack.providers.datatypes import *  # noqa: F403
-
+from llama_stack.apis.agents import AgentConfig, Turn
+from llama_stack.apis.inference import SamplingParams, UserMessage
+from llama_stack.providers.datatypes import Api
 from llama_stack.providers.utils.kvstore import kvstore_impl, SqliteKVStoreConfig
 from .fixtures import pick_inference_model
 

llama_stack/providers/tests/datasetio/test_datasetio.py

@@ -4,16 +4,17 @@
 # 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
-from llama_stack.apis.common.type_system import *  # noqa: F403
-from llama_stack.apis.datasetio import *  # noqa: F403
-from llama_stack.distribution.datatypes import *  # noqa: F403
 import base64
 import mimetypes
+import os
 from pathlib import Path
 
+import pytest
+
+from llama_stack.apis.common.content_types import URL
+from llama_stack.apis.common.type_system import ChatCompletionInputType, StringType
+from llama_stack.apis.datasets import Datasets
+
 # How to run this test:
 #
 # pytest llama_stack/providers/tests/datasetio/test_datasetio.py
@@ -37,9 +38,15 @@ def data_url_from_file(file_path: str) -> str:
 
 
 async def register_dataset(
-    datasets_impl: Datasets, for_generation=False, dataset_id="test_dataset"
+    datasets_impl: Datasets,
+    for_generation=False,
+    for_rag=False,
+    dataset_id="test_dataset",
 ):
-    test_file = Path(os.path.abspath(__file__)).parent / "test_dataset.csv"
+    if for_rag:
+        test_file = Path(os.path.abspath(__file__)).parent / "test_rag_dataset.csv"
+    else:
+        test_file = Path(os.path.abspath(__file__)).parent / "test_dataset.csv"
     test_url = data_url_from_file(str(test_file))
 
     if for_generation:
@@ -48,6 +55,13 @@ async def register_dataset(
             "input_query": StringType(),
             "chat_completion_input": ChatCompletionInputType(),
         }
+    elif for_rag:
+        dataset_schema = {
+            "expected_answer": StringType(),
+            "input_query": StringType(),
+            "generated_answer": StringType(),
+            "context": StringType(),
+        }
     else:
         dataset_schema = {
             "expected_answer": StringType(),

llama_stack/providers/tests/datasetio/test_rag_dataset.csv

@@ -0,0 +1,6 @@
+input_query,context,generated_answer,expected_answer
+What is the capital of France?,"France is a country in Western Europe with a population of about 67 million people. Its capital city has been a major European cultural center since the 17th century and is known for landmarks like the Eiffel Tower and the Louvre Museum.",London,Paris
+Who is the CEO of Meta?,"Meta Platforms, formerly known as Facebook, is one of the world's largest technology companies. Founded by Mark Zuckerberg in 2004, the company has expanded to include platforms like Instagram, WhatsApp, and virtual reality technologies.",Mark Zuckerberg,Mark Zuckerberg
+What is the largest planet in our solar system?,"The solar system consists of eight planets orbiting around the Sun. These planets, in order from the Sun, are Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune. Gas giants are significantly larger than terrestrial planets.",Jupiter,Jupiter
+What is the smallest country in the world?,"Independent city-states and micronations are among the world's smallest sovereign territories. Some notable examples include Monaco, San Marino, and Vatican City, which is an enclave within Rome, Italy.",China,Vatican City
+What is the currency of Japan?,"Japan is an island country in East Asia with a rich cultural heritage and one of the world's largest economies. Its financial system has been established since the Meiji period, with its modern currency being introduced in 1871.",Yen,Yen

llama_stack/providers/tests/eval/test_eval.py

@@ -7,8 +7,7 @@
 
 import pytest
 
-from llama_models.llama3.api import SamplingParams, URL
-
+from llama_stack.apis.common.content_types import URL
 from llama_stack.apis.common.type_system import ChatCompletionInputType, StringType
 
 from llama_stack.apis.eval.eval import (
@@ -16,6 +15,7 @@ from llama_stack.apis.eval.eval import (
     BenchmarkEvalTaskConfig,
     ModelCandidate,
 )
+from llama_stack.apis.inference import SamplingParams
 from llama_stack.apis.scoring_functions import LLMAsJudgeScoringFnParams
 from llama_stack.distribution.datatypes import Api
 from llama_stack.providers.tests.datasetio.test_datasetio import register_dataset

llama_stack/providers/tests/inference/fixtures.py

@@ -20,6 +20,7 @@ from llama_stack.providers.remote.inference.bedrock import BedrockConfig
 
 from llama_stack.providers.remote.inference.cerebras import CerebrasImplConfig
 from llama_stack.providers.remote.inference.fireworks import FireworksImplConfig
+from llama_stack.providers.remote.inference.groq import GroqConfig
 from llama_stack.providers.remote.inference.nvidia import NVIDIAConfig
 from llama_stack.providers.remote.inference.ollama import OllamaImplConfig
 from llama_stack.providers.remote.inference.tgi import TGIImplConfig
@@ -172,6 +173,22 @@ def inference_together() -> ProviderFixture:
     )
 
 
+@pytest.fixture(scope="session")
+def inference_groq() -> ProviderFixture:
+    return ProviderFixture(
+        providers=[
+            Provider(
+                provider_id="groq",
+                provider_type="remote::groq",
+                config=GroqConfig().model_dump(),
+            )
+        ],
+        provider_data=dict(
+            groq_api_key=get_env_or_fail("GROQ_API_KEY"),
+        ),
+    )
+
+
 @pytest.fixture(scope="session")
 def inference_bedrock() -> ProviderFixture:
     return ProviderFixture(
@@ -258,6 +275,7 @@ INFERENCE_FIXTURES = [
     "fireworks",
     "together",
     "vllm",
+    "groq",
     "vllm_remote",
     "remote",
     "bedrock",

llama_stack/providers/tests/inference/groq/test_groq_utils.py

@@ -0,0 +1,271 @@
+# 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 groq.types.chat.chat_completion import ChatCompletion, Choice
+from groq.types.chat.chat_completion_chunk import (
+    ChatCompletionChunk,
+    Choice as StreamChoice,
+    ChoiceDelta,
+)
+from groq.types.chat.chat_completion_message import ChatCompletionMessage
+
+from llama_stack.apis.inference import (
+    ChatCompletionRequest,
+    ChatCompletionResponseEventType,
+    CompletionMessage,
+    StopReason,
+    SystemMessage,
+    UserMessage,
+)
+from llama_stack.providers.remote.inference.groq.groq_utils import (
+    convert_chat_completion_request,
+    convert_chat_completion_response,
+    convert_chat_completion_response_stream,
+)
+
+
+class TestConvertChatCompletionRequest:
+    def test_sets_model(self):
+        request = self._dummy_chat_completion_request()
+        request.model = "Llama-3.2-3B"
+
+        converted = convert_chat_completion_request(request)
+
+        assert converted["model"] == "Llama-3.2-3B"
+
+    def test_converts_user_message(self):
+        request = self._dummy_chat_completion_request()
+        request.messages = [UserMessage(content="Hello World")]
+
+        converted = convert_chat_completion_request(request)
+
+        assert converted["messages"] == [
+            {"role": "user", "content": "Hello World"},
+        ]
+
+    def test_converts_system_message(self):
+        request = self._dummy_chat_completion_request()
+        request.messages = [SystemMessage(content="You are a helpful assistant.")]
+
+        converted = convert_chat_completion_request(request)
+
+        assert converted["messages"] == [
+            {"role": "system", "content": "You are a helpful assistant."},
+        ]
+
+    def test_converts_completion_message(self):
+        request = self._dummy_chat_completion_request()
+        request.messages = [
+            UserMessage(content="Hello World"),
+            CompletionMessage(
+                content="Hello World! How can I help you today?",
+                stop_reason=StopReason.end_of_message,
+            ),
+        ]
+
+        converted = convert_chat_completion_request(request)
+
+        assert converted["messages"] == [
+            {"role": "user", "content": "Hello World"},
+            {"role": "assistant", "content": "Hello World! How can I help you today?"},
+        ]
+
+    def test_does_not_include_logprobs(self):
+        request = self._dummy_chat_completion_request()
+        request.logprobs = True
+
+        with pytest.warns(Warning) as warnings:
+            converted = convert_chat_completion_request(request)
+
+        assert "logprobs are not supported yet" in warnings[0].message.args[0]
+        assert converted.get("logprobs") is None
+
+    def test_does_not_include_response_format(self):
+        request = self._dummy_chat_completion_request()
+        request.response_format = {
+            "type": "json_object",
+            "json_schema": {
+                "type": "object",
+                "properties": {
+                    "name": {"type": "string"},
+                    "age": {"type": "number"},
+                },
+            },
+        }
+
+        with pytest.warns(Warning) as warnings:
+            converted = convert_chat_completion_request(request)
+
+        assert "response_format is not supported yet" in warnings[0].message.args[0]
+        assert converted.get("response_format") is None
+
+    def test_does_not_include_repetition_penalty(self):
+        request = self._dummy_chat_completion_request()
+        request.sampling_params.repetition_penalty = 1.5
+
+        with pytest.warns(Warning) as warnings:
+            converted = convert_chat_completion_request(request)
+
+        assert "repetition_penalty is not supported" in warnings[0].message.args[0]
+        assert converted.get("repetition_penalty") is None
+        assert converted.get("frequency_penalty") is None
+
+    def test_includes_stream(self):
+        request = self._dummy_chat_completion_request()
+        request.stream = True
+
+        converted = convert_chat_completion_request(request)
+
+        assert converted["stream"] is True
+
+    def test_if_max_tokens_is_0_then_it_is_not_included(self):
+        request = self._dummy_chat_completion_request()
+        # 0 is the default value for max_tokens
+        # So we assume that if it's 0, the user didn't set it
+        request.sampling_params.max_tokens = 0
+
+        converted = convert_chat_completion_request(request)
+
+        assert converted.get("max_tokens") is None
+
+    def test_includes_max_tokens_if_set(self):
+        request = self._dummy_chat_completion_request()
+        request.sampling_params.max_tokens = 100
+
+        converted = convert_chat_completion_request(request)
+
+        assert converted["max_tokens"] == 100
+
+    def _dummy_chat_completion_request(self):
+        return ChatCompletionRequest(
+            model="Llama-3.2-3B",
+            messages=[UserMessage(content="Hello World")],
+        )
+
+    def test_includes_temperature(self):
+        request = self._dummy_chat_completion_request()
+        request.sampling_params.temperature = 0.5
+
+        converted = convert_chat_completion_request(request)
+
+        assert converted["temperature"] == 0.5
+
+    def test_includes_top_p(self):
+        request = self._dummy_chat_completion_request()
+        request.sampling_params.top_p = 0.95
+
+        converted = convert_chat_completion_request(request)
+
+        assert converted["top_p"] == 0.95
+
+
+class TestConvertNonStreamChatCompletionResponse:
+    def test_returns_response(self):
+        response = self._dummy_chat_completion_response()
+        response.choices[0].message.content = "Hello World"
+
+        converted = convert_chat_completion_response(response)
+
+        assert converted.completion_message.content == "Hello World"
+
+    def test_maps_stop_to_end_of_message(self):
+        response = self._dummy_chat_completion_response()
+        response.choices[0].finish_reason = "stop"
+
+        converted = convert_chat_completion_response(response)
+
+        assert converted.completion_message.stop_reason == StopReason.end_of_turn
+
+    def test_maps_length_to_end_of_message(self):
+        response = self._dummy_chat_completion_response()
+        response.choices[0].finish_reason = "length"
+
+        converted = convert_chat_completion_response(response)
+
+        assert converted.completion_message.stop_reason == StopReason.out_of_tokens
+
+    def _dummy_chat_completion_response(self):
+        return ChatCompletion(
+            id="chatcmpl-123",
+            model="Llama-3.2-3B",
+            choices=[
+                Choice(
+                    index=0,
+                    message=ChatCompletionMessage(
+                        role="assistant", content="Hello World"
+                    ),
+                    finish_reason="stop",
+                )
+            ],
+            created=1729382400,
+            object="chat.completion",
+        )
+
+
+class TestConvertStreamChatCompletionResponse:
+    @pytest.mark.asyncio
+    async def test_returns_stream(self):
+        def chat_completion_stream():
+            messages = ["Hello ", "World ", " !"]
+            for i, message in enumerate(messages):
+                chunk = self._dummy_chat_completion_chunk()
+                chunk.choices[0].delta.content = message
+                if i == len(messages) - 1:
+                    chunk.choices[0].finish_reason = "stop"
+                else:
+                    chunk.choices[0].finish_reason = None
+                yield chunk
+
+            chunk = self._dummy_chat_completion_chunk()
+            chunk.choices[0].delta.content = None
+            chunk.choices[0].finish_reason = "stop"
+            yield chunk
+
+        stream = chat_completion_stream()
+        converted = convert_chat_completion_response_stream(stream)
+
+        iter = converted.__aiter__()
+        chunk = await iter.__anext__()
+        assert chunk.event.event_type == ChatCompletionResponseEventType.start
+        assert chunk.event.delta == "Hello "
+
+        chunk = await iter.__anext__()
+        assert chunk.event.event_type == ChatCompletionResponseEventType.progress
+        assert chunk.event.delta == "World "
+
+        chunk = await iter.__anext__()
+        assert chunk.event.event_type == ChatCompletionResponseEventType.progress
+        assert chunk.event.delta == " !"
+
+        # Dummy chunk to ensure the last chunk is really the end of the stream
+        # This one technically maps to Groq's final "stop" chunk
+        chunk = await iter.__anext__()
+        assert chunk.event.event_type == ChatCompletionResponseEventType.progress
+        assert chunk.event.delta == ""
+
+        chunk = await iter.__anext__()
+        assert chunk.event.event_type == ChatCompletionResponseEventType.complete
+        assert chunk.event.delta == ""
+        assert chunk.event.stop_reason == StopReason.end_of_turn
+
+        with pytest.raises(StopAsyncIteration):
+            await iter.__anext__()
+
+    def _dummy_chat_completion_chunk(self):
+        return ChatCompletionChunk(
+            id="chatcmpl-123",
+            model="Llama-3.2-3B",
+            choices=[
+                StreamChoice(
+                    index=0,
+                    delta=ChoiceDelta(role="assistant", content="Hello World"),
+                )
+            ],
+            created=1729382400,
+            object="chat.completion.chunk",
+            x_groq=None,
+        )

llama_stack/providers/tests/inference/groq/test_init.py

@@ -0,0 +1,29 @@
+# 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 llama_stack.apis.inference import Inference
+from llama_stack.providers.remote.inference.groq import get_adapter_impl
+from llama_stack.providers.remote.inference.groq.config import GroqConfig
+from llama_stack.providers.remote.inference.groq.groq import GroqInferenceAdapter
+
+from llama_stack.providers.remote.inference.ollama import OllamaImplConfig
+
+
+class TestGroqInit:
+    @pytest.mark.asyncio
+    async def test_raises_runtime_error_if_config_is_not_groq_config(self):
+        config = OllamaImplConfig(model="llama3.1-8b-8192")
+
+        with pytest.raises(RuntimeError):
+            await get_adapter_impl(config, None)
+
+    @pytest.mark.asyncio
+    async def test_returns_groq_adapter(self):
+        config = GroqConfig()
+        adapter = await get_adapter_impl(config, None)
+        assert type(adapter) is GroqInferenceAdapter
+        assert isinstance(adapter, Inference)

llama_stack/providers/tests/inference/test_prompt_adapter.py

@@ -6,8 +6,14 @@
 
 import unittest
 
-from llama_models.llama3.api import *  # noqa: F403
-from llama_stack.apis.inference.inference import *  # noqa: F403
+from llama_models.llama3.api.datatypes import (
+    BuiltinTool,
+    ToolDefinition,
+    ToolParamDefinition,
+    ToolPromptFormat,
+)
+
+from llama_stack.apis.inference import ChatCompletionRequest, SystemMessage, UserMessage
 from llama_stack.providers.utils.inference.prompt_adapter import (
     chat_completion_request_to_messages,
 )
@@ -24,7 +30,7 @@ class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
                 UserMessage(content=content),
             ],
         )
-        messages = chat_completion_request_to_messages(request)
+        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)
@@ -41,7 +47,7 @@ class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
                 ToolDefinition(tool_name=BuiltinTool.brave_search),
             ],
         )
-        messages = chat_completion_request_to_messages(request)
+        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)
@@ -69,7 +75,7 @@ class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
             ],
             tool_prompt_format=ToolPromptFormat.json,
         )
-        messages = chat_completion_request_to_messages(request)
+        messages = chat_completion_request_to_messages(request, MODEL)
         self.assertEqual(len(messages), 3)
         self.assertTrue("Environment: ipython" in messages[0].content)
 
@@ -99,7 +105,7 @@ class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
                 ),
             ],
         )
-        messages = chat_completion_request_to_messages(request)
+        messages = chat_completion_request_to_messages(request, MODEL)
         self.assertEqual(len(messages), 3)
 
         self.assertTrue("Environment: ipython" in messages[0].content)
@@ -121,7 +127,7 @@ class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
                 ToolDefinition(tool_name=BuiltinTool.code_interpreter),
             ],
         )
-        messages = chat_completion_request_to_messages(request)
+        messages = chat_completion_request_to_messages(request, MODEL)
         self.assertEqual(len(messages), 2, messages)
         self.assertTrue(messages[0].content.endswith(system_prompt))
 

llama_stack/providers/tests/inference/test_text_inference.py

@@ -7,13 +7,32 @@
 
 import pytest
 
+from llama_models.llama3.api.datatypes import (
+    SamplingParams,
+    StopReason,
+    ToolCall,
+    ToolDefinition,
+    ToolParamDefinition,
+    ToolPromptFormat,
+)
+
 from pydantic import BaseModel, ValidationError
 
-from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.apis.inference import *  # noqa: F403
-
-from llama_stack.distribution.datatypes import *  # noqa: F403
-
+from llama_stack.apis.inference import (
+    ChatCompletionResponse,
+    ChatCompletionResponseEventType,
+    ChatCompletionResponseStreamChunk,
+    CompletionResponse,
+    CompletionResponseStreamChunk,
+    JsonSchemaResponseFormat,
+    LogProbConfig,
+    SystemMessage,
+    ToolCallDelta,
+    ToolCallParseStatus,
+    ToolChoice,
+    UserMessage,
+)
+from llama_stack.apis.models import Model
 from .utils import group_chunks
 
 
@@ -193,6 +212,7 @@ class TestInference:
         provider = inference_impl.routing_table.get_provider_impl(inference_model)
         if provider.__provider_spec__.provider_type not in (
             "inline::meta-reference",
+            "remote::ollama",
             "remote::tgi",
             "remote::together",
             "remote::fireworks",
@@ -255,6 +275,7 @@ class TestInference:
         provider = inference_impl.routing_table.get_provider_impl(inference_model)
         if provider.__provider_spec__.provider_type not in (
             "inline::meta-reference",
+            "remote::ollama",
             "remote::fireworks",
             "remote::tgi",
             "remote::together",
@@ -352,6 +373,14 @@ class TestInference:
         sample_messages,
         sample_tool_definition,
     ):
+        inference_impl, _ = inference_stack
+        provider = inference_impl.routing_table.get_provider_impl(inference_model)
+        if provider.__provider_spec__.provider_type in ("remote::groq",):
+            pytest.skip(
+                provider.__provider_spec__.provider_type
+                + " doesn't support tool calling yet"
+            )
+
         inference_impl, _ = inference_stack
         messages = sample_messages + [
             UserMessage(
@@ -392,6 +421,13 @@ class TestInference:
         sample_tool_definition,
     ):
         inference_impl, _ = inference_stack
+        provider = inference_impl.routing_table.get_provider_impl(inference_model)
+        if provider.__provider_spec__.provider_type in ("remote::groq",):
+            pytest.skip(
+                provider.__provider_spec__.provider_type
+                + " doesn't support tool calling yet"
+            )
+
         messages = sample_messages + [
             UserMessage(
                 content="What's the weather like in San Francisco?",

llama_stack/providers/tests/inference/test_vision_inference.py

@@ -8,11 +8,16 @@ from pathlib import Path
 
 import pytest
 
-
-from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.apis.inference import *  # noqa: F403
 from llama_stack.apis.common.content_types import ImageContentItem, TextContentItem, URL
 
+from llama_stack.apis.inference import (
+    ChatCompletionResponse,
+    ChatCompletionResponseEventType,
+    ChatCompletionResponseStreamChunk,
+    SamplingParams,
+    UserMessage,
+)
+
 from .utils import group_chunks
 
 THIS_DIR = Path(__file__).parent

llama_stack/providers/tests/memory/fixtures.py

@@ -10,8 +10,7 @@ import tempfile
 import pytest
 import pytest_asyncio
 
-from llama_stack.apis.inference import ModelInput, ModelType
-
+from llama_stack.apis.models import ModelInput, ModelType
 from llama_stack.distribution.datatypes import Api, Provider
 from llama_stack.providers.inline.memory.chroma import ChromaInlineImplConfig
 from llama_stack.providers.inline.memory.faiss import FaissImplConfig
@@ -19,7 +18,7 @@ from llama_stack.providers.remote.memory.chroma import ChromaRemoteImplConfig
 from llama_stack.providers.remote.memory.pgvector import PGVectorConfig
 from llama_stack.providers.remote.memory.weaviate import WeaviateConfig
 from llama_stack.providers.tests.resolver import construct_stack_for_test
-from llama_stack.providers.utils.kvstore import SqliteKVStoreConfig
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
 from ..conftest import ProviderFixture, remote_stack_fixture
 from ..env import get_env_or_fail
 

llama_stack/providers/tests/memory/test_memory.py

@@ -8,14 +8,18 @@ import uuid
 
 import pytest
 
-from llama_stack.apis.memory import *  # noqa: F403
-from llama_stack.distribution.datatypes import *  # noqa: F403
-from llama_stack.apis.memory_banks.memory_banks import VectorMemoryBankParams
+from llama_stack.apis.memory import MemoryBankDocument, QueryDocumentsResponse
+
+from llama_stack.apis.memory_banks import (
+    MemoryBank,
+    MemoryBanks,
+    VectorMemoryBankParams,
+)
 
 # How to run this test:
 #
 # pytest llama_stack/providers/tests/memory/test_memory.py
-#   -m "meta_reference"
+#   -m "sentence_transformers" --env EMBEDDING_DIMENSION=384
 #   -v -s --tb=short --disable-warnings
 
 

llama_stack/providers/tests/post_training/fixtures.py

@@ -7,8 +7,9 @@
 import pytest
 import pytest_asyncio
 
-from llama_stack.apis.common.type_system import *  # noqa: F403
 from llama_stack.apis.common.content_types import URL
+
+from llama_stack.apis.common.type_system import StringType
 from llama_stack.apis.datasets import DatasetInput
 from llama_stack.apis.models import ModelInput
 

llama_stack/providers/tests/post_training/test_post_training.py

@@ -4,9 +4,18 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 import pytest
-from llama_stack.apis.common.type_system import *  # noqa: F403
-from llama_stack.apis.post_training import *  # noqa: F403
-from llama_stack.distribution.datatypes import *  # noqa: F403
+
+from llama_stack.apis.common.type_system import JobStatus
+from llama_stack.apis.post_training import (
+    Checkpoint,
+    DataConfig,
+    LoraFinetuningConfig,
+    OptimizerConfig,
+    PostTrainingJob,
+    PostTrainingJobArtifactsResponse,
+    PostTrainingJobStatusResponse,
+    TrainingConfig,
+)
 
 # How to run this test:
 #

llama_stack/providers/tests/resolver.py

@@ -8,14 +8,24 @@ import json
 import tempfile
 from typing import Any, Dict, List, Optional
 
-from llama_stack.distribution.datatypes import *  # noqa: F403
+from pydantic import BaseModel
+
+from llama_stack.apis.datasets import DatasetInput
+from llama_stack.apis.eval_tasks import EvalTaskInput
+from llama_stack.apis.memory_banks import MemoryBankInput
+from llama_stack.apis.models import ModelInput
+from llama_stack.apis.scoring_functions import ScoringFnInput
+from llama_stack.apis.shields import ShieldInput
+
 from llama_stack.distribution.build import print_pip_install_help
 from llama_stack.distribution.configure import parse_and_maybe_upgrade_config
+from llama_stack.distribution.datatypes import Provider, StackRunConfig
 from llama_stack.distribution.distribution import get_provider_registry
 from llama_stack.distribution.request_headers import set_request_provider_data
 from llama_stack.distribution.resolver import resolve_remote_stack_impls
 from llama_stack.distribution.stack import construct_stack
-from llama_stack.providers.utils.kvstore import SqliteKVStoreConfig
+from llama_stack.providers.datatypes import Api, RemoteProviderConfig
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
 
 
 class TestStack(BaseModel):

llama_stack/providers/tests/safety/test_safety.py

@@ -6,11 +6,9 @@
 
 import pytest
 
-from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.apis.safety import *  # noqa: F403
-
-from llama_stack.distribution.datatypes import *  # noqa: F403
 from llama_stack.apis.inference import UserMessage
+from llama_stack.apis.safety import ViolationLevel
+from llama_stack.apis.shields import Shield
 
 # How to run this test:
 #

llama_stack/providers/tests/scoring/test_scoring.py

@@ -60,7 +60,7 @@ class TestScoring:
                 f"{provider_id} provider does not support scoring without params"
             )
 
-        await register_dataset(datasets_impl)
+        await register_dataset(datasets_impl, for_rag=True)
         response = await datasets_impl.list_datasets()
         assert len(response) == 1
 
@@ -112,7 +112,7 @@ class TestScoring:
             scoring_stack[Api.datasets],
             scoring_stack[Api.models],
         )
-        await register_dataset(datasets_impl)
+        await register_dataset(datasets_impl, for_rag=True)
         response = await datasets_impl.list_datasets()
         assert len(response) == 1
 
@@ -173,7 +173,7 @@ class TestScoring:
             scoring_stack[Api.datasets],
             scoring_stack[Api.models],
         )
-        await register_dataset(datasets_impl)
+        await register_dataset(datasets_impl, for_rag=True)
         rows = await datasetio_impl.get_rows_paginated(
             dataset_id="test_dataset",
             rows_in_page=3,
@@ -197,7 +197,7 @@ class TestScoring:
                     judge_score_regexes=[r"Score: (\d+)"],
                     aggregation_functions=aggr_fns,
                 )
-            elif x.provider_id == "basic":
+            elif x.provider_id == "basic" or x.provider_id == "braintrust":
                 if "regex_parser" in x.identifier:
                     scoring_functions[x.identifier] = RegexParserScoringFnParams(
                         aggregation_functions=aggr_fns,