mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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>
This commit is contained in:
kenneth-berri
GitHub
parent
f0021859a7
commit
effcb0705f
17 changed files with 365 additions and 27 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue