feat(api)!: support extra_body to embeddings and vector_stores APIs (#3794)

Applies the same pattern from
https://github.com/llamastack/llama-stack/pull/3777 to embeddings and
vector_stores.create() endpoints.

This should _not_ be a breaking change since (a) our tests were already
using the `extra_body` parameter when passing in to the backend (b) but
the backend probably wasn't extracting the parameters correctly. This PR
will fix that.

Updated APIs: `openai_embeddings(), openai_create_vector_store(),
openai_create_vector_store_file_batch()`

llama_stack/apis/inference/inference.py

@@ -1140,6 +1140,25 @@ class OpenAIChatCompletionRequestWithExtraBody(BaseModel, extra="allow"):
     user: str | None = None
 
 
+# extra_body can be accessed via .model_extra
+@json_schema_type
+class OpenAIEmbeddingsRequestWithExtraBody(BaseModel, extra="allow"):
+    """Request parameters for OpenAI-compatible embeddings endpoint.
+
+    :param model: The identifier of the model to use. The model must be an embedding model registered with Llama Stack and available via the /models endpoint.
+    :param input: Input text to embed, encoded as a string or array of strings. To embed multiple inputs in a single request, pass an array of strings.
+    :param encoding_format: (Optional) The format to return the embeddings in. Can be either "float" or "base64". Defaults to "float".
+    :param dimensions: (Optional) The number of dimensions the resulting output embeddings should have. Only supported in text-embedding-3 and later models.
+    :param user: (Optional) A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
+    """
+
+    model: str
+    input: str | list[str]
+    encoding_format: str | None = "float"
+    dimensions: int | None = None
+    user: str | None = None
+
+
 @runtime_checkable
 @trace_protocol
 class InferenceProvider(Protocol):
@@ -1200,21 +1219,11 @@ class InferenceProvider(Protocol):
     @webmethod(route="/embeddings", method="POST", level=LLAMA_STACK_API_V1)
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: Annotated[OpenAIEmbeddingsRequestWithExtraBody, Body(...)],
     ) -> OpenAIEmbeddingsResponse:
         """Create embeddings.
 
         Generate OpenAI-compatible embeddings for the given input using the specified model.
-
-        :param model: The identifier of the model to use. The model must be an embedding model registered with Llama Stack and available via the /models endpoint.
-        :param input: Input text to embed, encoded as a string or array of strings. To embed multiple inputs in a single request, pass an array of strings.
-        :param encoding_format: (Optional) The format to return the embeddings in. Can be either "float" or "base64". Defaults to "float".
-        :param dimensions: (Optional) The number of dimensions the resulting output embeddings should have. Only supported in text-embedding-3 and later models.
-        :param user: (Optional) A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
         :returns: An OpenAIEmbeddingsResponse containing the embeddings.
         """
         ...

llama_stack/apis/vector_io/vector_io.py

@@ -11,6 +11,7 @@
 import uuid
 from typing import Annotated, Any, Literal, Protocol, runtime_checkable
 
+from fastapi import Body
 from pydantic import BaseModel, Field
 
 from llama_stack.apis.inference import InterleavedContent
@@ -466,6 +467,40 @@ class VectorStoreFilesListInBatchResponse(BaseModel):
     has_more: bool = False
 
 
+# extra_body can be accessed via .model_extra
+@json_schema_type
+class OpenAICreateVectorStoreRequestWithExtraBody(BaseModel, extra="allow"):
+    """Request to create a vector store with extra_body support.
+
+    :param name: (Optional) A name for the vector store
+    :param file_ids: List of file IDs to include in the vector store
+    :param expires_after: (Optional) Expiration policy for the vector store
+    :param chunking_strategy: (Optional) Strategy for splitting files into chunks
+    :param metadata: Set of key-value pairs that can be attached to the vector store
+    """
+
+    name: str | None = None
+    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
+
+
+# extra_body can be accessed via .model_extra
+@json_schema_type
+class OpenAICreateVectorStoreFileBatchRequestWithExtraBody(BaseModel, extra="allow"):
+    """Request to create a vector store file batch with extra_body support.
+
+    :param file_ids: A list of File IDs that the vector store should use
+    :param attributes: (Optional) Key-value attributes to store with the files
+    :param chunking_strategy: (Optional) The chunking strategy used to chunk the file(s). Defaults to auto
+    """
+
+    file_ids: list[str]
+    attributes: dict[str, Any] | None = None
+    chunking_strategy: VectorStoreChunkingStrategy | None = None
+
+
 class VectorDBStore(Protocol):
     def get_vector_db(self, vector_db_id: str) -> VectorDB | None: ...
 
