Merge branch 'main' into feature/dpo-training

llama_stack/providers/datatypes.py

@@ -43,10 +43,24 @@ class ModelsProtocolPrivate(Protocol):
        -> Provider uses provider-model-id for inference
     """
 
+    # this should be called `on_model_register` or something like that.
+    # the provider should _not_ be able to change the object in this
+    # callback
     async def register_model(self, model: Model) -> Model: ...
 
     async def unregister_model(self, model_id: str) -> None: ...
 
+    # the Stack router will query each provider for their list of models
+    # if a `refresh_interval_seconds` is provided, this method will be called
+    # periodically to refresh the list of models
+    #
+    # NOTE: each model returned will be registered with the model registry. this means
+    # a callback to the `register_model()` method will be made. this is duplicative and
+    # may be removed in the future.
+    async def list_models(self) -> list[Model] | None: ...
+
+    async def should_refresh_models(self) -> bool: ...
+
 
 class ShieldsProtocolPrivate(Protocol):
     async def register_shield(self, shield: Shield) -> None: ...
@@ -104,6 +118,19 @@ class ProviderSpec(BaseModel):
         description="If this provider is deprecated and does NOT work, specify the error message here",
     )
 
+    module: str | None = Field(
+        default=None,
+        description="""
+ Fully-qualified name of the module to import. The module is expected to have:
+
+  - `get_adapter_impl(config, deps)`: returns the adapter implementation
+
+  Example: `module: ramalama_stack`
+ """,
+    )
+
+    is_external: bool = Field(default=False, description="Notes whether this provider is an external provider.")
+
     # used internally by the resolver; this is a hack for now
     deps__: list[str] = Field(default_factory=list)
 
@@ -113,7 +140,7 @@ class ProviderSpec(BaseModel):
 
 
 class RoutingTable(Protocol):
-    def get_provider_impl(self, routing_key: str) -> Any: ...
+    async def get_provider_impl(self, routing_key: str) -> Any: ...
 
 
 # TODO: this can now be inlined into RemoteProviderSpec
@@ -124,7 +151,7 @@ class AdapterSpec(BaseModel):
         description="Unique identifier for this adapter",
     )
     module: str = Field(
-        ...,
+        default_factory=str,
         description="""
 Fully-qualified name of the module to import. The module is expected to have:
 
@@ -162,14 +189,7 @@ The container image to use for this implementation. If one is provided, pip_pack
 If a provider depends on other providers, the dependencies MUST NOT specify a container image.
 """,
     )
-    module: str = Field(
-        ...,
-        description="""
-Fully-qualified name of the module to import. The module is expected to have:
-
- - `get_provider_impl(config, deps)`: returns the local implementation
-""",
-    )
+    # module field is inherited from ProviderSpec
     provider_data_validator: str | None = Field(
         default=None,
     )
@@ -212,9 +232,7 @@ API responses, specify the adapter here.
     def container_image(self) -> str | None:
         return None
 
-    @property
-    def module(self) -> str:
-        return self.adapter.module
+    # module field is inherited from ProviderSpec
 
     @property
     def pip_packages(self) -> list[str]:
@@ -226,14 +244,19 @@ API responses, specify the adapter here.
 
 
 def remote_provider_spec(
-    api: Api, adapter: AdapterSpec, api_dependencies: list[Api] | None = None
+    api: Api,
+    adapter: AdapterSpec,
+    api_dependencies: list[Api] | None = None,
+    optional_api_dependencies: list[Api] | None = None,
 ) -> RemoteProviderSpec:
     return RemoteProviderSpec(
         api=api,
         provider_type=f"remote::{adapter.adapter_type}",
         config_class=adapter.config_class,
+        module=adapter.module,
         adapter=adapter,
         api_dependencies=api_dependencies or [],
+        optional_api_dependencies=optional_api_dependencies or [],
     )
 
 

llama_stack/providers/inline/agents/meta_reference/agent_instance.py

@@ -10,6 +10,7 @@ import re
 import secrets
 import string
 import uuid
+import warnings
 from collections.abc import AsyncGenerator
 from datetime import UTC, datetime
 
@@ -911,8 +912,16 @@ async def load_data_from_url(url: str) -> str:
 
 
 async def get_raw_document_text(document: Document) -> str:
-    if not document.mime_type.startswith("text/"):
+    # Handle deprecated text/yaml mime type with warning
+    if document.mime_type == "text/yaml":
+        warnings.warn(
+            "The 'text/yaml' MIME type is deprecated. Please use 'application/yaml' instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+    elif not (document.mime_type.startswith("text/") or document.mime_type == "application/yaml"):
         raise ValueError(f"Unexpected document mime type: {document.mime_type}")
+
     if isinstance(document.content, URL):
         return await load_data_from_url(document.content.uri)
     elif isinstance(document.content, str):

llama_stack/providers/inline/agents/meta_reference/persistence.py

@@ -128,6 +128,11 @@ class AgentPersistence:
             except Exception as e:
                 log.error(f"Error parsing turn: {e}")
                 continue
+
+        # The kvstore does not guarantee order, so we sort by started_at
+        # to ensure consistent ordering of turns.
+        turns.sort(key=lambda t: t.started_at)
+
         return turns
 
     async def get_session_turn(self, session_id: str, turn_id: str) -> Turn | None:

llama_stack/providers/inline/files/localfs/__init__.py

@@ -6,7 +6,7 @@
 
 from typing import Any
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.distribution.datatypes import AccessRule, Api
 
 from .config import LocalfsFilesImplConfig
 from .files import LocalfsFilesImpl
@@ -14,7 +14,7 @@ from .files import LocalfsFilesImpl
 __all__ = ["LocalfsFilesImpl", "LocalfsFilesImplConfig"]
 
 
-async def get_provider_impl(config: LocalfsFilesImplConfig, deps: dict[Api, Any]):
-    impl = LocalfsFilesImpl(config)
+async def get_provider_impl(config: LocalfsFilesImplConfig, deps: dict[Api, Any], policy: list[AccessRule]):
+    impl = LocalfsFilesImpl(config, policy)
     await impl.initialize()
     return impl

llama_stack/providers/inline/files/localfs/files.py

@@ -19,16 +19,19 @@ from llama_stack.apis.files import (
     OpenAIFileObject,
     OpenAIFilePurpose,
 )
+from llama_stack.distribution.datatypes import AccessRule
 from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
-from llama_stack.providers.utils.sqlstore.sqlstore import SqlStore, sqlstore_impl
+from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
 
 from .config import LocalfsFilesImplConfig
 
 
 class LocalfsFilesImpl(Files):
-    def __init__(self, config: LocalfsFilesImplConfig) -> None:
+    def __init__(self, config: LocalfsFilesImplConfig, policy: list[AccessRule]) -> None:
         self.config = config
-        self.sql_store: SqlStore | None = None
+        self.policy = policy
+        self.sql_store: AuthorizedSqlStore | None = None
 
     async def initialize(self) -> None:
         """Initialize the files provider by setting up storage directory and metadata database."""
@@ -37,7 +40,7 @@ class LocalfsFilesImpl(Files):
         storage_path.mkdir(parents=True, exist_ok=True)
 
         # Initialize SQL store for metadata
-        self.sql_store = sqlstore_impl(self.config.metadata_store)
+        self.sql_store = AuthorizedSqlStore(sqlstore_impl(self.config.metadata_store))
         await self.sql_store.create_table(
             "openai_files",
             {
@@ -126,6 +129,7 @@ class LocalfsFilesImpl(Files):
 
         paginated_result = await self.sql_store.fetch_all(
             table="openai_files",
+            policy=self.policy,
             where=where_conditions if where_conditions else None,
             order_by=[("created_at", order.value)],
             cursor=("id", after) if after else None,
@@ -156,7 +160,7 @@ class LocalfsFilesImpl(Files):
         if not self.sql_store:
             raise RuntimeError("Files provider not initialized")
 
-        row = await self.sql_store.fetch_one("openai_files", where={"id": file_id})
+        row = await self.sql_store.fetch_one("openai_files", policy=self.policy, where={"id": file_id})
         if not row:
             raise ValueError(f"File with id {file_id} not found")
 
@@ -174,7 +178,7 @@ class LocalfsFilesImpl(Files):
         if not self.sql_store:
             raise RuntimeError("Files provider not initialized")
 
-        row = await self.sql_store.fetch_one("openai_files", where={"id": file_id})
+        row = await self.sql_store.fetch_one("openai_files", policy=self.policy, where={"id": file_id})
         if not row:
             raise ValueError(f"File with id {file_id} not found")
 
@@ -197,7 +201,7 @@ class LocalfsFilesImpl(Files):
             raise RuntimeError("Files provider not initialized")
 
         # Get file metadata
-        row = await self.sql_store.fetch_one("openai_files", where={"id": file_id})
+        row = await self.sql_store.fetch_one("openai_files", policy=self.policy, where={"id": file_id})
         if not row:
             raise ValueError(f"File with id {file_id} not found")
 

llama_stack/providers/inline/inference/meta_reference/inference.py

@@ -102,6 +102,12 @@ class MetaReferenceInferenceImpl(
         if self.config.create_distributed_process_group:
             self.generator.stop()
 
+    async def should_refresh_models(self) -> bool:
+        return False
+
+    async def list_models(self) -> list[Model] | None:
+        return None
+
     async def unregister_model(self, model_id: str) -> None:
         pass
 

llama_stack/providers/inline/inference/sentence_transformers/sentence_transformers.py

@@ -20,6 +20,7 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
+from llama_stack.apis.models import ModelType
 from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.embedding_mixin import (
     SentenceTransformerEmbeddingMixin,
@@ -41,6 +42,8 @@ class SentenceTransformersInferenceImpl(
     InferenceProvider,
     ModelsProtocolPrivate,
 ):
+    __provider_id__: str
+
     def __init__(self, config: SentenceTransformersInferenceConfig) -> None:
         self.config = config
 
@@ -50,6 +53,22 @@ class SentenceTransformersInferenceImpl(
     async def shutdown(self) -> None:
         pass
 
+    async def should_refresh_models(self) -> bool:
+        return False
+
+    async def list_models(self) -> list[Model] | None:
+        return [
+            Model(
+                identifier="all-MiniLM-L6-v2",
+                provider_resource_id="all-MiniLM-L6-v2",
+                provider_id=self.__provider_id__,
+                metadata={
+                    "embedding_dimension": 384,
+                },
+                model_type=ModelType.embedding,
+            ),
+        ]
+
     async def register_model(self, model: Model) -> Model:
         return model
 

llama_stack/providers/inline/inference/vllm/__init__.py

@@ -1,17 +0,0 @@
-# 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 typing import Any
-
-from .config import VLLMConfig
-
-
-async def get_provider_impl(config: VLLMConfig, _deps: dict[str, Any]):
-    from .vllm import VLLMInferenceImpl
-
-    impl = VLLMInferenceImpl(config)
-    await impl.initialize()
-    return impl

llama_stack/providers/inline/inference/vllm/config.py

@@ -1,53 +0,0 @@
-# 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 typing import Any
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-@json_schema_type
-class VLLMConfig(BaseModel):
-    """Configuration for the vLLM inference provider.
-
-    Note that the model name is no longer part of this static configuration.
-    You can bind an instance of this provider to a specific model with the
-    ``models.register()`` API call."""
-
-    tensor_parallel_size: int = Field(
-        default=1,
-        description="Number of tensor parallel replicas (number of GPUs to use).",
-    )
-    max_tokens: int = Field(
-        default=4096,
-        description="Maximum number of tokens to generate.",
-    )
-    max_model_len: int = Field(default=4096, description="Maximum context length to use during serving.")
-    max_num_seqs: int = Field(default=4, description="Maximum parallel batch size for generation.")
-    enforce_eager: bool = Field(
-        default=False,
-        description="Whether to use eager mode for inference (otherwise cuda graphs are used).",
-    )
-    gpu_memory_utilization: float = Field(
-        default=0.3,
-        description=(
-            "How much GPU memory will be allocated when this provider has finished "
-            "loading, including memory that was already allocated before loading."
-        ),
-    )
-
-    @classmethod
-    def sample_run_config(cls, **kwargs: Any) -> dict[str, Any]:
-        return {
-            "tensor_parallel_size": "${env.TENSOR_PARALLEL_SIZE:=1}",
-            "max_tokens": "${env.MAX_TOKENS:=4096}",
-            "max_model_len": "${env.MAX_MODEL_LEN:=4096}",
-            "max_num_seqs": "${env.MAX_NUM_SEQS:=4}",
-            "enforce_eager": "${env.ENFORCE_EAGER:=False}",
-            "gpu_memory_utilization": "${env.GPU_MEMORY_UTILIZATION:=0.3}",
-        }

llama_stack/providers/inline/inference/vllm/openai_utils.py

@@ -1,170 +0,0 @@
-# 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 vllm
-
-from llama_stack.apis.inference import (
-    ChatCompletionRequest,
-    GrammarResponseFormat,
-    JsonSchemaResponseFormat,
-    Message,
-    ToolChoice,
-    ToolDefinition,
-    UserMessage,
-)
-from llama_stack.models.llama.datatypes import BuiltinTool
-from llama_stack.providers.utils.inference.openai_compat import (
-    convert_message_to_openai_dict,
-    get_sampling_options,
-)
-
-###############################################################################
-# This file contains OpenAI compatibility code that is currently only used
-# by the inline vLLM connector. Some or all of this code may be moved to a
-# central location at a later date.
-
-
-def _merge_context_into_content(message: Message) -> Message:  # type: ignore
-    """
-    Merge the ``context`` field of a Llama Stack ``Message`` object into
-    the content field for compabilitiy with OpenAI-style APIs.
-
-    Generates a content string that emulates the current behavior
-    of ``llama_models.llama3.api.chat_format.encode_message()``.
-
-    :param message: Message that may include ``context`` field
-
-    :returns: A version of ``message`` with any context merged into the
-     ``content`` field.
-    """
-    if not isinstance(message, UserMessage):  # Separate type check for linter
-        return message
-    if message.context is None:
-        return message
-    return UserMessage(
-        role=message.role,
-        # Emumate llama_models.llama3.api.chat_format.encode_message()
-        content=message.content + "\n\n" + message.context,
-        context=None,
-    )
-
-
-def _llama_stack_tools_to_openai_tools(
-    tools: list[ToolDefinition] | None = None,
-) -> list[vllm.entrypoints.openai.protocol.ChatCompletionToolsParam]:
-    """
-    Convert the list of available tools from Llama Stack's format to vLLM's
-    version of OpenAI's format.
-    """
-    if tools is None:
-        return []
-
-    result = []
-    for t in tools:
-        if isinstance(t.tool_name, BuiltinTool):
-            raise NotImplementedError("Built-in tools not yet implemented")
-        if t.parameters is None:
-            parameters = None
-        else:  # if t.parameters is not None
-            # Convert the "required" flags to a list of required params
-            required_params = [k for k, v in t.parameters.items() if v.required]
-            parameters = {
-                "type": "object",  # Mystery value that shows up in OpenAI docs
-                "properties": {
-                    k: {"type": v.param_type, "description": v.description} for k, v in t.parameters.items()
-                },
-                "required": required_params,
-            }
-
-        function_def = vllm.entrypoints.openai.protocol.FunctionDefinition(
-            name=t.tool_name, description=t.description, parameters=parameters
-        )
-
-        # Every tool definition is double-boxed in a ChatCompletionToolsParam
-        result.append(vllm.entrypoints.openai.protocol.ChatCompletionToolsParam(function=function_def))
-    return result
-
-
-async def llama_stack_chat_completion_to_openai_chat_completion_dict(
-    request: ChatCompletionRequest,
-) -> dict:
-    """
-    Convert a chat completion request in Llama Stack format into an
-    equivalent set of arguments to pass to an OpenAI-compatible
-    chat completions API.
-
-    :param request: Bundled request parameters in Llama Stack format.
-
-    :returns: Dictionary of key-value pairs to use as an initializer
-     for a dataclass or to be converted directly to JSON and sent
-     over the wire.
-    """
-
-    converted_messages = [
-        # This mystery async call makes the parent function also be async
-        await convert_message_to_openai_dict(_merge_context_into_content(m), download=True)
-        for m in request.messages
-    ]
-    converted_tools = _llama_stack_tools_to_openai_tools(request.tools)
-
-    # Llama will try to use built-in tools with no tool catalog, so don't enable
-    # tool choice unless at least one tool is enabled.
-    converted_tool_choice = "none"
-    if (
-        request.tool_config is not None
-        and request.tool_config.tool_choice == ToolChoice.auto
-        and request.tools is not None
-        and len(request.tools) > 0
-    ):
-        converted_tool_choice = "auto"
-
-    # TODO: Figure out what to do with the tool_prompt_format argument.
-    #  Other connectors appear to drop it quietly.
-
-    # Use Llama Stack shared code to translate sampling parameters.
-    sampling_options = get_sampling_options(request.sampling_params)
-
-    # get_sampling_options() translates repetition penalties to an option that
-    # OpenAI's APIs don't know about.
-    # vLLM's OpenAI-compatible API also handles repetition penalties wrong.
-    # For now, translate repetition penalties into a format that vLLM's broken
-    # API will handle correctly. Two wrongs make a right...
-    if "repeat_penalty" in sampling_options:
-        del sampling_options["repeat_penalty"]
-    if request.sampling_params.repetition_penalty is not None and request.sampling_params.repetition_penalty != 1.0:
-        sampling_options["repetition_penalty"] = request.sampling_params.repetition_penalty
-
-    # Convert a single response format into four different parameters, per
-    # the OpenAI spec
-    guided_decoding_options = dict()
-    if request.response_format is None:
-        # Use defaults
-        pass
-    elif isinstance(request.response_format, JsonSchemaResponseFormat):
-        guided_decoding_options["guided_json"] = request.response_format.json_schema
-    elif isinstance(request.response_format, GrammarResponseFormat):
-        guided_decoding_options["guided_grammar"] = request.response_format.bnf
-    else:
-        raise TypeError(f"ResponseFormat object is of unexpected subtype '{type(request.response_format)}'")
-
-    logprob_options = dict()
-    if request.logprobs is not None:
-        logprob_options["logprobs"] = request.logprobs.top_k
-
-    # Marshall together all the arguments for a ChatCompletionRequest
-    request_options = {
-        "model": request.model,
-        "messages": converted_messages,
-        "tools": converted_tools,
-        "tool_choice": converted_tool_choice,
-        "stream": request.stream,
-        **sampling_options,
-        **guided_decoding_options,
-        **logprob_options,
-    }
-
-    return request_options

llama_stack/providers/inline/inference/vllm/vllm.py

@@ -1,811 +0,0 @@
-# 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 json
-import re
-import uuid
-from collections.abc import AsyncGenerator, AsyncIterator
-
-# These vLLM modules contain names that overlap with Llama Stack names, so we import
-# fully-qualified names
-import vllm.entrypoints.openai.protocol
-import vllm.sampling_params
-from vllm.engine.arg_utils import AsyncEngineArgs
-from vllm.engine.async_llm_engine import AsyncLLMEngine
-from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
-from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels
-
-from llama_stack.apis.common.content_types import (
-    InterleavedContent,
-    InterleavedContentItem,
-    TextDelta,
-    ToolCallDelta,
-)
-from llama_stack.apis.inference import (
-    ChatCompletionRequest,
-    ChatCompletionResponse,
-    ChatCompletionResponseEvent,
-    ChatCompletionResponseEventType,
-    ChatCompletionResponseStreamChunk,
-    CompletionMessage,
-    CompletionResponse,
-    CompletionResponseStreamChunk,
-    EmbeddingsResponse,
-    EmbeddingTaskType,
-    GrammarResponseFormat,
-    Inference,
-    JsonSchemaResponseFormat,
-    LogProbConfig,
-    Message,
-    OpenAIEmbeddingsResponse,
-    ResponseFormat,
-    SamplingParams,
-    TextTruncation,
-    TokenLogProbs,
-    ToolChoice,
-    ToolConfig,
-    TopKSamplingStrategy,
-    TopPSamplingStrategy,
-)
-from llama_stack.apis.models import Model
-from llama_stack.log import get_logger
-from llama_stack.models.llama import sku_list
-from llama_stack.models.llama.datatypes import (
-    StopReason,
-    ToolCall,
-    ToolDefinition,
-    ToolPromptFormat,
-)
-from llama_stack.models.llama.llama3.chat_format import ChatFormat
-from llama_stack.models.llama.llama3.tokenizer import Tokenizer
-from llama_stack.providers.remote.inference.vllm.vllm import build_hf_repo_model_entries
-from llama_stack.providers.utils.inference.model_registry import (
-    ModelRegistryHelper,
-    ModelsProtocolPrivate,
-)
-from llama_stack.providers.utils.inference.openai_compat import (
-    OpenAIChatCompletionToLlamaStackMixin,
-    OpenAICompatCompletionChoice,
-    OpenAICompatCompletionResponse,
-    OpenAICompletionToLlamaStackMixin,
-    get_stop_reason,
-    process_chat_completion_stream_response,
-)
-from llama_stack.providers.utils.inference.prompt_adapter import (
-    chat_completion_request_to_prompt,
-)
-
-from .config import VLLMConfig
-from .openai_utils import llama_stack_chat_completion_to_openai_chat_completion_dict
-
-# Map from Hugging Face model architecture name to appropriate tool parser.
-# See vllm.entrypoints.openai.tool_parsers.ToolParserManager.tool_parsers for the full list of
-# available parsers.
-# TODO: Expand this list
-CONFIG_TYPE_TO_TOOL_PARSER = {
-    "GraniteConfig": "granite",
-    "MllamaConfig": "llama3_json",
-    "LlamaConfig": "llama3_json",
-}
-DEFAULT_TOOL_PARSER = "pythonic"
-
-
-logger = get_logger(__name__, category="inference")
-
-
-def _random_uuid_str() -> str:
-    return str(uuid.uuid4().hex)
-
-
-def _response_format_to_guided_decoding_params(
-    response_format: ResponseFormat | None,  # type: ignore
-) -> vllm.sampling_params.GuidedDecodingParams:
-    """
-    Translate constrained decoding parameters from Llama Stack's format to vLLM's format.
-
-    :param response_format: Llama Stack version of constrained decoding info. Can be ``None``,
-     indicating no constraints.
-    :returns: The equivalent dataclass object for the low-level inference layer of vLLM.
-    """
-    if response_format is None:
-        # As of vLLM 0.6.3, the default constructor for GuidedDecodingParams() returns an invalid
-        # value that crashes the executor on some code paths. Use ``None`` instead.
-        return None
-
-    # Llama Stack currently implements fewer types of constrained decoding than vLLM does.
-    # Translate the types that exist and detect if Llama Stack adds new ones.
-    if isinstance(response_format, JsonSchemaResponseFormat):
-        return vllm.sampling_params.GuidedDecodingParams(json=response_format.json_schema)
-    elif isinstance(response_format, GrammarResponseFormat):
-        # BNF grammar.
-        # Llama Stack uses the parse tree of the grammar, while vLLM uses the string
-        # representation of the grammar.
-        raise TypeError(
-            "Constrained decoding with BNF grammars is not currently implemented, because the "
-            "reference implementation does not implement it."
-        )
-    else:
-        raise TypeError(f"ResponseFormat object is of unexpected subtype '{type(response_format)}'")
-
-
-def _convert_sampling_params(
-    sampling_params: SamplingParams | None,
-    response_format: ResponseFormat | None,  # type: ignore
-    log_prob_config: LogProbConfig | None,
-) -> vllm.SamplingParams:
-    """Convert sampling and constrained decoding configuration from Llama Stack's format to vLLM's
-    format."""
-    # In the absence of provided config values, use Llama Stack defaults as encoded in the Llama
-    # Stack dataclasses. These defaults are different from vLLM's defaults.
-    if sampling_params is None:
-        sampling_params = SamplingParams()
-    if log_prob_config is None:
-        log_prob_config = LogProbConfig()
-
-    if isinstance(sampling_params.strategy, TopKSamplingStrategy):
-        if sampling_params.strategy.top_k == 0:
-            # vLLM treats "k" differently for top-k sampling
-            vllm_top_k = -1
-        else:
-            vllm_top_k = sampling_params.strategy.top_k
-    else:
-        vllm_top_k = -1
-
-    if isinstance(sampling_params.strategy, TopPSamplingStrategy):
-        vllm_top_p = sampling_params.strategy.top_p
-        # Llama Stack only allows temperature with top-P.
-        vllm_temperature = sampling_params.strategy.temperature
-    else:
-        vllm_top_p = 1.0
-        vllm_temperature = 0.0
-
-    # vLLM allows top-p and top-k at the same time.
-    vllm_sampling_params = vllm.SamplingParams.from_optional(
-        max_tokens=(None if sampling_params.max_tokens == 0 else sampling_params.max_tokens),
-        temperature=vllm_temperature,
-        top_p=vllm_top_p,
-        top_k=vllm_top_k,
-        repetition_penalty=sampling_params.repetition_penalty,
-        guided_decoding=_response_format_to_guided_decoding_params(response_format),
-        logprobs=log_prob_config.top_k,
-    )
-    return vllm_sampling_params
-
-
-class VLLMInferenceImpl(
-    Inference,
-    OpenAIChatCompletionToLlamaStackMixin,
-    OpenAICompletionToLlamaStackMixin,
-    ModelsProtocolPrivate,
-):
-    """
-    vLLM-based inference model adapter for Llama Stack with support for multiple models.
-
-    Requires the configuration parameters documented in the :class:`VllmConfig2` class.
-    """
-
-    config: VLLMConfig
-    register_helper: ModelRegistryHelper
-    model_ids: set[str]
-    resolved_model_id: str | None
-    engine: AsyncLLMEngine | None
-    chat: OpenAIServingChat | None
-    is_meta_llama_model: bool
-
-    def __init__(self, config: VLLMConfig):
-        self.config = config
-        logger.info(f"Config is: {self.config}")
-
-        self.register_helper = ModelRegistryHelper(build_hf_repo_model_entries())
-        self.formatter = ChatFormat(Tokenizer.get_instance())
-
-        # The following are initialized when paths are bound to this provider
-        self.resolved_model_id = None
-        self.model_ids = set()
-        self.engine = None
-        self.chat = None
-        self.is_meta_llama_model = False
-
-    ###########################################################################
-    # METHODS INHERITED FROM IMPLICIT BASE CLASS.
-    # TODO: Make this class inherit from the new base class ProviderBase once that class exists.
-
-    async def initialize(self) -> None:
-        """
-        Callback that is invoked through many levels of indirection during provider class
-        instantiation, sometime after when __init__() is called and before any model registration
-        methods or methods connected to a REST API are called.
-
-        It's not clear what assumptions the class can make about the platform's initialization
-        state here that can't be made during __init__(), and vLLM can't be started until we know
-        what model it's supposed to be serving, so nothing happens here currently.
-        """
-        pass
-
-    async def shutdown(self) -> None:
-        logger.info(f"Shutting down inline vLLM inference provider {self}.")
-        if self.engine is not None:
-            self.engine.shutdown_background_loop()
-            self.engine = None
-            self.chat = None
-            self.model_ids = set()
-            self.resolved_model_id = None
-
-    ###########################################################################
-    # METHODS INHERITED FROM ModelsProtocolPrivate INTERFACE
-
-    # Note that the return type of the superclass method is WRONG
-    async def register_model(self, model: Model) -> Model:
-        """
-        Callback that is called when the server associates an inference endpoint with an
-        inference provider.
-
-        :param model: Object that encapsulates parameters necessary for identifying a specific
-         LLM.
-
-        :returns: The input ``Model`` object. It may or may not be permissible to change fields
-         before returning this object.
-        """
-        logger.debug(f"In register_model({model})")
-
-        # First attempt to interpret the model coordinates as a Llama model name
-        resolved_llama_model = sku_list.resolve_model(model.provider_model_id)
-        if resolved_llama_model is not None:
-            # Load from Hugging Face repo into default local cache dir
-            model_id_for_vllm = resolved_llama_model.huggingface_repo
-
-            # Detect a genuine Meta Llama model to trigger Meta-specific preprocessing.
-            # Don't set self.is_meta_llama_model until we actually load the model.
-            is_meta_llama_model = True
-        else:  # if resolved_llama_model is None
-            # Not a Llama model name. Pass the model id through to vLLM's loader
-            model_id_for_vllm = model.provider_model_id
-            is_meta_llama_model = False
-
-        if self.resolved_model_id is not None:
-            if model_id_for_vllm != self.resolved_model_id:
-                raise ValueError(
-                    f"Attempted to serve two LLMs (ids '{self.resolved_model_id}') and "
-                    f"'{model_id_for_vllm}') from one copy of provider '{self}'. Use multiple "
-                    f"copies of the provider instead."
-                )
-            else:
-                # Model already loaded
-                logger.info(
-                    f"Requested id {model} resolves to {model_id_for_vllm}, which is already loaded. Continuing."
-                )
-                self.model_ids.add(model.model_id)
-                return model
-
-        logger.info(f"Requested id {model} resolves to {model_id_for_vllm}. Loading {model_id_for_vllm}.")
-        if is_meta_llama_model:
-            logger.info(f"Model {model_id_for_vllm} is a Meta Llama model.")
-        self.is_meta_llama_model = is_meta_llama_model
-
-        # If we get here, this is the first time registering a model.
-        # Preload so that the first inference request won't time out.
-        engine_args = AsyncEngineArgs(
-            model=model_id_for_vllm,
-            tokenizer=model_id_for_vllm,
-            tensor_parallel_size=self.config.tensor_parallel_size,
-            enforce_eager=self.config.enforce_eager,
-            gpu_memory_utilization=self.config.gpu_memory_utilization,
-            max_num_seqs=self.config.max_num_seqs,
-            max_model_len=self.config.max_model_len,
-        )
-        self.engine = AsyncLLMEngine.from_engine_args(engine_args)
-
-        # vLLM currently requires the user to specify the tool parser manually. To choose a tool
-        # parser, we need to determine what model architecture is being used. For now, we infer
-        # that information from what config class the model uses.
-        low_level_model_config = self.engine.engine.get_model_config()
-        hf_config = low_level_model_config.hf_config
-        hf_config_class_name = hf_config.__class__.__name__
-        if hf_config_class_name in CONFIG_TYPE_TO_TOOL_PARSER:
-            tool_parser = CONFIG_TYPE_TO_TOOL_PARSER[hf_config_class_name]
-        else:
-            # No info -- choose a default so we can at least attempt tool
-            # use.
-            tool_parser = DEFAULT_TOOL_PARSER
-        logger.debug(f"{hf_config_class_name=}")
-        logger.debug(f"{tool_parser=}")
-
-        # Wrap the lower-level engine in an OpenAI-compatible chat API
-        model_config = await self.engine.get_model_config()
-        self.chat = OpenAIServingChat(
-            engine_client=self.engine,
-            model_config=model_config,
-            models=OpenAIServingModels(
-                engine_client=self.engine,
-                model_config=model_config,
-                base_model_paths=[
-                    # The layer below us will only see resolved model IDs
-                    BaseModelPath(model_id_for_vllm, model_id_for_vllm)
-                ],
-            ),
-            response_role="assistant",
-            request_logger=None,  # Use default logging
-            chat_template=None,  # Use default template from model checkpoint
-            enable_auto_tools=True,
-            tool_parser=tool_parser,
-            chat_template_content_format="auto",
-        )
-        self.resolved_model_id = model_id_for_vllm
-        self.model_ids.add(model.model_id)
-
-        logger.info(f"Finished preloading model: {model_id_for_vllm}")
-
-        return model
-
-    async def unregister_model(self, model_id: str) -> None:
-        """
-        Callback that is called when the server removes an inference endpoint from an inference
-        provider.
-
-        :param model_id: The same external ID that the higher layers of the stack previously passed
-        to :func:`register_model()`
-        """
-        if model_id not in self.model_ids:
-            raise ValueError(
-                f"Attempted to unregister model ID '{model_id}', but that ID is not registered to this provider."
-            )
-        self.model_ids.remove(model_id)
-
-        if len(self.model_ids) == 0:
-            # Last model was just unregistered. Shut down the connection to vLLM and free up
-            # resources.
-            # Note that this operation may cause in-flight chat completion requests on the
-            # now-unregistered model to return errors.
-            self.resolved_model_id = None
-            self.chat = None
-            self.engine.shutdown_background_loop()
-            self.engine = None
-
-    ###########################################################################
-    # METHODS INHERITED FROM Inference INTERFACE
-
-    async def completion(
-        self,
-        model_id: str,
-        content: InterleavedContent,
-        sampling_params: SamplingParams | None = None,
-        response_format: ResponseFormat | None = None,
-        stream: bool | None = False,
-        logprobs: LogProbConfig | None = None,
-    ) -> CompletionResponse | AsyncIterator[CompletionResponseStreamChunk]:
-        if model_id not in self.model_ids:
-            raise ValueError(
-                f"This adapter is not registered to model id '{model_id}'. Registered IDs are: {self.model_ids}"
-            )
-        if not isinstance(content, str):
-            raise NotImplementedError("Multimodal input not currently supported")
-        if sampling_params is None:
-            sampling_params = SamplingParams()
-
-        converted_sampling_params = _convert_sampling_params(sampling_params, response_format, logprobs)
-
-        logger.debug(f"{converted_sampling_params=}")
-
-        if stream:
-            return self._streaming_completion(content, converted_sampling_params)
-        else:
-            streaming_result = None
-            async for _ in self._streaming_completion(content, converted_sampling_params):
-                pass
-            return CompletionResponse(
-                content=streaming_result.delta,
-                stop_reason=streaming_result.stop_reason,
-                logprobs=streaming_result.logprobs,
-            )
-
-    async def embeddings(
-        self,
-        model_id: str,
-        contents: list[str] | list[InterleavedContentItem],
-        text_truncation: TextTruncation | None = TextTruncation.none,
-        output_dimension: int | None = None,
-        task_type: EmbeddingTaskType | None = None,
-    ) -> EmbeddingsResponse:
-        raise NotImplementedError()
-
-    async def openai_embeddings(
-        self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
-    ) -> OpenAIEmbeddingsResponse:
-        raise NotImplementedError()
-
-    async def chat_completion(
-        self,
-        model_id: str,
-        messages: list[Message],  # type: ignore
-        sampling_params: SamplingParams | None = None,
-        response_format: ResponseFormat | None = None,  # type: ignore
-        tools: list[ToolDefinition] | None = None,
-        tool_choice: ToolChoice | None = ToolChoice.auto,
-        tool_prompt_format: ToolPromptFormat | None = None,
-        stream: bool | None = False,
-        logprobs: LogProbConfig | None = None,
-        tool_config: ToolConfig | None = None,
-    ) -> ChatCompletionResponse | ChatCompletionResponseStreamChunk:
-        sampling_params = sampling_params or SamplingParams()
-        if model_id not in self.model_ids:
-            raise ValueError(
-                f"This adapter is not registered to model id '{model_id}'. Registered IDs are: {self.model_ids}"
-            )
-
-        # Convert to Llama Stack internal format for consistency
-        request = ChatCompletionRequest(
-            model=self.resolved_model_id,
-            messages=messages,
-            sampling_params=sampling_params,
-            response_format=response_format,
-            tools=tools,
-            tool_choice=tool_choice,
-            tool_prompt_format=tool_prompt_format,
-            stream=stream,
-            logprobs=logprobs,
-        )
-
-        if self.is_meta_llama_model:
-            # Bypass vLLM chat templating layer for Meta Llama models, because the
-            # templating layer in Llama Stack currently produces better results.
-            logger.debug(
-                f"Routing {self.resolved_model_id} chat completion through "
-                f"Llama Stack's templating layer instead of vLLM's."
-            )
-            return await self._chat_completion_for_meta_llama(request)
-
-        logger.debug(f"{self.resolved_model_id} is not a Meta Llama model")
-
-        # Arguments to the vLLM call must be packaged as a ChatCompletionRequest dataclass.
-        # Note that this dataclass has the same name as a similar dataclass in Llama Stack.
-        request_options = await llama_stack_chat_completion_to_openai_chat_completion_dict(request)
-        chat_completion_request = vllm.entrypoints.openai.protocol.ChatCompletionRequest(**request_options)
-
-        logger.debug(f"Converted request: {chat_completion_request}")
-
-        vllm_result = await self.chat.create_chat_completion(chat_completion_request)
-        logger.debug(f"Result from vLLM: {vllm_result}")
-        if isinstance(vllm_result, vllm.entrypoints.openai.protocol.ErrorResponse):
-            raise ValueError(f"Error from vLLM layer: {vllm_result}")
-
-        # Return type depends on "stream" argument
-        if stream:
-            if not isinstance(vllm_result, AsyncGenerator):
-                raise TypeError(f"Unexpected result type {type(vllm_result)} for streaming inference call")
-            # vLLM client returns a stream of strings, which need to be parsed.
-            # Stream comes in the form of an async generator.
-            return self._convert_streaming_results(vllm_result)
-        else:
-            if not isinstance(vllm_result, vllm.entrypoints.openai.protocol.ChatCompletionResponse):
-                raise TypeError(f"Unexpected result type {type(vllm_result)} for non-streaming inference call")
-            return self._convert_non_streaming_results(vllm_result)
-
-    ###########################################################################
-    # INTERNAL METHODS
-
-    async def _streaming_completion(
-        self, content: str, sampling_params: vllm.SamplingParams
-    ) -> AsyncIterator[CompletionResponseStreamChunk]:
-        """Internal implementation of :func:`completion()` API for the streaming case. Assumes
-        that arguments have been validated upstream.
-
-        :param content: Must be a string
-        :param sampling_params: Paramters from  public API's ``response_format``
-         and ``sampling_params`` arguments, converted to VLLM format
-        """
-        # We run agains the vLLM generate() call directly instead of using the OpenAI-compatible
-        # layer, because doing so simplifies the code here.
-
-        # The vLLM engine requires a unique identifier for each call to generate()
-        request_id = _random_uuid_str()
-
-        # The vLLM generate() API is streaming-only and returns an async generator.
-        # The generator returns objects of type vllm.RequestOutput.
-        results_generator = self.engine.generate(content, sampling_params, request_id)
-
-        # Need to know the model's EOS token ID for the conversion code below.
-        # AsyncLLMEngine is a wrapper around LLMEngine, and the tokenizer is only available if
-        # we drill down to the LLMEngine inside the AsyncLLMEngine.
-        # Similarly, the tokenizer in an LLMEngine is a wrapper around a BaseTokenizerGroup,
-        # and we need to drill down to the Hugging Face tokenizer inside the BaseTokenizerGroup.
-        llm_engine = self.engine.engine
-        tokenizer_group = llm_engine.tokenizer
-        eos_token_id = tokenizer_group.tokenizer.eos_token_id
-
-        request_output: vllm.RequestOutput = None
-        async for request_output in results_generator:
-            # Check for weird inference failures
-            if request_output.outputs is None or len(request_output.outputs) == 0:
-                # This case also should never happen
-                raise ValueError("Inference produced empty result")
-
-            # If we get here, then request_output contains the final output of the generate() call.
-            # The result may include multiple alternate outputs, but Llama Stack APIs only allow
-            # us to return one.
-            output: vllm.CompletionOutput = request_output.outputs[0]
-            completion_string = output.text
-
-            # Convert logprobs from vLLM's format to Llama Stack's format
-            logprobs = [
-                TokenLogProbs(logprobs_by_token={v.decoded_token: v.logprob for _, v in logprob_dict.items()})
-                for logprob_dict in output.logprobs
-            ]
-
-            # The final output chunk should be labeled with the reason that the overall generate()
-            # call completed.
-            logger.debug(f"{output.stop_reason=}; {type(output.stop_reason)=}")
-            if output.stop_reason is None:
-                stop_reason = None  # Still going
-            elif output.stop_reason == "stop":
-                stop_reason = StopReason.end_of_turn
-            elif output.stop_reason == "length":
-                stop_reason = StopReason.out_of_tokens
-            elif isinstance(output.stop_reason, int):
-                # If the model config specifies multiple end-of-sequence tokens, then vLLM
-                # will return the token ID of the EOS token in the stop_reason field.
-                stop_reason = StopReason.end_of_turn
-            else:
-                raise ValueError(f"Unrecognized stop reason '{output.stop_reason}'")
-
-            # vLLM's protocol outputs the stop token, then sets end of message on the next step for
-            # some reason.
-            if request_output.outputs[-1].token_ids[-1] == eos_token_id:
-                stop_reason = StopReason.end_of_message
-
-            yield CompletionResponseStreamChunk(delta=completion_string, stop_reason=stop_reason, logprobs=logprobs)
-
-        # Llama Stack requires that the last chunk have a stop reason, but vLLM doesn't always
-        # provide one if it runs out of tokens.
-        if stop_reason is None:
-            yield CompletionResponseStreamChunk(
-                delta=completion_string,
-                stop_reason=StopReason.out_of_tokens,
-                logprobs=logprobs,
-            )
-
-    def _convert_non_streaming_results(
-        self, vllm_result: vllm.entrypoints.openai.protocol.ChatCompletionResponse
-    ) -> ChatCompletionResponse:
-        """
-        Subroutine to convert the non-streaming output of vLLM's OpenAI-compatible API into an
-        equivalent Llama Stack object.
-
-        The result from vLLM's non-streaming API is a dataclass with the same name as the Llama
-        Stack ChatCompletionResponse dataclass, but with more and different field names. We ignore
-        the fields that aren't currently present in the Llama Stack dataclass.
-        """
-
-        # There may be multiple responses, but we can only pass through the first one.
-        if len(vllm_result.choices) == 0:
-            raise ValueError("Don't know how to convert response object without any responses")
-        vllm_message = vllm_result.choices[0].message
-        vllm_finish_reason = vllm_result.choices[0].finish_reason
-
-        converted_message = CompletionMessage(
-            role=vllm_message.role,
-            # Llama Stack API won't accept None for content field.
-            content=("" if vllm_message.content is None else vllm_message.content),
-            stop_reason=get_stop_reason(vllm_finish_reason),
-            tool_calls=[
-                ToolCall(
-                    call_id=t.id,
-                    tool_name=t.function.name,
-                    # vLLM function args come back as a string. Llama Stack expects JSON.
-                    arguments=json.loads(t.function.arguments),
-                    arguments_json=t.function.arguments,
-                )
-                for t in vllm_message.tool_calls
-            ],
-        )
-
-        # TODO: Convert logprobs
-
-        logger.debug(f"Converted message: {converted_message}")
-
-        return ChatCompletionResponse(
-            completion_message=converted_message,
-        )
-
-    async def _chat_completion_for_meta_llama(
-        self, request: ChatCompletionRequest
-    ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionResponseStreamChunk]:
-        """
-        Subroutine that routes chat completions for Meta Llama models through Llama Stack's
-        chat template instead of using vLLM's version of that template. The Llama Stack version
-        of the chat template currently produces more reliable outputs.
-
-        Once vLLM's support for Meta Llama models has matured more, we should consider routing
-        Meta Llama requests through the vLLM chat completions API instead of using this method.
-        """
-        formatter = ChatFormat(Tokenizer.get_instance())
-
-        # Note that this function call modifies `request` in place.
-        prompt = await chat_completion_request_to_prompt(request, self.resolved_model_id)
-
-        model_id = list(self.model_ids)[0]  # Any model ID will do here
-        completion_response_or_iterator = await self.completion(
-            model_id=model_id,
-            content=prompt,
-            sampling_params=request.sampling_params,
-            response_format=request.response_format,
-            stream=request.stream,
-            logprobs=request.logprobs,
-        )
-
-        if request.stream:
-            if not isinstance(completion_response_or_iterator, AsyncIterator):
-                raise TypeError(
-                    f"Received unexpected result type {type(completion_response_or_iterator)}for streaming request."
-                )
-            return self._chat_completion_for_meta_llama_streaming(completion_response_or_iterator, request)
-
-        # elsif not request.stream:
-        if not isinstance(completion_response_or_iterator, CompletionResponse):
-            raise TypeError(
-                f"Received unexpected result type {type(completion_response_or_iterator)}for non-streaming request."
-            )
-        completion_response: CompletionResponse = completion_response_or_iterator
-        raw_message = formatter.decode_assistant_message_from_content(
-            completion_response.content, completion_response.stop_reason
-        )
-        return ChatCompletionResponse(
-            completion_message=CompletionMessage(
-                content=raw_message.content,
-                stop_reason=raw_message.stop_reason,
-                tool_calls=raw_message.tool_calls,
-            ),
-            logprobs=completion_response.logprobs,
-        )
-
-    async def _chat_completion_for_meta_llama_streaming(
-        self, results_iterator: AsyncIterator, request: ChatCompletionRequest
-    ) -> AsyncIterator:
-        """
-        Code from :func:`_chat_completion_for_meta_llama()` that needs to be a separate
-        method to keep asyncio happy.
-        """
-
-        # Convert to OpenAI format, then use shared code to convert to Llama Stack format.
-        async def _generate_and_convert_to_openai_compat():
-            chunk: CompletionResponseStreamChunk  # Make Pylance happy
-            last_text_len = 0
-            async for chunk in results_iterator:
-                if chunk.stop_reason == StopReason.end_of_turn:
-                    finish_reason = "stop"
-                elif chunk.stop_reason == StopReason.end_of_message:
-                    finish_reason = "eos"
-                elif chunk.stop_reason == StopReason.out_of_tokens:
-                    finish_reason = "length"
-                else:
-                    finish_reason = None
-
-                # Convert delta back to an actual delta
-                text_delta = chunk.delta[last_text_len:]
-                last_text_len = len(chunk.delta)
-
-                logger.debug(f"{text_delta=}; {finish_reason=}")
-
-                yield OpenAICompatCompletionResponse(
-                    choices=[OpenAICompatCompletionChoice(finish_reason=finish_reason, text=text_delta)]
-                )
-
-        stream = _generate_and_convert_to_openai_compat()
-        async for chunk in process_chat_completion_stream_response(stream, request):
-            logger.debug(f"Returning chunk: {chunk}")
-            yield chunk
-
-    async def _convert_streaming_results(self, vllm_result: AsyncIterator) -> AsyncIterator:
-        """
-        Subroutine that wraps the streaming outputs of vLLM's OpenAI-compatible
-        API into a second async iterator that returns Llama Stack objects.
-
-        :param vllm_result: Stream of strings that need to be parsed
-        """
-        # Tool calls come in pieces, but Llama Stack expects them in bigger chunks. We build up
-        # those chunks and output them at the end.
-        # This data structure holds the current set of partial tool calls.
-        index_to_tool_call: dict[int, dict] = dict()
-
-        # The Llama Stack event stream must always start with a start event. Use an empty one to
-        # simplify logic below
-        yield ChatCompletionResponseStreamChunk(
-            event=ChatCompletionResponseEvent(
-                event_type=ChatCompletionResponseEventType.start,
-                delta=TextDelta(text=""),
-                stop_reason=None,
-            )
-        )
-
-        converted_stop_reason = None
-        async for chunk_str in vllm_result:
-            # Due to OpenAI compatibility, each event in the stream will start with "data: " and
-            # end with "\n\n".
-            _prefix = "data: "
-            _suffix = "\n\n"
-            if not chunk_str.startswith(_prefix) or not chunk_str.endswith(_suffix):
-                raise ValueError(f"Can't parse result string from vLLM: '{re.escape(chunk_str)}'")
-
-            # In between the "data: " and newlines is an event record
-            data_str = chunk_str[len(_prefix) : -len(_suffix)]
-
-            # The end of the stream is indicated with "[DONE]"
-            if data_str == "[DONE]":
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.complete,
-                        delta=TextDelta(text=""),
-                        stop_reason=converted_stop_reason,
-                    )
-                )
-                return
-
-            # Anything that is not "[DONE]" should be a JSON record
-            parsed_chunk = json.loads(data_str)
-
-            logger.debug(f"Parsed JSON event to:\n{json.dumps(parsed_chunk, indent=2)}")
-
-            # The result may contain multiple completions, but Llama Stack APIs only support
-            # returning one.
-            first_choice = parsed_chunk["choices"][0]
-            converted_stop_reason = get_stop_reason(first_choice["finish_reason"])
-            delta_record = first_choice["delta"]
-
-            if "content" in delta_record:
-                # Text delta
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=TextDelta(text=delta_record["content"]),
-                        stop_reason=converted_stop_reason,
-                    )
-                )
-            elif "tool_calls" in delta_record:
-                # Tool call(s). Llama Stack APIs do not have a clear way to return partial tool
-                # calls, so buffer until we get a "tool calls" stop reason
-                for tc in delta_record["tool_calls"]:
-                    index = tc["index"]
-                    if index not in index_to_tool_call:
-                        # First time this tool call is showing up
-                        index_to_tool_call[index] = dict()
-                    tool_call = index_to_tool_call[index]
-                    if "id" in tc:
-                        tool_call["call_id"] = tc["id"]
-                    if "function" in tc:
-                        if "name" in tc["function"]:
-                            tool_call["tool_name"] = tc["function"]["name"]
-                        if "arguments" in tc["function"]:
-                            # Arguments comes in as pieces of a string
-                            if "arguments_str" not in tool_call:
-                                tool_call["arguments_str"] = ""
-                            tool_call["arguments_str"] += tc["function"]["arguments"]
-            else:
-                raise ValueError(f"Don't know how to parse event delta: {delta_record}")
-
-            if first_choice["finish_reason"] == "tool_calls":
-                # Special OpenAI code for "tool calls complete".
-                # Output the buffered tool calls. Llama Stack requires a separate event per tool
-                # call.
-                for tool_call_record in index_to_tool_call.values():
-                    # Arguments come in as a string. Parse the completed string.
-                    tool_call_record["arguments"] = json.loads(tool_call_record["arguments_str"])
-                    del tool_call_record["arguments_str"]
-
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=ToolCallDelta(tool_call=tool_call_record, parse_status="succeeded"),
-                            stop_reason=converted_stop_reason,
-                        )
-                    )
-
-        # If we get here, we've lost the connection with the vLLM event stream before it ended
-        # normally.
-        raise ValueError("vLLM event stream ended without [DONE] message.")

llama_stack/providers/inline/post_training/huggingface/recipes/finetune_single_device.py

@@ -319,7 +319,7 @@ class HFFinetuningSingleDevice:
             use_cpu=True if device.type == "cpu" and not torch.backends.mps.is_available() else False,
             save_strategy=save_strategy,
             report_to="none",
-            max_seq_length=provider_config.max_seq_length,
+            max_length=provider_config.max_seq_length,
             gradient_accumulation_steps=config.gradient_accumulation_steps,
             gradient_checkpointing=provider_config.gradient_checkpointing,
             learning_rate=lr,

llama_stack/providers/inline/safety/llama_guard/llama_guard.py

@@ -146,9 +146,9 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
         pass
 
     async def register_shield(self, shield: Shield) -> None:
-        # Allow any model to be registered as a shield
-        # The model will be validated during runtime when making inference calls
-        pass
+        model_id = shield.provider_resource_id
+        if not model_id:
+            raise ValueError("Llama Guard shield must have a model id")
 
     async def run_shield(
         self,

llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py

@@ -11,19 +11,9 @@ from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanProcessor
 from opentelemetry.trace.status import StatusCode
 
-# Colors for console output
-COLORS = {
-    "reset": "\033[0m",
-    "bold": "\033[1m",
-    "dim": "\033[2m",
-    "red": "\033[31m",
-    "green": "\033[32m",
-    "yellow": "\033[33m",
-    "blue": "\033[34m",
-    "magenta": "\033[35m",
-    "cyan": "\033[36m",
-    "white": "\033[37m",
-}
+from llama_stack.log import get_logger
+
+logger = get_logger(name="console_span_processor", category="telemetry")
 
 
 class ConsoleSpanProcessor(SpanProcessor):
@@ -35,34 +25,21 @@ class ConsoleSpanProcessor(SpanProcessor):
             return
 
         timestamp = datetime.fromtimestamp(span.start_time / 1e9, tz=UTC).strftime("%H:%M:%S.%f")[:-3]
-
-        print(
-            f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
-            f"{COLORS['magenta']}[START]{COLORS['reset']} "
-            f"{COLORS['dim']}{span.name}{COLORS['reset']}"
-        )
+        logger.info(f"[dim]{timestamp}[/dim] [bold magenta][START][/bold magenta] [dim]{span.name}[/dim]")
 
     def on_end(self, span: ReadableSpan) -> None:
         if span.attributes and span.attributes.get("__autotraced__"):
             return
 
         timestamp = datetime.fromtimestamp(span.end_time / 1e9, tz=UTC).strftime("%H:%M:%S.%f")[:-3]
-
-        span_context = (
-            f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
-            f"{COLORS['magenta']}[END]{COLORS['reset']} "
-            f"{COLORS['dim']}{span.name}{COLORS['reset']}"
-        )
-
+        span_context = f"[dim]{timestamp}[/dim] [bold magenta][END][/bold magenta] [dim]{span.name}[/dim]"
         if span.status.status_code == StatusCode.ERROR:
-            span_context += f"{COLORS['reset']} {COLORS['red']}[ERROR]{COLORS['reset']}"
+            span_context += " [bold red][ERROR][/bold red]"
         elif span.status.status_code != StatusCode.UNSET:
-            span_context += f"{COLORS['reset']} [{span.status.status_code}]"
-
+            span_context += f" [{span.status.status_code}]"
         duration_ms = (span.end_time - span.start_time) / 1e6
-        span_context += f"{COLORS['reset']} ({duration_ms:.2f}ms)"
-
-        print(span_context)
+        span_context += f" ({duration_ms:.2f}ms)"
+        logger.info(span_context)
 
         if self.print_attributes and span.attributes:
             for key, value in span.attributes.items():
@@ -71,31 +48,26 @@ class ConsoleSpanProcessor(SpanProcessor):
                 str_value = str(value)
                 if len(str_value) > 1000:
                     str_value = str_value[:997] + "..."
-                print(f"    {COLORS['dim']}{key}: {str_value}{COLORS['reset']}")
+                logger.info(f"    [dim]{key}[/dim]: {str_value}")
 
         for event in span.events:
             event_time = datetime.fromtimestamp(event.timestamp / 1e9, tz=UTC).strftime("%H:%M:%S.%f")[:-3]
-
             severity = event.attributes.get("severity", "info")
             message = event.attributes.get("message", event.name)
-            if isinstance(message, dict | list):
+            if isinstance(message, dict) or isinstance(message, list):
                 message = json.dumps(message, indent=2)
-
-            severity_colors = {
-                "error": f"{COLORS['bold']}{COLORS['red']}",
-                "warn": f"{COLORS['bold']}{COLORS['yellow']}",
-                "info": COLORS["white"],
-                "debug": COLORS["dim"],
-            }
-            msg_color = severity_colors.get(severity, COLORS["white"])
-
-            print(f" {event_time} {msg_color}[{severity.upper()}] {message}{COLORS['reset']}")
-
+            severity_color = {
+                "error": "red",
+                "warn": "yellow",
+                "info": "white",
+                "debug": "dim",
+            }.get(severity, "white")
+            logger.info(f" {event_time} [bold {severity_color}][{severity.upper()}][/bold {severity_color}] {message}")
             if event.attributes:
                 for key, value in event.attributes.items():
                     if key.startswith("__") or key in ["message", "severity"]:
                         continue
-                    print(f"   {COLORS['dim']}{key}: {value}{COLORS['reset']}")
+                    logger.info(f"/r[dim]{key}[/dim]: {value}")
 
     def shutdown(self) -> None:
         """Shutdown the processor."""

llama_stack/providers/inline/vector_io/chroma/__init__.py

@@ -16,6 +16,6 @@ async def get_provider_impl(config: ChromaVectorIOConfig, deps: dict[Api, Any]):
         ChromaVectorIOAdapter,
     )
 
-    impl = ChromaVectorIOAdapter(config, deps[Api.inference])
+    impl = ChromaVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
     await impl.initialize()
     return impl

llama_stack/providers/inline/vector_io/chroma/config.py

@@ -6,12 +6,25 @@
 
 from typing import Any
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+
+from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig
+from llama_stack.schema_utils import json_schema_type
 
 
+@json_schema_type
 class ChromaVectorIOConfig(BaseModel):
     db_path: str
+    kvstore: KVStoreConfig = Field(description="Config for KV store backend")
 
     @classmethod
-    def sample_run_config(cls, db_path: str = "${env.CHROMADB_PATH}", **kwargs: Any) -> dict[str, Any]:
-        return {"db_path": db_path}
+    def sample_run_config(
+        cls, __distro_dir__: str, db_path: str = "${env.CHROMADB_PATH}", **kwargs: Any
+    ) -> dict[str, Any]:
+        return {
+            "db_path": db_path,
+            "kvstore": SqliteKVStoreConfig.sample_run_config(
+                __distro_dir__=__distro_dir__,
+                db_name="chroma_inline_registry.db",
+            ),
+        }

llama_stack/providers/inline/vector_io/faiss/faiss.py

@@ -55,6 +55,11 @@ class FaissIndex(EmbeddingIndex):
         self.kvstore = kvstore
         self.bank_id = bank_id
 
+        # A list of chunk id's in the same order as they are in the index,
+        # must be updated when chunks are added or removed
+        self.chunk_id_lock = asyncio.Lock()
+        self.chunk_ids: list[Any] = []
+
     @classmethod
     async def create(cls, dimension: int, kvstore: KVStore | None = None, bank_id: str | None = None):
         instance = cls(dimension, kvstore, bank_id)
@@ -75,6 +80,7 @@ class FaissIndex(EmbeddingIndex):
             buffer = io.BytesIO(base64.b64decode(data["faiss_index"]))
             try:
                 self.index = faiss.deserialize_index(np.load(buffer, allow_pickle=False))
+                self.chunk_ids = [chunk.chunk_id for chunk in self.chunk_by_index.values()]
             except Exception as e:
                 logger.debug(e, exc_info=True)
                 raise ValueError(
@@ -114,11 +120,33 @@ class FaissIndex(EmbeddingIndex):
         for i, chunk in enumerate(chunks):
             self.chunk_by_index[indexlen + i] = chunk
 
-        self.index.add(np.array(embeddings).astype(np.float32))
+        async with self.chunk_id_lock:
+            self.index.add(np.array(embeddings).astype(np.float32))
+            self.chunk_ids.extend([chunk.chunk_id for chunk in chunks])
 
         # Save updated index
         await self._save_index()
 
+    async def delete_chunk(self, chunk_id: str) -> None:
+        if chunk_id not in self.chunk_ids:
+            return
+
+        async with self.chunk_id_lock:
+            index = self.chunk_ids.index(chunk_id)
+            self.index.remove_ids(np.array([index]))
+
+            new_chunk_by_index = {}
+            for idx, chunk in self.chunk_by_index.items():
+                # Shift all chunks after the removed chunk to the left
+                if idx > index:
+                    new_chunk_by_index[idx - 1] = chunk
+                else:
+                    new_chunk_by_index[idx] = chunk
+            self.chunk_by_index = new_chunk_by_index
+            self.chunk_ids.pop(index)
+
+        await self._save_index()
+
     async def query_vector(
         self,
         embedding: NDArray,
@@ -261,47 +289,8 @@ class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPr
 
         return await index.query_chunks(query, params)
 
-    async def _save_openai_vector_store_file(
-        self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
-    ) -> None:
-        """Save vector store file data to kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
-        await self.kvstore.set(key=key, value=json.dumps(file_info))
-        content_key = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}"
-        await self.kvstore.set(key=content_key, value=json.dumps(file_contents))
-
-    async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
-        """Load vector store file metadata from kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
-        stored_data = await self.kvstore.get(key)
-        return json.loads(stored_data) if stored_data else {}
-
-    async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
-        """Load vector store file contents from kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}"
-        stored_data = await self.kvstore.get(key)
-        return json.loads(stored_data) if stored_data else []
-
-    async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
-        """Update vector store file metadata in kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
-        await self.kvstore.set(key=key, value=json.dumps(file_info))
-
-    async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
-        """Delete vector store data from kvstore."""
-        assert self.kvstore is not None
-
-        keys_to_delete = [
-            f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}",
-            f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}",
-        ]
-        for key in keys_to_delete:
-            try:
-                await self.kvstore.delete(key)
-            except Exception as e:
-                logger.warning(f"Failed to delete key {key}: {e}")
-                continue
+    async def delete_chunks(self, store_id: str, chunk_ids: list[str]) -> None:
+        """Delete a chunk from a faiss index"""
+        faiss_index = self.cache[store_id].index
+        for chunk_id in chunk_ids:
+            await faiss_index.delete_chunk(chunk_id)

llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import json
 import logging
 import re
 import sqlite3
@@ -426,6 +425,35 @@ class SQLiteVecIndex(EmbeddingIndex):
 
         return QueryChunksResponse(chunks=chunks, scores=scores)
 
+    async def delete_chunk(self, chunk_id: str) -> None:
+        """Remove a chunk from the SQLite vector store."""
+
+        def _delete_chunk():
+            connection = _create_sqlite_connection(self.db_path)
+            cur = connection.cursor()
+            try:
+                cur.execute("BEGIN TRANSACTION")
+
+                # Delete from metadata table
+                cur.execute(f"DELETE FROM {self.metadata_table} WHERE id = ?", (chunk_id,))
+
+                # Delete from vector table
+                cur.execute(f"DELETE FROM {self.vector_table} WHERE id = ?", (chunk_id,))
+
+                # Delete from FTS table
+                cur.execute(f"DELETE FROM {self.fts_table} WHERE id = ?", (chunk_id,))
+
+                connection.commit()
+            except Exception as e:
+                connection.rollback()
+                logger.error(f"Error deleting chunk {chunk_id}: {e}")
+                raise
+            finally:
+                cur.close()
+                connection.close()
+
+        await asyncio.to_thread(_delete_chunk)
+
 
 class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPrivate):
     """
@@ -506,140 +534,6 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc
         await self.cache[vector_db_id].index.delete()
         del self.cache[vector_db_id]
 
-    async def _save_openai_vector_store_file(
-        self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
-    ) -> None:
-        """Save vector store file metadata to SQLite database."""
-
-        def _create_or_store():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                # Create a table to persist OpenAI vector store files.
-                cur.execute("""
-                    CREATE TABLE IF NOT EXISTS openai_vector_store_files (
-                        store_id TEXT,
-                        file_id TEXT,
-                        metadata TEXT,
-                        PRIMARY KEY (store_id, file_id)
-                    );
-                """)
-                cur.execute("""
-                    CREATE TABLE IF NOT EXISTS openai_vector_store_files_contents (
-                        store_id TEXT,
-                        file_id TEXT,
-                        contents TEXT,
-                        PRIMARY KEY (store_id, file_id)
-                    );
-                """)
-                connection.commit()
-                cur.execute(
-                    "INSERT OR REPLACE INTO openai_vector_store_files (store_id, file_id, metadata) VALUES (?, ?, ?)",
-                    (store_id, file_id, json.dumps(file_info)),
-                )
-                cur.execute(
-                    "INSERT OR REPLACE INTO openai_vector_store_files_contents (store_id, file_id, contents) VALUES (?, ?, ?)",
-                    (store_id, file_id, json.dumps(file_contents)),
-                )
-                connection.commit()
-            except Exception as e:
-                logger.error(f"Error saving openai vector store file {store_id} {file_id}: {e}")
-                raise
-            finally:
-                cur.close()
-                connection.close()
-
-        try:
-            await asyncio.to_thread(_create_or_store)
-        except Exception as e:
-            logger.error(f"Error saving openai vector store file {store_id} {file_id}: {e}")
-            raise
-
-    async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
-        """Load vector store file metadata from SQLite database."""
-
-        def _load():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "SELECT metadata FROM openai_vector_store_files WHERE store_id = ? AND file_id = ?",
-                    (store_id, file_id),
-                )
-                row = cur.fetchone()
-                if row is None:
-                    return None
-                (metadata,) = row
-                return metadata
-            finally:
-                cur.close()
-                connection.close()
-
-        stored_data = await asyncio.to_thread(_load)
-        return json.loads(stored_data) if stored_data else {}
-
-    async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
-        """Load vector store file contents from SQLite database."""
-
-        def _load():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "SELECT contents FROM openai_vector_store_files_contents WHERE store_id = ? AND file_id = ?",
-                    (store_id, file_id),
-                )
-                row = cur.fetchone()
-                if row is None:
-                    return None
-                (contents,) = row
-                return contents
-            finally:
-                cur.close()
-                connection.close()
-
-        stored_contents = await asyncio.to_thread(_load)
-        return json.loads(stored_contents) if stored_contents else []
-
-    async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
-        """Update vector store file metadata in SQLite database."""
-
-        def _update():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "UPDATE openai_vector_store_files SET metadata = ? WHERE store_id = ? AND file_id = ?",
-                    (json.dumps(file_info), store_id, file_id),
-                )
-                connection.commit()
-            finally:
-                cur.close()
-                connection.close()
-
-        await asyncio.to_thread(_update)
-
-    async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
-        """Delete vector store file metadata from SQLite database."""
-
-        def _delete():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "DELETE FROM openai_vector_store_files WHERE store_id = ? AND file_id = ?", (store_id, file_id)
-                )
-                cur.execute(
-                    "DELETE FROM openai_vector_store_files_contents WHERE store_id = ? AND file_id = ?",
-                    (store_id, file_id),
-                )
-                connection.commit()
-            finally:
-                cur.close()
-                connection.close()
-
-        await asyncio.to_thread(_delete)
-
     async def insert_chunks(self, vector_db_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None:
         index = await self._get_and_cache_vector_db_index(vector_db_id)
         if not index:
@@ -655,3 +549,13 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc
         if not index:
             raise ValueError(f"Vector DB {vector_db_id} not found")
         return await index.query_chunks(query, params)
+
+    async def delete_chunks(self, store_id: str, chunk_ids: list[str]) -> None:
+        """Delete a chunk from a sqlite_vec index."""
+        index = await self._get_and_cache_vector_db_index(store_id)
+        if not index:
+            raise ValueError(f"Vector DB {store_id} not found")
+
+        for chunk_id in chunk_ids:
+            # Use the index's delete_chunk method
+            await index.index.delete_chunk(chunk_id)

llama_stack/providers/registry/inference.py

@@ -37,16 +37,6 @@ def available_providers() -> list[ProviderSpec]:
             config_class="llama_stack.providers.inline.inference.meta_reference.MetaReferenceInferenceConfig",
             description="Meta's reference implementation of inference with support for various model formats and optimization techniques.",
         ),
-        InlineProviderSpec(
-            api=Api.inference,
-            provider_type="inline::vllm",
-            pip_packages=[
-                "vllm",
-            ],
-            module="llama_stack.providers.inline.inference.vllm",
-            config_class="llama_stack.providers.inline.inference.vllm.VLLMConfig",
-            description="vLLM inference provider for high-performance model serving with PagedAttention and continuous batching.",
-        ),
         InlineProviderSpec(
             api=Api.inference,
             provider_type="inline::sentence-transformers",
@@ -234,17 +224,6 @@ def available_providers() -> list[ProviderSpec]:
                 description="Groq inference provider for ultra-fast inference using Groq's LPU technology.",
             ),
         ),
-        remote_provider_spec(
-            api=Api.inference,
-            adapter=AdapterSpec(
-                adapter_type="fireworks-openai-compat",
-                pip_packages=["litellm"],
-                module="llama_stack.providers.remote.inference.fireworks_openai_compat",
-                config_class="llama_stack.providers.remote.inference.fireworks_openai_compat.config.FireworksCompatConfig",
-                provider_data_validator="llama_stack.providers.remote.inference.fireworks_openai_compat.config.FireworksProviderDataValidator",
-                description="Fireworks AI OpenAI-compatible provider for using Fireworks models with OpenAI API format.",
-            ),
-        ),
         remote_provider_spec(
             api=Api.inference,
             adapter=AdapterSpec(
@@ -256,50 +235,6 @@ def available_providers() -> list[ProviderSpec]:
                 description="Llama OpenAI-compatible provider for using Llama models with OpenAI API format.",
             ),
         ),
