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/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(