@@ -516,25 +551,11 @@ class VectorIO(Protocol):
     @webmethod(route="/vector_stores", method="POST", level=LLAMA_STACK_API_V1)
     async def openai_create_vector_store(
         self,
-        name: str | None = None,
-        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,
+        params: Annotated[OpenAICreateVectorStoreRequestWithExtraBody, Body(...)],
     ) -> VectorStoreObject:
         """Creates a vector store.
 
-        :param name: A name for the vector store.
-        :param file_ids: A list of File IDs that the vector store should use. Useful for tools like `file_search` that can access files.
-        :param expires_after: The expiration policy for a vector store.
-        :param chunking_strategy: The chunking strategy used to chunk the file(s). If not set, will use the `auto` strategy.
-        :param metadata: Set of 16 key-value pairs that can be attached to an object.
-        :param embedding_model: The embedding model to use for this vector store.
-        :param embedding_dimension: The dimension of the embedding vectors (default: 384).
-        :param provider_id: The ID of the provider to use for this vector store.
+        Generate an OpenAI-compatible vector store with the given parameters.
         :returns: A VectorStoreObject representing the created vector store.
         """
         ...
@@ -827,16 +848,12 @@ class VectorIO(Protocol):
     async def openai_create_vector_store_file_batch(
         self,
         vector_store_id: str,
-        file_ids: list[str],
-        attributes: dict[str, Any] | None = None,
-        chunking_strategy: VectorStoreChunkingStrategy | None = None,
+        params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)],
     ) -> VectorStoreFileBatchObject:
         """Create a vector store file batch.
 
+        Generate an OpenAI-compatible vector store file batch for the given vector store.
         :param vector_store_id: The ID of the vector store to create the file batch for.
-        :param file_ids: A list of File IDs that the vector store should use.
-        :param attributes: (Optional) Key-value attributes to store with the files.
-        :param chunking_strategy: (Optional) The chunking strategy used to chunk the file(s). Defaults to auto.
         :returns: A VectorStoreFileBatchObject representing the created file batch.
         """
         ...

llama_stack/core/library_client.py

@@ -513,6 +513,14 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
         # Strip NOT_GIVENs to use the defaults in signature
         body = {k: v for k, v in body.items() if v is not NOT_GIVEN}
 
+        # Check if there's an unwrapped body parameter among multiple parameters
+        # (e.g., path param + body param like: vector_store_id: str, params: Annotated[Model, Body(...)])
+        unwrapped_body_param = None
+        for param in params_list:
+            if is_unwrapped_body_param(param.annotation):
+                unwrapped_body_param = param
+                break
+
         # Convert parameters to Pydantic models where needed
         converted_body = {}
         for param_name, param in sig.parameters.items():
@@ -522,5 +530,11 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
                     converted_body[param_name] = value
                 else:
                     converted_body[param_name] = convert_to_pydantic(param.annotation, value)
+            elif unwrapped_body_param and param.name == unwrapped_body_param.name:
+                # This is the unwrapped body param - construct it from remaining body keys
+                base_type = get_args(param.annotation)[0]
+                # Extract only the keys that aren't already used by other params
+                remaining_keys = {k: v for k, v in body.items() if k not in converted_body}
+                converted_body[param.name] = base_type(**remaining_keys)
 
         return converted_body

llama_stack/core/routers/inference.py

@@ -40,6 +40,7 @@ from llama_stack.apis.inference import (
     OpenAICompletion,
     OpenAICompletionRequestWithExtraBody,
     OpenAICompletionWithInputMessages,
+    OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
     OpenAIMessageParam,
     Order,
@@ -279,26 +280,18 @@ class InferenceRouter(Inference):
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: Annotated[OpenAIEmbeddingsRequestWithExtraBody, Body(...)],
     ) -> OpenAIEmbeddingsResponse:
         logger.debug(
-            f"InferenceRouter.openai_embeddings: {model=}, input_type={type(input)}, {encoding_format=}, {dimensions=}",
-        )
-        model_obj = await self._get_model(model, ModelType.embedding)
-        params = dict(
-            model=model_obj.identifier,
-            input=input,
-            encoding_format=encoding_format,
-            dimensions=dimensions,
-            user=user,
+            f"InferenceRouter.openai_embeddings: model={params.model}, input_type={type(params.input)}, encoding_format={params.encoding_format}, dimensions={params.dimensions}",
         )
+        model_obj = await self._get_model(params.model, ModelType.embedding)
+
+        # Update model to use resolved identifier
+        params.model = model_obj.identifier
 
         provider = await self.routing_table.get_provider_impl(model_obj.identifier)