-        remote_provider_spec(
-            api=Api.inference,
-            adapter=AdapterSpec(
-                adapter_type="together-openai-compat",
-                pip_packages=["litellm"],
-                module="llama_stack.providers.remote.inference.together_openai_compat",
-                config_class="llama_stack.providers.remote.inference.together_openai_compat.config.TogetherCompatConfig",
-                provider_data_validator="llama_stack.providers.remote.inference.together_openai_compat.config.TogetherProviderDataValidator",
-                description="Together AI OpenAI-compatible provider for using Together models with OpenAI API format.",
-            ),
-        ),
-        remote_provider_spec(
-            api=Api.inference,
-            adapter=AdapterSpec(
-                adapter_type="groq-openai-compat",
-                pip_packages=["litellm"],
-                module="llama_stack.providers.remote.inference.groq_openai_compat",
-                config_class="llama_stack.providers.remote.inference.groq_openai_compat.config.GroqCompatConfig",
-                provider_data_validator="llama_stack.providers.remote.inference.groq_openai_compat.config.GroqProviderDataValidator",
-                description="Groq OpenAI-compatible provider for using Groq models with OpenAI API format.",
-            ),
-        ),
-        remote_provider_spec(
-            api=Api.inference,
-            adapter=AdapterSpec(
-                adapter_type="sambanova-openai-compat",
-                pip_packages=["litellm"],
-                module="llama_stack.providers.remote.inference.sambanova_openai_compat",
-                config_class="llama_stack.providers.remote.inference.sambanova_openai_compat.config.SambaNovaCompatConfig",
-                provider_data_validator="llama_stack.providers.remote.inference.sambanova_openai_compat.config.SambaNovaProviderDataValidator",
-                description="SambaNova OpenAI-compatible provider for using SambaNova models with OpenAI API format.",
-            ),
-        ),
-        remote_provider_spec(
-            api=Api.inference,
-            adapter=AdapterSpec(
-                adapter_type="cerebras-openai-compat",
-                pip_packages=["litellm"],
-                module="llama_stack.providers.remote.inference.cerebras_openai_compat",
-                config_class="llama_stack.providers.remote.inference.cerebras_openai_compat.config.CerebrasCompatConfig",
-                provider_data_validator="llama_stack.providers.remote.inference.cerebras_openai_compat.config.CerebrasProviderDataValidator",
-                description="Cerebras OpenAI-compatible provider for using Cerebras models with OpenAI API format.",
-            ),
-        ),
         remote_provider_spec(
             api=Api.inference,
             adapter=AdapterSpec(

llama_stack/providers/registry/vector_io.py

@@ -395,7 +395,7 @@ That means you'll get fast and efficient vector retrieval.
 To use PGVector in your Llama Stack project, follow these steps:
 
 1. Install the necessary dependencies.
-2. Configure your Llama Stack project to use Faiss.
+2. Configure your Llama Stack project to use pgvector. (e.g. remote::pgvector).
 3. Start storing and querying vectors.
 
 ## Installation
@@ -410,6 +410,7 @@ See [PGVector's documentation](https://github.com/pgvector/pgvector) for more de
 """,
             ),
             api_dependencies=[Api.inference],
+            optional_api_dependencies=[Api.files],
         ),
         remote_provider_spec(
             Api.vector_io,

llama_stack/providers/remote/inference/anthropic/anthropic.py

@@ -15,6 +15,7 @@ class AnthropicInferenceAdapter(LiteLLMOpenAIMixin):
         LiteLLMOpenAIMixin.__init__(
             self,
             MODEL_ENTRIES,
+            litellm_provider_name="anthropic",
             api_key_from_config=config.api_key,
             provider_data_api_key_field="anthropic_api_key",
         )

llama_stack/providers/remote/inference/anthropic/config.py

@@ -26,7 +26,7 @@ class AnthropicConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.ANTHROPIC_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.ANTHROPIC_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "api_key": api_key,
         }

llama_stack/providers/remote/inference/anthropic/models.py

@@ -10,9 +10,9 @@ from llama_stack.providers.utils.inference.model_registry import (
 )
 
 LLM_MODEL_IDS = [
-    "anthropic/claude-3-5-sonnet-latest",
-    "anthropic/claude-3-7-sonnet-latest",
-    "anthropic/claude-3-5-haiku-latest",
+    "claude-3-5-sonnet-latest",
+    "claude-3-7-sonnet-latest",
+    "claude-3-5-haiku-latest",
 ]
 
 SAFETY_MODELS_ENTRIES = []
@@ -21,17 +21,17 @@ MODEL_ENTRIES = (
     [ProviderModelEntry(provider_model_id=m) for m in LLM_MODEL_IDS]
     + [
         ProviderModelEntry(
-            provider_model_id="anthropic/voyage-3",
+            provider_model_id="voyage-3",
             model_type=ModelType.embedding,
             metadata={"embedding_dimension": 1024, "context_length": 32000},
         ),
         ProviderModelEntry(
-            provider_model_id="anthropic/voyage-3-lite",
+            provider_model_id="voyage-3-lite",
             model_type=ModelType.embedding,
             metadata={"embedding_dimension": 512, "context_length": 32000},
         ),
         ProviderModelEntry(
-            provider_model_id="anthropic/voyage-code-3",
+            provider_model_id="voyage-code-3",
             model_type=ModelType.embedding,
             metadata={"embedding_dimension": 1024, "context_length": 32000},
         ),

llama_stack/providers/remote/inference/bedrock/bedrock.py

@@ -63,18 +63,20 @@ class BedrockInferenceAdapter(
     def __init__(self, config: BedrockConfig) -> None:
         ModelRegistryHelper.__init__(self, MODEL_ENTRIES)
         self._config = config
-
-        self._client = create_bedrock_client(config)
+        self._client = None
 
     @property
     def client(self) -> BaseClient:
+        if self._client is None:
+            self._client = create_bedrock_client(self._config)
         return self._client
 
     async def initialize(self) -> None:
         pass
 
     async def shutdown(self) -> None:
-        self.client.close()
+        if self._client is not None:
+            self._client.close()
 
     async def completion(
         self,

llama_stack/providers/remote/inference/cerebras/cerebras.py

@@ -65,6 +65,7 @@ class CerebrasInferenceAdapter(
         )
         self.config = config
 
+        # TODO: make this use provider data, etc. like other providers
         self.client = AsyncCerebras(
             base_url=self.config.base_url,
             api_key=self.config.api_key.get_secret_value(),

llama_stack/providers/remote/inference/cerebras/config.py

@@ -26,7 +26,7 @@ class CerebrasImplConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.CEREBRAS_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.CEREBRAS_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "base_url": DEFAULT_BASE_URL,
             "api_key": api_key,

llama_stack/providers/remote/inference/cerebras_openai_compat/__init__.py

@@ -1,17 +0,0 @@
-# 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 llama_stack.apis.inference import InferenceProvider
-
-from .config import CerebrasCompatConfig
-
-
-async def get_adapter_impl(config: CerebrasCompatConfig, _deps) -> InferenceProvider:
-    # import dynamically so the import is used only when it is needed
-    from .cerebras import CerebrasCompatInferenceAdapter
-
-    adapter = CerebrasCompatInferenceAdapter(config)
-    return adapter

llama_stack/providers/remote/inference/cerebras_openai_compat/cerebras.py

@@ -1,30 +0,0 @@
-# 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 llama_stack.providers.remote.inference.cerebras_openai_compat.config import CerebrasCompatConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-
-from ..cerebras.models import MODEL_ENTRIES
-
-
-class CerebrasCompatInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: CerebrasCompatConfig
-
-    def __init__(self, config: CerebrasCompatConfig):
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="cerebras_api_key",
-            openai_compat_api_base=config.openai_compat_api_base,
-        )
-        self.config = config
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

llama_stack/providers/remote/inference/cerebras_openai_compat/config.py

@@ -1,38 +0,0 @@
-# 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 typing import Any
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-class CerebrasProviderDataValidator(BaseModel):
-    cerebras_api_key: str | None = Field(
-        default=None,
-        description="API key for Cerebras models",
-    )
-
-
-@json_schema_type
-class CerebrasCompatConfig(BaseModel):
-    api_key: str | None = Field(
-        default=None,
-        description="The Cerebras API key",
-    )
-
-    openai_compat_api_base: str = Field(
-        default="https://api.cerebras.ai/v1",
-        description="The URL for the Cerebras API server",
-    )
-
-    @classmethod
-    def sample_run_config(cls, api_key: str = "${env.CEREBRAS_API_KEY}", **kwargs) -> dict[str, Any]:
-        return {
-            "openai_compat_api_base": "https://api.cerebras.ai/v1",
-            "api_key": api_key,
-        }

llama_stack/providers/remote/inference/databricks/config.py

@@ -25,8 +25,8 @@ class DatabricksImplConfig(BaseModel):
     @classmethod
     def sample_run_config(
         cls,
-        url: str = "${env.DATABRICKS_URL}",
-        api_token: str = "${env.DATABRICKS_API_TOKEN}",
+        url: str = "${env.DATABRICKS_URL:=}",
+        api_token: str = "${env.DATABRICKS_API_TOKEN:=}",
         **kwargs: Any,
     ) -> dict[str, Any]:
         return {

llama_stack/providers/remote/inference/fireworks/config.py

@@ -6,13 +6,14 @@
 
 from typing import Any
 
-from pydantic import BaseModel, Field, SecretStr
+from pydantic import Field, SecretStr
 
+from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
 from llama_stack.schema_utils import json_schema_type
 
 
 @json_schema_type
-class FireworksImplConfig(BaseModel):
+class FireworksImplConfig(RemoteInferenceProviderConfig):
     url: str = Field(
         default="https://api.fireworks.ai/inference/v1",
         description="The URL for the Fireworks server",
@@ -23,7 +24,7 @@ class FireworksImplConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.FIREWORKS_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.FIREWORKS_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "url": "https://api.fireworks.ai/inference/v1",
             "api_key": api_key,

llama_stack/providers/remote/inference/fireworks/fireworks.py

@@ -70,7 +70,7 @@ logger = get_logger(name=__name__, category="inference")
 
 class FireworksInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProviderData):
     def __init__(self, config: FireworksImplConfig) -> None:
-        ModelRegistryHelper.__init__(self, MODEL_ENTRIES)
+        ModelRegistryHelper.__init__(self, MODEL_ENTRIES, config.allowed_models)
         self.config = config
 
     async def initialize(self) -> None:

llama_stack/providers/remote/inference/fireworks_openai_compat/__init__.py

@@ -1,17 +0,0 @@
-# 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 llama_stack.apis.inference import InferenceProvider
-
-from .config import FireworksCompatConfig
-
-
-async def get_adapter_impl(config: FireworksCompatConfig, _deps) -> InferenceProvider:
-    # import dynamically so the import is used only when it is needed
-    from .fireworks import FireworksCompatInferenceAdapter
-
-    adapter = FireworksCompatInferenceAdapter(config)
-    return adapter

llama_stack/providers/remote/inference/fireworks_openai_compat/config.py

@@ -1,38 +0,0 @@
-# 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 typing import Any
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-class FireworksProviderDataValidator(BaseModel):
-    fireworks_api_key: str | None = Field(
-        default=None,
-        description="API key for Fireworks models",
-    )
-
-
-@json_schema_type
-class FireworksCompatConfig(BaseModel):
-    api_key: str | None = Field(
-        default=None,
-        description="The Fireworks API key",
-    )
-
-    openai_compat_api_base: str = Field(
-        default="https://api.fireworks.ai/inference/v1",
-        description="The URL for the Fireworks API server",
-    )
-
-    @classmethod
-    def sample_run_config(cls, api_key: str = "${env.FIREWORKS_API_KEY}", **kwargs) -> dict[str, Any]:
-        return {
-            "openai_compat_api_base": "https://api.fireworks.ai/inference/v1",
-            "api_key": api_key,
-        }

llama_stack/providers/remote/inference/fireworks_openai_compat/fireworks.py

@@ -1,30 +0,0 @@
-# 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 llama_stack.providers.remote.inference.fireworks_openai_compat.config import FireworksCompatConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-
-from ..fireworks.models import MODEL_ENTRIES
-
-
-class FireworksCompatInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: FireworksCompatConfig
-
-    def __init__(self, config: FireworksCompatConfig):
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="fireworks_api_key",
-            openai_compat_api_base=config.openai_compat_api_base,
-        )
-        self.config = config
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

llama_stack/providers/remote/inference/gemini/config.py

@@ -26,7 +26,7 @@ class GeminiConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.GEMINI_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.GEMINI_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "api_key": api_key,
         }

llama_stack/providers/remote/inference/gemini/gemini.py

@@ -15,6 +15,7 @@ class GeminiInferenceAdapter(LiteLLMOpenAIMixin):
         LiteLLMOpenAIMixin.__init__(
             self,
             MODEL_ENTRIES,
+            litellm_provider_name="gemini",
             api_key_from_config=config.api_key,
             provider_data_api_key_field="gemini_api_key",
         )

llama_stack/providers/remote/inference/gemini/models.py

@@ -10,11 +10,11 @@ from llama_stack.providers.utils.inference.model_registry import (
 )
 
 LLM_MODEL_IDS = [
-    "gemini/gemini-1.5-flash",
-    "gemini/gemini-1.5-pro",
-    "gemini/gemini-2.0-flash",
-    "gemini/gemini-2.5-flash",
-    "gemini/gemini-2.5-pro",
+    "gemini-1.5-flash",
+    "gemini-1.5-pro",
+    "gemini-2.0-flash",
+    "gemini-2.5-flash",
+    "gemini-2.5-pro",
 ]
 
 SAFETY_MODELS_ENTRIES = []
@@ -23,7 +23,7 @@ MODEL_ENTRIES = (
     [ProviderModelEntry(provider_model_id=m) for m in LLM_MODEL_IDS]
     + [
         ProviderModelEntry(
-            provider_model_id="gemini/text-embedding-004",
+            provider_model_id="text-embedding-004",
             model_type=ModelType.embedding,
             metadata={"embedding_dimension": 768, "context_length": 2048},
         ),

llama_stack/providers/remote/inference/groq/config.py

@@ -32,7 +32,7 @@ class GroqConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.GROQ_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.GROQ_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "url": "https://api.groq.com",
             "api_key": api_key,

llama_stack/providers/remote/inference/groq/groq.py

@@ -34,6 +34,7 @@ class GroqInferenceAdapter(LiteLLMOpenAIMixin):
         LiteLLMOpenAIMixin.__init__(
             self,
             model_entries=MODEL_ENTRIES,
+            litellm_provider_name="groq",
             api_key_from_config=config.api_key,
             provider_data_api_key_field="groq_api_key",
         )
@@ -96,7 +97,7 @@ class GroqInferenceAdapter(LiteLLMOpenAIMixin):
             tool_choice = "required"
 
         params = await prepare_openai_completion_params(
-            model=model_obj.provider_resource_id.replace("groq/", ""),
+            model=model_obj.provider_resource_id,
             messages=messages,
             frequency_penalty=frequency_penalty,
             function_call=function_call,

llama_stack/providers/remote/inference/groq/models.py

@@ -14,19 +14,19 @@ SAFETY_MODELS_ENTRIES = []
 
 MODEL_ENTRIES = [
     build_hf_repo_model_entry(
-        "groq/llama3-8b-8192",
+        "llama3-8b-8192",
         CoreModelId.llama3_1_8b_instruct.value,
     ),
     build_model_entry(
-        "groq/llama-3.1-8b-instant",
+        "llama-3.1-8b-instant",
         CoreModelId.llama3_1_8b_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "groq/llama3-70b-8192",
+        "llama3-70b-8192",
         CoreModelId.llama3_70b_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "groq/llama-3.3-70b-versatile",
+        "llama-3.3-70b-versatile",
         CoreModelId.llama3_3_70b_instruct.value,
     ),
     # Groq only contains a preview version for llama-3.2-3b
@@ -34,23 +34,15 @@ MODEL_ENTRIES = [
     # to pass the test fixture
     # TODO(aidand): Replace this with a stable model once Groq supports it
     build_hf_repo_model_entry(
-        "groq/llama-3.2-3b-preview",
+        "llama-3.2-3b-preview",
         CoreModelId.llama3_2_3b_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "groq/llama-4-scout-17b-16e-instruct",
+        "meta-llama/llama-4-scout-17b-16e-instruct",
         CoreModelId.llama4_scout_17b_16e_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "groq/meta-llama/llama-4-scout-17b-16e-instruct",
-        CoreModelId.llama4_scout_17b_16e_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "groq/llama-4-maverick-17b-128e-instruct",
-        CoreModelId.llama4_maverick_17b_128e_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "groq/meta-llama/llama-4-maverick-17b-128e-instruct",
+        "meta-llama/llama-4-maverick-17b-128e-instruct",
         CoreModelId.llama4_maverick_17b_128e_instruct.value,
     ),
 ] + SAFETY_MODELS_ENTRIES

llama_stack/providers/remote/inference/groq_openai_compat/__init__.py

@@ -1,17 +0,0 @@
-# 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 llama_stack.apis.inference import InferenceProvider
-
-from .config import GroqCompatConfig
-
-
-async def get_adapter_impl(config: GroqCompatConfig, _deps) -> InferenceProvider:
-    # import dynamically so the import is used only when it is needed
-    from .groq import GroqCompatInferenceAdapter
-
-    adapter = GroqCompatInferenceAdapter(config)
-    return adapter

llama_stack/providers/remote/inference/groq_openai_compat/config.py

@@ -1,38 +0,0 @@
-# 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 typing import Any
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-class GroqProviderDataValidator(BaseModel):
-    groq_api_key: str | None = Field(
-        default=None,
-        description="API key for Groq models",
-    )
-
-
-@json_schema_type
-class GroqCompatConfig(BaseModel):
-    api_key: str | None = Field(
-        default=None,
-        description="The Groq API key",
-    )
-
-    openai_compat_api_base: str = Field(
-        default="https://api.groq.com/openai/v1",
-        description="The URL for the Groq API server",
-    )
-
-    @classmethod
-    def sample_run_config(cls, api_key: str = "${env.GROQ_API_KEY}", **kwargs) -> dict[str, Any]:
-        return {
-            "openai_compat_api_base": "https://api.groq.com/openai/v1",
-            "api_key": api_key,
-        }

llama_stack/providers/remote/inference/groq_openai_compat/groq.py

@@ -1,30 +0,0 @@
-# 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 llama_stack.providers.remote.inference.groq_openai_compat.config import GroqCompatConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-
-from ..groq.models import MODEL_ENTRIES
-
-
-class GroqCompatInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: GroqCompatConfig
-
-    def __init__(self, config: GroqCompatConfig):
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="groq_api_key",
-            openai_compat_api_base=config.openai_compat_api_base,
-        )
-        self.config = config
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

llama_stack/providers/remote/inference/llama_openai_compat/llama.py

@@ -5,55 +5,53 @@
 # the root directory of this source tree.
 import logging
 
-from llama_api_client import AsyncLlamaAPIClient, NotFoundError
-
 from llama_stack.providers.remote.inference.llama_openai_compat.config import LlamaCompatConfig
 from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
+from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .models import MODEL_ENTRIES
 
 logger = logging.getLogger(__name__)
 
 
-class LlamaCompatInferenceAdapter(LiteLLMOpenAIMixin):
+class LlamaCompatInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin):
+    """
+    Llama API Inference Adapter for Llama Stack.
+
+    Note: The inheritance order is important here. OpenAIMixin must come before
+    LiteLLMOpenAIMixin to ensure that OpenAIMixin.check_model_availability()
+    is used instead of ModelRegistryHelper.check_model_availability().
+
+    - OpenAIMixin.check_model_availability() queries the Llama API to check if a model exists
+    - ModelRegistryHelper.check_model_availability() (inherited by LiteLLMOpenAIMixin) just returns False and shows a warning
+    """
+
     _config: LlamaCompatConfig
 
     def __init__(self, config: LlamaCompatConfig):
         LiteLLMOpenAIMixin.__init__(
             self,
             model_entries=MODEL_ENTRIES,
+            litellm_provider_name="meta_llama",
             api_key_from_config=config.api_key,
             provider_data_api_key_field="llama_api_key",
             openai_compat_api_base=config.openai_compat_api_base,
         )
         self.config = config
 
-    async def check_model_availability(self, model: str) -> bool:
+    # Delegate the client data handling get_api_key method to LiteLLMOpenAIMixin
+    get_api_key = LiteLLMOpenAIMixin.get_api_key
+
+    def get_base_url(self) -> str:
         """
-        Check if a specific model is available from Llama API.
+        Get the base URL for OpenAI mixin.
 
-        :param model: The model identifier to check.
-        :return: True if the model is available dynamically, False otherwise.
+        :return: The Llama API base URL
         """
-        try:
-            llama_api_client = self._get_llama_api_client()
-            retrieved_model = await llama_api_client.models.retrieve(model)
-            logger.info(f"Model {retrieved_model.id} is available from Llama API")
-            return True
-
-        except NotFoundError:
-            logger.error(f"Model {model} is not available from Llama API")
-            return False
-
-        except Exception as e:
-            logger.error(f"Failed to check model availability from Llama API: {e}")
-            return False
+        return self.config.openai_compat_api_base
 
     async def initialize(self):
         await super().initialize()
 
     async def shutdown(self):
         await super().shutdown()
-
-    def _get_llama_api_client(self) -> AsyncLlamaAPIClient:
-        return AsyncLlamaAPIClient(api_key=self.get_api_key(), base_url=self.config.openai_compat_api_base)

llama_stack/providers/remote/inference/nvidia/nvidia.py

@@ -7,9 +7,8 @@
 import logging
 import warnings
 from collections.abc import AsyncIterator
-from typing import Any
 
-from openai import APIConnectionError, AsyncOpenAI, BadRequestError, NotFoundError
+from openai import APIConnectionError, BadRequestError
 
 from llama_stack.apis.common.content_types import (
     InterleavedContent,
@@ -28,12 +27,6 @@ from llama_stack.apis.inference import (
     Inference,
     LogProbConfig,
     Message,
-    OpenAIChatCompletion,
-    OpenAIChatCompletionChunk,
-    OpenAICompletion,
-    OpenAIEmbeddingsResponse,
-    OpenAIMessageParam,
-    OpenAIResponseFormatParam,
     ResponseFormat,
     SamplingParams,
     TextTruncation,
@@ -47,8 +40,8 @@ from llama_stack.providers.utils.inference.model_registry import (
 from llama_stack.providers.utils.inference.openai_compat import (
     convert_openai_chat_completion_choice,
     convert_openai_chat_completion_stream,
-    prepare_openai_completion_params,
 )
+from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 from llama_stack.providers.utils.inference.prompt_adapter import content_has_media
 
 from . import NVIDIAConfig
@@ -64,7 +57,20 @@ from .utils import _is_nvidia_hosted
 logger = logging.getLogger(__name__)
 
 
-class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
+class NVIDIAInferenceAdapter(OpenAIMixin, Inference, ModelRegistryHelper):
+    """
+    NVIDIA Inference Adapter for Llama Stack.
+
+    Note: The inheritance order is important here. OpenAIMixin must come before
+    ModelRegistryHelper to ensure that OpenAIMixin.check_model_availability()
+    is used instead of ModelRegistryHelper.check_model_availability(). It also
+    must come before Inference to ensure that OpenAIMixin methods are available
+    in the Inference interface.
+
+    - OpenAIMixin.check_model_availability() queries the NVIDIA API to check if a model exists
+    - ModelRegistryHelper.check_model_availability() just returns False and shows a warning
+    """
+
     def __init__(self, config: NVIDIAConfig) -> None:
         # TODO(mf): filter by available models
         ModelRegistryHelper.__init__(self, model_entries=MODEL_ENTRIES)
@@ -88,45 +94,21 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
 
         self._config = config
 
-    async def check_model_availability(self, model: str) -> bool:
+    def get_api_key(self) -> str:
         """
-        Check if a specific model is available.
+        Get the API key for OpenAI mixin.
 
-        :param model: The model identifier to check.
-        :return: True if the model is available dynamically, False otherwise.
+        :return: The NVIDIA API key
         """
-        try:
-            await self._client.models.retrieve(model)
-            return True
-        except NotFoundError:
-            logger.error(f"Model {model} is not available")
-        except Exception as e:
-            logger.error(f"Failed to check model availability: {e}")
-        return False
+        return self._config.api_key.get_secret_value() if self._config.api_key else "NO KEY"
 
-    @property
-    def _client(self) -> AsyncOpenAI:
+    def get_base_url(self) -> str:
         """
-        Returns an OpenAI client for the configured NVIDIA API endpoint.
+        Get the base URL for OpenAI mixin.
 
-        :return: An OpenAI client
+        :return: The NVIDIA API base URL
         """
-
-        base_url = f"{self._config.url}/v1" if self._config.append_api_version else self._config.url
-
-        return AsyncOpenAI(
-            base_url=base_url,
-            api_key=(self._config.api_key.get_secret_value() if self._config.api_key else "NO KEY"),
-            timeout=self._config.timeout,
-        )
-
-    async def _get_provider_model_id(self, model_id: str) -> str:
-        if not self.model_store:
-            raise RuntimeError("Model store is not set")
-        model = await self.model_store.get_model(model_id)
-        if model is None:
-            raise ValueError(f"Model {model_id} is unknown")
-        return model.provider_model_id
+        return f"{self._config.url}/v1" if self._config.append_api_version else self._config.url
 
     async def completion(
         self,
@@ -160,7 +142,7 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
         )
 
         try:
-            response = await self._client.completions.create(**request)
+            response = await self.client.completions.create(**request)
         except APIConnectionError as e:
             raise ConnectionError(f"Failed to connect to NVIDIA NIM at {self._config.url}: {e}") from e
 
@@ -213,7 +195,7 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
             extra_body["input_type"] = task_type_options[task_type]
 
         try:
-            response = await self._client.embeddings.create(
+            response = await self.client.embeddings.create(
                 model=provider_model_id,
                 input=input,
                 extra_body=extra_body,
@@ -228,16 +210,6 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
         #
         return EmbeddingsResponse(embeddings=[embedding.embedding for embedding in response.data])
 
-    async def openai_embeddings(
-        self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
-    ) -> OpenAIEmbeddingsResponse:
-        raise NotImplementedError()
-
     async def chat_completion(
         self,
         model_id: str,
@@ -274,7 +246,7 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
         )
 
         try:
-            response = await self._client.chat.completions.create(**request)
+            response = await self.client.chat.completions.create(**request)
         except APIConnectionError as e:
             raise ConnectionError(f"Failed to connect to NVIDIA NIM at {self._config.url}: {e}") from e
 
@@ -283,112 +255,3 @@ class NVIDIAInferenceAdapter(Inference, ModelRegistryHelper):
         else:
             # we pass n=1 to get only one completion
             return convert_openai_chat_completion_choice(response.choices[0])
-
-    async def openai_completion(
-        self,
-        model: str,
-        prompt: str | list[str] | list[int] | list[list[int]],
-        best_of: int | None = None,
-        echo: bool | None = None,
-        frequency_penalty: float | None = None,
-        logit_bias: dict[str, float] | None = None,
-        logprobs: bool | None = None,
-        max_tokens: int | None = None,
-        n: int | None = None,
-        presence_penalty: float | None = None,
-        seed: int | None = None,
-        stop: str | list[str] | None = None,
-        stream: bool | None = None,
-        stream_options: dict[str, Any] | None = None,
-        temperature: float | None = None,
-        top_p: float | None = None,
-        user: str | None = None,
-        guided_choice: list[str] | None = None,
-        prompt_logprobs: int | None = None,
-        suffix: str | None = None,
-    ) -> OpenAICompletion:
-        provider_model_id = await self._get_provider_model_id(model)
-
-        params = await prepare_openai_completion_params(
-            model=provider_model_id,
-            prompt=prompt,
-            best_of=best_of,
-            echo=echo,
-            frequency_penalty=frequency_penalty,
-            logit_bias=logit_bias,
-            logprobs=logprobs,
-            max_tokens=max_tokens,
-            n=n,
-            presence_penalty=presence_penalty,
-            seed=seed,
-            stop=stop,
-            stream=stream,
-            stream_options=stream_options,
-            temperature=temperature,
-            top_p=top_p,
-            user=user,
-        )
-
-        try:
-            return await self._client.completions.create(**params)
-        except APIConnectionError as e:
-            raise ConnectionError(f"Failed to connect to NVIDIA NIM at {self._config.url}: {e}") from e
-
-    async def openai_chat_completion(
-        self,
-        model: str,
-        messages: list[OpenAIMessageParam],
-        frequency_penalty: float | None = None,
-        function_call: str | dict[str, Any] | None = None,
-        functions: list[dict[str, Any]] | None = None,
-        logit_bias: dict[str, float] | None = None,
-        logprobs: bool | None = None,
-        max_completion_tokens: int | None = None,
-        max_tokens: int | None = None,
-        n: int | None = None,
-        parallel_tool_calls: bool | None = None,
-        presence_penalty: float | None = None,
-        response_format: OpenAIResponseFormatParam | None = None,
-        seed: int | None = None,
-        stop: str | list[str] | None = None,
-        stream: bool | None = None,
-        stream_options: dict[str, Any] | None = None,
-        temperature: float | None = None,
-        tool_choice: str | dict[str, Any] | None = None,
-        tools: list[dict[str, Any]] | None = None,
-        top_logprobs: int | None = None,
-        top_p: float | None = None,
-        user: str | None = None,
-    ) -> OpenAIChatCompletion | AsyncIterator[OpenAIChatCompletionChunk]:
-        provider_model_id = await self._get_provider_model_id(model)
-
-        params = await prepare_openai_completion_params(
-            model=provider_model_id,
-            messages=messages,
-            frequency_penalty=frequency_penalty,
-            function_call=function_call,
-            functions=functions,
-            logit_bias=logit_bias,
-            logprobs=logprobs,
-            max_completion_tokens=max_completion_tokens,
-            max_tokens=max_tokens,
-            n=n,
-            parallel_tool_calls=parallel_tool_calls,
-            presence_penalty=presence_penalty,
-            response_format=response_format,
-            seed=seed,
-            stop=stop,
-            stream=stream,
-            stream_options=stream_options,
-            temperature=temperature,
-            tool_choice=tool_choice,
-            tools=tools,
-            top_logprobs=top_logprobs,
-            top_p=top_p,
-            user=user,
-        )
-
-        try:
-            return await self._client.chat.completions.create(**params)
-        except APIConnectionError as e:
-            raise ConnectionError(f"Failed to connect to NVIDIA NIM at {self._config.url}: {e}") from e

llama_stack/providers/remote/inference/ollama/config.py

@@ -6,13 +6,17 @@
 
 from typing import Any
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 DEFAULT_OLLAMA_URL = "http://localhost:11434"
 
 
 class OllamaImplConfig(BaseModel):
     url: str = DEFAULT_OLLAMA_URL
+    refresh_models: bool = Field(
+        default=False,
+        description="Whether to refresh models periodically",
+    )
 
     @classmethod
     def sample_run_config(cls, url: str = "${env.OLLAMA_URL:=http://localhost:11434}", **kwargs) -> dict[str, Any]:

llama_stack/providers/remote/inference/ollama/ollama.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 
+import asyncio
 import base64
 import uuid
 from collections.abc import AsyncGenerator, AsyncIterator
@@ -91,23 +92,92 @@ class OllamaInferenceAdapter(
     InferenceProvider,
     ModelsProtocolPrivate,
 ):
+    # automatically set by the resolver when instantiating the provider
+    __provider_id__: str
+
     def __init__(self, config: OllamaImplConfig) -> None:
         self.register_helper = ModelRegistryHelper(MODEL_ENTRIES)
-        self.url = config.url
+        self.config = config
+        self._clients: dict[asyncio.AbstractEventLoop, AsyncClient] = {}
+        self._openai_client = None
 
     @property
     def client(self) -> AsyncClient:
-        return AsyncClient(host=self.url)
+        # ollama client attaches itself to the current event loop (sadly?)
+        loop = asyncio.get_running_loop()
+        if loop not in self._clients:
+            self._clients[loop] = AsyncClient(host=self.config.url)
+        return self._clients[loop]
 
     @property
     def openai_client(self) -> AsyncOpenAI:
-        return AsyncOpenAI(base_url=f"{self.url}/v1", api_key="ollama")
+        if self._openai_client is None:
+            self._openai_client = AsyncOpenAI(base_url=f"{self.config.url}/v1", api_key="ollama")
+        return self._openai_client
 
     async def initialize(self) -> None:
-        logger.debug(f"checking connectivity to Ollama at `{self.url}`...")
+        logger.info(f"checking connectivity to Ollama at `{self.config.url}`...")
         health_response = await self.health()
         if health_response["status"] == HealthStatus.ERROR:
-            raise RuntimeError("Ollama Server is not running, start it using `ollama serve` in a separate terminal")
+            logger.warning(
+                "Ollama Server is not running, make sure to start it using `ollama serve` in a separate terminal"
+            )
+
+    async def should_refresh_models(self) -> bool:
+        return self.config.refresh_models
+
+    async def list_models(self) -> list[Model] | None:
+        provider_id = self.__provider_id__
+        response = await self.client.list()
+
+        # always add the two embedding models which can be pulled on demand
+        models = [
+            Model(
+                identifier="all-minilm:l6-v2",
+                provider_resource_id="all-minilm:l6-v2",
+                provider_id=provider_id,
+                metadata={
+                    "embedding_dimension": 384,
+                    "context_length": 512,
+                },
+                model_type=ModelType.embedding,
+            ),
+            # add all-minilm alias
+            Model(
+                identifier="all-minilm",
+                provider_resource_id="all-minilm:l6-v2",
+                provider_id=provider_id,
+                metadata={
+                    "embedding_dimension": 384,
+                    "context_length": 512,
+                },
+                model_type=ModelType.embedding,
+            ),
+            Model(
+                identifier="nomic-embed-text",
+                provider_resource_id="nomic-embed-text",
+                provider_id=provider_id,
+                metadata={
+                    "embedding_dimension": 768,
+                    "context_length": 8192,
+                },
+                model_type=ModelType.embedding,
+            ),
+        ]
+        for m in response.models:
+            # kill embedding models since we don't know dimensions for them
+            if "bert" in m.details.family:
+                continue
+            models.append(
+                Model(
+                    identifier=m.model,
+                    provider_resource_id=m.model,
+                    provider_id=provider_id,
+                    metadata={},
+                    model_type=ModelType.llm,
+                )
+            )
+        return models
 
     async def health(self) -> HealthResponse:
         """
@@ -124,7 +194,7 @@ class OllamaInferenceAdapter(
             return HealthResponse(status=HealthStatus.ERROR, message=f"Health check failed: {str(e)}")
 
     async def shutdown(self) -> None:
-        pass
+        self._clients.clear()
 
     async def unregister_model(self, model_id: str) -> None:
         pass
@@ -350,12 +420,7 @@ class OllamaInferenceAdapter(
         except ValueError:
             pass  # Ignore statically unknown model, will check live listing
 
-        if model.provider_resource_id is None:
-            raise ValueError("Model provider_resource_id cannot be None")
-
         if model.model_type == ModelType.embedding:
-            logger.info(f"Pulling embedding model `{model.provider_resource_id}` if necessary...")
-            # TODO: you should pull here only if the model is not found in a list
             response = await self.client.list()
             if model.provider_resource_id not in [m.model for m in response.models]:
                 await self.client.pull(model.provider_resource_id)
@@ -365,9 +430,9 @@ class OllamaInferenceAdapter(
         #  - models not currently running are run by the ollama server as needed
         response = await self.client.list()
         available_models = [m.model for m in response.models]
-        provider_resource_id = self.register_helper.get_provider_model_id(model.provider_resource_id)
-        if provider_resource_id is None:
-            provider_resource_id = model.provider_resource_id
+
+        provider_resource_id = model.provider_resource_id
+        assert provider_resource_id is not None  # mypy
         if provider_resource_id not in available_models:
             available_models_latest = [m.model.split(":latest")[0] for m in response.models]
             if provider_resource_id in available_models_latest:
@@ -375,7 +440,9 @@ class OllamaInferenceAdapter(
                     f"Imprecise provider resource id was used but 'latest' is available in Ollama - using '{model.provider_resource_id}:latest'"
                 )
                 return model
-            raise UnsupportedModelError(model.provider_resource_id, available_models)
+            raise UnsupportedModelError(provider_resource_id, available_models)
+
+        # mutating this should be considered an anti-pattern
         model.provider_resource_id = provider_resource_id
 
         return model

llama_stack/providers/remote/inference/openai/config.py

@@ -24,9 +24,19 @@ class OpenAIConfig(BaseModel):
         default=None,
         description="API key for OpenAI models",
     )
+    base_url: str = Field(
+        default="https://api.openai.com/v1",
+        description="Base URL for OpenAI API",
+    )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.OPENAI_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(
+        cls,
+        api_key: str = "${env.OPENAI_API_KEY:=}",
+        base_url: str = "${env.OPENAI_BASE_URL:=https://api.openai.com/v1}",
+        **kwargs,
+    ) -> dict[str, Any]:
         return {
             "api_key": api_key,
+            "base_url": base_url,
         }

llama_stack/providers/remote/inference/openai/models.py

@@ -12,11 +12,6 @@ from llama_stack.providers.utils.inference.model_registry import (
 )
 
 LLM_MODEL_IDS = [
-    # the models w/ "openai/" prefix are the litellm specific model names.
-    # they should be deprecated in favor of the canonical openai model names.
-    "openai/gpt-4o",
-    "openai/gpt-4o-mini",
-    "openai/chatgpt-4o-latest",
     "gpt-3.5-turbo-0125",
     "gpt-3.5-turbo",
     "gpt-3.5-turbo-instruct",
@@ -43,8 +38,6 @@ class EmbeddingModelInfo:
 
 
 EMBEDDING_MODEL_IDS: dict[str, EmbeddingModelInfo] = {
-    "openai/text-embedding-3-small": EmbeddingModelInfo(1536, 8192),
-    "openai/text-embedding-3-large": EmbeddingModelInfo(3072, 8192),
     "text-embedding-3-small": EmbeddingModelInfo(1536, 8192),
     "text-embedding-3-large": EmbeddingModelInfo(3072, 8192),
 }

llama_stack/providers/remote/inference/openai/openai.py

@@ -5,23 +5,9 @@
 # the root directory of this source tree.
 
 import logging
-from collections.abc import AsyncIterator
-from typing import Any
 
-from openai import AsyncOpenAI, NotFoundError
-
-from llama_stack.apis.inference import (
-    OpenAIChatCompletion,
-    OpenAIChatCompletionChunk,
-    OpenAICompletion,
-    OpenAIEmbeddingData,
-    OpenAIEmbeddingsResponse,
-    OpenAIEmbeddingUsage,
-    OpenAIMessageParam,
-    OpenAIResponseFormatParam,
-)
 from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-from llama_stack.providers.utils.inference.openai_compat import prepare_openai_completion_params
+from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .config import OpenAIConfig
 from .models import MODEL_ENTRIES
@@ -30,7 +16,7 @@ logger = logging.getLogger(__name__)
 
 
 #
-# This OpenAI adapter implements Inference methods using two clients -
+# This OpenAI adapter implements Inference methods using two mixins -
 #
 # | Inference Method           | Implementation Source    |
 # |----------------------------|--------------------------|
@@ -39,15 +25,27 @@ logger = logging.getLogger(__name__)
 # | embedding                  | LiteLLMOpenAIMixin       |
 # | batch_completion           | LiteLLMOpenAIMixin       |
 # | batch_chat_completion      | LiteLLMOpenAIMixin       |
-# | openai_completion          | AsyncOpenAI              |
-# | openai_chat_completion     | AsyncOpenAI              |
-# | openai_embeddings          | AsyncOpenAI              |
+# | openai_completion          | OpenAIMixin              |
+# | openai_chat_completion     | OpenAIMixin              |
+# | openai_embeddings          | OpenAIMixin              |
 #
-class OpenAIInferenceAdapter(LiteLLMOpenAIMixin):
+class OpenAIInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin):
+    """
+    OpenAI Inference Adapter for Llama Stack.
+
+    Note: The inheritance order is important here. OpenAIMixin must come before
+    LiteLLMOpenAIMixin to ensure that OpenAIMixin.check_model_availability()
+    is used instead of ModelRegistryHelper.check_model_availability().
+
+    - OpenAIMixin.check_model_availability() queries the OpenAI API to check if a model exists
+    - ModelRegistryHelper.check_model_availability() (inherited by LiteLLMOpenAIMixin) just returns False and shows a warning
+    """
+
     def __init__(self, config: OpenAIConfig) -> None:
         LiteLLMOpenAIMixin.__init__(
             self,
             MODEL_ENTRIES,
+            litellm_provider_name="openai",
             api_key_from_config=config.api_key,
             provider_data_api_key_field="openai_api_key",
         )
@@ -60,191 +58,19 @@ class OpenAIInferenceAdapter(LiteLLMOpenAIMixin):
         # litellm specific model names, an abstraction leak.
         self.is_openai_compat = True
 
-    async def check_model_availability(self, model: str) -> bool:
+    # Delegate the client data handling get_api_key method to LiteLLMOpenAIMixin
+    get_api_key = LiteLLMOpenAIMixin.get_api_key
+
+    def get_base_url(self) -> str:
         """
-        Check if a specific model is available from OpenAI.
+        Get the OpenAI API base URL.
 
-        :param model: The model identifier to check.
-        :return: True if the model is available dynamically, False otherwise.
+        Returns the OpenAI API base URL from the configuration.
         """
-        try:
-            openai_client = self._get_openai_client()
-            retrieved_model = await openai_client.models.retrieve(model)
-            logger.info(f"Model {retrieved_model.id} is available from OpenAI")
-            return True
-
-        except NotFoundError:
-            logger.error(f"Model {model} is not available from OpenAI")
-            return False
-
-        except Exception as e:
-            logger.error(f"Failed to check model availability from OpenAI: {e}")
-            return False
+        return self.config.base_url
 
     async def initialize(self) -> None:
         await super().initialize()
 
     async def shutdown(self) -> None:
         await super().shutdown()
-
-    def _get_openai_client(self) -> AsyncOpenAI:
-        return AsyncOpenAI(
-            api_key=self.get_api_key(),
-        )
-
-    async def openai_completion(
-        self,
-        model: str,
-        prompt: str | list[str] | list[int] | list[list[int]],
-        best_of: int | None = None,
-        echo: bool | None = None,
-        frequency_penalty: float | None = None,
-        logit_bias: dict[str, float] | None = None,
-        logprobs: bool | None = None,
-        max_tokens: int | None = None,
-        n: int | None = None,
-        presence_penalty: float | None = None,
-        seed: int | None = None,
-        stop: str | list[str] | None = None,
-        stream: bool | None = None,
-        stream_options: dict[str, Any] | None = None,
-        temperature: float | None = None,
-        top_p: float | None = None,
-        user: str | None = None,
-        guided_choice: list[str] | None = None,
-        prompt_logprobs: int | None = None,
-        suffix: str | None = None,
-    ) -> OpenAICompletion:
-        if guided_choice is not None:
-            logging.warning("guided_choice is not supported by the OpenAI API. Ignoring.")
-        if prompt_logprobs is not None:
-            logging.warning("prompt_logprobs is not supported by the OpenAI API. Ignoring.")
-
-        model_id = (await self.model_store.get_model(model)).provider_resource_id
-        if model_id.startswith("openai/"):
-            model_id = model_id[len("openai/") :]
-        params = await prepare_openai_completion_params(
-            model=model_id,
-            prompt=prompt,
-            best_of=best_of,
-            echo=echo,
-            frequency_penalty=frequency_penalty,
-            logit_bias=logit_bias,
-            logprobs=logprobs,
-            max_tokens=max_tokens,
-            n=n,
-            presence_penalty=presence_penalty,
-            seed=seed,
-            stop=stop,
-            stream=stream,
-            stream_options=stream_options,
-            temperature=temperature,
-            top_p=top_p,
-            user=user,
-            suffix=suffix,
-        )
-        return await self._get_openai_client().completions.create(**params)
-
-    async def openai_chat_completion(
-        self,
-        model: str,
-        messages: list[OpenAIMessageParam],
-        frequency_penalty: float | None = None,
-        function_call: str | dict[str, Any] | None = None,
-        functions: list[dict[str, Any]] | None = None,
-        logit_bias: dict[str, float] | None = None,
-        logprobs: bool | None = None,
-        max_completion_tokens: int | None = None,
-        max_tokens: int | None = None,
-        n: int | None = None,
-        parallel_tool_calls: bool | None = None,
-        presence_penalty: float | None = None,
-        response_format: OpenAIResponseFormatParam | None = None,
-        seed: int | None = None,
-        stop: str | list[str] | None = None,
-        stream: bool | None = None,
-        stream_options: dict[str, Any] | None = None,
-        temperature: float | None = None,
-        tool_choice: str | dict[str, Any] | None = None,
-        tools: list[dict[str, Any]] | None = None,
-        top_logprobs: int | None = None,
-        top_p: float | None = None,
-        user: str | None = None,
-    ) -> OpenAIChatCompletion | AsyncIterator[OpenAIChatCompletionChunk]:
-        model_id = (await self.model_store.get_model(model)).provider_resource_id
-        if model_id.startswith("openai/"):
-            model_id = model_id[len("openai/") :]
-        params = await prepare_openai_completion_params(
-            model=model_id,
-            messages=messages,
-            frequency_penalty=frequency_penalty,
-            function_call=function_call,
-            functions=functions,
-            logit_bias=logit_bias,
-            logprobs=logprobs,
-            max_completion_tokens=max_completion_tokens,
-            max_tokens=max_tokens,
-            n=n,
-            parallel_tool_calls=parallel_tool_calls,
-            presence_penalty=presence_penalty,
-            response_format=response_format,
-            seed=seed,
-            stop=stop,
-            stream=stream,
-            stream_options=stream_options,
-            temperature=temperature,
-            tool_choice=tool_choice,
-            tools=tools,
-            top_logprobs=top_logprobs,
-            top_p=top_p,
-            user=user,
-        )
-        return await self._get_openai_client().chat.completions.create(**params)
-
-    async def openai_embeddings(
-        self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
-    ) -> OpenAIEmbeddingsResponse:
-        model_id = (await self.model_store.get_model(model)).provider_resource_id
-        if model_id.startswith("openai/"):
-            model_id = model_id[len("openai/") :]
-
-        # Prepare parameters for OpenAI embeddings API
-        params = {
-            "model": model_id,
-            "input": input,
-        }
-
-        if encoding_format is not None:
-            params["encoding_format"] = encoding_format
-        if dimensions is not None:
-            params["dimensions"] = dimensions
-        if user is not None:
-            params["user"] = user
-
-        # Call OpenAI embeddings API
-        response = await self._get_openai_client().embeddings.create(**params)
-
-        data = []
-        for i, embedding_data in enumerate(response.data):
-            data.append(
-                OpenAIEmbeddingData(
-                    embedding=embedding_data.embedding,
-                    index=i,
-                )
-            )
-
-        usage = OpenAIEmbeddingUsage(
-            prompt_tokens=response.usage.prompt_tokens,
-            total_tokens=response.usage.total_tokens,
-        )
-
-        return OpenAIEmbeddingsResponse(
-            data=data,
-            model=response.model,
-            usage=usage,
-        )

llama_stack/providers/remote/inference/sambanova/config.py

@@ -30,7 +30,7 @@ class SambaNovaImplConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.SAMBANOVA_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.SAMBANOVA_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "url": "https://api.sambanova.ai/v1",
             "api_key": api_key,

llama_stack/providers/remote/inference/sambanova/models.py

@@ -9,49 +9,20 @@ from llama_stack.providers.utils.inference.model_registry import (
     build_hf_repo_model_entry,
 )
 
-SAFETY_MODELS_ENTRIES = [
-    build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-Guard-3-8B",
-        CoreModelId.llama_guard_3_8b.value,
-    ),
-]
+SAFETY_MODELS_ENTRIES = []
 
 
 MODEL_ENTRIES = [
     build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-3.1-8B-Instruct",
+        "Meta-Llama-3.1-8B-Instruct",
         CoreModelId.llama3_1_8b_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-3.1-405B-Instruct",
-        CoreModelId.llama3_1_405b_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-3.2-1B-Instruct",
-        CoreModelId.llama3_2_1b_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-3.2-3B-Instruct",
-        CoreModelId.llama3_2_3b_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Meta-Llama-3.3-70B-Instruct",
+        "Meta-Llama-3.3-70B-Instruct",
         CoreModelId.llama3_3_70b_instruct.value,
     ),
     build_hf_repo_model_entry(
-        "sambanova/Llama-3.2-11B-Vision-Instruct",
-        CoreModelId.llama3_2_11b_vision_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Llama-3.2-90B-Vision-Instruct",
-        CoreModelId.llama3_2_90b_vision_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Llama-4-Scout-17B-16E-Instruct",
-        CoreModelId.llama4_scout_17b_16e_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "sambanova/Llama-4-Maverick-17B-128E-Instruct",
+        "Llama-4-Maverick-17B-128E-Instruct",
         CoreModelId.llama4_maverick_17b_128e_instruct.value,
     ),
 ] + SAFETY_MODELS_ENTRIES

llama_stack/providers/remote/inference/sambanova/sambanova.py

@@ -182,6 +182,7 @@ class SambaNovaInferenceAdapter(LiteLLMOpenAIMixin):
         LiteLLMOpenAIMixin.__init__(
             self,
             model_entries=MODEL_ENTRIES,
+            litellm_provider_name="sambanova",
             api_key_from_config=self.config.api_key.get_secret_value() if self.config.api_key else None,
             provider_data_api_key_field="sambanova_api_key",
         )

llama_stack/providers/remote/inference/sambanova_openai_compat/__init__.py

@@ -1,17 +0,0 @@
-# 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 llama_stack.apis.inference import InferenceProvider
-
-from .config import SambaNovaCompatConfig
-
-
-async def get_adapter_impl(config: SambaNovaCompatConfig, _deps) -> InferenceProvider:
-    # import dynamically so the import is used only when it is needed
-    from .sambanova import SambaNovaCompatInferenceAdapter
-
-    adapter = SambaNovaCompatInferenceAdapter(config)
-    return adapter

llama_stack/providers/remote/inference/sambanova_openai_compat/config.py

@@ -1,38 +0,0 @@
-# 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 typing import Any
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-class SambaNovaProviderDataValidator(BaseModel):
-    sambanova_api_key: str | None = Field(
-        default=None,
-        description="API key for SambaNova models",
-    )
-
-
-@json_schema_type
-class SambaNovaCompatConfig(BaseModel):
-    api_key: str | None = Field(
-        default=None,
-        description="The SambaNova API key",
-    )
-
-    openai_compat_api_base: str = Field(
-        default="https://api.sambanova.ai/v1",
-        description="The URL for the SambaNova API server",
-    )
-
-    @classmethod
-    def sample_run_config(cls, api_key: str = "${env.SAMBANOVA_API_KEY}", **kwargs) -> dict[str, Any]:
-        return {
-            "openai_compat_api_base": "https://api.sambanova.ai/v1",
-            "api_key": api_key,
-        }

llama_stack/providers/remote/inference/sambanova_openai_compat/sambanova.py

@@ -1,30 +0,0 @@
-# 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 llama_stack.providers.remote.inference.sambanova_openai_compat.config import SambaNovaCompatConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-
-from ..sambanova.models import MODEL_ENTRIES
-
-
-class SambaNovaCompatInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: SambaNovaCompatConfig
-
-    def __init__(self, config: SambaNovaCompatConfig):
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="sambanova_api_key",
-            openai_compat_api_base=config.openai_compat_api_base,
-        )
-        self.config = config
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

llama_stack/providers/remote/inference/tgi/config.py

@@ -19,7 +19,7 @@ class TGIImplConfig(BaseModel):
     @classmethod
     def sample_run_config(
         cls,
-        url: str = "${env.TGI_URL}",
+        url: str = "${env.TGI_URL:=}",
         **kwargs,
     ):
         return {

llama_stack/providers/remote/inference/tgi/tgi.py

@@ -305,6 +305,8 @@ class _HfAdapter(
 
 class TGIAdapter(_HfAdapter):
     async def initialize(self, config: TGIImplConfig) -> None:
+        if not config.url:
+            raise ValueError("You must provide a URL in run.yaml (or via the TGI_URL environment variable) to use TGI.")
         log.info(f"Initializing TGI client with url={config.url}")
         self.client = AsyncInferenceClient(
             model=config.url,

llama_stack/providers/remote/inference/together/config.py

@@ -6,13 +6,14 @@
 
 from typing import Any
 
-from pydantic import BaseModel, Field, SecretStr
+from pydantic import Field, SecretStr
 
+from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
 from llama_stack.schema_utils import json_schema_type
 
 
 @json_schema_type
-class TogetherImplConfig(BaseModel):
+class TogetherImplConfig(RemoteInferenceProviderConfig):
     url: str = Field(
         default="https://api.together.xyz/v1",
         description="The URL for the Together AI server",
@@ -26,5 +27,5 @@ class TogetherImplConfig(BaseModel):
     def sample_run_config(cls, **kwargs) -> dict[str, Any]:
         return {
             "url": "https://api.together.xyz/v1",
-            "api_key": "${env.TOGETHER_API_KEY}",
+            "api_key": "${env.TOGETHER_API_KEY:=}",
         }

llama_stack/providers/remote/inference/together/models.py

@@ -69,15 +69,9 @@ MODEL_ENTRIES = [
     build_hf_repo_model_entry(
         "meta-llama/Llama-4-Scout-17B-16E-Instruct",
         CoreModelId.llama4_scout_17b_16e_instruct.value,
-        additional_aliases=[
-            "together/meta-llama/Llama-4-Scout-17B-16E-Instruct",
-        ],
     ),
     build_hf_repo_model_entry(
         "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
         CoreModelId.llama4_maverick_17b_128e_instruct.value,
-        additional_aliases=[
-            "together/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
-        ],
     ),
 ] + SAFETY_MODELS_ENTRIES

llama_stack/providers/remote/inference/together/together.py

@@ -66,7 +66,7 @@ logger = get_logger(name=__name__, category="inference")
 
 class TogetherInferenceAdapter(ModelRegistryHelper, Inference, NeedsRequestProviderData):
     def __init__(self, config: TogetherImplConfig) -> None:
-        ModelRegistryHelper.__init__(self, MODEL_ENTRIES)
+        ModelRegistryHelper.__init__(self, MODEL_ENTRIES, config.allowed_models)
         self.config = config
 
     async def initialize(self) -> None:

llama_stack/providers/remote/inference/together_openai_compat/__init__.py

@@ -1,17 +0,0 @@
-# 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 llama_stack.apis.inference import InferenceProvider
-
-from .config import TogetherCompatConfig
-
-
-async def get_adapter_impl(config: TogetherCompatConfig, _deps) -> InferenceProvider:
-    # import dynamically so the import is used only when it is needed
-    from .together import TogetherCompatInferenceAdapter
-
-    adapter = TogetherCompatInferenceAdapter(config)
-    return adapter

llama_stack/providers/remote/inference/together_openai_compat/config.py

@@ -1,38 +0,0 @@
-# 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 typing import Any
-
-from pydantic import BaseModel, Field
-
-from llama_stack.schema_utils import json_schema_type
-
-
-class TogetherProviderDataValidator(BaseModel):
-    together_api_key: str | None = Field(
-        default=None,
-        description="API key for Together models",
-    )
-
-
-@json_schema_type
-class TogetherCompatConfig(BaseModel):
-    api_key: str | None = Field(
-        default=None,
-        description="The Together API key",
-    )
-
-    openai_compat_api_base: str = Field(
-        default="https://api.together.xyz/v1",
-        description="The URL for the Together API server",
-    )
-
-    @classmethod
-    def sample_run_config(cls, api_key: str = "${env.TOGETHER_API_KEY}", **kwargs) -> dict[str, Any]:
-        return {
-            "openai_compat_api_base": "https://api.together.xyz/v1",
-            "api_key": api_key,
-        }

llama_stack/providers/remote/inference/together_openai_compat/together.py

@@ -1,30 +0,0 @@
-# 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 llama_stack.providers.remote.inference.together_openai_compat.config import TogetherCompatConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-
-from ..together.models import MODEL_ENTRIES
-
-
-class TogetherCompatInferenceAdapter(LiteLLMOpenAIMixin):
-    _config: TogetherCompatConfig
-
-    def __init__(self, config: TogetherCompatConfig):
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="together_api_key",
-            openai_compat_api_base=config.openai_compat_api_base,
-        )
-        self.config = config
-
-    async def initialize(self):
-        await super().initialize()
-
-    async def shutdown(self):
-        await super().shutdown()

llama_stack/providers/remote/inference/vllm/config.py

@@ -29,6 +29,10 @@ class VLLMInferenceAdapterConfig(BaseModel):
         default=True,
         description="Whether to verify TLS certificates. Can be a boolean or a path to a CA certificate file.",
     )
+    refresh_models: bool = Field(
+        default=False,
+        description="Whether to refresh models periodically",
+    )
 
     @field_validator("tls_verify")
     @classmethod
@@ -46,7 +50,7 @@ class VLLMInferenceAdapterConfig(BaseModel):
     @classmethod
     def sample_run_config(
         cls,
-        url: str = "${env.VLLM_URL}",
+        url: str = "${env.VLLM_URL:=}",
         **kwargs,
     ):
         return {

llama_stack/providers/remote/inference/vllm/vllm.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 import json
-import logging
 from collections.abc import AsyncGenerator, AsyncIterator
 from typing import Any
 
@@ -38,6 +37,7 @@ from llama_stack.apis.inference import (
     JsonSchemaResponseFormat,
     LogProbConfig,
     Message,
+    ModelStore,
     OpenAIChatCompletion,
     OpenAICompletion,
     OpenAIEmbeddingData,
@@ -54,6 +54,7 @@ from llama_stack.apis.inference import (
     ToolPromptFormat,
 )
 from llama_stack.apis.models import Model, ModelType
+from llama_stack.log import get_logger
 from llama_stack.models.llama.datatypes import BuiltinTool, StopReason, ToolCall
 from llama_stack.models.llama.sku_list import all_registered_models
 from llama_stack.providers.datatypes import (
@@ -84,7 +85,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 
 from .config import VLLMInferenceAdapterConfig
 
-log = logging.getLogger(__name__)
+log = get_logger(name=__name__, category="inference")
 
 
 def build_hf_repo_model_entries():
@@ -288,13 +289,40 @@ async def _process_vllm_chat_completion_stream_response(
 
 
 class VLLMInferenceAdapter(Inference, ModelsProtocolPrivate):
+    # automatically set by the resolver when instantiating the provider
+    __provider_id__: str
+    model_store: ModelStore | None = None
+
     def __init__(self, config: VLLMInferenceAdapterConfig) -> None:
         self.register_helper = ModelRegistryHelper(build_hf_repo_model_entries())
         self.config = config
         self.client = None
 
     async def initialize(self) -> None:
-        pass
+        if not self.config.url:
+            raise ValueError(
+                "You must provide a URL in run.yaml (or via the VLLM_URL environment variable) to use vLLM."
+            )
+
+    async def should_refresh_models(self) -> bool:
+        return self.config.refresh_models
+
+    async def list_models(self) -> list[Model] | None:
+        self._lazy_initialize_client()
+        assert self.client is not None  # mypy
+        models = []
+        async for m in self.client.models.list():
+            model_type = ModelType.llm  # unclear how to determine embedding vs. llm models
+            models.append(
+                Model(
+                    identifier=m.id,
+                    provider_resource_id=m.id,
+                    provider_id=self.__provider_id__,
+                    metadata={},
+                    model_type=model_type,
+                )
+            )
+        return models
 
     async def shutdown(self) -> None:
         pass

llama_stack/providers/remote/safety/sambanova/config.py

@@ -30,7 +30,7 @@ class SambaNovaSafetyConfig(BaseModel):
     )
 
     @classmethod
-    def sample_run_config(cls, api_key: str = "${env.SAMBANOVA_API_KEY}", **kwargs) -> dict[str, Any]:
+    def sample_run_config(cls, api_key: str = "${env.SAMBANOVA_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
             "url": "https://api.sambanova.ai/v1",
             "api_key": api_key,

llama_stack/providers/remote/vector_io/chroma/__init__.py

@@ -12,6 +12,6 @@ from .config import ChromaVectorIOConfig
 async def get_adapter_impl(config: ChromaVectorIOConfig, deps: dict[Api, ProviderSpec]):
     from .chroma import ChromaVectorIOAdapter
 
-    impl = ChromaVectorIOAdapter(config, deps[Api.inference])
+    impl = ChromaVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
     await impl.initialize()
     return impl

llama_stack/providers/remote/vector_io/chroma/chroma.py

@@ -12,25 +12,19 @@ from urllib.parse import urlparse
 import chromadb
 from numpy.typing import NDArray
 
+from llama_stack.apis.files import Files
 from llama_stack.apis.inference import InterleavedContent
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.apis.vector_io import (
     Chunk,
     QueryChunksResponse,
-    SearchRankingOptions,
     VectorIO,
-    VectorStoreChunkingStrategy,
-    VectorStoreDeleteResponse,
-    VectorStoreFileContentsResponse,
-    VectorStoreFileObject,
-    VectorStoreFileStatus,
-    VectorStoreListFilesResponse,
-    VectorStoreListResponse,
-    VectorStoreObject,
-    VectorStoreSearchResponsePage,
 )
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig
+from llama_stack.providers.utils.kvstore import kvstore_impl
+from llama_stack.providers.utils.kvstore.api import KVStore
+from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import (
     EmbeddingIndex,
     VectorDBWithIndex,
@@ -42,6 +36,13 @@ log = logging.getLogger(__name__)
 
 ChromaClientType = chromadb.api.AsyncClientAPI | chromadb.api.ClientAPI
 
+VERSION = "v3"
+VECTOR_DBS_PREFIX = f"vector_dbs:chroma:{VERSION}::"
+VECTOR_INDEX_PREFIX = f"vector_index:chroma:{VERSION}::"
+OPENAI_VECTOR_STORES_PREFIX = f"openai_vector_stores:chroma:{VERSION}::"
+OPENAI_VECTOR_STORES_FILES_PREFIX = f"openai_vector_stores_files:chroma:{VERSION}::"
+OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX = f"openai_vector_stores_files_contents:chroma:{VERSION}::"
+
 
 # this is a helper to allow us to use async and non-async chroma clients interchangeably
 async def maybe_await(result):
@@ -51,16 +52,20 @@ async def maybe_await(result):
 
 
 class ChromaIndex(EmbeddingIndex):
-    def __init__(self, client: ChromaClientType, collection):
+    def __init__(self, client: ChromaClientType, collection, kvstore: KVStore | None = None):
         self.client = client
         self.collection = collection
+        self.kvstore = kvstore
+
+    async def initialize(self):
+        pass
 
     async def add_chunks(self, chunks: list[Chunk], embeddings: NDArray):
         assert len(chunks) == len(embeddings), (
             f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
         )
 
-        ids = [f"{c.metadata['document_id']}:chunk-{i}" for i, c in enumerate(chunks)]
+        ids = [f"{c.metadata.get('document_id', '')}:{c.chunk_id}" for c in chunks]
         await maybe_await(
             self.collection.add(
                 documents=[chunk.model_dump_json() for chunk in chunks],
@@ -110,6 +115,9 @@ class ChromaIndex(EmbeddingIndex):
     ) -> QueryChunksResponse:
         raise NotImplementedError("Keyword search is not supported in Chroma")
 
+    async def delete_chunk(self, chunk_id: str) -> None:
+        raise NotImplementedError("delete_chunk is not supported in Chroma")
+
     async def query_hybrid(
         self,
         embedding: NDArray,
@@ -122,24 +130,26 @@ class ChromaIndex(EmbeddingIndex):
         raise NotImplementedError("Hybrid search is not supported in Chroma")
 
 
-class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
+class ChromaVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPrivate):
     def __init__(
         self,
         config: RemoteChromaVectorIOConfig | InlineChromaVectorIOConfig,
         inference_api: Api.inference,
+        files_api: Files | None,
     ) -> None:
         log.info(f"Initializing ChromaVectorIOAdapter with url: {config}")
         self.config = config
         self.inference_api = inference_api
-
         self.client = None
         self.cache = {}
+        self.kvstore: KVStore | None = None
+        self.vector_db_store = None
 
     async def initialize(self) -> None:
-        if isinstance(self.config, RemoteChromaVectorIOConfig):
-            if not self.config.url:
-                raise ValueError("URL is a required parameter for the remote Chroma provider's config")
+        self.kvstore = await kvstore_impl(self.config.kvstore)
+        self.vector_db_store = self.kvstore
 
+        if isinstance(self.config, RemoteChromaVectorIOConfig):
             log.info(f"Connecting to Chroma server at: {self.config.url}")
             url = self.config.url.rstrip("/")
             parsed = urlparse(url)
@@ -151,6 +161,7 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         else:
             log.info(f"Connecting to Chroma local db at: {self.config.db_path}")
             self.client = chromadb.PersistentClient(path=self.config.db_path)
+        self.openai_vector_stores = await self._load_openai_vector_stores()
 
     async def shutdown(self) -> None:
         pass
@@ -170,6 +181,10 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         )
 
     async def unregister_vector_db(self, vector_db_id: str) -> None:
+        if vector_db_id not in self.cache:
+            log.warning(f"Vector DB {vector_db_id} not found")
+            return
+
         await self.cache[vector_db_id].index.delete()
         del self.cache[vector_db_id]
 
@@ -180,6 +195,8 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         ttl_seconds: int | None = None,
     ) -> None:
         index = await self._get_and_cache_vector_db_index(vector_db_id)
+        if index is None:
+            raise ValueError(f"Vector DB {vector_db_id} not found in Chroma")
 
         await index.insert_chunks(chunks)
 
@@ -191,6 +208,9 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
     ) -> QueryChunksResponse:
         index = await self._get_and_cache_vector_db_index(vector_db_id)
 
+        if index is None:
+            raise ValueError(f"Vector DB {vector_db_id} not found in Chroma")
+
         return await index.query_chunks(query, params)
 
     async def _get_and_cache_vector_db_index(self, vector_db_id: str) -> VectorDBWithIndex:
@@ -207,106 +227,5 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         self.cache[vector_db_id] = index
         return index
 
-    async def openai_create_vector_store(
-        self,
-        name: str,
-        file_ids: list[str] | None = None,
-        expires_after: dict[str, Any] | None = None,
-        chunking_strategy: dict[str, Any] | None = None,
-        metadata: dict[str, Any] | None = None,
-        embedding_model: str | None = None,
-        embedding_dimension: int | None = 384,
-        provider_id: str | None = None,
-    ) -> VectorStoreObject:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_list_vector_stores(
-        self,
-        limit: int | None = 20,
-        order: str | None = "desc",
-        after: str | None = None,
-        before: str | None = None,
-    ) -> VectorStoreListResponse:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_retrieve_vector_store(
-        self,
-        vector_store_id: str,
-    ) -> VectorStoreObject:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_update_vector_store(
-        self,
-        vector_store_id: str,
-        name: str | None = None,
-        expires_after: dict[str, Any] | None = None,
-        metadata: dict[str, Any] | None = None,
-    ) -> VectorStoreObject:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_delete_vector_store(
-        self,
-        vector_store_id: str,
-    ) -> VectorStoreDeleteResponse:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_search_vector_store(
-        self,
-        vector_store_id: str,
-        query: str | list[str],
-        filters: dict[str, Any] | None = None,
-        max_num_results: int | None = 10,
-        ranking_options: SearchRankingOptions | None = None,
-        rewrite_query: bool | None = False,
-        search_mode: str | None = "vector",
-    ) -> VectorStoreSearchResponsePage:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_attach_file_to_vector_store(
-        self,
-        vector_store_id: str,
-        file_id: str,
-        attributes: dict[str, Any] | None = None,
-        chunking_strategy: VectorStoreChunkingStrategy | None = None,
-    ) -> VectorStoreFileObject:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_list_files_in_vector_store(
-        self,
-        vector_store_id: str,
-        limit: int | None = 20,
-        order: str | None = "desc",
-        after: str | None = None,
-        before: str | None = None,
-        filter: VectorStoreFileStatus | None = None,
-    ) -> VectorStoreListFilesResponse:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_retrieve_vector_store_file(
-        self,
-        vector_store_id: str,
-        file_id: str,
-    ) -> VectorStoreFileObject:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_retrieve_vector_store_file_contents(
-        self,
-        vector_store_id: str,
-        file_id: str,
-    ) -> VectorStoreFileContentsResponse:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_update_vector_store_file(
-        self,
-        vector_store_id: str,
-        file_id: str,
-        attributes: dict[str, Any] | None = None,
-    ) -> VectorStoreFileObject:
-        raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")
-
-    async def openai_delete_vector_store_file(
-        self,
-        vector_store_id: str,
-        file_id: str,
-    ) -> VectorStoreFileObject:
+    async def delete_chunks(self, store_id: str, chunk_ids: list[str]) -> None:
         raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")

llama_stack/providers/remote/vector_io/chroma/config.py

@@ -6,12 +6,23 @@
 
 from typing import Any
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+
+from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig
+from llama_stack.schema_utils import json_schema_type
 
 
+@json_schema_type
 class ChromaVectorIOConfig(BaseModel):
     url: str | None
+    kvstore: KVStoreConfig = Field(description="Config for KV store backend")
 
     @classmethod
-    def sample_run_config(cls, url: str = "${env.CHROMADB_URL}", **kwargs: Any) -> dict[str, Any]:
-        return {"url": url}
+    def sample_run_config(cls, __distro_dir__: str, url: str = "${env.CHROMADB_URL}", **kwargs: Any) -> dict[str, Any]:
+        return {
+            "url": url,
+            "kvstore": SqliteKVStoreConfig.sample_run_config(
+                __distro_dir__=__distro_dir__,
+                db_name="chroma_remote_registry.db",
+            ),
+        }

llama_stack/providers/remote/vector_io/milvus/milvus.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import json
 import logging
 import os
 import re
@@ -248,6 +247,16 @@ class MilvusIndex(EmbeddingIndex):
     ) -> QueryChunksResponse:
         raise NotImplementedError("Hybrid search is not supported in Milvus")
 
+    async def delete_chunk(self, chunk_id: str) -> None:
+        """Remove a chunk from the Milvus collection."""
+        try:
+            await asyncio.to_thread(
+                self.client.delete, collection_name=self.collection_name, filter=f'chunk_id == "{chunk_id}"'
+            )
+        except Exception as e:
+            logger.error(f"Error deleting chunk {chunk_id} from Milvus collection {self.collection_name}: {e}")
+            raise
+
 
 class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPrivate):
     def __init__(
@@ -371,185 +380,12 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
 
         return await index.query_chunks(query, params)
 
-    async def _save_openai_vector_store_file(
-        self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
-    ) -> None:
-        """Save vector store file metadata to Milvus database."""
-        if store_id not in self.openai_vector_stores:
-            store_info = await self._load_openai_vector_stores(store_id)
-            if not store_info:
-                logger.error(f"OpenAI vector store {store_id} not found")
-                raise ValueError(f"No vector store found with id {store_id}")
+    async def delete_chunks(self, store_id: str, chunk_ids: list[str]) -> None:
+        """Delete a chunk from a milvus vector store."""
+        index = await self._get_and_cache_vector_db_index(store_id)
+        if not index:
+            raise ValueError(f"Vector DB {store_id} not found")
 
-        try:
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files"):
-                file_schema = MilvusClient.create_schema(
-                    auto_id=False,
-                    enable_dynamic_field=True,
-                    description="Metadata for OpenAI vector store files",
-                )
-                file_schema.add_field(
-                    field_name="store_file_id", datatype=DataType.VARCHAR, is_primary=True, max_length=512
-                )
-                file_schema.add_field(field_name="store_id", datatype=DataType.VARCHAR, max_length=512)
-                file_schema.add_field(field_name="file_id", datatype=DataType.VARCHAR, max_length=512)
-                file_schema.add_field(field_name="file_info", datatype=DataType.VARCHAR, max_length=65535)
-
-                await asyncio.to_thread(
-                    self.client.create_collection,
-                    collection_name="openai_vector_store_files",
-                    schema=file_schema,
-                )
-
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files_contents"):
-                content_schema = MilvusClient.create_schema(
-                    auto_id=False,
-                    enable_dynamic_field=True,
-                    description="Contents for OpenAI vector store files",
-                )
-                content_schema.add_field(
-                    field_name="chunk_id", datatype=DataType.VARCHAR, is_primary=True, max_length=1024
-                )
-                content_schema.add_field(field_name="store_file_id", datatype=DataType.VARCHAR, max_length=1024)
-                content_schema.add_field(field_name="store_id", datatype=DataType.VARCHAR, max_length=512)
-                content_schema.add_field(field_name="file_id", datatype=DataType.VARCHAR, max_length=512)
-                content_schema.add_field(field_name="content", datatype=DataType.VARCHAR, max_length=65535)
-
-                await asyncio.to_thread(
-                    self.client.create_collection,
-                    collection_name="openai_vector_store_files_contents",
-                    schema=content_schema,
-                )
-
-            file_data = [
-                {
-                    "store_file_id": f"{store_id}_{file_id}",
-                    "store_id": store_id,
-                    "file_id": file_id,
-                    "file_info": json.dumps(file_info),
-                }
-            ]
-            await asyncio.to_thread(
-                self.client.upsert,
-                collection_name="openai_vector_store_files",
-                data=file_data,
-            )
-
-            # Save file contents
-            contents_data = [
-                {
-                    "chunk_id": content.get("chunk_metadata").get("chunk_id"),
-                    "store_file_id": f"{store_id}_{file_id}",
-                    "store_id": store_id,
-                    "file_id": file_id,
-                    "content": json.dumps(content),
-                }
-                for content in file_contents
-            ]
-            await asyncio.to_thread(
-                self.client.upsert,
-                collection_name="openai_vector_store_files_contents",
-                data=contents_data,
-            )
-
-        except Exception as e:
-            logger.error(f"Error saving openai vector store file {file_id} for store {store_id}: {e}")
-
-    async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
-        """Load vector store file metadata from Milvus database."""
-        try:
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files"):
-                return {}
-
-            query_filter = f"store_file_id == '{store_id}_{file_id}'"
-            results = await asyncio.to_thread(
-                self.client.query,
-                collection_name="openai_vector_store_files",
-                filter=query_filter,
-                output_fields=["file_info"],
-            )
-
-            if results:
-                try:
-                    return json.loads(results[0]["file_info"])
-                except json.JSONDecodeError as e:
-                    logger.error(f"Failed to decode file_info for store {store_id}, file {file_id}: {e}")
-                    return {}
-            return {}
-        except Exception as e:
-            logger.error(f"Error loading openai vector store file {file_id} for store {store_id}: {e}")
-            return {}
-
-    async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
-        """Update vector store file metadata in Milvus database."""
-        try:
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files"):
-                return
-
-            file_data = [
-                {
-                    "store_file_id": f"{store_id}_{file_id}",
-                    "store_id": store_id,
-                    "file_id": file_id,
-                    "file_info": json.dumps(file_info),
-                }
-            ]
-            await asyncio.to_thread(
-                self.client.upsert,
-                collection_name="openai_vector_store_files",
-                data=file_data,
-            )
-        except Exception as e:
-            logger.error(f"Error updating openai vector store file {file_id} for store {store_id}: {e}")
-            raise
-
-    async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
-        """Load vector store file contents from Milvus database."""
-        try:
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files_contents"):
-                return []
-
-            query_filter = (
-                f"store_id == '{store_id}' AND file_id == '{file_id}' AND store_file_id == '{store_id}_{file_id}'"
-            )
-            results = await asyncio.to_thread(
-                self.client.query,
-                collection_name="openai_vector_store_files_contents",
-                filter=query_filter,
-                output_fields=["chunk_id", "store_id", "file_id", "content"],
-            )
-
-            contents = []
-            for result in results:
-                try:
-                    content = json.loads(result["content"])
-                    contents.append(content)
-                except json.JSONDecodeError as e:
-                    logger.error(f"Failed to decode content for store {store_id}, file {file_id}: {e}")
-            return contents
-        except Exception as e:
-            logger.error(f"Error loading openai vector store file contents for {file_id} in store {store_id}: {e}")
-            return []
-
-    async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
-        """Delete vector store file metadata from Milvus database."""
-        try:
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files"):
-                return
-
-            query_filter = f"store_file_id in ['{store_id}_{file_id}']"
-            await asyncio.to_thread(
-                self.client.delete,
-                collection_name="openai_vector_store_files",
-                filter=query_filter,
-            )
-            if await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files_contents"):
-                await asyncio.to_thread(
-                    self.client.delete,
-                    collection_name="openai_vector_store_files_contents",
-                    filter=query_filter,
-                )
-
-        except Exception as e:
-            logger.error(f"Error deleting openai vector store file {file_id} for store {store_id}: {e}")
-            raise
+        for chunk_id in chunk_ids:
+            # Use the index's delete_chunk method
+            await index.index.delete_chunk(chunk_id)

llama_stack/providers/remote/vector_io/pgvector/__init__.py

@@ -12,6 +12,6 @@ from .config import PGVectorVectorIOConfig
 async def get_adapter_impl(config: PGVectorVectorIOConfig, deps: dict[Api, ProviderSpec]):
     from .pgvector import PGVectorVectorIOAdapter
 
-    impl = PGVectorVectorIOAdapter(config, deps[Api.inference])
+    impl = PGVectorVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files, None))
     await impl.initialize()
     return impl

llama_stack/providers/remote/vector_io/pgvector/pgvector.py

@@ -99,7 +99,7 @@ class PGVectorIndex(EmbeddingIndex):
         for i, chunk in enumerate(chunks):
             values.append(
                 (
-                    f"{chunk.metadata['document_id']}:chunk-{i}",
+                    f"{chunk.chunk_id}",
                     Json(chunk.model_dump()),
                     embeddings[i].tolist(),
                 )
@@ -159,6 +159,11 @@ class PGVectorIndex(EmbeddingIndex):
         with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
             cur.execute(f"DROP TABLE IF EXISTS {self.table_name}")
 
+    async def delete_chunk(self, chunk_id: str) -> None:
+        """Remove a chunk from the PostgreSQL table."""
+        with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
+            cur.execute(f"DELETE FROM {self.table_name} WHERE id = %s", (chunk_id,))
+
 
 class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPrivate):
     def __init__(
@@ -266,124 +271,12 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
         self.cache[vector_db_id] = VectorDBWithIndex(vector_db, index, self.inference_api)
         return self.cache[vector_db_id]
 
-    # OpenAI Vector Stores File operations are not supported in PGVector
-    async def _save_openai_vector_store_file(
-        self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
-    ) -> None:
-        """Save vector store file metadata to Postgres database."""
-        if self.conn is None:
-            raise RuntimeError("PostgreSQL connection is not initialized")
-        try:
-            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-                cur.execute(
-                    """
-                    CREATE TABLE IF NOT EXISTS openai_vector_store_files (
-                        store_id TEXT,
-                        file_id TEXT,
-                        metadata JSONB,
-                        PRIMARY KEY (store_id, file_id)
-                    )
-                    """
-                )
-                cur.execute(
-                    """
-                    CREATE TABLE IF NOT EXISTS openai_vector_store_files_contents (
-                        store_id TEXT,
-                        file_id TEXT,
-                        contents JSONB,
-                        PRIMARY KEY (store_id, file_id)
-                    )
-                    """
-                )
-                # Insert file metadata
-                files_query = sql.SQL(
-                    """
-                    INSERT INTO openai_vector_store_files (store_id, file_id, metadata)
-                    VALUES %s
-                    ON CONFLICT (store_id, file_id) DO UPDATE SET metadata = EXCLUDED.metadata
-                    """
-                )
-                files_values = [(store_id, file_id, Json(file_info))]
-                execute_values(cur, files_query, files_values, template="(%s, %s, %s)")
-                # Insert file contents
-                contents_query = sql.SQL(
-                    """
-                    INSERT INTO openai_vector_store_files_contents (store_id, file_id, contents)
-                    VALUES %s
-                    ON CONFLICT (store_id, file_id) DO UPDATE SET contents = EXCLUDED.contents
-                    """
-                )
-                contents_values = [(store_id, file_id, Json(file_contents))]
-                execute_values(cur, contents_query, contents_values, template="(%s, %s, %s)")
-        except Exception as e:
-            log.error(f"Error saving openai vector store file {file_id} for store {store_id}: {e}")
-            raise
+    async def delete_chunks(self, store_id: str, chunk_ids: list[str]) -> None:
+        """Delete a chunk from a PostgreSQL vector store."""
+        index = await self._get_and_cache_vector_db_index(store_id)
+        if not index:
+            raise ValueError(f"Vector DB {store_id} not found")
 
-    async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
-        """Load vector store file metadata from Postgres database."""
-        if self.conn is None:
-            raise RuntimeError("PostgreSQL connection is not initialized")
-        try:
-            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-                cur.execute(
-                    "SELECT metadata FROM openai_vector_store_files WHERE store_id = %s AND file_id = %s",
-                    (store_id, file_id),
-                )
-                row = cur.fetchone()
-                return row[0] if row and row[0] is not None else {}
-        except Exception as e:
-            log.error(f"Error loading openai vector store file {file_id} for store {store_id}: {e}")
-            return {}
-
-    async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
-        """Load vector store file contents from Postgres database."""
-        if self.conn is None:
-            raise RuntimeError("PostgreSQL connection is not initialized")
-        try:
-            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-                cur.execute(
-                    "SELECT contents FROM openai_vector_store_files_contents WHERE store_id = %s AND file_id = %s",
-                    (store_id, file_id),
-                )
-                row = cur.fetchone()
-                return row[0] if row and row[0] is not None else []
-        except Exception as e:
-            log.error(f"Error loading openai vector store file contents for {file_id} in store {store_id}: {e}")
-            return []
-
-    async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
-        """Update vector store file metadata in Postgres database."""
-        if self.conn is None:
-            raise RuntimeError("PostgreSQL connection is not initialized")
-        try:
-            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-                query = sql.SQL(
-                    """
-                    INSERT INTO openai_vector_store_files (store_id, file_id, metadata)
-                    VALUES %s
-                    ON CONFLICT (store_id, file_id) DO UPDATE SET metadata = EXCLUDED.metadata
-                    """
-                )
-                values = [(store_id, file_id, Json(file_info))]
-                execute_values(cur, query, values, template="(%s, %s, %s)")
-        except Exception as e:
-            log.error(f"Error updating openai vector store file {file_id} for store {store_id}: {e}")
-            raise
-
-    async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
-        """Delete vector store file metadata from Postgres database."""
-        if self.conn is None:
-            raise RuntimeError("PostgreSQL connection is not initialized")
-        try:
-            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-                cur.execute(
-                    "DELETE FROM openai_vector_store_files WHERE store_id = %s AND file_id = %s",
-                    (store_id, file_id),
-                )
-                cur.execute(
-                    "DELETE FROM openai_vector_store_files_contents WHERE store_id = %s AND file_id = %s",
-                    (store_id, file_id),
-                )
-        except Exception as e:
-            log.error(f"Error deleting openai vector store file {file_id} for store {store_id}: {e}")
-            raise
+        for chunk_id in chunk_ids:
+            # Use the index's delete_chunk method
+            await index.index.delete_chunk(chunk_id)

llama_stack/providers/remote/vector_io/qdrant/qdrant.py

@@ -82,6 +82,9 @@ class QdrantIndex(EmbeddingIndex):
 
         await self.client.upsert(collection_name=self.collection_name, points=points)
 
+    async def delete_chunk(self, chunk_id: str) -> None:
+        raise NotImplementedError("delete_chunk is not supported in qdrant")
+
     async def query_vector(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
         results = (
             await self.client.query_points(
@@ -307,3 +310,6 @@ class QdrantVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         file_id: str,
     ) -> VectorStoreFileObject:
         raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
+
+    async def delete_chunks(self, store_id: str, chunk_ids: list[str]) -> None:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")

llama_stack/providers/remote/vector_io/weaviate/weaviate.py

@@ -66,6 +66,9 @@ class WeaviateIndex(EmbeddingIndex):
         # TODO: make this async friendly
         collection.data.insert_many(data_objects)
 
+    async def delete_chunk(self, chunk_id: str) -> None:
+        raise NotImplementedError("delete_chunk is not supported in Chroma")
+
     async def query_vector(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
         collection = self.client.collections.get(self.collection_name)
 
@@ -264,3 +267,6 @@ class WeaviateVectorIOAdapter(
 
     async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
         raise NotImplementedError("OpenAI Vector Stores API is not supported in Weaviate")
+
+    async def delete_chunks(self, store_id: str, chunk_ids: list[str]) -> None:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in Weaviate")

llama_stack/providers/utils/inference/embedding_mixin.py

@@ -88,7 +88,7 @@ class SentenceTransformerEmbeddingMixin:
         usage = OpenAIEmbeddingUsage(prompt_tokens=-1, total_tokens=-1)
         return OpenAIEmbeddingsResponse(
             data=data,
-            model=model_obj.provider_resource_id,
+            model=model,
             usage=usage,
         )
 

llama_stack/providers/utils/inference/litellm_openai_mixin.py

@@ -68,11 +68,23 @@ class LiteLLMOpenAIMixin(
     def __init__(
         self,
         model_entries,
+        litellm_provider_name: str,
         api_key_from_config: str | None,
         provider_data_api_key_field: str,
         openai_compat_api_base: str | None = None,
     ):
+        """
+        Initialize the LiteLLMOpenAIMixin.
+
+        :param model_entries: The model entries to register.
+        :param api_key_from_config: The API key to use from the config.
+        :param provider_data_api_key_field: The field in the provider data that contains the API key.
+        :param litellm_provider_name: The name of the provider, used for model lookups.
+        :param openai_compat_api_base: The base URL for OpenAI compatibility, or None if not using OpenAI compatibility.
+        """
         ModelRegistryHelper.__init__(self, model_entries)
+
+        self.litellm_provider_name = litellm_provider_name
         self.api_key_from_config = api_key_from_config
         self.provider_data_api_key_field = provider_data_api_key_field
         self.api_base = openai_compat_api_base
@@ -91,7 +103,11 @@ class LiteLLMOpenAIMixin(
     def get_litellm_model_name(self, model_id: str) -> str:
         # users may be using openai/ prefix in their model names. the openai/models.py did this by default.
         # model_id.startswith("openai/") is for backwards compatibility.
-        return "openai/" + model_id if self.is_openai_compat and not model_id.startswith("openai/") else model_id
+        return (
+            f"{self.litellm_provider_name}/{model_id}"
+            if self.is_openai_compat and not model_id.startswith(self.litellm_provider_name)
+            else model_id
+        )
 
     async def completion(
         self,
@@ -421,3 +437,17 @@ class LiteLLMOpenAIMixin(
         logprobs: LogProbConfig | None = None,
     ):
         raise NotImplementedError("Batch chat completion is not supported for OpenAI Compat")
+
+    async def check_model_availability(self, model: str) -> bool:
+        """
+        Check if a specific model is available via LiteLLM for the current
+        provider (self.litellm_provider_name).
+
+        :param model: The model identifier to check.
+        :return: True if the model is available dynamically, False otherwise.
+        """
+        if self.litellm_provider_name not in litellm.models_by_provider:
+            logger.error(f"Provider {self.litellm_provider_name} is not registered in litellm.")
+            return False
+
+        return model in litellm.models_by_provider[self.litellm_provider_name]

llama_stack/providers/utils/inference/model_registry.py

@@ -10,12 +10,22 @@ from pydantic import BaseModel, Field
 
 from llama_stack.apis.common.errors import UnsupportedModelError
 from llama_stack.apis.models import ModelType
+from llama_stack.log import get_logger
 from llama_stack.models.llama.sku_list import all_registered_models
 from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
 from llama_stack.providers.utils.inference import (
     ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR,
 )
 
+logger = get_logger(name=__name__, category="core")
+
+
+class RemoteInferenceProviderConfig(BaseModel):
+    allowed_models: list[str] | None = Field(
+        default=None,
+        description="List of models that should be registered with the model registry. If None, all models are allowed.",
+    )
+
 
 # TODO: this class is more confusing than useful right now. We need to make it
 # more closer to the Model class.
@@ -40,7 +50,8 @@ def build_hf_repo_model_entry(
     additional_aliases: list[str] | None = None,
 ) -> ProviderModelEntry:
     aliases = [
-        get_huggingface_repo(model_descriptor),
+        # NOTE: avoid HF aliases because they _cannot_ be unique across providers
+        # get_huggingface_repo(model_descriptor),
     ]
     if additional_aliases:
         aliases.extend(additional_aliases)
@@ -62,7 +73,12 @@ def build_model_entry(provider_model_id: str, model_descriptor: str) -> Provider
 
 
 class ModelRegistryHelper(ModelsProtocolPrivate):
-    def __init__(self, model_entries: list[ProviderModelEntry]):
+    __provider_id__: str
+
+    def __init__(self, model_entries: list[ProviderModelEntry], allowed_models: list[str] | None = None):
+        self.model_entries = model_entries
+        self.allowed_models = allowed_models
+
         self.alias_to_provider_id_map = {}
         self.provider_id_to_llama_model_map = {}
         for entry in model_entries:
@@ -76,6 +92,27 @@ class ModelRegistryHelper(ModelsProtocolPrivate):
                 self.alias_to_provider_id_map[entry.llama_model] = entry.provider_model_id
                 self.provider_id_to_llama_model_map[entry.provider_model_id] = entry.llama_model
 
+    async def list_models(self) -> list[Model] | None:
+        models = []
+        for entry in self.model_entries:
+            ids = [entry.provider_model_id] + entry.aliases
+            for id in ids:
+                if self.allowed_models and id not in self.allowed_models:
+                    continue
+                models.append(
+                    Model(
+                        identifier=id,
+                        provider_resource_id=entry.provider_model_id,
+                        model_type=ModelType.llm,
+                        metadata=entry.metadata,
+                        provider_id=self.__provider_id__,
+                    )
+                )
+        return models
+
+    async def should_refresh_models(self) -> bool:
+        return False
+
     def get_provider_model_id(self, identifier: str) -> str | None:
         return self.alias_to_provider_id_map.get(identifier, None)
 
@@ -98,6 +135,9 @@ class ModelRegistryHelper(ModelsProtocolPrivate):
         :param model: The model identifier to check.
         :return: True if the model is available dynamically, False otherwise.
         """
+        logger.info(
+            f"check_model_availability is not implemented for {self.__class__.__name__}. Returning False by default."
+        )
         return False
 
     async def register_model(self, model: Model) -> Model:
@@ -148,8 +188,8 @@ class ModelRegistryHelper(ModelsProtocolPrivate):
         return model
 
     async def unregister_model(self, model_id: str) -> None:
-        # TODO: should we block unregistering base supported provider model IDs?
-        if model_id not in self.alias_to_provider_id_map:
-            raise ValueError(f"Model id '{model_id}' is not registered.")
-
-        del self.alias_to_provider_id_map[model_id]
+        # model_id is the identifier, not the provider_resource_id
+        # unfortunately, this ID can be of the form provider_id/model_id which
+        # we never registered. TODO: fix this by significantly rewriting
+        # registration and registry helper
+        pass

llama_stack/providers/utils/inference/openai_mixin.py

@@ -0,0 +1,272 @@
+# 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 abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from typing import Any
+
+import openai
+from openai import NOT_GIVEN, AsyncOpenAI
+
+from llama_stack.apis.inference import (
+    Model,
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAICompletion,
+    OpenAIEmbeddingData,
+    OpenAIEmbeddingsResponse,
+    OpenAIEmbeddingUsage,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
+)
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.inference.openai_compat import prepare_openai_completion_params
+
+logger = get_logger(name=__name__, category="core")
+
+
+class OpenAIMixin(ABC):
+    """
+    Mixin class that provides OpenAI-specific functionality for inference providers.
+    This class handles direct OpenAI API calls using the AsyncOpenAI client.
+
+    This is an abstract base class that requires child classes to implement:
+    - get_api_key(): Method to retrieve the API key
+    - get_base_url(): Method to retrieve the OpenAI-compatible API base URL
+
+    Expected Dependencies:
+    - self.model_store: Injected by the Llama Stack distribution system at runtime.
+      This provides model registry functionality for looking up registered models.
+      The model_store is set in routing_tables/common.py during provider initialization.
+    """
+
+    @abstractmethod
+    def get_api_key(self) -> str:
+        """
+        Get the API key.
+
+        This method must be implemented by child classes to provide the API key
+        for authenticating with the OpenAI API or compatible endpoints.
+
+        :return: The API key as a string
+        """
+        pass
+
+    @abstractmethod
+    def get_base_url(self) -> str:
+        """
+        Get the OpenAI-compatible API base URL.
+
+        This method must be implemented by child classes to provide the base URL
+        for the OpenAI API or compatible endpoints (e.g., "https://api.openai.com/v1").
+
+        :return: The base URL as a string
+        """
+        pass
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        """
+        Get an AsyncOpenAI client instance.
+
+        Uses the abstract methods get_api_key() and get_base_url() which must be
+        implemented by child classes.
+        """
+        return AsyncOpenAI(
+            api_key=self.get_api_key(),
+            base_url=self.get_base_url(),
+        )
+
+    async def _get_provider_model_id(self, model: str) -> str:
+        """
+        Get the provider-specific model ID from the model store.
+
+        This is a utility method that looks up the registered model and returns
+        the provider_resource_id that should be used for actual API calls.
+
+        :param model: The registered model name/identifier
+        :return: The provider-specific model ID (e.g., "gpt-4")
+        """
+        # Look up the registered model to get the provider-specific model ID
+        # self.model_store is injected by the distribution system at runtime
+        model_obj: Model = await self.model_store.get_model(model)  # type: ignore[attr-defined]
+        # provider_resource_id is str | None, but we expect it to be str for OpenAI calls
+        if model_obj.provider_resource_id is None:
+            raise ValueError(f"Model {model} has no provider_resource_id")
+        return model_obj.provider_resource_id
+
+    async def openai_completion(
+        self,
+        model: str,
+        prompt: str | list[str] | list[int] | list[list[int]],
+        best_of: int | None = None,
+        echo: bool | None = None,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[str, float] | None = None,
+        logprobs: bool | None = None,
+        max_tokens: int | None = None,
+        n: int | None = None,
+        presence_penalty: float | None = None,
+        seed: int | None = None,
+        stop: str | list[str] | None = None,
+        stream: bool | None = None,
+        stream_options: dict[str, Any] | None = None,
+        temperature: float | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+        guided_choice: list[str] | None = None,
+        prompt_logprobs: int | None = None,
+        suffix: str | None = None,
+    ) -> OpenAICompletion:
+        """
+        Direct OpenAI completion API call.
+        """
+        if guided_choice is not None:
+            logger.warning("guided_choice is not supported by the OpenAI API. Ignoring.")
+        if prompt_logprobs is not None:
+            logger.warning("prompt_logprobs is not supported by the OpenAI API. Ignoring.")
+
+        # TODO: fix openai_completion to return type compatible with OpenAI's API response
+        return await self.client.completions.create(  # type: ignore[no-any-return]
+            **await prepare_openai_completion_params(
+                model=await self._get_provider_model_id(model),
+                prompt=prompt,
+                best_of=best_of,
+                echo=echo,
+                frequency_penalty=frequency_penalty,
+                logit_bias=logit_bias,
+                logprobs=logprobs,
+                max_tokens=max_tokens,
+                n=n,
+                presence_penalty=presence_penalty,
+                seed=seed,
+                stop=stop,
+                stream=stream,
+                stream_options=stream_options,
+                temperature=temperature,
+                top_p=top_p,
+                user=user,
+                suffix=suffix,
+            )
+        )
+
+    async def openai_chat_completion(
+        self,
+        model: str,
+        messages: list[OpenAIMessageParam],
+        frequency_penalty: float | None = None,
+        function_call: str | dict[str, Any] | None = None,
+        functions: list[dict[str, Any]] | None = None,
+        logit_bias: dict[str, float] | None = None,
+        logprobs: bool | None = None,
+        max_completion_tokens: int | None = None,
+        max_tokens: int | None = None,
+        n: int | None = None,
+        parallel_tool_calls: bool | None = None,
+        presence_penalty: float | None = None,
+        response_format: OpenAIResponseFormatParam | None = None,
+        seed: int | None = None,
+        stop: str | list[str] | None = None,
+        stream: bool | None = None,
+        stream_options: dict[str, Any] | None = None,
+        temperature: float | None = None,
+        tool_choice: str | dict[str, Any] | None = None,
+        tools: list[dict[str, Any]] | None = None,
+        top_logprobs: int | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+    ) -> OpenAIChatCompletion | AsyncIterator[OpenAIChatCompletionChunk]:
+        """
+        Direct OpenAI chat completion API call.
+        """
+        # Type ignore because return types are compatible
+        return await self.client.chat.completions.create(  # type: ignore[no-any-return]
+            **await prepare_openai_completion_params(
+                model=await self._get_provider_model_id(model),
+                messages=messages,
+                frequency_penalty=frequency_penalty,
+                function_call=function_call,
+                functions=functions,
+                logit_bias=logit_bias,
+                logprobs=logprobs,
+                max_completion_tokens=max_completion_tokens,
+                max_tokens=max_tokens,
+                n=n,
+                parallel_tool_calls=parallel_tool_calls,
+                presence_penalty=presence_penalty,
+                response_format=response_format,
+                seed=seed,
+                stop=stop,
+                stream=stream,
+                stream_options=stream_options,
+                temperature=temperature,
+                tool_choice=tool_choice,
+                tools=tools,
+                top_logprobs=top_logprobs,
+                top_p=top_p,
+                user=user,
+            )
+        )
+
+    async def openai_embeddings(
+        self,
+        model: str,
+        input: str | list[str],
+        encoding_format: str | None = "float",
+        dimensions: int | None = None,
+        user: str | None = None,
+    ) -> OpenAIEmbeddingsResponse:
+        """
+        Direct OpenAI embeddings API call.
+        """
+        # Call OpenAI embeddings API with properly typed parameters
+        response = await self.client.embeddings.create(
+            model=await self._get_provider_model_id(model),
+            input=input,
+            encoding_format=encoding_format if encoding_format is not None else NOT_GIVEN,
+            dimensions=dimensions if dimensions is not None else NOT_GIVEN,
+            user=user if user is not None else NOT_GIVEN,
+        )
+
+        data = []
+        for i, embedding_data in enumerate(response.data):
+            data.append(
+                OpenAIEmbeddingData(
+                    embedding=embedding_data.embedding,
+                    index=i,
+                )
+            )
+
+        usage = OpenAIEmbeddingUsage(
+            prompt_tokens=response.usage.prompt_tokens,
+            total_tokens=response.usage.total_tokens,
+        )
+
+        return OpenAIEmbeddingsResponse(
+            data=data,
+            model=response.model,
+            usage=usage,
+        )
+
+    async def check_model_availability(self, model: str) -> bool:
+        """
+        Check if a specific model is available from OpenAI.
+
+        :param model: The model identifier to check.
+        :return: True if the model is available dynamically, False otherwise.
+        """
+        try:
+            # Direct model lookup - returns model or raises NotFoundError
+            await self.client.models.retrieve(model)
+            return True
+        except openai.NotFoundError:
+            # Model doesn't exist - this is expected for unavailable models
+            pass
+        except Exception as e:
+            # All other errors (auth, rate limit, network, etc.)
+            logger.warning(f"Failed to check model availability for {model}: {e}")
+
+        return False

llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -66,7 +66,7 @@ class OpenAIVectorStoreMixin(ABC):
 
     async def _save_openai_vector_store(self, store_id: str, store_info: dict[str, Any]) -> None:
         """Save vector store metadata to persistent storage."""
-        assert self.kvstore is not None
+        assert self.kvstore
         key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
         await self.kvstore.set(key=key, value=json.dumps(store_info))
         # update in-memory cache
@@ -74,7 +74,7 @@ class OpenAIVectorStoreMixin(ABC):
 
     async def _load_openai_vector_stores(self) -> dict[str, dict[str, Any]]:
         """Load all vector store metadata from persistent storage."""
-        assert self.kvstore is not None
+        assert self.kvstore
         start_key = OPENAI_VECTOR_STORES_PREFIX
         end_key = f"{OPENAI_VECTOR_STORES_PREFIX}\xff"
         stored_data = await self.kvstore.values_in_range(start_key, end_key)
@@ -87,7 +87,7 @@ class OpenAIVectorStoreMixin(ABC):
 
     async def _update_openai_vector_store(self, store_id: str, store_info: dict[str, Any]) -> None:
         """Update vector store metadata in persistent storage."""
-        assert self.kvstore is not None
+        assert self.kvstore
         key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
         await self.kvstore.set(key=key, value=json.dumps(store_info))
         # update in-memory cache
@@ -95,37 +95,66 @@ class OpenAIVectorStoreMixin(ABC):
 
     async def _delete_openai_vector_store_from_storage(self, store_id: str) -> None:
         """Delete vector store metadata from persistent storage."""
-        assert self.kvstore is not None
+        assert self.kvstore
         key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
         await self.kvstore.delete(key)
         # remove from in-memory cache
         self.openai_vector_stores.pop(store_id, None)
 
-    @abstractmethod
     async def _save_openai_vector_store_file(
         self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
     ) -> None:
         """Save vector store file metadata to persistent storage."""
-        pass
+        assert self.kvstore
+        meta_key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
+        await self.kvstore.set(key=meta_key, value=json.dumps(file_info))
+        contents_prefix = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}:"
+        for idx, chunk in enumerate(file_contents):
+            await self.kvstore.set(key=f"{contents_prefix}{idx}", value=json.dumps(chunk))
 
-    @abstractmethod
     async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
         """Load vector store file metadata from persistent storage."""
-        pass
+        assert self.kvstore
+        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
+        stored_data = await self.kvstore.get(key)
+        return json.loads(stored_data) if stored_data else {}
 
-    @abstractmethod
     async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
         """Load vector store file contents from persistent storage."""
-        pass
+        assert self.kvstore
+        prefix = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}:"
+        end_key = f"{prefix}\xff"
+        raw_items = await self.kvstore.values_in_range(prefix, end_key)
+        return [json.loads(item) for item in raw_items]
 
-    @abstractmethod
     async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
         """Update vector store file metadata in persistent storage."""
-        pass
+        assert self.kvstore
+        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
+        await self.kvstore.set(key=key, value=json.dumps(file_info))
 
-    @abstractmethod
     async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
         """Delete vector store file metadata from persistent storage."""
+        assert self.kvstore
+
+        meta_key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
+        await self.kvstore.delete(meta_key)
+
+        contents_prefix = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}:"
+        end_key = f"{contents_prefix}\xff"
+        # load all stored chunk values (values_in_range is implemented by all backends)
+        raw_items = await self.kvstore.values_in_range(contents_prefix, end_key)
+        # delete each chunk by its index suffix
+        for idx in range(len(raw_items)):
+            await self.kvstore.delete(f"{contents_prefix}{idx}")
+
+    async def initialize_openai_vector_stores(self) -> None:
+        """Load existing OpenAI vector stores into the in-memory cache."""
+        self.openai_vector_stores = await self._load_openai_vector_stores()
+
+    @abstractmethod
+    async def delete_chunks(self, store_id: str, chunk_ids: list[str]) -> None:
+        """Delete a chunk from a vector store."""
         pass
 
     @abstractmethod
@@ -138,10 +167,6 @@ class OpenAIVectorStoreMixin(ABC):
         """Unregister a vector database (provider-specific implementation)."""
         pass
 
-    async def initialize_openai_vector_stores(self) -> None:
-        """Load existing OpenAI vector stores into the in-memory cache."""
-        self.openai_vector_stores = await self._load_openai_vector_stores()
-
     @abstractmethod
     async def insert_chunks(
         self,
@@ -161,7 +186,7 @@ class OpenAIVectorStoreMixin(ABC):
 
     async def openai_create_vector_store(
         self,
-        name: str,
+        name: str | None = None,
         file_ids: list[str] | None = None,
         expires_after: dict[str, Any] | None = None,
         chunking_strategy: dict[str, Any] | None = None,
@@ -743,17 +768,15 @@ class OpenAIVectorStoreMixin(ABC):
         if vector_store_id not in self.openai_vector_stores:
             raise ValueError(f"Vector store {vector_store_id} not found")
 
+        dict_chunks = await self._load_openai_vector_store_file_contents(vector_store_id, file_id)
+        chunks = [Chunk.model_validate(c) for c in dict_chunks]
+        await self.delete_chunks(vector_store_id, [str(c.chunk_id) for c in chunks if c.chunk_id])
+
         store_info = self.openai_vector_stores[vector_store_id].copy()
 
         file = await self.openai_retrieve_vector_store_file(vector_store_id, file_id)
         await self._delete_openai_vector_store_file_from_storage(vector_store_id, file_id)
 
-        # TODO: We need to actually delete the embeddings from the underlying vector store...
-        # Also uncomment the corresponding integration test marked as xfail
-        #
-        # test_openai_vector_store_delete_file_removes_from_vector_store in
-        # tests/integration/vector_io/test_openai_vector_stores.py
-
         # Update in-memory cache
         store_info["file_ids"].remove(file_id)
         store_info["file_counts"][file.status] -= 1

llama_stack/providers/utils/memory/vector_store.py

@@ -231,6 +231,10 @@ class EmbeddingIndex(ABC):
     async def add_chunks(self, chunks: list[Chunk], embeddings: NDArray):
         raise NotImplementedError()
 
+    @abstractmethod
+    async def delete_chunk(self, chunk_id: str):
+        raise NotImplementedError()
+
     @abstractmethod
     async def query_vector(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
         raise NotImplementedError()

llama_stack/providers/utils/telemetry/sqlite_trace_store.py

@@ -83,6 +83,7 @@ class SQLiteTraceStore(TraceStore):
             )
             SELECT DISTINCT trace_id, root_span_id, start_time, end_time
             FROM filtered_traces
+            WHERE root_span_id IS NOT NULL
             LIMIT {limit} OFFSET {offset}
         """
 
@@ -166,7 +167,11 @@ class SQLiteTraceStore(TraceStore):
                 return spans_by_id
 
     async def get_trace(self, trace_id: str) -> Trace:
-        query = "SELECT * FROM traces WHERE trace_id = ?"
+        query = """
+            SELECT *
+            FROM traces t
+            WHERE t.trace_id = ?
+        """
         async with aiosqlite.connect(self.conn_string) as conn:
             conn.row_factory = aiosqlite.Row
             async with conn.execute(query, (trace_id,)) as cursor:

llama_stack/providers/utils/tools/mcp.py

@@ -4,13 +4,16 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
+from enum import Enum
 from typing import Any, cast
 
 import httpx
-from mcp import ClientSession
+from mcp import ClientSession, McpError
 from mcp import types as mcp_types
 from mcp.client.sse import sse_client
+from mcp.client.streamable_http import streamablehttp_client
 
 from llama_stack.apis.common.content_types import ImageContentItem, InterleavedContentItem, TextContentItem
 from llama_stack.apis.tools import (
@@ -21,31 +24,61 @@ from llama_stack.apis.tools import (
 )
 from llama_stack.distribution.datatypes import AuthenticationRequiredError
 from llama_stack.log import get_logger
+from llama_stack.providers.utils.tools.ttl_dict import TTLDict
 
 logger = get_logger(__name__, category="tools")
 
+protocol_cache = TTLDict(ttl_seconds=3600)
+
+
+class MCPProtol(Enum):
+    UNKNOWN = 0
+    STREAMABLE_HTTP = 1
+    SSE = 2
+
 
 @asynccontextmanager
-async def sse_client_wrapper(endpoint: str, headers: dict[str, str]):
-    try:
-        async with sse_client(endpoint, headers=headers) as streams:
-            async with ClientSession(*streams) as session:
-                await session.initialize()
-                yield session
-    except* httpx.HTTPStatusError as eg:
-        for exc in eg.exceptions:
-            # mypy does not currently narrow the type of `eg.exceptions` based on the `except*` filter,
-            # so we explicitly cast each item to httpx.HTTPStatusError. This is safe because
-            # `except* httpx.HTTPStatusError` guarantees all exceptions in `eg.exceptions` are of that type.
-            err = cast(httpx.HTTPStatusError, exc)
-            if err.response.status_code == 401:
-                raise AuthenticationRequiredError(exc) from exc
-        raise
+async def client_wrapper(endpoint: str, headers: dict[str, str]) -> AsyncGenerator[ClientSession, Any]:
+    # we use a ttl'd dict to cache the happy path protocol for each endpoint
+    # but, we always fall back to trying the other protocol if we cannot initialize the session
+    connection_strategies = [MCPProtol.STREAMABLE_HTTP, MCPProtol.SSE]
+    mcp_protocol = protocol_cache.get(endpoint, default=MCPProtol.UNKNOWN)
+    if mcp_protocol == MCPProtol.SSE:
+        connection_strategies = [MCPProtol.SSE, MCPProtol.STREAMABLE_HTTP]
+
+    for i, strategy in enumerate(connection_strategies):
+        try:
+            client = streamablehttp_client
+            if strategy == MCPProtol.SSE:
+                client = sse_client
+            async with client(endpoint, headers=headers) as client_streams:
+                async with ClientSession(read_stream=client_streams[0], write_stream=client_streams[1]) as session:
+                    await session.initialize()
+                    protocol_cache[endpoint] = strategy
+                    yield session
+                    return
+        except* httpx.HTTPStatusError as eg:
+            for exc in eg.exceptions:
+                # mypy does not currently narrow the type of `eg.exceptions` based on the `except*` filter,
+                # so we explicitly cast each item to httpx.HTTPStatusError. This is safe because
+                # `except* httpx.HTTPStatusError` guarantees all exceptions in `eg.exceptions` are of that type.
+                err = cast(httpx.HTTPStatusError, exc)
+                if err.response.status_code == 401:
+                    raise AuthenticationRequiredError(exc) from exc
+            if i == len(connection_strategies) - 1:
+                raise
+        except* McpError:
+            if i < len(connection_strategies) - 1:
+                logger.warning(
+                    f"failed to connect via {strategy.name}, falling back to {connection_strategies[i + 1].name}"
+                )
+            else:
+                raise
 
 
 async def list_mcp_tools(endpoint: str, headers: dict[str, str]) -> ListToolDefsResponse:
     tools = []
-    async with sse_client_wrapper(endpoint, headers) as session:
+    async with client_wrapper(endpoint, headers) as session:
         tools_result = await session.list_tools()
         for tool in tools_result.tools:
             parameters = []
@@ -73,7 +106,7 @@ async def list_mcp_tools(endpoint: str, headers: dict[str, str]) -> ListToolDefs
 async def invoke_mcp_tool(
     endpoint: str, headers: dict[str, str], tool_name: str, kwargs: dict[str, Any]
 ) -> ToolInvocationResult:
-    async with sse_client_wrapper(endpoint, headers) as session:
+    async with client_wrapper(endpoint, headers) as session:
         result = await session.call_tool(tool_name, kwargs)
 
         content: list[InterleavedContentItem] = []

llama_stack/providers/utils/tools/ttl_dict.py

@@ -0,0 +1,70 @@
+# 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 time
+from threading import RLock
+from typing import Any
+
+
+class TTLDict(dict):
+    """
+    A dictionary with a ttl for each item
+    """
+
+    def __init__(self, ttl_seconds: float, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.ttl_seconds = ttl_seconds
+        self._expires: dict[Any, Any] = {}  # expires holds when an item will expire
+        self._lock = RLock()
+
+        if args or kwargs:
+            for k, v in self.items():
+                self.__setitem__(k, v)
+
+    def __delitem__(self, key):
+        with self._lock:
+            del self._expires[key]
+            super().__delitem__(key)
+
+    def __setitem__(self, key, value):
+        with self._lock:
+            self._expires[key] = time.monotonic() + self.ttl_seconds
+            super().__setitem__(key, value)
+
+    def _is_expired(self, key):
+        if key not in self._expires:
+            return False
+        return time.monotonic() > self._expires[key]
+
+    def __getitem__(self, key):
+        with self._lock:
+            if self._is_expired(key):
+                del self._expires[key]
+                super().__delitem__(key)
+                raise KeyError(f"{key} has expired and was removed")
+
+            return super().__getitem__(key)
+
+    def get(self, key, default=None):
+        try:
+            return self[key]
+        except KeyError:
+            return default
+
+    def __contains__(self, key):
+        try:
+            _ = self[key]
+            return True
+        except KeyError:
+            return False
+
+    def __repr__(self):
+        with self._lock:
+            for key in self.keys():
+                if self._is_expired(key):
+                    del self._expires[key]
+                    super().__delitem__(key)
+            return f"TTLDict({self.ttl_seconds}, {super().__repr__()})"