phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-13 16:42:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

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

Browse source 
Applies the same pattern from #3777 to embeddings and vector_stores.create() endpoints.

Breaking change: Method signatures now accept a single params object with Pydantic extra="allow" instead of individual parameters. Provider-specific params can be passed via extra_body and accessed through params.model_extra.

Updated APIs: openai_embeddings(), openai_create_vector_store(), openai_create_vector_store_file_batch()
This commit is contained in:
Ashwin Bharambe 2025-10-11 15:27:47 -07:00
parent cfd2e303db
commit 74e2976c1e
20 changed files with 364 additions and 297 deletions
Download patch file Download diff file Expand all files Collapse all files

7
llama_stack/providers/remote/inference/bedrock/bedrock.py
View file

@ -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()

11
llama_stack/providers/remote/inference/cerebras/cerebras.py
View file

@ -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()

7
llama_stack/providers/remote/inference/llama_openai_compat/llama.py
View file

@ -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()

17
llama_stack/providers/remote/inference/nvidia/nvidia.py
View file

@ -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,
)

7
llama_stack/providers/remote/inference/passthrough/passthrough.py
View file

@ -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()

11
llama_stack/providers/remote/inference/tgi/tgi.py
View file

@ -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()

23
llama_stack/providers/remote/inference/together/together.py
View file

@ -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)

15
llama_stack/providers/utils/inference/embedding_mixin.py
View file

@ -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,
)

15
llama_stack/providers/utils/inference/litellm_openai_mixin.py
View file

@ -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"],

31
llama_stack/providers/utils/inference/openai_mixin.py
View file

@ -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,
)

60
llama_stack/providers/utils/memory/openai_vector_store_mixin.py
View file

@ -19,6 +19,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,39 +342,37 @@ 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: OpenAICreateVectorStoreRequestWithExtraBody,
) -> VectorStoreObject:
"""Creates a vector store."""
created_at = int(time.time())
# Extract provider_vector_db_id from extra_body if present
provider_vector_db_id = None
if params.model_extra and "provider_vector_db_id" in params.model_extra:
provider_vector_db_id = params.model_extra["provider_vector_db_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()}")
if provider_id is None:
if params.provider_id is None:
raise ValueError("Provider ID is required")
if embedding_model is None:
if params.embedding_model is None:
raise ValueError("Embedding model is required")
# Embedding dimension is required (defaulted to 384 if not provided)
if embedding_dimension is None:
if params.embedding_dimension is None:
raise ValueError("Embedding dimension is required")
# Register the VectorDB backing this vector store
vector_db = VectorDB(
identifier=vector_db_id,
embedding_dimension=embedding_dimension,
embedding_model=embedding_model,
provider_id=provider_id,
embedding_dimension=params.embedding_dimension,
embedding_model=params.embedding_model,
provider_id=params.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,21 +391,21 @@ 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 {}
if provider_id:
metadata["provider_id"] = provider_id
metadata = params.metadata or {}
if params.provider_id:
metadata["provider_id"] = params.provider_id
if provider_vector_db_id:
metadata["provider_vector_db_id"] = provider_vector_db_id
store_info["metadata"] = metadata
@ -417,7 +417,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 +976,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: OpenAICreateVectorStoreFileBatchRequestWithExtraBody,
) -> 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 +994,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 +1009,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,
}