diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index b928f33e961..05275d58c7d 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -445,7 +445,7 @@ class Cache: 1. If caching groups are set, then return the caching group as the model https://docs.litellm.ai/docs/routing#caching-across-model-groups 2. Else if a model_group is set, then return the model_group as the model. This is used for all requests sent through the litellm.Router() - 3. Else use the `model` passed in kwargs + 3. Else use the `model` passed in kwargs, plus `model_id` when set (e.g. a SageMaker inference component) """ litellm_params: Final[dict] = kwargs.get("litellm_params", {}) or {} metadata_sources: Final[tuple[dict, ...]] = ( @@ -460,7 +460,10 @@ class Cache: caching_group: Final = next( (group for source in metadata_sources if (group := self._get_caching_group(source, model_group))), None ) - return caching_group or model_group or kwargs["model"] + model_id_suffix: Final = ( + f" model_id: {kwargs['model_id']}" if "model_id" in kwargs and kwargs["model_id"] is not None else "" + ) + return caching_group or model_group or f"{kwargs['model']}{model_id_suffix}" def _get_caching_group(self, metadata: dict, model_group: str | None) -> str | None: caching_groups: Final[list | None] = metadata.get("caching_groups", []) diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 7772ca6869a..579309035f0 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -4,7 +4,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping from typing import Final import httpx -from pydantic import ConfigDict, RootModel, TypeAdapter, with_config +from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, ValidationError, with_config from typing_extensions import ReadOnly, Required, TypedDict import litellm @@ -13,7 +13,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import GenericStreamingChunk as GChunk from litellm.types.utils import StreamingChatCompletionChunk -_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) @@ -54,7 +54,9 @@ def _load_sagemaker_response_stream_shape(): from botocore.model import ServiceModel loader: Final = Loader() - service_dict: Final = _JSON_OBJECT.validate_python(loader.load_service_model("sagemaker-runtime", "service-2")) + service_dict: Final = STR_KEYED_MAPPING.validate_python( + loader.load_service_model("sagemaker-runtime", "service-2") + ) return ServiceModel(service_dict).shape_for("InvokeEndpointWithResponseStreamOutput") except Exception as e: verbose_logger.warning( @@ -76,6 +78,38 @@ def get_sagemaker_response_stream_shape(): return _load_sagemaker_response_stream_shape() +def with_inference_component_hint(error_message: str) -> str: + if "Inference Component Name header is required" in error_message: + return f"{error_message}\n pass in via `litellm.embedding(..., model_id={{InferenceComponentName}})`" + if "Inference Component Name header is not allowed" in error_message: + return f"{error_message}\n remove `model_id` from this deployment, the endpoint has no inference components" + return error_message + + +class _ClientErrorBody(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + message: str | None = Field(default=None, alias="Message") + + +class _ClientErrorMetadata(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + http_status_code: int = Field(default=500, alias="HTTPStatusCode") + + +class _ClientErrorResponse(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + error: _ClientErrorBody = Field(default_factory=_ClientErrorBody, alias="Error") + response_metadata: _ClientErrorMetadata = Field(default_factory=_ClientErrorMetadata, alias="ResponseMetadata") + + +def client_error_details(error: Exception) -> tuple[int, str]: + try: + response: Final = _ClientErrorResponse.model_validate(getattr(error, "response", None)) + except ValidationError: + return 500, str(error) + return response.response_metadata.http_status_code, response.error.message or str(error) + + class SagemakerError(BaseLLMException): def __init__( self, diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 069518021fa..b285200ef15 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -4,6 +4,7 @@ from copy import deepcopy from typing import Final, cast import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -24,12 +25,26 @@ from litellm.utils import ( get_secret, ) -from ..common_utils import AWSEventStreamDecoder, SagemakerError +from ..common_utils import ( + STR_KEYED_MAPPING, + AWSEventStreamDecoder, + SagemakerError, + client_error_details, + with_inference_component_hint, +) from ..embedding.transformation import SagemakerEmbeddingConfig from .transformation import SagemakerConfig sagemaker_config: Final = SagemakerConfig() + +class _InvokeEndpointArgs(TypedDict): + EndpointName: ReadOnly[str] + ContentType: ReadOnly[str] + Body: ReadOnly[bytes] + CustomAttributes: ReadOnly[str] + + """ SAGEMAKER AUTH Keys/Vars os.environ['AWS_ACCESS_KEY_ID'] = "" @@ -288,10 +303,7 @@ class SagemakerLLM(BaseAWSLLM): raise e except Exception as e: verbose_logger.error("Sagemaker error %s", str(e)) - status_code: Final = getattr(e, "response", {}).get("ResponseMetadata", {}).get("HTTPStatusCode", 500) - error_message = getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) - if "Inference Component Name header is required" in error_message: - error_message += "\n pass in via `litellm.completion(..., model_id={InferenceComponentName})`" + status_code, error_message = client_error_details(e) raise SagemakerError(status_code=status_code, message=error_message) return sagemaker_config.transform_response( @@ -512,10 +524,7 @@ class SagemakerLLM(BaseAWSLLM): ) raise e except Exception as e: - error_message = f"{e}" - if "Inference Component Name header is required" in error_message: - error_message += "\n pass in via `litellm.completion(..., model_id={InferenceComponentName})`" - raise SagemakerError(status_code=500, message=error_message) + raise SagemakerError(status_code=500, message=str(e)) return sagemaker_config.transform_response( model=model, raw_response=response, @@ -574,16 +583,29 @@ class SagemakerLLM(BaseAWSLLM): #### EMBEDDING LOGIC # Transform request based on model type provider_config: Final = SagemakerEmbeddingConfig.get_model_config(model) - request_data: Final = provider_config.transform_embedding_request(model, input, optional_params, {}) + endpoint_name: Final = SagemakerEmbeddingConfig.get_endpoint_name(model) + embedding_params: Final = STR_KEYED_MAPPING.validate_python(optional_params) + 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, {}) data: Final = json.dumps(request_data).encode("utf-8") + invoke_args: Final = _InvokeEndpointArgs( + EndpointName=endpoint_name, + ContentType="application/json", + Body=data, + CustomAttributes="accept_eula=true", + ) ## LOGGING request_str: Final = f""" response = client.invoke_endpoint( - EndpointName={model}, + EndpointName={endpoint_name}, ContentType="application/json", Body=f"{data!r}", # Use !r for safe representation CustomAttributes="accept_eula=true", + InferenceComponentName={inference_component_name}, )""" logging_obj.pre_call( input=input, @@ -592,16 +614,14 @@ class SagemakerLLM(BaseAWSLLM): ) ## EMBEDDING CALL try: - response = client.invoke_endpoint( - EndpointName=model, - ContentType="application/json", - Body=data, - CustomAttributes="accept_eula=true", + response = ( + client.invoke_endpoint(**invoke_args, InferenceComponentName=inference_component_name) + if inference_component_name + else client.invoke_endpoint(**invoke_args) ) except Exception as e: - status_code: Final = getattr(e, "response", {}).get("ResponseMetadata", {}).get("HTTPStatusCode", 500) - error_message: Final = getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) - raise SagemakerError(status_code=status_code, message=error_message) + status_code, error_message = client_error_details(e) + raise SagemakerError(status_code=status_code, message=with_inference_component_hint(error_message)) response = json.loads(response["Body"].read().decode("utf8")) ## LOGGING @@ -628,7 +648,7 @@ class SagemakerLLM(BaseAWSLLM): # Use the request_data that was already transformed above return provider_config.transform_embedding_response( - model=model, + model=endpoint_name, raw_response=mock_response, model_response=model_response, logging_obj=logging_obj, diff --git a/litellm/llms/sagemaker/embedding/openai_transformation.py b/litellm/llms/sagemaker/embedding/openai_transformation.py new file mode 100644 index 00000000000..a3f0ebfbf59 --- /dev/null +++ b/litellm/llms/sagemaker/embedding/openai_transformation.py @@ -0,0 +1,146 @@ +""" +Translate from OpenAI's `/v1/embeddings` to SageMaker's `/invoke` + +For endpoints whose container already serves the OpenAI embeddings API, opted into per +deployment with the `openai/` route (`sagemaker/openai/`): the request +goes out as `{"input": [...]}` plus the caller's extra body fields, and the response is +read as the OpenAI embeddings object (`data[].embedding`, `usage`). +""" + +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import AllEmbeddingInputValues + +from httpx._models import Headers, Response +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import EmbeddingResponse, Usage + +from ..common_utils import STR_KEYED_MAPPING, SagemakerError + +SUPPORTED_OPENAI_PARAMS: Final = ("dimensions", "encoding_format", "user") + + +class OpenAIEmbeddingUsage(BaseModel): + model_config = ConfigDict(frozen=True) + prompt_tokens: int = 0 + total_tokens: int | None = None + + +class OpenAIEmbeddingRow(BaseModel): + model_config = ConfigDict(frozen=True) + embedding: tuple[float, ...] | str + index: int | None = None + + +class OpenAIEmbeddingsPayload(BaseModel): + model_config = ConfigDict(frozen=True) + data: tuple[OpenAIEmbeddingRow, ...] + model: str | None = None + usage: OpenAIEmbeddingUsage | None = None + + +class SagemakerOpenAIEmbeddingConfig(BaseEmbeddingConfig): + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseConfig returns a list + return list(SUPPORTED_OPENAI_PARAMS) + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: callers keep adding to the mapped params + mapped: Final = {k: v for k, v in non_default_params.items() if k in SUPPORTED_OPENAI_PARAMS} + return {**optional_params, **mapped} + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str] | Headers, # mutable-ok: SagemakerError takes the same header types + ) -> BaseLLMException: + return SagemakerError(message=error_message, status_code=status_code, headers=headers) + + def transform_embedding_request( + self, + model: str, + input: "AllEmbeddingInputValues", + optional_params: Mapping[str, object], + headers: Mapping[str, str], + ) -> dict[str, object]: # mutable-ok: BaseEmbeddingConfig returns the request body as a dict + extra_body: Final = optional_params.get("extra_body") + extra_fields: Final[Mapping[str, object]] = ( + STR_KEYED_MAPPING.validate_python(extra_body) if isinstance(extra_body, dict) else {} + ) + top_level_params: Final = {k: v for k, v in optional_params.items() if k != "extra_body"} + return {"input": input, **top_level_params, **extra_fields} + + def transform_embedding_response( + self, + model: str, + raw_response: Response, + model_response: "EmbeddingResponse", + logging_obj: "LiteLLMLoggingObj", + api_key: str | None = None, + request_data: Mapping[str, object] | None = None, + optional_params: Mapping[str, object] | None = None, + litellm_params: Mapping[str, object] | None = None, + ) -> "EmbeddingResponse": + try: + response_data: Final = raw_response.json() + except ValueError as e: + raise SagemakerError( + message=f"Failed to parse response: {e}", + status_code=raw_response.status_code, + ) + + try: + payload: Final = OpenAIEmbeddingsPayload.model_validate(response_data) + except ValidationError as e: + raise SagemakerError( + status_code=500, + message=( + "Unexpected response format. Expected an OpenAI embeddings object with a 'data' list of " + f"{{'embedding': [...]}} rows: {e}" + ), + ) + + prompt_tokens: Final = payload.usage.prompt_tokens if payload.usage else 0 + total_tokens: Final = ( + payload.usage.total_tokens if payload.usage and payload.usage.total_tokens is not None else prompt_tokens + ) + + return model_response.model_copy( + update={ + "object": "list", + "data": [ + { + "object": "embedding", + "index": row.index if row.index is not None else idx, + "embedding": row.embedding if isinstance(row.embedding, str) else list(row.embedding), + } + for idx, row in enumerate(payload.data) + ], + "model": payload.model or model, + "usage": Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=total_tokens), + } + ) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseConfig returns the request headers as a dict + return {"Content-Type": "application/json"} diff --git a/litellm/llms/sagemaker/embedding/transformation.py b/litellm/llms/sagemaker/embedding/transformation.py index 97940929b09..1395f9b98ec 100644 --- a/litellm/llms/sagemaker/embedding/transformation.py +++ b/litellm/llms/sagemaker/embedding/transformation.py @@ -19,6 +19,20 @@ from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import SagemakerError from .cohere_transformation import SagemakerCohereEmbeddingConfig +from .openai_transformation import SagemakerOpenAIEmbeddingConfig + +OPENAI_COMPATIBLE_ROUTE: Final = "openai/" + + +def _model_specific_config(model: str) -> BaseEmbeddingConfig | None: + if model.startswith(OPENAI_COMPATIBLE_ROUTE): + return SagemakerOpenAIEmbeddingConfig() + model_lower: Final = model.lower() + if "voyage" in model_lower: + return VoyageEmbeddingConfig() + if "cohere" in model_lower: + return SagemakerCohereEmbeddingConfig() + return None class SagemakerEmbeddingConfig(BaseEmbeddingConfig): @@ -35,25 +49,22 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): Factory method to get the appropriate embedding config based on model type Args: - model: The model name + model: The model name, optionally prefixed with the `openai/` route for + endpoints whose container serves the OpenAI embeddings API Returns: Appropriate embedding config instance """ - model_lower: Final = model.lower() - if "voyage" in model_lower: - return VoyageEmbeddingConfig() - if "cohere" in model_lower: - return SagemakerCohereEmbeddingConfig() - return cls() + specific_config: Final = _model_specific_config(model) + return cls() if specific_config is None else specific_config + + @staticmethod + def get_endpoint_name(model: str) -> str: + return model.removeprefix(OPENAI_COMPATIBLE_ROUTE) def get_supported_openai_params(self, model: str) -> list[str]: - model_lower: Final = model.lower() - if "voyage" in model_lower: - return VoyageEmbeddingConfig().get_supported_openai_params(model) - if "cohere" in model_lower: - return SagemakerCohereEmbeddingConfig().get_supported_openai_params(model) - return [] + specific_config: Final = _model_specific_config(model) + return [] if specific_config is None else specific_config.get_supported_openai_params(model) def map_openai_params( self, diff --git a/tests/integration/_support/sagemaker_embedding.py b/tests/integration/_support/sagemaker_embedding.py new file mode 100644 index 00000000000..67dda14f664 --- /dev/null +++ b/tests/integration/_support/sagemaker_embedding.py @@ -0,0 +1,210 @@ +import base64 +import json +import struct +from collections.abc import Mapping +from typing import Final +from urllib.parse import unquote + +from integration._support.wire import Reply, Request +from pydantic import JsonValue, TypeAdapter + +ACCESS_KEY: Final = "AKIAINTEGRATION000004" +SECRET_KEY: Final = "synthetic-secret-key-for-testing" +REGION: Final = "us-east-1" +COMPONENT: Final = "integration-embed-component" +SERVED_MODEL: Final = "integration-org/served-embedding-model" +VOYAGE_SERVED_MODEL: Final = "voyage-served-model" +ENDPOINT_PREFIX: Final = "integration-embed-" +COMPONENT_HEADER: Final = "x-amzn-sagemaker-inference-component" +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +NOT_JSON_BODY: Final = b"upstream gateway error" +EXPECTED_SHAPE: Final = "Expected an OpenAI embeddings object with a 'data' list of {'embedding': [...]} rows" +NON_STRING_MODEL_ID: Final = "model_id must be the inference component name as a string" +IC_REQUIRED: Final = ( + "Inference Component Name header is required for endpoints to which you plan to deploy inference components. " + "Please include Inference Component Name header or consider using SageMaker models." +) +IC_NOT_ALLOWED: Final = ( + "Inference Component Name header is not allowed for endpoints to which you dont plan to deploy inference " + "components. Please remove the Inference Component Name header and try again." +) +PASS_IN_HINT: Final = "pass in via `litellm.embedding(..., model_id={InferenceComponentName})`" +REMOVE_HINT: Final = "remove `model_id` from this deployment, the endpoint has no inference components" + + +def endpoint(scenario: str) -> str: + return f"{ENDPOINT_PREFIX}{scenario}" + + +def openai_route(scenario: str) -> str: + return f"sagemaker/openai/{endpoint(scenario)}" + + +def plain_route(scenario: str) -> str: + return f"sagemaker/{endpoint(scenario)}" + + +def invocations_target(name: str) -> str: + return f"/endpoints/{name}/invocations" + + +def vector(index: int) -> list[float]: + return [0.25, -0.5, 1.0 + index] + + +def inputs(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + carried: Final = next((body[key] for key in ("input", "inputs", "texts") if key in body), []) + return tuple(carried) if isinstance(carried, list) else (carried,) + + +def embedding_value(index: int, encoding_format: JsonValue) -> JsonValue: + floats: Final = vector(index) + if encoding_format == "base64": + return base64.b64encode(struct.pack(f"<{len(floats)}f", *floats)).decode() + return floats + + +def openai_rows(count: int, encoding_format: JsonValue, indexed: bool) -> list[JsonValue]: + return [ + { + "object": "embedding", + **({"index": index} if indexed else {}), + "embedding": embedding_value(index, encoding_format), + } + for index in range(count) + ] + + +def _openai_usage(scenario: str) -> dict[str, JsonValue]: + if scenario == "nousage": + return {} + if scenario == "promptonly": + return {"usage": {"prompt_tokens": 5}} + return {"usage": {"prompt_tokens": 7, "total_tokens": 7}} + + +def openai_payload(body: Mapping[str, JsonValue], scenario: str = "openai") -> dict[str, JsonValue]: + return { + "object": "list", + "data": openai_rows(len(inputs(body)), body.get("encoding_format"), indexed=scenario != "noindex"), + **({} if scenario == "nomodel" else {"model": SERVED_MODEL}), + **_openai_usage(scenario), + } + + +def json_reply(payload: JsonValue) -> Reply: + return Reply(body=json.dumps(payload).encode()) + + +def aws_error(status: int, code: str, message: str, **fields: JsonValue) -> Reply: + return Reply( + status=status, + body=json.dumps({"__type": code, "Message": message, **fields}).encode(), + headers={"x-amzn-ErrorType": code}, + ) + + +def endpoint_name_of(request: Request) -> str: + parts: Final = request.target.split("/") + return unquote(parts[2]) if len(parts) > 2 else "" + + +def scenario_of(name: str) -> str: + if not name.startswith(ENDPOINT_PREFIX): + return "" + return name.removeprefix(ENDPOINT_PREFIX).split("-")[0] + + +def respond(request: Request) -> Reply: + name: Final = endpoint_name_of(request) + body: Final = JSON_OBJECT.validate_json(request.body) + match scenario_of(name): + case "openai" | "nousage" | "nomodel" | "noindex" | "promptonly": + return json_reply(openai_payload(body, scenario_of(name))) + case "hf": + return json_reply([vector(index) for index in range(len(inputs(body)))]) + case "icrequired": + if COMPONENT_HEADER not in request.headers: + return aws_error(400, "ValidationError", IC_REQUIRED) + return json_reply(openai_payload(body)) + case "noic": + if COMPONENT_HEADER in request.headers: + return aws_error(400, "ValidationError", IC_NOT_ALLOWED) + return json_reply([vector(index) for index in range(len(inputs(body)))]) + case "hfdict": + return json_reply({"embedding": [vector(index) for index in range(len(inputs(body)))]}) + case "voyage": + return json_reply( + { + "object": "list", + "data": openai_rows(len(inputs(body)), None, indexed=True), + "model": VOYAGE_SERVED_MODEL, + "usage": {"total_tokens": 9}, + } + ) + case "cohere": + return json_reply( + { + "id": "integration-cohere-response", + "texts": list(inputs(body)), + "embeddings": [vector(index) for index in range(len(inputs(body)))], + "meta": {"api_version": {"version": "1"}, "billed_units": {"input_tokens": 6}}, + "response_type": "embeddings_floats", + } + ) + case "datadict": + return json_reply({"object": "list", "data": {"embedding": vector(0)}}) + case "hfshape": + return json_reply([vector(0)]) + case "notjson": + return Reply(body=NOT_JSON_BODY, content_type="text/html") + case "err400": + return aws_error( + 400, "ValidationError", "1 validation error detected: Value at 'body' failed to satisfy constraint" + ) + case "err424": + return aws_error( + 424, + "ModelError", + 'Received client error (400) from model with message "{\\"error\\":\\"unexpected field inputs\\"}".', + OriginalStatusCode=400, + OriginalMessage='{"error":"unexpected field inputs"}', + LogStreamArn="arn:aws:logs:us-east-1:000000000000:log-group:/aws/sagemaker/Endpoints/integration", + ) + case "err429": + return aws_error(429, "ModelNotReadyException", "Model is not ready for inference yet") + case "err503": + return aws_error(503, "ServiceUnavailable", "The endpoint is scaling, try again later") + case "err500": + return aws_error(500, "InternalFailure", "An internal failure occurred") + case "drop": + return Reply(drop_connection=True) + case _: + return aws_error(400, "ValidationError", f"Endpoint {name} of account 000000000000 not found.") + + +def litellm_params(model: str, **parameters: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "aws_access_key_id": ACCESS_KEY, + "aws_secret_access_key": SECRET_KEY, + "aws_region_name": REGION, + **parameters, + } + + +def deployment( + name: str, model: str, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue +) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": litellm_params(model, **parameters), + **({} if model_info is None else {"model_info": dict(model_info)}), + } + + +def openai_deployment( + name: str, scenario: str, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue +) -> dict[str, JsonValue]: + defaults: Final[dict[str, JsonValue]] = {"model_id": COMPONENT, "extra_body": {"model": SERVED_MODEL}} + return deployment(name, openai_route(scenario), model_info=model_info, **{**defaults, **parameters}) diff --git a/tests/integration/providers/test_sagemaker_embedding_chaos.py b/tests/integration/providers/test_sagemaker_embedding_chaos.py new file mode 100644 index 00000000000..680ad3bfd2d --- /dev/null +++ b/tests/integration/providers/test_sagemaker_embedding_chaos.py @@ -0,0 +1,302 @@ +import asyncio +import re +import socket +import threading +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, string_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy, owned_proxy_process +from integration._support.sagemaker_embedding import ( + COMPONENT, + COMPONENT_HEADER, + JSON_OBJECT, + SERVED_MODEL, + json_reply, + openai_deployment, + openai_payload, + respond, +) +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark = pytest.mark.timeout(300) + +_MODEL: Final = "sm-openai-chaos" +_ALIASES: Final = ( + "/v1/embeddings", + "/embeddings", + f"/openai/deployments/{_MODEL}/embeddings", + f"/engines/{_MODEL}/embeddings", +) +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_ROWS_BY_CALL: Final = ( + "SELECT litellm_call_id, status FROM \"LiteLLM_SpendLogs\" WHERE litellm_call_id = ANY(string_to_array(%s, ','))" +) + + +@dataclass(frozen=True, slots=True) +class _Served: + text: str + status: int + body: str + call_id: str + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] + + +def _config(directory: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = [openai_deployment(_MODEL, "openai", model_info={"mode": "embedding"})] + path: Final = directory / "sagemaker-embedding-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _runtime(peer_url: str) -> dict[str, str]: + return {"AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME": peer_url, "AWS_MAX_ATTEMPTS": "1"} + + +@pytest.fixture(scope="module") +def rig() -> Iterator[Gateway]: + with gateway_from_environment() as gateway: + yield gateway + + +@pytest.fixture(scope="module") +def peer_port() -> int: + return _free_port() + + +@pytest.fixture(scope="module") +def proxy(rig: Gateway, peer_port: int, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("sagemaker-embedding-chaos") + runtime: Final = _runtime(f"http://127.0.0.1:{peer_port}") + with owned_proxy(rig, directory, runtime, config=_config(directory), workers=2) as owned: + yield owned + + +def _texts(count: int) -> tuple[str, ...]: + return tuple(f"integration embedding {uuid.uuid4().hex}" for _ in range(count)) + + +async def _send(client: httpx.AsyncClient, key: str, index: int, text: str) -> _Served: + path: Final = _ALIASES[index % len(_ALIASES)] + body: Final[dict[str, JsonValue]] = { + "input": text, + "cache": {"no-cache": True}, + **({"model": _MODEL} if path.endswith("/v1/embeddings") or path == "/embeddings" else {}), + } + response: Final = await client.post(path, json=body, headers={"Authorization": f"Bearer {key}"}) + return _Served(text, response.status_code, response.text, response.headers.get("x-litellm-call-id", "")) + + +async def _burst( + base_url: str, key: str, texts: tuple[str, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=90, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, index, text) for index, text in enumerate(texts)), + return_exceptions=tolerate_transport_errors, + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _assert_embedded(served: _Served) -> None: + assert served.status == 200, served.body + payload: Final = JSON_OBJECT.validate_json(served.body) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], served.body + + +def _assert_peer_saw(received: tuple[Request, ...], texts: tuple[str, ...]) -> None: + assert {request.headers.get(COMPONENT_HEADER) for request in received} == {COMPONENT}, received + bodies: Final = [JSON_OBJECT.validate_json(request.body) for request in received] + assert sorted(string_value(body["input"]) for body in bodies) == sorted(texts), bodies + assert {body["model"] for body in bodies} == {SERVED_MODEL}, bodies + + +def _spend_status_by_call(served: tuple[_Served, ...]) -> dict[str, JsonValue]: + wanted: Final = sorted(item.call_id for item in served) + rows: Final = eventually( + lambda: read_rows(_ROWS_BY_CALL, (",".join(wanted),)), lambda found: len(found) >= len(wanted), seconds=70 + ) + assert sorted(string_value(row["litellm_call_id"]) for row in rows) == wanted, rows + return {string_value(row["litellm_call_id"]): row["status"] for row in rows} + + +def _unhealthy_count(gateway: Gateway) -> JsonValue: + response: Final = gateway.request("GET", "/health", params={"model": _MODEL}) + return JSON_OBJECT.validate_json(response.content).get("unhealthy_count") + + +async def test_concurrent_burst_across_every_alias_reaches_the_component_and_is_logged_once_per_call( + proxy: Gateway, peer_port: int +) -> None: + texts: Final = _texts(24) + with wire_server(respond, port=peer_port) as wire: + served: Final = await _burst(str(proxy.client.base_url), proxy.key, texts) + assert len(served) == 24 + for item in served: + _assert_embedded(item) + _assert_peer_saw(wire.drain(), texts) + assert _spend_status_by_call(served) == {item.call_id: "success" for item in served} + + +async def test_peer_outage_fails_each_call_once_reports_unhealthy_and_recovers_on_restart( + proxy: Gateway, peer_port: int +) -> None: + before: Final = _texts(12) + during: Final = _texts(4) + after: Final = _texts(4) + base_url: Final = str(proxy.client.base_url) + with wire_server(respond, port=peer_port) as wire: + served: Final = await _burst(base_url, proxy.key, before) + assert len(served) == 12 + for item in served: + _assert_embedded(item) + _assert_peer_saw(wire.drain(), before) + refused: Final = await _burst(base_url, proxy.key, during) + assert len(refused) == 4 + for item in refused: + assert item.status == 503, item.body + assert "Could not connect to the endpoint URL" in item.body, item.body + assert await asyncio.to_thread(eventually, lambda: _unhealthy_count(proxy), lambda count: count == 1, 60) + with wire_server(respond, port=peer_port) as restarted: + recovered: Final = await _burst(base_url, proxy.key, after) + assert len(recovered) == 4 + for item in recovered: + _assert_embedded(item) + healthy: Final = [request for request in restarted.drain() if "test from litellm" not in request.body.decode()] + _assert_peer_saw(tuple(healthy), after) + assert _spend_status_by_call((*served, *refused, *recovered)) == { + **{item.call_id: "success" for item in served}, + **{item.call_id: "failure" for item in refused}, + **{item.call_id: "success" for item in recovered}, + } + + +def _slow(request: Request) -> Reply: + body: Final = json_reply(openai_payload(JSON_OBJECT.validate_json(request.body))).body + return Reply(chunks=(body[:8], body[8:]), pause_between_chunks=0.5) + + +async def test_slow_peer_answers_every_caller_whole_without_duplicate_rows(proxy: Gateway, peer_port: int) -> None: + texts: Final = _texts(10) + with wire_server(_slow, port=peer_port) as wire: + served: Final = await _burst(str(proxy.client.base_url), proxy.key, texts) + assert len(served) == 10 + for item in served: + _assert_embedded(item) + _assert_peer_saw(wire.drain(), texts) + assert _spend_status_by_call(served) == {item.call_id: "success" for item in served} + + +def _live_workers(log: Path) -> tuple[int, ...]: + started: Final = (int(pid) for pid in _STARTED_WORKER.findall(log.read_text())) + return tuple(pid for pid in started if psutil.pid_exists(pid)) + + +def _open_peer_connections(pid: int, peer_url: str) -> int: + port: Final = urlsplit(peer_url).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@dataclass(frozen=True, slots=True) +class _HeldPeer: + release: threading.Event + held: SimpleQueue[str] + + def respond(self, request: Request) -> Reply: + self.held.put(string_value(JSON_OBJECT.validate_json(request.body)["input"])) + assert self.release.wait(timeout=120), "The burst was never released" + return respond(request) + + +@pytest.mark.timeout(600) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_the_component(rig: Gateway, tmp_path: Path) -> None: + texts: Final = _texts(8) + peer: Final = _HeldPeer(threading.Event(), SimpleQueue()) + with wire_server(peer.respond) as wire: + with owned_proxy_process(rig, tmp_path, _runtime(wire.url), config=_config(tmp_path), workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually(lambda: _live_workers(owned.log), lambda pids: len(pids) == 2, seconds=30) + burst: Final = asyncio.create_task( + _burst(str(candidate.client.base_url), candidate.key, texts, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, peer.held.qsize, lambda size: size == len(texts), 120) + held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == len(texts), held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + peer.release.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_embedded(item) + (follow_up,) = _texts(1) + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, (follow_up,)) + _assert_embedded(answered) + _assert_peer_saw(wire.drain(), (*texts, follow_up)) + assert _spend_status_by_call((*served, answered)) == { + item.call_id: "success" for item in (*served, answered) + } + + +@pytest.mark.timeout(600) +async def test_proxy_sigterm_mid_burst_finishes_in_flight_calls_and_a_relaunch_serves_again( + rig: Gateway, tmp_path: Path +) -> None: + texts: Final = _texts(8) + peer: Final = _HeldPeer(threading.Event(), SimpleQueue()) + with wire_server(peer.respond) as wire: + first: Final = tmp_path / "first" + first.mkdir() + with owned_proxy_process(rig, first, _runtime(wire.url), config=_config(first), workers=2) as owned: + burst: Final = asyncio.create_task( + _burst(str(owned.gateway.client.base_url), owned.gateway.key, texts, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, peer.held.qsize, lambda size: size == len(texts), 120) + owned.begin_shutdown() + peer.release.set() + served: Final = await burst + exit_code: Final = await asyncio.to_thread( + eventually, owned.process.poll, lambda code: code is not None, 120 + ) + assert exit_code is not None + for item in served: + _assert_embedded(item) + _assert_peer_saw(wire.drain(), texts) + assert len(served) == len(texts), (len(served), len(texts)) + second: Final = tmp_path / "second" + second.mkdir() + with owned_proxy_process(rig, second, _runtime(wire.url), config=_config(second), workers=2) as relaunched: + (follow_up,) = _texts(1) + (answered,) = await _burst(str(relaunched.gateway.client.base_url), relaunched.gateway.key, (follow_up,)) + _assert_embedded(answered) + _assert_peer_saw(wire.drain(), (follow_up,)) + assert _spend_status_by_call((*served, answered)) == { + item.call_id: "success" for item in (*served, answered) + } diff --git a/tests/integration/providers/test_sagemaker_embedding_wire.py b/tests/integration/providers/test_sagemaker_embedding_wire.py new file mode 100644 index 00000000000..0ed38b1a768 --- /dev/null +++ b/tests/integration/providers/test_sagemaker_embedding_wire.py @@ -0,0 +1,742 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Final + +import httpx +import openai +import pytest +import yaml +from botocore.exceptions import ConnectionClosedError +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.sagemaker_embedding import ( + ACCESS_KEY, + COMPONENT, + COMPONENT_HEADER, + EXPECTED_SHAPE, + IC_NOT_ALLOWED, + IC_REQUIRED, + JSON_OBJECT, + NON_STRING_MODEL_ID, + NOT_JSON_BODY, + PASS_IN_HINT, + REMOVE_HINT, + SERVED_MODEL, + deployment, + embedding_value, + endpoint, + invocations_target, + openai_deployment, + openai_route, + plain_route, + respond, +) +from integration._support.wire import Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark = pytest.mark.timeout(240) + +_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} +_SPEND_ROW: Final = 'SELECT status, call_type, model_group FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s' + + +def _model_list(health_endpoint: str) -> list[JsonValue]: + return [ + openai_deployment("sm-openai", "openai"), + deployment( + "sm-openai-explicit", + f"openai/{endpoint('openai')}", + custom_llm_provider="sagemaker", + model_id=COMPONENT, + extra_body={"model": SERVED_MODEL}, + ), + deployment("sm-openai-bare", openai_route("openai")), + deployment("sm-openai-no-model-id", openai_route("openai"), extra_body={"model": SERVED_MODEL}), + openai_deployment("sm-openai-custom-key", "openai", input_type="query"), + deployment( + "sm-openai-pinned", + openai_route("openai"), + model_id=COMPONENT, + extra_body={"model": SERVED_MODEL, "encoding_format": "float"}, + ), + openai_deployment("sm-openai-model-id-int", "openai", model_id=7), + openai_deployment("sm-openai-model-id-list", "openai", model_id=[COMPONENT]), + openai_deployment("sm-openai-model-id-empty", "openai", model_id=""), + openai_deployment("sm-openai-model-id-null", "openai", model_id=None), + openai_deployment("sm-openai-model-id-long", "openai", model_id="c" * 5120), + openai_deployment("sm-openai-extra-body-string", "openai", extra_body="not-an-object"), + openai_deployment("sm-openai-extra-body-empty", "openai", extra_body={}), + openai_deployment("sm-openai-nousage", "nousage"), + openai_deployment("sm-openai-nomodel", "nomodel"), + openai_deployment("sm-openai-noindex", "noindex"), + openai_deployment("sm-openai-promptonly", "promptonly"), + openai_deployment("sm-openai-datadict", "datadict"), + openai_deployment("sm-openai-hfshape", "hfshape"), + openai_deployment("sm-openai-notjson", "notjson"), + openai_deployment("sm-openai-err400", "err400"), + openai_deployment("sm-openai-err424", "err424"), + openai_deployment("sm-openai-err429", "err429"), + openai_deployment("sm-openai-err503", "err503"), + openai_deployment("sm-openai-err500", "err500"), + openai_deployment("sm-openai-drop", "drop"), + deployment("sm-openai-unknown", "sagemaker/openai/unknown-embed-endpoint", model_id=COMPONENT), + openai_deployment("sm-openai-ic", "icrequired"), + deployment("sm-openai-ic-missing", openai_route("icrequired"), extra_body={"model": SERVED_MODEL}), + deployment("sm-plain-noic", plain_route("noic"), model_id=COMPONENT), + deployment("sm-plain-noic-control", plain_route("noic")), + deployment( + "sm-openai-health", + f"sagemaker/openai/{health_endpoint}", + model_id=COMPONENT, + extra_body={"model": SERVED_MODEL}, + model_info={"mode": "embedding"}, + ), + deployment("sm-plain", plain_route("hf"), model_id=COMPONENT), + deployment("sm-plain-control", plain_route("hf")), + deployment("sm-plain-drop-params", plain_route("hf"), model_id=COMPONENT, drop_params=True), + deployment("sm-plain-dict", plain_route("hfdict"), model_id=COMPONENT), + deployment("sm-plain-err424", plain_route("err424"), model_id=COMPONENT), + deployment("sm-voyage", plain_route("voyage"), model_id=COMPONENT), + deployment("sm-cohere", plain_route("cohere"), model_id=COMPONENT), + ] + + +def _config(directory: Path, health_endpoint: str) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = _model_list(health_endpoint) + path: Final = directory / "sagemaker-embedding.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def rig() -> Iterator[Gateway]: + with gateway_from_environment() as gateway: + yield gateway + + +@pytest.fixture(scope="module") +def wire() -> Iterator[Wire]: + with wire_server(respond) as served: + yield served + + +@pytest.fixture(autouse=True) +def _drained_wire(wire: Wire) -> None: + wire.drain() + + +@pytest.fixture(scope="module") +def health_endpoint() -> str: + return endpoint(f"openai-health-{uuid.uuid4().hex[:8]}") + + +@pytest.fixture(scope="module") +def proxy( + rig: Gateway, wire: Wire, health_endpoint: str, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("sagemaker-embedding") + overrides: Final = {"AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME": wire.url, "AWS_MAX_ATTEMPTS": "2"} + with owned_proxy(rig, directory, overrides, config=_config(directory, health_endpoint), workers=2) as owned: + yield owned + + +def _text() -> str: + return f"integration embedding {uuid.uuid4().hex}" + + +def _embed(proxy: Gateway, model: str, text: JsonValue, **fields: JsonValue) -> httpx.Response: + return proxy.request("POST", "/v1/embeddings", {"model": model, "input": text, "cache": _NO_CACHE, **fields}) + + +def _spend_row(response: httpx.Response) -> dict[str, JsonValue]: + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (call_id,)), lambda found: len(found) == 1, seconds=70) + return rows[0] + + +def _usage(payload: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + usage: Final = payload["usage"] + assert isinstance(usage, dict), payload + return tuple(usage[key] for key in ("prompt_tokens", "completion_tokens", "total_tokens")) + + +def _only_request(wire: Wire) -> Request: + received: Final = wire.drain() + assert len(received) == 1, [(request.method, request.target) for request in received] + return received[0] + + +def _body(request: Request) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_json(request.body) + + +def _forwarded_body(sent: httpx.Request, text: str) -> dict[str, JsonValue]: + asked: Final = JSON_OBJECT.validate_json(sent.content) + encoding: Final = {"encoding_format": asked["encoding_format"]} if "encoding_format" in asked else {} + return {"input": [text], **encoding, "model": SERVED_MODEL} + + +def _signed_headers(request: Request) -> tuple[str, ...]: + authorization: Final = request.headers["authorization"] + assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/"), authorization + signed: Final = next(part for part in authorization.split(", ") if part.startswith("SignedHeaders=")) + return tuple(signed.removeprefix("SignedHeaders=").split(";")) + + +def _sdk_base_url(proxy: Gateway) -> str: + return f"{str(proxy.client.base_url).rstrip('/')}/v1" + + +def _sdk(proxy: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=_sdk_base_url(proxy), api_key=proxy.key, max_retries=0, timeout=60) + + +def _async_sdk(proxy: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=_sdk_base_url(proxy), api_key=proxy.key, max_retries=0, timeout=60) + + +def test_openai_route_signs_model_id_as_the_inference_component_header_and_sends_the_openai_body( + proxy: Gateway, wire: Wire +) -> None: + texts: Final = [_text(), _text()] + response: Final = _embed(proxy, "sm-openai", texts) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [ + {"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}, + {"object": "embedding", "index": 1, "embedding": [0.25, -0.5, 2.0]}, + ], response.text + assert payload["model"] == "sm-openai", response.text + assert _usage(payload) == (7, 0, 7), response.text + request: Final = _only_request(wire) + assert (request.method, request.target) == ("POST", invocations_target(endpoint("openai"))), request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert request.headers["x-amzn-sagemaker-custom-attributes"] == "accept_eula=true", dict(request.headers) + assert request.headers["content-type"] == "application/json", dict(request.headers) + assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"] + assert _body(request) == {"input": texts, "model": SERVED_MODEL}, request.body + row: Final = _spend_row(response) + assert row == {"status": "success", "call_type": "aembedding", "model_group": "sm-openai"}, row + + +@pytest.mark.parametrize( + ("path", "model_in_body"), + [ + ("/v1/embeddings", True), + ("/embeddings", True), + ("/openai/deployments/sm-openai/embeddings", False), + ("/engines/sm-openai/embeddings", False), + ], +) +def test_every_embeddings_alias_reaches_the_inference_component( + proxy: Gateway, wire: Wire, path: str, model_in_body: bool +) -> None: + text: Final = _text() + body: Final[dict[str, JsonValue]] = { + "input": text, + "cache": _NO_CACHE, + **({"model": "sm-openai"} if model_in_body else {}), + } + response: Final = proxy.request("POST", path, body) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": text, "model": SERVED_MODEL}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_openai_sdk_sync_client_reads_the_vectors_in_the_encoding_it_asked_for(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + raw: Final = _sdk(proxy).embeddings.with_raw_response.create( + model="sm-openai", input=[text], extra_body={"cache": _NO_CACHE} + ) + created: Final = raw.parse() + assert [row.embedding for row in created.data] == [[0.25, -0.5, 1.0]], created + assert created.model == "sm-openai", created + assert (created.usage.prompt_tokens, created.usage.total_tokens) == (7, 7), created + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == _forwarded_body(raw.http_request, text), request.body + call_id: Final = raw.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (call_id,)), lambda found: len(found) == 1, seconds=70) + assert rows[0]["status"] == "success", rows + + +async def test_openai_sdk_async_client_reads_the_vectors_in_the_encoding_it_asked_for( + proxy: Gateway, wire: Wire +) -> None: + text: Final = _text() + raw: Final = await _async_sdk(proxy).embeddings.with_raw_response.create( + model="sm-openai", input=[text], extra_body={"cache": _NO_CACHE} + ) + created: Final = raw.parse() + assert [row.embedding for row in created.data] == [[0.25, -0.5, 1.0]], created + assert created.model == "sm-openai", created + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == _forwarded_body(raw.http_request, text), request.body + call_id: Final = raw.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (call_id,)), lambda found: len(found) == 1, seconds=70) + assert rows[0]["status"] == "success", rows + + +def test_explicit_custom_llm_provider_spelling_takes_the_openai_route(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-explicit", [text]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["model"] == "sm-openai-explicit", response.text + request: Final = _only_request(wire) + assert request.target == f"/endpoints/{endpoint('openai')}/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_plain_route_signs_model_id_as_the_header_and_keeps_the_hf_body(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain", [text]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + assert payload["model"] == "sm-plain", response.text + assert _usage(payload) == (3, 0, 3), response.text + request: Final = _only_request(wire) + assert request.target == f"/endpoints/{endpoint('hf')}/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"] + assert _body(request) == {"inputs": [text]}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_plain_route_without_model_id_sends_no_inference_component_header(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain-control", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert COMPONENT_HEADER not in _signed_headers(request), request.headers["authorization"] + assert _body(request) == {"inputs": [text]}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_plain_route_reads_the_wrapped_hf_embedding_key(proxy: Gateway, wire: Wire) -> None: + texts: Final = [_text(), _text()] + response: Final = _embed(proxy, "sm-plain-dict", texts) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [ + {"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}, + {"object": "embedding", "index": 1, "embedding": [0.25, -0.5, 2.0]}, + ], response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"inputs": texts}, request.body + + +def test_voyage_endpoint_keeps_the_voyage_body_and_signs_model_id_as_the_header(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-voyage", [text]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + assert payload["model"] == "sm-voyage", response.text + assert _usage(payload) == (9, 0, 9), response.text + request: Final = _only_request(wire) + assert request.target == f"/endpoints/{endpoint('voyage')}/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "model": endpoint("voyage")}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_cohere_endpoint_keeps_the_cohere_body_and_signs_model_id_as_the_header(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-cohere", [text]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + assert _usage(payload) == (6, 0, 6), response.text + request: Final = _only_request(wire) + assert request.target == f"/endpoints/{endpoint('cohere')}/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"texts": [text], "input_type": "search_document"}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_openai_route_forwards_dimensions_encoding_format_and_user_in_the_body(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed( + proxy, "sm-openai", [text], dimensions=3, encoding_format="float", user="integration-embedding-user" + ) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == { + "input": [text], + "dimensions": 3, + "encoding_format": "float", + "user": "integration-embedding-user", + "model": SERVED_MODEL, + }, request.body + + +def test_openai_route_without_model_id_or_extra_body_sends_the_bare_openai_body(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-bare", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert _body(request) == {"input": [text]}, request.body + + +def test_openai_route_sends_a_custom_deployment_key_at_the_top_level_of_the_body(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-custom-key", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "input_type": "query", "model": SERVED_MODEL}, request.body + + +def test_deployment_extra_body_wins_over_the_callers_encoding_format(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + created: Final = _sdk(proxy).embeddings.create( + model="sm-openai-pinned", input=[text], extra_body={"cache": _NO_CACHE} + ) + assert [row.embedding for row in created.data] == [[0.25, -0.5, 1.0]], created + request: Final = _only_request(wire) + assert _body(request) == {"input": [text], "encoding_format": "float", "model": SERVED_MODEL}, request.body + + +def test_openai_route_returns_the_base64_string_a_raw_caller_asked_for_unchanged(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai", [text], encoding_format="base64") + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": embedding_value(0, "base64")}], ( + response.text + ) + request: Final = _only_request(wire) + assert _body(request) == {"input": [text], "encoding_format": "base64", "model": SERVED_MODEL}, request.body + + +def test_openai_route_reports_zero_usage_when_the_container_sends_none(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-nousage", [_text()]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert _usage(payload) == (0, 0, 0), response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + assert _spend_row(response)["status"] == "success" + + +def test_openai_route_names_the_endpoint_when_the_container_sends_no_model(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-nomodel", [_text()]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["model"] == "sm-openai-nomodel", response.text + assert _usage(payload) == (7, 0, 7), response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + + +def test_openai_route_numbers_rows_the_container_left_unindexed(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-noindex", [_text(), _text()]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [ + {"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}, + {"object": "embedding", "index": 1, "embedding": [0.25, -0.5, 2.0]}, + ], response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + + +def test_openai_route_takes_total_tokens_from_prompt_tokens_when_the_container_omits_it( + proxy: Gateway, wire: Wire +) -> None: + response: Final = _embed(proxy, "sm-openai-promptonly", [_text()]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert _usage(payload) == (5, 0, 5), response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + + +@pytest.mark.parametrize("model", ["sm-openai-datadict", "sm-openai-hfshape"]) +def test_openai_route_rejects_a_container_body_that_is_not_an_openai_embeddings_object( + proxy: Gateway, wire: Wire, model: str +) -> None: + response: Final = _embed(proxy, model, [_text()]) + assert response.status_code == 503, response.text + assert EXPECTED_SHAPE in response.text, response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + assert _spend_row(response)["status"] == "failure" + + +def test_openai_route_reports_a_non_json_container_body_as_a_server_error(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-notjson", [_text()]) + with pytest.raises(json.JSONDecodeError) as decode_error: + json.loads(NOT_JSON_BODY) + assert response.status_code == 500, response.text + assert str(decode_error.value) in response.text, response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + assert _spend_row(response)["status"] == "failure" + + +def test_container_validation_error_reaches_the_caller_as_400(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-err400", [_text()]) + assert response.status_code == 400, response.text + assert "1 validation error detected" in response.text, response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + assert _spend_row(response)["status"] == "failure" + + +def test_container_model_error_reaches_the_caller_with_the_original_message(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-err424", [_text()]) + assert response.status_code == 424, response.text + assert "unexpected field inputs" in response.text, response.text + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT + assert _spend_row(response)["status"] == "failure" + + +def test_plain_route_body_rejected_by_an_openai_container_reaches_the_caller(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain-err424", [text]) + assert response.status_code == 424, response.text + assert "unexpected field inputs" in response.text, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"inputs": [text]}, request.body + assert _spend_row(response)["status"] == "failure" + + +@pytest.mark.parametrize( + ("model", "status", "message"), + [ + ("sm-openai-err429", 429, "Model is not ready for inference yet"), + ("sm-openai-err503", 503, "The endpoint is scaling, try again later"), + ("sm-openai-err500", 503, "An internal failure occurred"), + ], +) +def test_container_outage_errors_reach_the_caller_and_every_attempt_carries_the_header( + proxy: Gateway, wire: Wire, model: str, status: int, message: str +) -> None: + text: Final = _text() + response: Final = _embed(proxy, model, [text]) + assert response.status_code == status, response.text + assert message in response.text, response.text + attempts: Final = wire.drain() + assert len(attempts) >= 1, attempts + assert {attempt.headers[COMPONENT_HEADER] for attempt in attempts} == {COMPONENT}, attempts + assert {attempt.body for attempt in attempts} == {json.dumps({"input": [text], "model": SERVED_MODEL}).encode()} + assert _spend_row(response)["status"] == "failure" + + +def test_dropped_connection_reaches_the_caller_as_a_server_error(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-drop", [_text()]) + assert response.status_code == 503, response.text + assert ConnectionClosedError.fmt.partition(":")[0] in response.text, response.text + attempts: Final = wire.drain() + assert len(attempts) >= 1, attempts + assert {attempt.headers[COMPONENT_HEADER] for attempt in attempts} == {COMPONENT}, attempts + assert _spend_row(response)["status"] == "failure" + + +def test_inference_component_endpoint_answers_once_the_header_unlocks_it(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-ic", [text]) + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_inference_component_endpoint_without_model_id_tells_the_caller_how_to_pass_it( + proxy: Gateway, wire: Wire +) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-ic-missing", [text]) + assert response.status_code == 400, response.text + assert IC_REQUIRED in response.text, response.text + assert PASS_IN_HINT in response.text, response.text + request: Final = _only_request(wire) + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + assert _spend_row(response)["status"] == "failure" + + +def test_endpoint_without_inference_components_tells_the_deployment_to_drop_model_id( + proxy: Gateway, wire: Wire +) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain-noic", [text]) + assert response.status_code == 400, response.text + assert IC_NOT_ALLOWED in response.text, response.text + assert REMOVE_HINT in response.text, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"inputs": [text]}, request.body + assert _spend_row(response)["status"] == "failure" + + +def test_endpoint_without_inference_components_serves_a_deployment_without_model_id(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain-noic-control", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert _body(request) == {"inputs": [text]}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_unknown_endpoint_reaches_the_caller_as_400(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-openai-unknown", [_text()]) + assert response.status_code == 400, response.text + assert "Endpoint unknown-embed-endpoint of account 000000000000 not found." in response.text, response.text + request: Final = _only_request(wire) + assert request.target == "/endpoints/unknown-embed-endpoint/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _spend_row(response)["status"] == "failure" + + +def test_plain_route_rejects_an_unsupported_openai_param_without_drop_params(proxy: Gateway, wire: Wire) -> None: + response: Final = _embed(proxy, "sm-plain", [_text()], dimensions=3) + assert response.status_code == 400, response.text + assert "does not support parameters: {'dimensions': 3}" in response.text, response.text + assert "drop_params" in response.text, response.text + assert wire.drain() == (), response.text + assert _spend_row(response)["status"] == "failure" + + +def test_plain_route_drops_an_unsupported_openai_param_under_drop_params(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-plain-drop-params", [text], dimensions=3) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"inputs": [text]}, request.body + + +@pytest.mark.parametrize("model", ["sm-openai-model-id-int", "sm-openai-model-id-list"]) +def test_a_non_string_model_id_fails_before_any_request_leaves_the_proxy( + proxy: Gateway, wire: Wire, model: str +) -> None: + response: Final = _embed(proxy, model, [_text()]) + assert response.status_code == 400, response.text + assert NON_STRING_MODEL_ID in response.text, response.text + assert wire.drain() == (), response.text + assert _spend_row(response)["status"] == "failure" + + +@pytest.mark.parametrize("model", ["sm-openai-model-id-empty", "sm-openai-model-id-null"]) +def test_an_empty_or_null_model_id_sends_no_inference_component_header(proxy: Gateway, wire: Wire, model: str) -> None: + text: Final = _text() + response: Final = _embed(proxy, model, [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + + +def test_a_long_model_id_is_signed_and_forwarded_whole(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-model-id-long", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == "c" * 5120, len(request.headers[COMPONENT_HEADER]) + assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"] + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + + +@pytest.mark.parametrize("model", ["sm-openai-extra-body-string", "sm-openai-extra-body-empty"]) +def test_an_extra_body_that_is_not_a_populated_object_adds_nothing_to_the_body( + proxy: Gateway, wire: Wire, model: str +) -> None: + text: Final = _text() + response: Final = _embed(proxy, model, [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text]}, request.body + + +def test_a_caller_sent_model_id_overrides_the_deployments_inference_component(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai", [text], model_id="caller-chosen-component") + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == "caller-chosen-component", dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + + +def test_a_caller_sent_model_id_is_forwarded_when_the_deployment_has_none(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai-bare", [text], model_id="caller-chosen-component") + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == "caller-chosen-component", dict(request.headers) + assert _body(request) == {"input": [text]}, request.body + + +def test_a_caller_sent_non_string_model_id_fails_before_any_request_leaves_the_proxy( + proxy: Gateway, wire: Wire +) -> None: + response: Final = _embed(proxy, "sm-openai-bare", [_text()], model_id=["caller-chosen-component"]) + assert response.status_code == 400, response.text + assert NON_STRING_MODEL_ID in response.text, response.text + assert wire.drain() == (), response.text + + +def test_a_caller_sent_extra_body_replaces_the_deployments_extra_body(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + response: Final = _embed(proxy, "sm-openai", [text], extra_body={"truncate": "END"}) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "truncate": "END"}, request.body + + +def test_a_failed_call_does_not_poison_the_next_one(proxy: Gateway, wire: Wire) -> None: + failed: Final = _embed(proxy, "sm-openai-err503", [_text()]) + assert failed.status_code == 503, failed.text + wire.drain() + text: Final = _text() + response: Final = _embed(proxy, "sm-openai", [text]) + assert response.status_code == 200, response.text + request: Final = _only_request(wire) + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body + assert _spend_row(response)["status"] == "success" + + +def test_two_identical_uncached_requests_each_reach_the_container_and_log_once(proxy: Gateway, wire: Wire) -> None: + text: Final = _text() + first: Final = _embed(proxy, "sm-openai", [text]) + second: Final = _embed(proxy, "sm-openai", [text]) + assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text) + assert first.headers["x-litellm-call-id"] != second.headers["x-litellm-call-id"], (first.headers, second.headers) + requests: Final = wire.drain() + assert [_body(request) for request in requests] == [{"input": [text], "model": SERVED_MODEL}] * 2, requests + assert {request.headers[COMPONENT_HEADER] for request in requests} == {COMPONENT}, requests + assert (_spend_row(first)["status"], _spend_row(second)["status"]) == ("success", "success") + + +def test_health_check_reaches_the_container_with_the_inference_component_header( + proxy: Gateway, wire: Wire, health_endpoint: str +) -> None: + report: Final = proxy.get("/health", {"model": "sm-openai-health"}) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + request: Final = _only_request(wire) + assert request.target == f"/endpoints/{health_endpoint}/invocations", request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + body: Final = _body(request) + assert body["model"] == SERVED_MODEL, request.body + assert isinstance(body["input"], list) and len(body["input"]) == 1, request.body diff --git a/tests/integration/sdk/test_sagemaker_embedding_sdk.py b/tests/integration/sdk/test_sagemaker_embedding_sdk.py new file mode 100644 index 00000000000..6b70a27fbec --- /dev/null +++ b/tests/integration/sdk/test_sagemaker_embedding_sdk.py @@ -0,0 +1,171 @@ +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.sagemaker_embedding import ( + ACCESS_KEY, + COMPONENT, + COMPONENT_HEADER, + JSON_OBJECT, + NON_STRING_MODEL_ID, + REGION, + SECRET_KEY, + SERVED_MODEL, + endpoint, + invocations_target, + litellm_params, + openai_route, + plain_route, + respond, +) +from integration._support.wire import Request, Wire, wire_server +from pydantic import JsonValue + +import litellm + +_AWS: Final[dict[str, str]] = { + "aws_access_key_id": ACCESS_KEY, + "aws_secret_access_key": SECRET_KEY, + "aws_region_name": REGION, +} + + +@pytest.fixture(scope="module") +def wire() -> Iterator[Wire]: + with wire_server(respond) as served: + yield served + + +@pytest.fixture(autouse=True) +def _runtime_at_the_peer(wire: Wire, monkeypatch: pytest.MonkeyPatch) -> None: + wire.drain() + monkeypatch.setenv("AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME", wire.url) + monkeypatch.setenv("AWS_MAX_ATTEMPTS", "1") + + +def _text() -> str: + return f"integration embedding {uuid.uuid4().hex}" + + +def _only_request(wire: Wire) -> Request: + received: Final = wire.drain() + assert len(received) == 1, [(request.method, request.target) for request in received] + return received[0] + + +def _signed_headers(request: Request) -> tuple[str, ...]: + authorization: Final = request.headers["authorization"] + assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/"), authorization + signed: Final = next(part for part in authorization.split(", ") if part.startswith("SignedHeaders=")) + return tuple(signed.removeprefix("SignedHeaders=").split(";")) + + +def _assert_component_request(request: Request, scenario: str, body: dict[str, JsonValue]) -> None: + assert (request.method, request.target) == ("POST", invocations_target(endpoint(scenario))), request + assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers) + assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"] + assert JSON_OBJECT.validate_json(request.body) == body, request.body + + +def _vectors(response: litellm.EmbeddingResponse) -> list[JsonValue]: + return [row["embedding"] for row in response.data] + + +def test_embedding_signs_model_id_as_the_inference_component_and_reads_the_served_model(wire: Wire) -> None: + text: Final = _text() + response: Final = litellm.embedding( + model=openai_route("openai"), input=[text], model_id=COMPONENT, extra_body={"model": SERVED_MODEL}, **_AWS + ) + assert response.model == SERVED_MODEL, response + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + assert (response.usage.prompt_tokens, response.usage.total_tokens) == (7, 7), response.usage + _assert_component_request(_only_request(wire), "openai", {"input": [text], "model": SERVED_MODEL}) + + +async def test_aembedding_signs_model_id_as_the_inference_component_and_reads_the_served_model(wire: Wire) -> None: + texts: Final = [_text(), _text()] + response: Final = await litellm.aembedding( + model=openai_route("openai"), input=texts, model_id=COMPONENT, extra_body={"model": SERVED_MODEL}, **_AWS + ) + assert response.model == SERVED_MODEL, response + assert _vectors(response) == [[0.25, -0.5, 1.0], [0.25, -0.5, 2.0]], response + _assert_component_request(_only_request(wire), "openai", {"input": texts, "model": SERVED_MODEL}) + + +def test_embedding_names_the_endpoint_when_the_container_reports_no_model(wire: Wire) -> None: + text: Final = _text() + response: Final = litellm.embedding(model=openai_route("nomodel"), input=[text], model_id=COMPONENT, **_AWS) + assert response.model == endpoint("nomodel"), response + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + _assert_component_request(_only_request(wire), "nomodel", {"input": [text]}) + + +def test_plain_route_keeps_the_hf_body_and_signs_the_component_header(wire: Wire) -> None: + text: Final = _text() + response: Final = litellm.embedding(model=plain_route("hf"), input=[text], model_id=COMPONENT, **_AWS) + assert response.model == endpoint("hf"), response + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + assert response.usage.prompt_tokens == 3, response.usage + _assert_component_request(_only_request(wire), "hf", {"inputs": [text]}) + + +def test_plain_route_without_model_id_sends_no_component_header(wire: Wire) -> None: + text: Final = _text() + response: Final = litellm.embedding(model=plain_route("hf"), input=[text], **_AWS) + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + request: Final = _only_request(wire) + assert (request.method, request.target) == ("POST", invocations_target(endpoint("hf"))), request + assert COMPONENT_HEADER not in request.headers, dict(request.headers) + assert COMPONENT_HEADER not in _signed_headers(request), request.headers["authorization"] + assert JSON_OBJECT.validate_json(request.body) == {"inputs": [text]}, request.body + + +def test_explicit_provider_spelling_signs_model_id_as_the_inference_component(wire: Wire) -> None: + text: Final = _text() + response: Final = litellm.embedding( + model=f"openai/{endpoint('openai')}", + custom_llm_provider="sagemaker", + input=[text], + model_id=COMPONENT, + extra_body={"model": SERVED_MODEL}, + **_AWS, + ) + assert response.model == SERVED_MODEL, response + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + _assert_component_request(_only_request(wire), "openai", {"input": [text], "model": SERVED_MODEL}) + + +async def test_router_aembedding_forwards_the_deployment_model_id_as_the_component(wire: Wire) -> None: + text: Final = _text() + router: Final = litellm.Router( + model_list=[ + { + "model_name": "sm-openai-router", + "litellm_params": litellm_params( + openai_route("openai"), model_id=COMPONENT, extra_body={"model": SERVED_MODEL} + ), + } + ], + num_retries=0, + ) + response: Final = await router.aembedding(model="sm-openai-router", input=[text]) + assert response.model == SERVED_MODEL, response + assert _vectors(response) == [[0.25, -0.5, 1.0]], response + _assert_component_request(_only_request(wire), "openai", {"input": [text], "model": SERVED_MODEL}) + + +def test_embedding_rejects_a_non_string_model_id_before_calling_aws(wire: Wire) -> None: + with pytest.raises(litellm.BadRequestError) as caught: + litellm.embedding(model=openai_route("openai"), input=[_text()], model_id=[COMPONENT], **_AWS) + assert caught.value.status_code == 400, caught.value + assert NON_STRING_MODEL_ID in str(caught.value), caught.value + assert wire.drain() == (), "AWS was called with a list-valued model_id" + + +def test_embedding_surfaces_the_container_model_error_with_its_status(wire: Wire) -> None: + with pytest.raises(litellm.BadRequestError) as caught: + litellm.embedding(model=openai_route("err424"), input=[_text()], model_id=COMPONENT, **_AWS) + assert caught.value.status_code == 424, caught.value + assert "unexpected field inputs" in str(caught.value), caught.value + assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 0cad6e350e3..c28455f3aca 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -1548,3 +1548,23 @@ async def test_redis_caching_ttl_sadd( ) redis_client.expire.assert_awaited_once_with("cache-key", timedelta(seconds=120)) + + +def test_cache_key_separates_model_ids_unless_a_model_group_owns_the_key() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + variants: Final[tuple[dict[str, object], ...]] = ( + {"model_id": "component-a"}, + {"model_id": "component-b"}, + {}, + {"model_id": "component-a"}, + {"model_id": "component-a", "metadata": {"model_group": "embed"}}, + {"model_id": "component-b", "metadata": {"model_group": "embed"}}, + ) + key_a, key_b, bare, key_a_again, grouped_a, grouped_b = ( + cache.get_cache_key(model="sagemaker/openai/embed", input=["hello"], **extra) for extra in variants + ) + + assert key_a != key_b + assert key_a != bare + assert key_a == key_a_again + assert grouped_a == grouped_b diff --git a/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py b/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py index 51a2ccb5640..17fac935679 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py @@ -4,7 +4,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder +from botocore.exceptions import ClientError, ConnectionClosedError + +from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder, client_error_details from litellm.llms.sagemaker.completion.transformation import SagemakerConfig @@ -32,9 +34,7 @@ def test_sagemaker_response_stream_shape_lazy_loads_once(): import litellm.llms.sagemaker.common_utils as mod sentinel = MagicMock() - with patch.object( - mod, "_load_sagemaker_response_stream_shape", return_value=sentinel - ) as mock_load: + with patch.object(mod, "_load_sagemaker_response_stream_shape", return_value=sentinel) as mock_load: assert mod.get_sagemaker_response_stream_shape() is sentinel assert mod.get_sagemaker_response_stream_shape() is sentinel mock_load.assert_called_once() @@ -94,9 +94,7 @@ def test_sagemaker_response_stream_shape_is_structure_shape(): from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape shape = get_sagemaker_response_stream_shape() - assert ( - shape is not None - ), "get_sagemaker_response_stream_shape() is None — botocore may not be installed" + assert shape is not None, "get_sagemaker_response_stream_shape() is None — botocore may not be installed" assert isinstance(shape, StructureShape) assert shape.name == "InvokeEndpointWithResponseStreamOutput" @@ -274,3 +272,38 @@ class TestSagemakerTransform: # The function should properly map max_tokens if max_completion_tokens is not provided assert result == {"temperature": 0.7, "max_new_tokens": 200} + + +NULL_MESSAGE_ERROR = ClientError( + {"Error": {"Code": "InternalFailure", "Message": None}, "ResponseMetadata": {"HTTPStatusCode": 500}}, + "InvokeEndpoint", +) +DROPPED_CONNECTION_ERROR = ConnectionClosedError(endpoint_url="https://runtime.sagemaker.test/x") +HTTPX_STATUS_ERROR = httpx.HTTPStatusError( + "boom", request=httpx.Request("POST", "https://runtime.sagemaker.test/x"), response=httpx.Response(400) +) +PLAIN_ERROR = RuntimeError("plain failure") + + +@pytest.mark.parametrize( + ("error", "expected_status", "expected_message"), + [ + ( + ClientError( + { + "Error": {"Code": "ValidationError", "Message": "bad input"}, + "ResponseMetadata": {"HTTPStatusCode": 400}, + }, + "InvokeEndpoint", + ), + 400, + "bad input", + ), + (NULL_MESSAGE_ERROR, 500, str(NULL_MESSAGE_ERROR)), + (DROPPED_CONNECTION_ERROR, 500, str(DROPPED_CONNECTION_ERROR)), + (HTTPX_STATUS_ERROR, 500, str(HTTPX_STATUS_ERROR)), + (PLAIN_ERROR, 500, str(PLAIN_ERROR)), + ], +) +def test_client_error_details_reads_every_error_shape(error, expected_status, expected_message): + assert client_error_details(error) == (expected_status, expected_message) diff --git a/tests/unit/llms/sagemaker/test_sagemaker_embedding_openai.py b/tests/unit/llms/sagemaker/test_sagemaker_embedding_openai.py new file mode 100644 index 00000000000..aacbddf7d10 --- /dev/null +++ b/tests/unit/llms/sagemaker/test_sagemaker_embedding_openai.py @@ -0,0 +1,378 @@ +""" +Test cases for OpenAI-compatible SageMaker embedding endpoints and inference components + +Covers the `sagemaker/openai/` route (request and response in the OpenAI +embeddings shape) and `model_id` forwarded to `invoke_endpoint` as `InferenceComponentName` +for every SageMaker embedding config. +""" + +import asyncio +import json +from collections.abc import Mapping +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from botocore.exceptions import ClientError, ConnectionClosedError + +import litellm +from litellm import Router +from litellm.llms.sagemaker.common_utils import SagemakerError +from litellm.llms.sagemaker.completion.handler import SagemakerLLM +from litellm.llms.sagemaker.embedding.openai_transformation import ( + SagemakerOpenAIEmbeddingConfig, +) +from litellm.llms.sagemaker.embedding.transformation import SagemakerEmbeddingConfig +from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig +from litellm.types.utils import EmbeddingResponse + +OPENAI_RESPONSE = { + "object": "list", + "data": [ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}, + {"object": "embedding", "index": 1, "embedding": [0.4, 0.5, 0.6]}, + ], + "model": "served-embedding-model", + "usage": {"prompt_tokens": 11, "total_tokens": 11}, +} + + +def fake_sagemaker_session(response_body: object) -> tuple[MagicMock, MagicMock]: + client = MagicMock() + client.invoke_endpoint.return_value = { + "Body": MagicMock(read=MagicMock(return_value=json.dumps(response_body).encode())) + } + session = MagicMock() + session.client.return_value = client + return session, client + + +def invoke_kwargs(client: MagicMock) -> Mapping[str, object]: + client.invoke_endpoint.assert_called_once() + return client.invoke_endpoint.call_args.kwargs + + +class TestSagemakerEmbeddingFactoryOpenAIRoute: + def test_openai_route_selects_openai_config(self): + config = SagemakerEmbeddingConfig.get_model_config("openai/my-embed-endpoint") + + assert isinstance(config, SagemakerOpenAIEmbeddingConfig) + assert SagemakerEmbeddingConfig().get_supported_openai_params("openai/my-embed-endpoint") == [ + "dimensions", + "encoding_format", + "user", + ] + + def test_openai_route_wins_over_model_name_substrings(self): + assert isinstance( + SagemakerEmbeddingConfig.get_model_config("openai/voyage-endpoint"), + SagemakerOpenAIEmbeddingConfig, + ) + assert isinstance( + SagemakerEmbeddingConfig.get_model_config("openai/cohere-endpoint"), + SagemakerOpenAIEmbeddingConfig, + ) + + def test_plain_endpoint_names_keep_their_config(self): + assert type(SagemakerEmbeddingConfig.get_model_config("my-embed-endpoint")) is SagemakerEmbeddingConfig + assert isinstance(SagemakerEmbeddingConfig.get_model_config("voyage-endpoint"), VoyageEmbeddingConfig) + + def test_endpoint_name_strips_only_the_route(self): + assert SagemakerEmbeddingConfig.get_endpoint_name("openai/my-embed-endpoint") == "my-embed-endpoint" + assert SagemakerEmbeddingConfig.get_endpoint_name("my-embed-endpoint") == "my-embed-endpoint" + assert SagemakerEmbeddingConfig.get_endpoint_name("my-openai-endpoint") == "my-openai-endpoint" + + +class TestSagemakerOpenAIEmbeddingConfig: + def setup_method(self): + self.config = SagemakerOpenAIEmbeddingConfig() + + def test_request_is_openai_shaped_with_extra_body_flattened(self): + result = self.config.transform_embedding_request( + model="my-embed-endpoint", + input=["hello", "world"], + optional_params={"truncate": True, "extra_body": {"model": "served-embedding-model"}}, + headers={}, + ) + + assert result == {"input": ["hello", "world"], "truncate": True, "model": "served-embedding-model"} + + def test_request_without_extra_body_sends_input_only(self): + result = self.config.transform_embedding_request( + model="my-embed-endpoint", input=["hello"], optional_params={}, headers={} + ) + + assert result == {"input": ["hello"]} + + def test_openai_params_pass_through(self): + mapped = self.config.map_openai_params( + non_default_params={"dimensions": 256, "encoding_format": "float", "user": "u1", "stream": True}, + optional_params={}, + model="my-embed-endpoint", + drop_params=False, + ) + + assert mapped == {"dimensions": 256, "encoding_format": "float", "user": "u1"} + + def test_response_reads_data_rows_and_usage(self): + raw = httpx.Response(200, json=OPENAI_RESPONSE) + + result = self.config.transform_embedding_response( + model="my-embed-endpoint", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + ) + + assert result.object == "list" + assert result.data == [ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}, + {"object": "embedding", "index": 1, "embedding": [0.4, 0.5, 0.6]}, + ] + assert result.model == "served-embedding-model" + assert result.usage is not None + assert (result.usage.prompt_tokens, result.usage.total_tokens) == (11, 11) + + def test_response_without_usage_or_model_falls_back(self): + raw = httpx.Response(200, json={"data": [{"embedding": [0.9]}]}) + + result = self.config.transform_embedding_response( + model="my-embed-endpoint", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + ) + + assert result.data == [{"object": "embedding", "index": 0, "embedding": [0.9]}] + assert result.model == "my-embed-endpoint" + assert result.usage is not None + assert (result.usage.prompt_tokens, result.usage.total_tokens) == (0, 0) + + def test_response_keeps_base64_embedding_strings(self): + raw = httpx.Response( + 200, json={"data": [{"index": 0, "embedding": "AACAPwAAAEA="}], "usage": {"prompt_tokens": 2}} + ) + + result = self.config.transform_embedding_response( + model="my-embed-endpoint", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + ) + + assert result.data == [{"object": "embedding", "index": 0, "embedding": "AACAPwAAAEA="}] + assert result.usage is not None + assert (result.usage.prompt_tokens, result.usage.total_tokens) == (2, 2) + + @pytest.mark.parametrize( + "body", + [ + {"embedding": [[0.1, 0.2]]}, + [[0.1, 0.2]], + {"data": [{"index": 0}]}, + {"data": "not-a-list"}, + ], + ) + def test_response_rejects_non_openai_shapes(self, body): + raw = httpx.Response(200, json=body) + + with pytest.raises(SagemakerError) as exc_info: + self.config.transform_embedding_response( + model="my-embed-endpoint", + raw_response=raw, + model_response=EmbeddingResponse(), + logging_obj=MagicMock(), + ) + + assert exc_info.value.status_code == 500 + assert "'data'" in exc_info.value.message + + +class TestSagemakerEmbeddingInferenceComponents: + def setup_method(self): + self.sagemaker_llm = SagemakerLLM() + + def embed_with(self, session: MagicMock, model: str, optional_params: Mapping[str, object]) -> EmbeddingResponse: + with patch("boto3.Session", return_value=session): + return self.sagemaker_llm.embedding( + model=model, + input=["hello", "good morning"], + model_response=EmbeddingResponse(), + print_verbose=print, + encoding=None, + logging_obj=MagicMock(), + optional_params={"aws_region_name": "us-east-1", **optional_params}, + ) + + def run_embedding( + self, model: str, optional_params: Mapping[str, object], response_body: object + ) -> tuple[EmbeddingResponse, Mapping[str, object]]: + session, client = fake_sagemaker_session(response_body) + result = self.embed_with(session, model, optional_params) + return result, invoke_kwargs(client) + + def test_openai_route_forwards_component_and_keeps_it_out_of_the_body(self): + result, kwargs = self.run_embedding( + model="openai/my-embed-endpoint", + optional_params={"model_id": "my-component", "extra_body": {"model": "served-embedding-model"}}, + response_body=OPENAI_RESPONSE, + ) + + assert kwargs["EndpointName"] == "my-embed-endpoint" + assert kwargs["InferenceComponentName"] == "my-component" + assert json.loads(kwargs["Body"]) == {"input": ["hello", "good morning"], "model": "served-embedding-model"} + assert [row["embedding"] for row in result.data] == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + assert result.usage is not None + assert result.usage.prompt_tokens == 11 + + def test_hf_endpoint_forwards_component_and_keeps_its_body_shape(self): + result, kwargs = self.run_embedding( + model="my-embed-endpoint", + optional_params={"model_id": "my-component", "normalize": True}, + response_body={"embedding": [[0.1, 0.2], [0.3, 0.4]]}, + ) + + assert kwargs["EndpointName"] == "my-embed-endpoint" + assert kwargs["InferenceComponentName"] == "my-component" + assert json.loads(kwargs["Body"]) == {"inputs": ["hello", "good morning"], "normalize": True} + assert [row["embedding"] for row in result.data] == [[0.1, 0.2], [0.3, 0.4]] + + def test_no_component_without_model_id(self): + _, kwargs = self.run_embedding( + model="my-embed-endpoint", + optional_params={}, + response_body={"embedding": [[0.1, 0.2], [0.3, 0.4]]}, + ) + + assert "InferenceComponentName" not in kwargs + assert json.loads(kwargs["Body"]) == {"inputs": ["hello", "good morning"]} + + @pytest.mark.parametrize( + ("sagemaker_message", "expected_hint"), + [ + ( + "Inference Component Name header is not allowed for endpoints to which you dont plan to deploy" + " inference components. Please remove the Inference Component Name header and try again.", + "remove `model_id` from this deployment, the endpoint has no inference components", + ), + ( + "Inference Component Name header is required for this endpoint.", + "pass in via `litellm.embedding(..., model_id={InferenceComponentName})`", + ), + ('Received server error (500) from model with message "boom".', ""), + ], + ) + def test_component_errors_keep_status_and_name_model_id(self, sagemaker_message: str, expected_hint: str): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + client.invoke_endpoint.side_effect = ClientError( + { + "Error": {"Code": "ValidationError", "Message": sagemaker_message}, + "ResponseMetadata": {"HTTPStatusCode": 400}, + }, + "InvokeEndpoint", + ) + with pytest.raises(SagemakerError) as raised: + self.embed_with(session, "my-embed-endpoint", {"model_id": "stray-component"}) + + assert raised.value.status_code == 400 + assert raised.value.message == ( + sagemaker_message + "\n " + expected_hint if expected_hint else sagemaker_message + ) + + def test_null_error_message_surfaces_the_error_code(self): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + error = ClientError( + { + "Error": {"Code": "InternalFailure", "Message": None}, + "ResponseMetadata": {"HTTPStatusCode": 500}, + }, + "InvokeEndpoint", + ) + client.invoke_endpoint.side_effect = error + with pytest.raises(SagemakerError) as raised: + self.embed_with(session, "openai/my-embed-endpoint", {"model_id": "x" * 2048}) + + assert raised.value.status_code == 500 + assert raised.value.message == str(error) + + def test_dropped_connection_surfaces_as_a_server_error(self): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + error = ConnectionClosedError(endpoint_url="https://runtime.sagemaker.test/x") + client.invoke_endpoint.side_effect = error + with pytest.raises(SagemakerError) as raised: + self.embed_with(session, "openai/my-embed-endpoint", {"model_id": "my-component"}) + + assert raised.value.status_code == 500 + assert raised.value.message == str(error) + + def test_non_string_model_id_is_rejected_before_the_call(self): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + with pytest.raises(SagemakerError) as raised: + self.embed_with(session, "openai/my-embed-endpoint", {"model_id": ["my-component"]}) + + assert raised.value.status_code == 400 + assert raised.value.message == "model_id must be the inference component name as a string" + client.invoke_endpoint.assert_not_called() + + +SDK_KWARGS = { + "model": "sagemaker/openai/my-embed-endpoint", + "input": ["hello", "good morning"], + "model_id": "my-component", + "extra_body": {"model": "served-embedding-model"}, + "aws_region_name": "us-east-1", + "aws_access_key_id": "test-key", + "aws_secret_access_key": "test-secret", +} + +EXPECTED_BODY = {"input": ["hello", "good morning"], "model": "served-embedding-model"} + + +class TestSagemakerOpenAIEmbeddingSdkCalls: + def test_sdk_sync(self): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + with patch("boto3.Session", return_value=session): + result = litellm.embedding(**SDK_KWARGS) + + kwargs = invoke_kwargs(client) + assert (kwargs["EndpointName"], kwargs["InferenceComponentName"]) == ("my-embed-endpoint", "my-component") + assert json.loads(kwargs["Body"]) == EXPECTED_BODY + assert len(result.data) == 2 + assert result.usage is not None + assert result.usage.total_tokens == 11 + + def test_sdk_async(self): + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + with patch("boto3.Session", return_value=session): + result = asyncio.run(litellm.aembedding(**SDK_KWARGS)) + + kwargs = invoke_kwargs(client) + assert (kwargs["EndpointName"], kwargs["InferenceComponentName"]) == ("my-embed-endpoint", "my-component") + assert json.loads(kwargs["Body"]) == EXPECTED_BODY + assert len(result.data) == 2 + + def test_router_deployment_with_model_id(self): + router = Router( + model_list=[ + { + "model_name": "my-embeddings", + "litellm_params": { + "model": "sagemaker/openai/my-embed-endpoint", + "model_id": "my-component", + "extra_body": {"model": "served-embedding-model"}, + "aws_region_name": "us-east-1", + "aws_access_key_id": "test-key", + "aws_secret_access_key": "test-secret", + }, + "model_info": {"mode": "embedding"}, + } + ] + ) + session, client = fake_sagemaker_session(OPENAI_RESPONSE) + with patch("boto3.Session", return_value=session): + result = asyncio.run(router.aembedding(model="my-embeddings", input=["hello", "good morning"])) + + kwargs = invoke_kwargs(client) + assert (kwargs["EndpointName"], kwargs["InferenceComponentName"]) == ("my-embed-endpoint", "my-component") + assert json.loads(kwargs["Body"]) == EXPECTED_BODY + assert [row["index"] for row in result.data] == [0, 1]