-        return await provider.openai_embeddings(**params)
+        return await provider.openai_embeddings(params)
 
     async def list_chat_completions(
         self,

llama_stack/core/routers/vector_io.py

@@ -6,12 +6,16 @@
 
 import asyncio
 import uuid
-from typing import Any
+from typing import Annotated, Any
+
+from fastapi import Body
 
 from llama_stack.apis.common.content_types import InterleavedContent
 from llama_stack.apis.models import ModelType
 from llama_stack.apis.vector_io import (
     Chunk,
+    OpenAICreateVectorStoreFileBatchRequestWithExtraBody,
+    OpenAICreateVectorStoreRequestWithExtraBody,
     QueryChunksResponse,
     SearchRankingOptions,
     VectorIO,
@@ -120,18 +124,19 @@ class VectorIORouter(VectorIO):
     # OpenAI Vector Stores API endpoints
     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 = None,
-        provider_id: str | None = None,
+        params: Annotated[OpenAICreateVectorStoreRequestWithExtraBody, Body(...)],
     ) -> VectorStoreObject:
-        logger.debug(f"VectorIORouter.openai_create_vector_store: name={name}, provider_id={provider_id}")
+        # Extract llama-stack-specific parameters from extra_body
+        extra = params.model_extra or {}
+        embedding_model = extra.get("embedding_model")
+        embedding_dimension = extra.get("embedding_dimension", 384)
+        provider_id = extra.get("provider_id")
+
+        logger.debug(f"VectorIORouter.openai_create_vector_store: name={params.name}, provider_id={provider_id}")
 
         # If no embedding model is provided, use the first available one
+        # TODO: this branch will soon be deleted so you _must_ provide the embedding_model when
+        # creating a vector store
         if embedding_model is None:
             embedding_model_info = await self._get_first_embedding_model()
             if embedding_model_info is None:
@@ -146,20 +151,19 @@ class VectorIORouter(VectorIO):
             embedding_dimension=embedding_dimension,
             provider_id=provider_id,
             provider_vector_db_id=vector_db_id,
-            vector_db_name=name,
+            vector_db_name=params.name,
         )
         provider = await self.routing_table.get_provider_impl(registered_vector_db.identifier)
-        return await provider.openai_create_vector_store(
-            name=name,
-            file_ids=file_ids,
-            expires_after=expires_after,
-            chunking_strategy=chunking_strategy,
-            metadata=metadata,
-            embedding_model=embedding_model,
-            embedding_dimension=embedding_dimension,
-            provider_id=registered_vector_db.provider_id,
-            provider_vector_db_id=registered_vector_db.provider_resource_id,
-        )
+
+        # Update model_extra with registered values so provider uses the already-registered vector_db
+        if params.model_extra is None:
+            params.model_extra = {}
+        params.model_extra["provider_vector_db_id"] = registered_vector_db.provider_resource_id
+        params.model_extra["provider_id"] = registered_vector_db.provider_id
+        params.model_extra["embedding_model"] = embedding_model
+        params.model_extra["embedding_dimension"] = embedding_dimension
+
+        return await provider.openai_create_vector_store(params)
 
     async def openai_list_vector_stores(
         self,
@@ -219,7 +223,8 @@ class VectorIORouter(VectorIO):
         vector_store_id: str,
     ) -> VectorStoreObject:
         logger.debug(f"VectorIORouter.openai_retrieve_vector_store: {vector_store_id}")
-        return await self.routing_table.openai_retrieve_vector_store(vector_store_id)
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_retrieve_vector_store(vector_store_id)
 
     async def openai_update_vector_store(
         self,
@@ -229,7 +234,8 @@ class VectorIORouter(VectorIO):
         metadata: dict[str, Any] | None = None,
     ) -> VectorStoreObject:
         logger.debug(f"VectorIORouter.openai_update_vector_store: {vector_store_id}")
