From effcb0705fbb2216ab36211b5ea074f457673910 Mon Sep 17 00:00:00 2001 From: kenneth-berri Date: Thu, 12 Nov 2026 15:41:50 -0800 Subject: [PATCH] fix(embeddings): unwrap extra_body for providers with their own embedding handlers (#51335) * fix(embeddings): unwrap extra_body for providers with their own embedding handlers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(embeddings): unwrap extra_body in get_optional_params_embeddings and keep lint gates flat Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(embeddings): let mapped params win over extra_body and drop the unused merge helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(embeddings): merge extra_body into the provider body after each native embedding transform Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(embeddings): let the transformed body win over extra_body so input and mapped params cannot be replaced Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(embeddings): sign the Bedrock body that is sent and align the shared handler so extra_body never replaces the input Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(embeddings): keep the shared handler's extra_body precedence, the router batch limit relies on extra_body.input Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/base_llm/embedding/transformation.py | 22 +++- litellm/llms/bedrock/embed/embedding.py | 25 ++-- litellm/llms/nvidia_nim/embed.py | 10 +- litellm/llms/ollama/completion/handler.py | 6 +- litellm/llms/sagemaker/completion/handler.py | 8 +- .../batch_embed_content_transformation.py | 9 +- .../multimodal_embeddings/transformation.py | 7 +- .../vertex_embeddings/transformation.py | 17 ++- litellm/utils.py | 10 +- .../test_provider_embedding_translation.py | 1 + .../bedrock/embed/test_bedrock_embedding.py | 119 ++++++++++++++++++ tests/unit/llms/nvidia_nim/test_nvidia_nim.py | 25 ++++ .../unit/llms/ollama/test_ollama_embedding.py | 18 +++ .../test_sagemaker_embedding_voyage.py | 37 ++++++ ...test_batch_embed_content_transformation.py | 29 +++++ ..._ai_multimodal_embedding_transformation.py | 32 +++++ .../test_embedding_handler.py | 17 +++ 17 files changed, 365 insertions(+), 27 deletions(-) diff --git a/litellm/llms/base_llm/embedding/transformation.py b/litellm/llms/base_llm/embedding/transformation.py index 46ac3ccf433..d25a0c3d524 100644 --- a/litellm/llms/base_llm/embedding/transformation.py +++ b/litellm/llms/base_llm/embedding/transformation.py @@ -1,7 +1,9 @@ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues @@ -16,6 +18,24 @@ else: LiteLLMLoggingObj = Any +_BODY_FIELDS: Final = TypeAdapter(dict[str, object]) + + +def extra_body_fields(optional_params: Mapping[str, object] | None) -> Mapping[str, object]: + """The caller's ``extra_body`` fields, merged under the provider body after the transform so they add fields but never + replace the input or a mapped param.""" + extra_body: Final = (optional_params or {}).get("extra_body") + return _BODY_FIELDS.validate_python(extra_body) if isinstance(extra_body, Mapping) else {} + + +def split_extra_body( + optional_params: Mapping[str, object] | None, +) -> tuple[Mapping[str, object], Mapping[str, object]]: + """The params without ``extra_body``, and the ``extra_body`` fields to merge into the provider body.""" + params: Final = _BODY_FIELDS.validate_python(optional_params or {}) + return {k: v for k, v in params.items() if k != "extra_body"}, extra_body_fields(params) + + class BaseEmbeddingConfig(BaseConfig, ABC): @abstractmethod def transform_embedding_request( diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index e28e39ef2d4..0743bd6e7e7 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -6,6 +6,7 @@ import copy import json import urllib.parse from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, get_args, overload import httpx @@ -15,6 +16,7 @@ import litellm from litellm.constants import BEDROCK_EMBEDDING_PROVIDERS_LITERAL from litellm.litellm_core_utils.hidden_params import set_hidden_param from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.embedding.transformation import extra_body_fields from litellm.llms.cohere.embed.handler import embedding as cohere_embedding from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -103,7 +105,7 @@ class BedrockEmbedding(BaseAWSLLM): region_from_model: str | None = None, ) -> tuple[Credentials | None, str]: auth_params: Final = pop_aws_auth_params(optional_params) - explicit_region_name: Final = optional_params.pop("aws_region_name", None) + explicit_region_name: Final = _OPTIONAL_STR.validate_python(optional_params.pop("aws_region_name", None)) aws_region_name = region_from_model if explicit_region_name is None else explicit_region_name ### SET REGION NAME ### @@ -291,9 +293,10 @@ class BedrockEmbedding(BaseAWSLLM): provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL, api_key: str | None = None, is_async_invoke: bool | None = False, + extra_body: Mapping[str, object] = MappingProxyType({}), ): responses: Final[list[dict]] = [] - for data in batch_data: + for data in ({**extra_body, **item} for item in batch_data): headers = {"Content-Type": "application/json"} if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} @@ -360,9 +363,10 @@ class BedrockEmbedding(BaseAWSLLM): provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL, api_key: str | None = None, is_async_invoke: bool | None = False, + extra_body: Mapping[str, object] = MappingProxyType({}), ): responses: Final[list[dict]] = [] - for data in batch_data: + for data in ({**extra_body, **item} for item in batch_data): headers = {"Content-Type": "application/json"} if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} @@ -462,9 +466,12 @@ class BedrockEmbedding(BaseAWSLLM): ) inference_params = copy.deepcopy(optional_params) inference_params = { - k: v for k, v in inference_params.items() if k.lower() not in self.aws_authentication_params + k: v + for k, v in inference_params.items() + if k.lower() not in self.aws_authentication_params and k != "extra_body" } inference_params.pop("user", None) # make sure user is not passed in for bedrock call + extra_body: Final = extra_body_fields(optional_params) data: CohereEmbeddingRequest | None = None batch_data: list | None = None @@ -535,6 +542,8 @@ class BedrockEmbedding(BaseAWSLLM): ) batch_data.append(nova_request) + request_data: Final = None if data is None else {**extra_body, **data} + ### SET RUNTIME ENDPOINT ### endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( api_base=api_base, @@ -552,6 +561,7 @@ class BedrockEmbedding(BaseAWSLLM): client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), timeout=timeout, batch_data=batch_data, + extra_body=extra_body, credentials=credentials, extra_headers=extra_headers, endpoint_url=endpoint_url, @@ -566,6 +576,7 @@ class BedrockEmbedding(BaseAWSLLM): client=(client if client is not None and isinstance(client, HTTPHandler) else None), timeout=timeout, batch_data=batch_data, + extra_body=extra_body, credentials=credentials, extra_headers=extra_headers, endpoint_url=endpoint_url, @@ -579,7 +590,7 @@ class BedrockEmbedding(BaseAWSLLM): if returned_response is None: raise Exception("Unable to map Bedrock request to provider") return returned_response - elif data is None: + elif request_data is None: raise Exception("Unable to map Bedrock request to provider") headers = {"Content-Type": "application/json"} @@ -591,7 +602,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_region_name=aws_region_name, extra_headers=extra_headers, endpoint_url=endpoint_url, - data=json.dumps(data), + data=json.dumps(request_data), headers=headers, api_key=api_key, ) @@ -607,7 +618,7 @@ class BedrockEmbedding(BaseAWSLLM): logging_obj=logging_obj, optional_params=optional_params, encoding=encoding, - data=data, + data=request_data, complete_api_base=prepped.url, api_key=None, aembedding=aembedding, diff --git a/litellm/llms/nvidia_nim/embed.py b/litellm/llms/nvidia_nim/embed.py index d2ef446299e..7368831f7a8 100644 --- a/litellm/llms/nvidia_nim/embed.py +++ b/litellm/llms/nvidia_nim/embed.py @@ -66,7 +66,7 @@ class NvidiaNimEmbeddingConfig: self, non_default_params: dict, optional_params: dict, - kwargs: dict | None = None, + kwargs: Mapping[str, object] | None = None, ): if "extra_body" not in optional_params: optional_params["extra_body"] = {} @@ -80,5 +80,11 @@ class NvidiaNimEmbeddingConfig: if kwargs is not None: # pass kwargs in extra_body - optional_params["extra_body"].update(kwargs) + nested_extra_body: Final = kwargs.get("extra_body") + optional_params["extra_body"].update( + { + **{k: v for k, v in kwargs.items() if k != "extra_body"}, + **(nested_extra_body if isinstance(nested_extra_body, Mapping) else {}), + } + ) return optional_params diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index e1fc9a8d262..fd874246980 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -14,6 +14,7 @@ from typing_extensions import NotRequired, ReadOnly import litellm from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.llms.base_llm.embedding.transformation import split_extra_body from litellm.types.mcp import has_header from litellm.types.utils import EmbeddingResponse @@ -45,14 +46,15 @@ def _prepare_ollama_embedding_payload( data: Final[dict[str, object]] = {"model": model, "input": prompts} special_optional_params: Final = ["truncate", "options", "keep_alive", "dimensions"] - for k, v in optional_params.items(): + body_params, extra_body = split_extra_body(optional_params) + for k, v in body_params.items(): if k in special_optional_params: data[k] = v else: data.setdefault("options", {}) if isinstance(data["options"], dict): data["options"].update({k: v}) - return data + return {**extra_body, **data} def _process_ollama_embedding_response( diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index a852ae8b1a1..c20f0df6768 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.optional_imports import ensure_optional_import +from litellm.llms.base_llm.embedding.transformation import extra_body_fields from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, pop_aws_auth_params from litellm.llms.bedrock.web_identity_session_policy import SAGEMAKER_SESSION_POLICY, WebIdentitySessionPolicy from litellm.llms.custom_httpx.http_handler import ( @@ -594,8 +595,11 @@ class SagemakerLLM(BaseAWSLLM): inference_component_name: Final = embedding_params.get("model_id") if inference_component_name is not None and not isinstance(inference_component_name, str): raise SagemakerError(status_code=400, message="model_id must be the inference component name as a string") - body_params: Final = {k: v for k, v in embedding_params.items() if k != "model_id"} - request_data: Final = provider_config.transform_embedding_request(endpoint_name, input, body_params, {}) + body_params: Final = {k: v for k, v in embedding_params.items() if k not in ("model_id", "extra_body")} + request_data: Final = { + **extra_body_fields(embedding_params), + **provider_config.transform_embedding_request(endpoint_name, input, body_params, {}), + } data: Final = json.dumps(request_data).encode("utf-8") invoke_args: Final = _InvokeEndpointArgs( EndpointName=endpoint_name, diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 96a785c06c9..b27003369f7 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -10,6 +10,7 @@ from typing import Annotated, Final, Literal, cast from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from litellm.exceptions import BadRequestError +from litellm.llms.base_llm.embedding.transformation import extra_body_fields from litellm.llms.vertex_ai.common_utils import GEMINI_FILES_API_URI_PREFIX, gemini_video_metadata_from_openai from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.vertex_ai import ( @@ -414,6 +415,7 @@ def transform_openai_input_gemini_content( gemini_model_name: Final = f"models/{model}" gemini_params: Final = _filter_embed_params(optional_params) + extra_body: Final = extra_body_fields(optional_params) input_list: Final = _top_level_elements(input) requests: Final[list[EmbedContentRequest]] = [] @@ -432,7 +434,11 @@ def transform_openai_input_gemini_content( content=ContentType(parts=parts), **gemini_params, ) - requests.append(request) + requests.append( + cast( # cast-ok: extra_body carries raw Gemini fields the TypedDict does not list + EmbedContentRequest, {**extra_body, **request} + ) + ) return VertexAIBatchEmbeddingsRequestBody(requests=requests) @@ -472,6 +478,7 @@ def transform_openai_input_gemini_embed_content( parts.append(_build_part_for_input(element, resolved_files=resolved_files, drop_params=drop_params)) request_body: Final[dict] = { + **extra_body_fields(optional_params), "content": ContentType(parts=parts), **gemini_params, } diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index c282ae22ea6..7c276b43dc7 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -21,7 +21,7 @@ from litellm.types.utils import ( ) from litellm.utils import count_characters, is_base64_encoded -from ...base_llm.embedding.transformation import BaseEmbeddingConfig +from ...base_llm.embedding.transformation import BaseEmbeddingConfig, extra_body_fields from ..common_utils import VertexAIError @@ -189,8 +189,11 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): request_data["instances"] = [vertex_request_instance] + extra_body: Final = extra_body_fields(optional_params) if "outputDimensionality" in optional_params: - request_data["parameters"] = {"dimension": optional_params["outputDimensionality"]} + request_data["parameters"] = {**extra_body, "dimension": optional_params["outputDimensionality"]} + elif extra_body: + request_data["parameters"] = dict(extra_body) return cast(dict, request_data) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 96fccf007a5..5d99ce6d90f 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -1,10 +1,11 @@ import types from collections.abc import Mapping -from typing import Final, Literal +from typing import Final, Literal, cast # noqa: TID251 # extra_body carries raw fields the TypedDict lacks from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict +from litellm.llms.base_llm.embedding.transformation import split_extra_body from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.utils import EmbeddingResponse, PromptTokensDetailsWrapper, Usage @@ -164,7 +165,12 @@ class VertexAITextEmbeddingConfig(LiteLLMBaseModel): vertex_text_embedding_input_list.append(embedding_input) vertex_request["instances"] = vertex_text_embedding_input_list - vertex_request["parameters"] = EmbeddingParameters(**optional_params) + body_params, extra_body = split_extra_body(optional_params) + vertex_request["parameters"] = ( + cast( # cast-ok: extra_body carries raw Vertex fields the TypedDict does not list + EmbeddingParameters, {**extra_body, **body_params} + ) + ) if labels: vertex_request["labels"] = labels @@ -205,7 +211,12 @@ class VertexAITextEmbeddingConfig(LiteLLMBaseModel): vertex_text_embedding_input_list.append(embedding_input) vertex_request["instances"] = vertex_text_embedding_input_list - vertex_request["parameters"] = TextEmbeddingFineTunedParameters(**optional_params) + body_params, extra_body = split_extra_body(optional_params) + vertex_request["parameters"] = ( + cast( # cast-ok: extra_body carries raw Vertex fields the TypedDict does not list + TextEmbeddingFineTunedParameters, {**extra_body, **body_params} + ) + ) # Remove 'shared_session' from parameters if present if vertex_request["parameters"] is not None and "shared_session" in vertex_request["parameters"]: del vertex_request["parameters"]["shared_session"] diff --git a/litellm/utils.py b/litellm/utils.py index 67b7b4f3d1e..118f89766a9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5024,9 +5024,7 @@ def get_optional_params_embeddings( optional_params["dimensions"] = non_default_params.pop("dimensions") if len(non_default_params.keys()) > 0: if litellm.drop_params is True or drop_params is True: # drop the unsupported non-default values - keys = list(non_default_params.keys()) - for k in keys: - non_default_params.pop(k, None) + non_default_params.clear() else: raise UnsupportedParamsError( status_code=500, @@ -5039,9 +5037,7 @@ def get_optional_params_embeddings( ): if len(non_default_params.keys()) > 0: if litellm.drop_params is True or drop_params is True: # drop the unsupported non-default values - keys = list(non_default_params.keys()) - for k in keys: - non_default_params.pop(k, None) + non_default_params.clear() else: raise UnsupportedParamsError( status_code=500, @@ -5060,7 +5056,7 @@ def get_optional_params_embeddings( additional_drop_params=kwargs.get("additional_drop_params", None), ) - if "extra_body" in final_params and len(final_params["extra_body"]) == 0: + if "extra_body" in final_params and not final_params["extra_body"]: final_params.pop("extra_body", None) return final_params diff --git a/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py b/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py index 24157263daa..316fdad5493 100644 --- a/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py +++ b/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py @@ -106,3 +106,4 @@ def test_image_embedding(case: _Case, respx_mock: MockRouter) -> None: else: assert body["inputImage"] == png_b64 assert response.data[0]["embedding"] == [0.11, 0.22, 0.33] + diff --git a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py index 904ca3bb8ca..86e56307210 100644 --- a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py @@ -2130,3 +2130,122 @@ def test_bedrock_embedding_region_bug_reproduction(): img_base_64 = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABkBAMAAACCzIhnAAAAG1BMVEURAAD///+ln5/h39/Dv79qX18uHx+If39MPz9oMSdmAAAACXBIWXMAAA7EAAAOxAGVKw4bAAABB0lEQVRYhe2SzWrEIBCAh2A0jxEs4j6GLDS9hqWmV5Flt0cJS+lRwv742DXpEjY1kOZW6HwHFZnPmVEBEARBEARB/jd0KYA/bcUYbPrRLh6amXHJ/K+ypMoyUaGthILzw0l+xI0jsO7ZcmCcm4ILd+QuVYgpHOmDmz6jBeJImdcUCmeBqQpuqRIbVmQsLCrAalrGpfoEqEogqbLTWuXCPCo+Ki1XGqgQ+jVVuhB8bOaHkvmYuzm/b0KYLWwoK58oFqi6XfxQ4Uz7d6WeKpna6ytUs5e8betMcqAv5YPC5EZB2Lm9FIn0/VP6R58+/GEY1X1egVoZ/3bt/EqF6malgSAIgiDIH+QL41409QMY0LMAAAAASUVORK5CYII=" + + +@pytest.mark.parametrize( + ("model_id", "embed_response", "extra_body", "expected_body"), + [ + ( + "amazon.titan-embed-text-v2:0", + titan_embedding_response, + {"normalize": False, "dimensions": 256}, + {"inputText": test_input, "normalize": False, "dimensions": 256}, + ), + ( + "cohere.embed-english-v3", + cohere_embedding_response, + {"truncate": "END", "embedding_types": ["float"]}, + {"texts": [test_input], "truncate": "END", "embedding_types": ["float"]}, + ), + ], +) +def test_bedrock_embedding_sends_extra_body_fields_in_the_request_body( + respx_mock: respx.MockRouter, + model_id: str, + embed_response: dict[str, object], + extra_body: dict[str, object], + expected_body: dict[str, object], +) -> None: + route = respx_mock.post(f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model_id}/invoke").mock( + return_value=httpx.Response(200, json=embed_response) + ) + + litellm.embedding( + model=f"bedrock/{model_id}", + input=[test_input], + api_key="test-bearer-token", + aws_region_name="us-east-1", + extra_body=extra_body, + ) + + body = json.loads(route.calls.last.request.read()) + assert "extra_body" not in body + assert expected_body.items() <= body.items() + + +def test_bedrock_nova_embedding_sends_extra_body_fields_in_the_request_body(respx_mock: respx.MockRouter) -> None: + model_id: Final = "amazon.nova-2-multimodal-embeddings-v1:0" + route = respx_mock.post(f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model_id}/invoke").mock( + return_value=httpx.Response(200, json={"embeddings": [{"embeddingType": "TEXT", "embedding": [0.1, 0.2]}]}) + ) + + litellm.embedding( + model=f"bedrock/{model_id}", + input=[test_input], + api_key="test-bearer-token", + aws_region_name="us-east-1", + extra_body={"segmentationConfig": {"durationSeconds": 10}}, + ) + + body = json.loads(route.calls.last.request.read()) + assert "extra_body" not in body + assert body["segmentationConfig"] == {"durationSeconds": 10} + + +@pytest.mark.parametrize( + ("extra_body", "field", "expected"), + [ + ({"dimensions": 512}, "dimensions", 256), + ({"inputText": "text the guardrails never saw"}, "inputText", test_input), + ], + ids=["mapped_param", "input"], +) +def test_bedrock_embedding_extra_body_cannot_replace_the_input_or_a_mapped_param( + respx_mock: respx.MockRouter, extra_body: dict[str, object], field: str, expected: object +) -> None: + route = respx_mock.post( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-embed-text-v2:0/invoke" + ).mock(return_value=httpx.Response(200, json=titan_embedding_response)) + + litellm.embedding( + model="bedrock/amazon.titan-embed-text-v2:0", + input=[test_input], + api_key="test-bearer-token", + aws_region_name="us-east-1", + dimensions=256, + extra_body=extra_body, + ) + + assert json.loads(route.calls.last.request.read())[field] == expected + + +def test_bedrock_embedding_signs_the_body_it_sends(respx_mock: respx.MockRouter) -> None: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + from botocore.credentials import Credentials + + route = respx_mock.post( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-embed-text-v2:0/invoke" + ).mock(return_value=httpx.Response(200, json=titan_embedding_response)) + + litellm.embedding( + model="bedrock/amazon.titan-embed-text-v2:0", + input=[test_input], + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_region_name="us-east-1", + dimensions=256, + extra_body={"dimensions": 512}, + ) + + sent: Final = route.calls.last.request + replay: Final = AWSRequest( + method="POST", url=str(sent.url), data=sent.content, headers={"Content-Type": "application/json"} + ) + replay.context["timestamp"] = sent.headers["x-amz-date"] + SigV4Auth( + Credentials("AKIAIOSFODNN7EXAMPLE", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), "bedrock", "us-east-1" + ).add_auth(replay) + + assert json.loads(sent.content)["dimensions"] == 256 + assert replay.headers["Authorization"] == sent.headers["authorization"] diff --git a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py index 3633ac97265..04bd26d53ef 100644 --- a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py +++ b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py @@ -237,3 +237,28 @@ def test_completion_nvidia_nim(): assert request_body["model"] == "databricks/dbrx-instruct" assert request_body["frequency_penalty"] == 0.1 assert request_body["presence_penalty"] == 0.5 + + +def test_embedding_nvidia_nim_sends_extra_body_fields_without_an_extra_body_key(): + from openai import OpenAI + + transport: Final = CapturingTransport( + CreateEmbeddingResponse( + object="list", + data=(Embedding(object="embedding", index=0, embedding=(0.1, 0.2, 0.3)),), + model="nvidia/nv-embedqa-e5-v5", + usage=EmbeddingUsage(prompt_tokens=6, total_tokens=6), + ) + ) + client: Final = OpenAI(api_key="fake-api-key", http_client=httpx.Client(transport=transport)) + + litellm.embedding( + model="nvidia_nim/nvidia/nv-embedqa-e5-v5", + input="What is the meaning of life?", + extra_body={"input_type": "query", "truncate": "END"}, + client=client, + ) + + request_body: Final = transport.request_bodies[0] + assert "extra_body" not in request_body + assert (request_body["input_type"], request_body["truncate"]) == ("query", "END") diff --git a/tests/unit/llms/ollama/test_ollama_embedding.py b/tests/unit/llms/ollama/test_ollama_embedding.py index b32bb9c4415..abbe449bf4a 100644 --- a/tests/unit/llms/ollama/test_ollama_embedding.py +++ b/tests/unit/llms/ollama/test_ollama_embedding.py @@ -348,3 +348,21 @@ async def test_sdk_embedding_rejects_non_string_batch(embedding_route: respx.Rou litellm.embedding(model="ollama/nomic-embed-text", input=[123], api_base="http://ollama.test") assert not embedding_route.called + + +def test_ollama_embeddings_sends_extra_body_fields_at_the_documented_places(embedding_route: respx.Route) -> None: + response: Final = litellm.embedding( + model="ollama/nomic-embed-text", + input=["hello", "world"], + api_base="http://ollama.test", + extra_body={"truncate": False, "keep_alive": "5m", "options": {"num_ctx": 2048}}, + ) + + assert json.loads(embedding_route.calls.last.request.content) == { + "model": "nomic-embed-text", + "input": ["hello", "world"], + "truncate": False, + "keep_alive": "5m", + "options": {"num_ctx": 2048}, + } + assert response.data[0]["embedding"] == [0.1, 0.2] diff --git a/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py b/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py index 3951b17db92..088718b9c60 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py @@ -626,3 +626,40 @@ class TestTEIEmbeddingResponse: if __name__ == "__main__": pytest.main([__file__]) + + +def test_voyage_endpoint_sends_extra_body_fields_in_the_invoke_body(): + client = MagicMock() + client.invoke_endpoint.return_value = { + "Body": MagicMock( + read=MagicMock( + return_value=json.dumps( + { + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "usage": {"total_tokens": 3}, + } + ).encode() + ) + ) + } + session = MagicMock() + session.client.return_value = client + + with patch("boto3.Session", return_value=session): + response = embedding( + model="sagemaker/voyage-3-5-endpoint", + input=["hello"], + aws_region_name="us-east-1", + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + extra_body={"input_type": "query", "truncation": False, "output_dimension": 512}, + ) + + assert json.loads(client.invoke_endpoint.call_args.kwargs["Body"]) == { + "input": ["hello"], + "model": "voyage-3-5-endpoint", + "input_type": "query", + "truncation": False, + "output_dimension": 512, + } + assert response.data[0]["embedding"] == [0.1, 0.2] diff --git a/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index d7d921ab35c..06306981daf 100644 --- a/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -652,3 +652,32 @@ def test_gemini_embedding(respx_mock: respx.MockRouter) -> None: expected_prompt_tokens: Final = litellm.token_counter(model="gemini/gemini-embedding-001", text="Hello, world!") assert response.usage.prompt_tokens == expected_prompt_tokens assert response.usage.total_tokens == expected_prompt_tokens + + +def test_gemini_embedding_sends_extra_body_fields_on_each_request(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents" + ).mock(return_value=httpx.Response(200, json={"embeddings": [{"values": [0.1, 0.2]}]})) + + litellm.embedding( + model="gemini/gemini-embedding-001", + input="Hello, world!", + api_key="gemini-test-key", + extra_body={ + "title": "Greeting", + "taskType": "RETRIEVAL_DOCUMENT", + "outputDimensionality": 256, + "documentOcr": True, + }, + ) + + assert json.loads(route.calls.last.request.content)["requests"] == [ + { + "model": "models/gemini-embedding-001", + "content": {"parts": [{"text": "Hello, world!"}]}, + "title": "Greeting", + "taskType": "RETRIEVAL_DOCUMENT", + "outputDimensionality": 256, + "documentOcr": True, + } + ] diff --git a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py index edb6e889814..123d4128679 100644 --- a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py +++ b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py @@ -1,8 +1,13 @@ import json from unittest.mock import AsyncMock, MagicMock, patch +import google.auth +import google.auth.credentials import httpx import pytest +import respx + +import litellm from litellm.llms.vertex_ai.multimodal_embeddings.transformation import ( @@ -174,3 +179,30 @@ class TestVertexMultimodalEmbedding: headers={}, ) assert "parameters" not in request + + +class _StaticGoogleCredentials(google.auth.credentials.Credentials): + def refresh(self, request: object) -> None: + self.token = "test-token" + + +def test_extra_body_fields_reach_the_multimodal_predict_body( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(google.auth, "default", lambda scopes: (_StaticGoogleCredentials(), "test-project")) + route = respx_mock.post(host="us-central1-aiplatform.googleapis.com").mock( + return_value=httpx.Response(200, json={"predictions": [{"textEmbedding": [0.1, 0.2]}]}) + ) + + litellm.embedding( + model="vertex_ai/multimodalembedding@001", + input="Hello world", + vertex_project="test-project", + vertex_location="us-central1", + extra_body={"dimension": 128}, + ) + + assert json.loads(route.calls.last.request.content) == { + "instances": [{"text": "Hello world"}], + "parameters": {"dimension": 128}, + } diff --git a/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py b/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py index 1058d4fb82d..0e8d49b01e5 100644 --- a/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py +++ b/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py @@ -305,3 +305,20 @@ async def test_vertex_embedding_usage_carries_the_billable_character_count_googl 8, character_count, ) + + +def test_vertex_embedding_call_sends_extra_body_fields_as_predict_parameters(predict_endpoint: respx.Route) -> None: + predict_endpoint.mock(return_value=_predict_response({"predictions": [_gecko_prediction([0.1, 2], 3)]})) + + litellm.embedding( + model=f"vertex_ai/{GECKO}", + input=["hello"], + vertex_project=PUBLIC_CALL_PROJECT, + vertex_location="us-central1", + extra_body={"autoTruncate": False}, + ) + + body: Final = json.loads(predict_endpoint.calls.last.request.content) + assert body["instances"] == [{"content": "hello"}] + assert "extra_body" not in body["parameters"] + assert body["parameters"]["autoTruncate"] is False