-        return await self.routing_table.openai_update_vector_store(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_update_vector_store(
             vector_store_id=vector_store_id,
             name=name,
             expires_after=expires_after,
@@ -241,7 +247,8 @@ class VectorIORouter(VectorIO):
         vector_store_id: str,
     ) -> VectorStoreDeleteResponse:
         logger.debug(f"VectorIORouter.openai_delete_vector_store: {vector_store_id}")
-        return await self.routing_table.openai_delete_vector_store(vector_store_id)
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_delete_vector_store(vector_store_id)
 
     async def openai_search_vector_store(
         self,
@@ -254,7 +261,8 @@ class VectorIORouter(VectorIO):
         search_mode: str | None = "vector",
     ) -> VectorStoreSearchResponsePage:
         logger.debug(f"VectorIORouter.openai_search_vector_store: {vector_store_id}")
-        return await self.routing_table.openai_search_vector_store(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_search_vector_store(
             vector_store_id=vector_store_id,
             query=query,
             filters=filters,
@@ -272,7 +280,8 @@ class VectorIORouter(VectorIO):
         chunking_strategy: VectorStoreChunkingStrategy | None = None,
     ) -> VectorStoreFileObject:
         logger.debug(f"VectorIORouter.openai_attach_file_to_vector_store: {vector_store_id}, {file_id}")
-        return await self.routing_table.openai_attach_file_to_vector_store(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_attach_file_to_vector_store(
             vector_store_id=vector_store_id,
             file_id=file_id,
             attributes=attributes,
@@ -289,7 +298,8 @@ class VectorIORouter(VectorIO):
         filter: VectorStoreFileStatus | None = None,
     ) -> list[VectorStoreFileObject]:
         logger.debug(f"VectorIORouter.openai_list_files_in_vector_store: {vector_store_id}")
-        return await self.routing_table.openai_list_files_in_vector_store(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_list_files_in_vector_store(
             vector_store_id=vector_store_id,
             limit=limit,
             order=order,
@@ -304,7 +314,8 @@ class VectorIORouter(VectorIO):
         file_id: str,
     ) -> VectorStoreFileObject:
         logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file: {vector_store_id}, {file_id}")
-        return await self.routing_table.openai_retrieve_vector_store_file(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_retrieve_vector_store_file(
             vector_store_id=vector_store_id,
             file_id=file_id,
         )
@@ -315,7 +326,8 @@ class VectorIORouter(VectorIO):
         file_id: str,
     ) -> VectorStoreFileContentsResponse:
         logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_contents: {vector_store_id}, {file_id}")
-        return await self.routing_table.openai_retrieve_vector_store_file_contents(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_retrieve_vector_store_file_contents(
             vector_store_id=vector_store_id,
             file_id=file_id,
         )
@@ -327,7 +339,8 @@ class VectorIORouter(VectorIO):
         attributes: dict[str, Any],
     ) -> VectorStoreFileObject:
         logger.debug(f"VectorIORouter.openai_update_vector_store_file: {vector_store_id}, {file_id}")
-        return await self.routing_table.openai_update_vector_store_file(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_update_vector_store_file(
             vector_store_id=vector_store_id,
             file_id=file_id,
             attributes=attributes,
@@ -339,7 +352,8 @@ class VectorIORouter(VectorIO):
         file_id: str,
     ) -> VectorStoreFileDeleteResponse:
         logger.debug(f"VectorIORouter.openai_delete_vector_store_file: {vector_store_id}, {file_id}")
-        return await self.routing_table.openai_delete_vector_store_file(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_delete_vector_store_file(
             vector_store_id=vector_store_id,
             file_id=file_id,
         )
@@ -370,17 +384,13 @@ class VectorIORouter(VectorIO):
     async def openai_create_vector_store_file_batch(
         self,
         vector_store_id: str,
-        file_ids: list[str],
-        attributes: dict[str, Any] | None = None,
-        chunking_strategy: VectorStoreChunkingStrategy | None = None,
+        params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)],
     ) -> VectorStoreFileBatchObject:
-        logger.debug(f"VectorIORouter.openai_create_vector_store_file_batch: {vector_store_id}, {len(file_ids)} files")
-        return await self.routing_table.openai_create_vector_store_file_batch(
-            vector_store_id=vector_store_id,
-            file_ids=file_ids,
-            attributes=attributes,
-            chunking_strategy=chunking_strategy,
+        logger.debug(
+            f"VectorIORouter.openai_create_vector_store_file_batch: {vector_store_id}, {len(params.file_ids)} files"
         )
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_create_vector_store_file_batch(vector_store_id, params)
 
     async def openai_retrieve_vector_store_file_batch(
         self,
@@ -388,7 +398,8 @@ class VectorIORouter(VectorIO):
         vector_store_id: str,
     ) -> VectorStoreFileBatchObject:
         logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_batch: {batch_id}, {vector_store_id}")
-        return await self.routing_table.openai_retrieve_vector_store_file_batch(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_retrieve_vector_store_file_batch(
             batch_id=batch_id,
             vector_store_id=vector_store_id,
         )
@@ -404,7 +415,8 @@ class VectorIORouter(VectorIO):
         order: str | None = "desc",
     ) -> VectorStoreFilesListInBatchResponse:
         logger.debug(f"VectorIORouter.openai_list_files_in_vector_store_file_batch: {batch_id}, {vector_store_id}")
-        return await self.routing_table.openai_list_files_in_vector_store_file_batch(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_list_files_in_vector_store_file_batch(
             batch_id=batch_id,
             vector_store_id=vector_store_id,
             after=after,
@@ -420,7 +432,8 @@ class VectorIORouter(VectorIO):
         vector_store_id: str,
     ) -> VectorStoreFileBatchObject:
         logger.debug(f"VectorIORouter.openai_cancel_vector_store_file_batch: {batch_id}, {vector_store_id}")
-        return await self.routing_table.openai_cancel_vector_store_file_batch(
+        provider = await self.routing_table.get_provider_impl(vector_store_id)
+        return await provider.openai_cancel_vector_store_file_batch(
             batch_id=batch_id,
             vector_store_id=vector_store_id,
         )

llama_stack/providers/inline/batches/reference/batches.py

@@ -25,6 +25,7 @@ from llama_stack.apis.inference import (
     OpenAIChatCompletionRequestWithExtraBody,
     OpenAICompletionRequestWithExtraBody,
     OpenAIDeveloperMessageParam,
+    OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIMessageParam,
     OpenAISystemMessageParam,
     OpenAIToolMessageParam,
@@ -640,7 +641,9 @@ class ReferenceBatchesImpl(Batches):
                     },
                 }
             else:  # /v1/embeddings
-                embeddings_response = await self.inference_api.openai_embeddings(**request.body)
+                embeddings_response = await self.inference_api.openai_embeddings(
+                    OpenAIEmbeddingsRequestWithExtraBody(**request.body)
+                )
                 assert hasattr(embeddings_response, "model_dump_json"), (
                     "Embeddings response must have model_dump_json method"
                 )

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

@@ -14,6 +14,7 @@ from llama_stack.apis.inference import (
     Inference,
     OpenAIChatCompletionRequestWithExtraBody,
     OpenAICompletionRequestWithExtraBody,
+    OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
 )
 from llama_stack.apis.inference.inference import (
@@ -124,11 +125,7 @@ class BedrockInferenceAdapter(
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
         raise NotImplementedError()
 

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

@@ -6,7 +6,10 @@
 
 from urllib.parse import urljoin
 
-from llama_stack.apis.inference import OpenAIEmbeddingsResponse
+from llama_stack.apis.inference import (
+    OpenAIEmbeddingsRequestWithExtraBody,
+    OpenAIEmbeddingsResponse,
+)
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .config import CerebrasImplConfig
@@ -20,10 +23,6 @@ class CerebrasInferenceAdapter(OpenAIMixin):
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
         raise NotImplementedError()

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

@@ -7,6 +7,7 @@
 from llama_stack.apis.inference.inference import (
     OpenAICompletion,
     OpenAICompletionRequestWithExtraBody,
+    OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
 )
 from llama_stack.log import get_logger
@@ -40,10 +41,6 @@ class LlamaCompatInferenceAdapter(OpenAIMixin):
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
         raise NotImplementedError()

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

@@ -9,6 +9,7 @@ from openai import NOT_GIVEN
 
 from llama_stack.apis.inference import (
     OpenAIEmbeddingData,
+    OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
     OpenAIEmbeddingUsage,
 )
@@ -78,11 +79,7 @@ class NVIDIAInferenceAdapter(OpenAIMixin):
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
         """
         OpenAI-compatible embeddings for NVIDIA NIM.
@@ -99,11 +96,11 @@ class NVIDIAInferenceAdapter(OpenAIMixin):
         )
 
         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,
+            model=await self._get_provider_model_id(params.model),
+            input=params.input,
+            encoding_format=params.encoding_format if params.encoding_format is not None else NOT_GIVEN,
+            dimensions=params.dimensions if params.dimensions is not None else NOT_GIVEN,
+            user=params.user if params.user is not None else NOT_GIVEN,
             extra_body=extra_body,
         )
 

llama_stack/providers/remote/inference/passthrough/passthrough.py

@@ -16,6 +16,7 @@ from llama_stack.apis.inference import (
     OpenAIChatCompletionRequestWithExtraBody,
     OpenAICompletion,
     OpenAICompletionRequestWithExtraBody,
+    OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
 )
 from llama_stack.apis.models import Model
@@ -69,11 +70,7 @@ class PassthroughInferenceAdapter(Inference):
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
         raise NotImplementedError()
 

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

@@ -10,7 +10,10 @@ from collections.abc import Iterable
 from huggingface_hub import AsyncInferenceClient, HfApi
 from pydantic import SecretStr
 
-from llama_stack.apis.inference import OpenAIEmbeddingsResponse
+from llama_stack.apis.inference import (
+    OpenAIEmbeddingsRequestWithExtraBody,
+    OpenAIEmbeddingsResponse,
+)
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
@@ -40,11 +43,7 @@ class _HfAdapter(OpenAIMixin):
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
         raise NotImplementedError()
 

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

@@ -11,6 +11,7 @@ from together import AsyncTogether
 from together.constants import BASE_URL
 
 from llama_stack.apis.inference import (
+    OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
 )
 from llama_stack.apis.inference.inference import OpenAIEmbeddingUsage
@@ -62,11 +63,7 @@ class TogetherInferenceAdapter(OpenAIMixin, NeedsRequestProviderData):
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
         """
         Together's OpenAI-compatible embeddings endpoint is not compatible with
@@ -78,25 +75,27 @@ class TogetherInferenceAdapter(OpenAIMixin, NeedsRequestProviderData):
          - does not support dimensions param, returns 400 Unrecognized request arguments supplied: dimensions
         """
         # Together support ticket #13332 -> will not fix
-        if user is not None:
+        if params.user is not None:
             raise ValueError("Together's embeddings endpoint does not support user param.")
         # Together support ticket #13333 -> escalated
-        if dimensions is not None:
+        if params.dimensions is not None:
             raise ValueError("Together's embeddings endpoint does not support dimensions param.")
 
         response = await self.client.embeddings.create(
-            model=await self._get_provider_model_id(model),
-            input=input,
-            encoding_format=encoding_format,
+            model=await self._get_provider_model_id(params.model),
+            input=params.input,
+            encoding_format=params.encoding_format,
         )
 
-        response.model = model  # return the user the same model id they provided, avoid exposing the provider model id
+        response.model = (
+            params.model
+        )  # return the user the same model id they provided, avoid exposing the provider model id
 
         # Together support ticket #13330 -> escalated
         #  - togethercomputer/m2-bert-80M-32k-retrieval *does not* return usage information
         if not hasattr(response, "usage") or response.usage is None:
             logger.warning(
-                f"Together's embedding endpoint for {model} did not return usage information, substituting -1s."
+                f"Together's embedding endpoint for {params.model} did not return usage information, substituting -1s."
             )
             response.usage = OpenAIEmbeddingUsage(prompt_tokens=-1, total_tokens=-1)
 

llama_stack/providers/utils/inference/embedding_mixin.py

@@ -17,6 +17,7 @@ if TYPE_CHECKING:
 from llama_stack.apis.inference import (
     ModelStore,
     OpenAIEmbeddingData,
+    OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
     OpenAIEmbeddingUsage,
 )
@@ -32,26 +33,22 @@ class SentenceTransformerEmbeddingMixin:
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
         # Convert input to list format if it's a single string
-        input_list = [input] if isinstance(input, str) else input
+        input_list = [params.input] if isinstance(params.input, str) else params.input
         if not input_list:
             raise ValueError("Empty list not supported")
 
         # Get the model and generate embeddings
-        model_obj = await self.model_store.get_model(model)
+        model_obj = await self.model_store.get_model(params.model)
         embedding_model = await self._load_sentence_transformer_model(model_obj.provider_resource_id)
         embeddings = await asyncio.to_thread(embedding_model.encode, input_list, show_progress_bar=False)
 
         # Convert embeddings to the requested format
         data = []
         for i, embedding in enumerate(embeddings):
-            if encoding_format == "base64":
+            if params.encoding_format == "base64":
                 # Convert float array to base64 string
                 float_bytes = struct.pack(f"{len(embedding)}f", *embedding)
                 embedding_value = base64.b64encode(float_bytes).decode("ascii")
@@ -70,7 +67,7 @@ class SentenceTransformerEmbeddingMixin:
         usage = OpenAIEmbeddingUsage(prompt_tokens=-1, total_tokens=-1)
         return OpenAIEmbeddingsResponse(
             data=data,
-            model=model,
+            model=params.model,
             usage=usage,
         )
 

llama_stack/providers/utils/inference/litellm_openai_mixin.py

@@ -20,6 +20,7 @@ from llama_stack.apis.inference import (
     OpenAICompletion,
     OpenAICompletionRequestWithExtraBody,
     OpenAIEmbeddingData,
+    OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
     OpenAIEmbeddingUsage,
     ToolChoice,
@@ -189,16 +190,12 @@ class LiteLLMOpenAIMixin(
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
-        model_obj = await self.model_store.get_model(model)
+        model_obj = await self.model_store.get_model(params.model)
 
         # Convert input to list if it's a string
-        input_list = [input] if isinstance(input, str) else input
+        input_list = [params.input] if isinstance(params.input, str) else params.input
 
         # Call litellm embedding function
         # litellm.drop_params = True
@@ -207,11 +204,11 @@ class LiteLLMOpenAIMixin(
             input=input_list,
             api_key=self.get_api_key(),
             api_base=self.api_base,
-            dimensions=dimensions,
+            dimensions=params.dimensions,
         )
 
         # Convert response to OpenAI format
-        data = b64_encode_openai_embeddings_response(response.data, encoding_format)
+        data = b64_encode_openai_embeddings_response(response.data, params.encoding_format)
 
         usage = OpenAIEmbeddingUsage(
             prompt_tokens=response["usage"]["prompt_tokens"],

llama_stack/providers/utils/inference/openai_mixin.py

@@ -21,6 +21,7 @@ from llama_stack.apis.inference import (
     OpenAICompletion,
     OpenAICompletionRequestWithExtraBody,
     OpenAIEmbeddingData,
+    OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
     OpenAIEmbeddingUsage,
     OpenAIMessageParam,
@@ -316,23 +317,27 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
 
     async def openai_embeddings(
         self,
-        model: str,
-        input: str | list[str],
-        encoding_format: str | None = "float",
-        dimensions: int | None = None,
-        user: str | None = None,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
         """
         Direct OpenAI embeddings API call.
         """
+        # Prepare request parameters
+        request_params = {
+            "model": await self._get_provider_model_id(params.model),
+            "input": params.input,
+            "encoding_format": params.encoding_format if params.encoding_format is not None else NOT_GIVEN,
+            "dimensions": params.dimensions if params.dimensions is not None else NOT_GIVEN,
+            "user": params.user if params.user is not None else NOT_GIVEN,
+        }
+
+        # Add extra_body if present
+        extra_body = params.model_extra
+        if extra_body:
+            request_params["extra_body"] = extra_body
+
         # 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,
-        )
+        response = await self.client.embeddings.create(**request_params)
 
         data = []
         for i, embedding_data in enumerate(response.data):
@@ -350,7 +355,7 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
 
         return OpenAIEmbeddingsResponse(
             data=data,
-            model=model,
+            model=params.model,
             usage=usage,
         )
 

llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -10,8 +10,9 @@ import mimetypes
 import time
 import uuid
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Annotated, Any
 
+from fastapi import Body
 from pydantic import TypeAdapter
 
 from llama_stack.apis.common.errors import VectorStoreNotFoundError
@@ -19,6 +20,8 @@ from llama_stack.apis.files import Files, OpenAIFileObject
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.apis.vector_io import (
     Chunk,
+    OpenAICreateVectorStoreFileBatchRequestWithExtraBody,
+    OpenAICreateVectorStoreRequestWithExtraBody,
     QueryChunksResponse,
     SearchRankingOptions,
     VectorStoreChunkingStrategy,
@@ -340,18 +343,18 @@ class OpenAIVectorStoreMixin(ABC):
 
     async def openai_create_vector_store(
         self,
-        name: str | None = None,
-        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,
-        provider_vector_db_id: str | None = None,
+        params: Annotated[OpenAICreateVectorStoreRequestWithExtraBody, Body(...)],
     ) -> VectorStoreObject:
         """Creates a vector store."""
         created_at = int(time.time())
+
+        # Extract llama-stack-specific parameters from extra_body
+        extra = params.model_extra or {}
+        provider_vector_db_id = extra.get("provider_vector_db_id")
+        embedding_model = extra.get("embedding_model")
+        embedding_dimension = extra.get("embedding_dimension", 384)
+        provider_id = extra.get("provider_id")
+
         # Derive the canonical vector_db_id (allow override, else generate)
         vector_db_id = provider_vector_db_id or generate_object_id("vector_store", lambda: f"vs_{uuid.uuid4()}")
 
@@ -372,7 +375,7 @@ class OpenAIVectorStoreMixin(ABC):
             embedding_model=embedding_model,
             provider_id=provider_id,
             provider_resource_id=vector_db_id,
-            vector_db_name=name,
+            vector_db_name=params.name,
         )
         await self.register_vector_db(vector_db)
 
@@ -391,19 +394,19 @@ class OpenAIVectorStoreMixin(ABC):
             "id": vector_db_id,
             "object": "vector_store",
             "created_at": created_at,
-            "name": name,
+            "name": params.name,
             "usage_bytes": 0,
             "file_counts": file_counts.model_dump(),
             "status": status,
-            "expires_after": expires_after,
+            "expires_after": params.expires_after,
             "expires_at": None,
             "last_active_at": created_at,
             "file_ids": [],
-            "chunking_strategy": chunking_strategy,
+            "chunking_strategy": params.chunking_strategy,
         }
 
         # Add provider information to metadata if provided
-        metadata = metadata or {}
+        metadata = params.metadata or {}
         if provider_id:
             metadata["provider_id"] = provider_id
         if provider_vector_db_id:
@@ -417,7 +420,7 @@ class OpenAIVectorStoreMixin(ABC):
         self.openai_vector_stores[vector_db_id] = store_info
 
         # Now that our vector store is created, attach any files that were provided
-        file_ids = file_ids or []
+        file_ids = params.file_ids or []
         tasks = [self.openai_attach_file_to_vector_store(vector_db_id, file_id) for file_id in file_ids]
         await asyncio.gather(*tasks)
 
@@ -976,15 +979,13 @@ class OpenAIVectorStoreMixin(ABC):
     async def openai_create_vector_store_file_batch(
         self,
         vector_store_id: str,
-        file_ids: list[str],
-        attributes: dict[str, Any] | None = None,
-        chunking_strategy: VectorStoreChunkingStrategy | None = None,
+        params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)],
     ) -> VectorStoreFileBatchObject:
         """Create a vector store file batch."""
         if vector_store_id not in self.openai_vector_stores:
             raise VectorStoreNotFoundError(vector_store_id)
 
-        chunking_strategy = chunking_strategy or VectorStoreChunkingStrategyAuto()
+        chunking_strategy = params.chunking_strategy or VectorStoreChunkingStrategyAuto()
 
         created_at = int(time.time())
         batch_id = generate_object_id("vector_store_file_batch", lambda: f"batch_{uuid.uuid4()}")
@@ -996,8 +997,8 @@ class OpenAIVectorStoreMixin(ABC):
             completed=0,
             cancelled=0,
             failed=0,
-            in_progress=len(file_ids),
-            total=len(file_ids),
+            in_progress=len(params.file_ids),
+            total=len(params.file_ids),
         )
 
         # Create batch object immediately with in_progress status
@@ -1011,8 +1012,8 @@ class OpenAIVectorStoreMixin(ABC):
 
         batch_info = {
             **batch_object.model_dump(),
-            "file_ids": file_ids,
-            "attributes": attributes,
+            "file_ids": params.file_ids,
+            "attributes": params.attributes,
             "chunking_strategy": chunking_strategy.model_dump(),
             "expires_at": expires_at,
         }

llama_stack/providers/utils/memory/vector_store.py

@@ -21,6 +21,7 @@ from llama_stack.apis.common.content_types import (
     URL,
     InterleavedContent,
 )
+from llama_stack.apis.inference import OpenAIEmbeddingsRequestWithExtraBody
 from llama_stack.apis.tools import RAGDocument
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.apis.vector_io import Chunk, ChunkMetadata, QueryChunksResponse
@@ -274,10 +275,11 @@ class VectorDBWithIndex:
                 _validate_embedding(c.embedding, i, self.vector_db.embedding_dimension)
 
         if chunks_to_embed:
-            resp = await self.inference_api.openai_embeddings(
-                self.vector_db.embedding_model,
-                [c.content for c in chunks_to_embed],
+            params = OpenAIEmbeddingsRequestWithExtraBody(
+                model=self.vector_db.embedding_model,
+                input=[c.content for c in chunks_to_embed],
             )
+            resp = await self.inference_api.openai_embeddings(params)
             for c, data in zip(chunks_to_embed, resp.data, strict=False):
                 c.embedding = data.embedding
 
@@ -316,7 +318,11 @@ class VectorDBWithIndex:
         if mode == "keyword":
             return await self.index.query_keyword(query_string, k, score_threshold)
 
-        embeddings_response = await self.inference_api.openai_embeddings(self.vector_db.embedding_model, [query_string])
+        params = OpenAIEmbeddingsRequestWithExtraBody(
+            model=self.vector_db.embedding_model,
+            input=[query_string],
+        )
+        embeddings_response = await self.inference_api.openai_embeddings(params)
         query_vector = np.array(embeddings_response.data[0].embedding, dtype=np.float32)
         if mode == "hybrid":
             return await self.index.query_hybrid(