Spaces:
Sleeping
Sleeping
| """ | |
| Handles embedding calls to Bedrock's `/invoke` endpoint | |
| """ | |
| import copy | |
| import json | |
| from typing import Any, Callable, List, Optional, Tuple, Union | |
| import httpx | |
| import litellm | |
| from litellm.llms.cohere.embed.handler import embedding as cohere_embedding | |
| from litellm.llms.custom_httpx.http_handler import ( | |
| AsyncHTTPHandler, | |
| HTTPHandler, | |
| _get_httpx_client, | |
| get_async_httpx_client, | |
| ) | |
| from litellm.secret_managers.main import get_secret | |
| from litellm.types.llms.bedrock import AmazonEmbeddingRequest, CohereEmbeddingRequest | |
| from litellm.types.utils import EmbeddingResponse | |
| from ..base_aws_llm import BaseAWSLLM | |
| from ..common_utils import BedrockError | |
| from .amazon_titan_g1_transformation import AmazonTitanG1Config | |
| from .amazon_titan_multimodal_transformation import ( | |
| AmazonTitanMultimodalEmbeddingG1Config, | |
| ) | |
| from .amazon_titan_v2_transformation import AmazonTitanV2Config | |
| from .cohere_transformation import BedrockCohereEmbeddingConfig | |
| class BedrockEmbedding(BaseAWSLLM): | |
| def _load_credentials( | |
| self, | |
| optional_params: dict, | |
| ) -> Tuple[Any, str]: | |
| try: | |
| from botocore.credentials import Credentials | |
| except ImportError: | |
| raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") | |
| ## CREDENTIALS ## | |
| # pop aws_secret_access_key, aws_access_key_id, aws_session_token, aws_region_name from kwargs, since completion calls fail with them | |
| aws_secret_access_key = optional_params.pop("aws_secret_access_key", None) | |
| aws_access_key_id = optional_params.pop("aws_access_key_id", None) | |
| aws_session_token = optional_params.pop("aws_session_token", None) | |
| aws_region_name = optional_params.pop("aws_region_name", None) | |
| aws_role_name = optional_params.pop("aws_role_name", None) | |
| aws_session_name = optional_params.pop("aws_session_name", None) | |
| aws_profile_name = optional_params.pop("aws_profile_name", None) | |
| aws_web_identity_token = optional_params.pop("aws_web_identity_token", None) | |
| aws_sts_endpoint = optional_params.pop("aws_sts_endpoint", None) | |
| ### SET REGION NAME ### | |
| if aws_region_name is None: | |
| # check env # | |
| litellm_aws_region_name = get_secret("AWS_REGION_NAME", None) | |
| if litellm_aws_region_name is not None and isinstance( | |
| litellm_aws_region_name, str | |
| ): | |
| aws_region_name = litellm_aws_region_name | |
| standard_aws_region_name = get_secret("AWS_REGION", None) | |
| if standard_aws_region_name is not None and isinstance( | |
| standard_aws_region_name, str | |
| ): | |
| aws_region_name = standard_aws_region_name | |
| if aws_region_name is None: | |
| aws_region_name = "us-west-2" | |
| credentials: Credentials = self.get_credentials( | |
| aws_access_key_id=aws_access_key_id, | |
| aws_secret_access_key=aws_secret_access_key, | |
| aws_session_token=aws_session_token, | |
| aws_region_name=aws_region_name, | |
| aws_session_name=aws_session_name, | |
| aws_profile_name=aws_profile_name, | |
| aws_role_name=aws_role_name, | |
| aws_web_identity_token=aws_web_identity_token, | |
| aws_sts_endpoint=aws_sts_endpoint, | |
| ) | |
| return credentials, aws_region_name | |
| async def async_embeddings(self): | |
| pass | |
| def _make_sync_call( | |
| self, | |
| client: Optional[HTTPHandler], | |
| timeout: Optional[Union[float, httpx.Timeout]], | |
| api_base: str, | |
| headers: dict, | |
| data: dict, | |
| ) -> dict: | |
| if client is None or not isinstance(client, HTTPHandler): | |
| _params = {} | |
| if timeout is not None: | |
| if isinstance(timeout, float) or isinstance(timeout, int): | |
| timeout = httpx.Timeout(timeout) | |
| _params["timeout"] = timeout | |
| client = _get_httpx_client(_params) # type: ignore | |
| else: | |
| client = client | |
| try: | |
| response = client.post(url=api_base, headers=headers, data=json.dumps(data)) # type: ignore | |
| response.raise_for_status() | |
| except httpx.HTTPStatusError as err: | |
| error_code = err.response.status_code | |
| raise BedrockError(status_code=error_code, message=err.response.text) | |
| except httpx.TimeoutException: | |
| raise BedrockError(status_code=408, message="Timeout error occurred.") | |
| return response.json() | |
| async def _make_async_call( | |
| self, | |
| client: Optional[AsyncHTTPHandler], | |
| timeout: Optional[Union[float, httpx.Timeout]], | |
| api_base: str, | |
| headers: dict, | |
| data: dict, | |
| ) -> dict: | |
| if client is None or not isinstance(client, AsyncHTTPHandler): | |
| _params = {} | |
| if timeout is not None: | |
| if isinstance(timeout, float) or isinstance(timeout, int): | |
| timeout = httpx.Timeout(timeout) | |
| _params["timeout"] = timeout | |
| client = get_async_httpx_client( | |
| params=_params, llm_provider=litellm.LlmProviders.BEDROCK | |
| ) | |
| else: | |
| client = client | |
| try: | |
| response = await client.post(url=api_base, headers=headers, data=json.dumps(data)) # type: ignore | |
| response.raise_for_status() | |
| except httpx.HTTPStatusError as err: | |
| error_code = err.response.status_code | |
| raise BedrockError(status_code=error_code, message=err.response.text) | |
| except httpx.TimeoutException: | |
| raise BedrockError(status_code=408, message="Timeout error occurred.") | |
| return response.json() | |
| def _single_func_embeddings( | |
| self, | |
| client: Optional[HTTPHandler], | |
| timeout: Optional[Union[float, httpx.Timeout]], | |
| batch_data: List[dict], | |
| credentials: Any, | |
| extra_headers: Optional[dict], | |
| endpoint_url: str, | |
| aws_region_name: str, | |
| model: str, | |
| logging_obj: Any, | |
| ): | |
| try: | |
| from botocore.auth import SigV4Auth | |
| from botocore.awsrequest import AWSRequest | |
| except ImportError: | |
| raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") | |
| responses: List[dict] = [] | |
| for data in batch_data: | |
| sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) | |
| headers = {"Content-Type": "application/json"} | |
| if extra_headers is not None: | |
| headers = {"Content-Type": "application/json", **extra_headers} | |
| request = AWSRequest( | |
| method="POST", url=endpoint_url, data=json.dumps(data), headers=headers | |
| ) | |
| sigv4.add_auth(request) | |
| if ( | |
| extra_headers is not None and "Authorization" in extra_headers | |
| ): # prevent sigv4 from overwriting the auth header | |
| request.headers["Authorization"] = extra_headers["Authorization"] | |
| prepped = request.prepare() | |
| ## LOGGING | |
| logging_obj.pre_call( | |
| input=data, | |
| api_key="", | |
| additional_args={ | |
| "complete_input_dict": data, | |
| "api_base": prepped.url, | |
| "headers": prepped.headers, | |
| }, | |
| ) | |
| response = self._make_sync_call( | |
| client=client, | |
| timeout=timeout, | |
| api_base=prepped.url, | |
| headers=prepped.headers, # type: ignore | |
| data=data, | |
| ) | |
| ## LOGGING | |
| logging_obj.post_call( | |
| input=data, | |
| api_key="", | |
| original_response=response, | |
| additional_args={"complete_input_dict": data}, | |
| ) | |
| responses.append(response) | |
| returned_response: Optional[EmbeddingResponse] = None | |
| ## TRANSFORM RESPONSE ## | |
| if model == "amazon.titan-embed-image-v1": | |
| returned_response = ( | |
| AmazonTitanMultimodalEmbeddingG1Config()._transform_response( | |
| response_list=responses, model=model | |
| ) | |
| ) | |
| elif model == "amazon.titan-embed-text-v1": | |
| returned_response = AmazonTitanG1Config()._transform_response( | |
| response_list=responses, model=model | |
| ) | |
| elif model == "amazon.titan-embed-text-v2:0": | |
| returned_response = AmazonTitanV2Config()._transform_response( | |
| response_list=responses, model=model | |
| ) | |
| if returned_response is None: | |
| raise Exception( | |
| "Unable to map model response to known provider format. model={}".format( | |
| model | |
| ) | |
| ) | |
| return returned_response | |
| async def _async_single_func_embeddings( | |
| self, | |
| client: Optional[AsyncHTTPHandler], | |
| timeout: Optional[Union[float, httpx.Timeout]], | |
| batch_data: List[dict], | |
| credentials: Any, | |
| extra_headers: Optional[dict], | |
| endpoint_url: str, | |
| aws_region_name: str, | |
| model: str, | |
| logging_obj: Any, | |
| ): | |
| try: | |
| from botocore.auth import SigV4Auth | |
| from botocore.awsrequest import AWSRequest | |
| except ImportError: | |
| raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") | |
| responses: List[dict] = [] | |
| for data in batch_data: | |
| sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) | |
| headers = {"Content-Type": "application/json"} | |
| if extra_headers is not None: | |
| headers = {"Content-Type": "application/json", **extra_headers} | |
| request = AWSRequest( | |
| method="POST", url=endpoint_url, data=json.dumps(data), headers=headers | |
| ) | |
| sigv4.add_auth(request) | |
| if ( | |
| extra_headers is not None and "Authorization" in extra_headers | |
| ): # prevent sigv4 from overwriting the auth header | |
| request.headers["Authorization"] = extra_headers["Authorization"] | |
| prepped = request.prepare() | |
| ## LOGGING | |
| logging_obj.pre_call( | |
| input=data, | |
| api_key="", | |
| additional_args={ | |
| "complete_input_dict": data, | |
| "api_base": prepped.url, | |
| "headers": prepped.headers, | |
| }, | |
| ) | |
| response = await self._make_async_call( | |
| client=client, | |
| timeout=timeout, | |
| api_base=prepped.url, | |
| headers=prepped.headers, # type: ignore | |
| data=data, | |
| ) | |
| ## LOGGING | |
| logging_obj.post_call( | |
| input=data, | |
| api_key="", | |
| original_response=response, | |
| additional_args={"complete_input_dict": data}, | |
| ) | |
| responses.append(response) | |
| returned_response: Optional[EmbeddingResponse] = None | |
| ## TRANSFORM RESPONSE ## | |
| if model == "amazon.titan-embed-image-v1": | |
| returned_response = ( | |
| AmazonTitanMultimodalEmbeddingG1Config()._transform_response( | |
| response_list=responses, model=model | |
| ) | |
| ) | |
| elif model == "amazon.titan-embed-text-v1": | |
| returned_response = AmazonTitanG1Config()._transform_response( | |
| response_list=responses, model=model | |
| ) | |
| elif model == "amazon.titan-embed-text-v2:0": | |
| returned_response = AmazonTitanV2Config()._transform_response( | |
| response_list=responses, model=model | |
| ) | |
| if returned_response is None: | |
| raise Exception( | |
| "Unable to map model response to known provider format. model={}".format( | |
| model | |
| ) | |
| ) | |
| return returned_response | |
| def embeddings( | |
| self, | |
| model: str, | |
| input: List[str], | |
| api_base: Optional[str], | |
| model_response: EmbeddingResponse, | |
| print_verbose: Callable, | |
| encoding, | |
| logging_obj, | |
| client: Optional[Union[HTTPHandler, AsyncHTTPHandler]], | |
| timeout: Optional[Union[float, httpx.Timeout]], | |
| aembedding: Optional[bool], | |
| extra_headers: Optional[dict], | |
| optional_params: dict, | |
| litellm_params: dict, | |
| ) -> EmbeddingResponse: | |
| try: | |
| from botocore.auth import SigV4Auth | |
| from botocore.awsrequest import AWSRequest | |
| except ImportError: | |
| raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") | |
| credentials, aws_region_name = self._load_credentials(optional_params) | |
| ### TRANSFORMATION ### | |
| provider = model.split(".")[0] | |
| inference_params = copy.deepcopy(optional_params) | |
| inference_params = { | |
| k: v | |
| for k, v in inference_params.items() | |
| if k.lower() not in self.aws_authentication_params | |
| } | |
| inference_params.pop( | |
| "user", None | |
| ) # make sure user is not passed in for bedrock call | |
| modelId = ( | |
| optional_params.pop("model_id", None) or model | |
| ) # default to model if not passed | |
| data: Optional[CohereEmbeddingRequest] = None | |
| batch_data: Optional[List] = None | |
| if provider == "cohere": | |
| data = BedrockCohereEmbeddingConfig()._transform_request( | |
| model=model, input=input, inference_params=inference_params | |
| ) | |
| elif provider == "amazon" and model in [ | |
| "amazon.titan-embed-image-v1", | |
| "amazon.titan-embed-text-v1", | |
| "amazon.titan-embed-text-v2:0", | |
| ]: | |
| batch_data = [] | |
| for i in input: | |
| if model == "amazon.titan-embed-image-v1": | |
| transformed_request: ( | |
| AmazonEmbeddingRequest | |
| ) = AmazonTitanMultimodalEmbeddingG1Config()._transform_request( | |
| input=i, inference_params=inference_params | |
| ) | |
| elif model == "amazon.titan-embed-text-v1": | |
| transformed_request = AmazonTitanG1Config()._transform_request( | |
| input=i, inference_params=inference_params | |
| ) | |
| elif model == "amazon.titan-embed-text-v2:0": | |
| transformed_request = AmazonTitanV2Config()._transform_request( | |
| input=i, inference_params=inference_params | |
| ) | |
| else: | |
| raise Exception( | |
| "Unmapped model. Received={}. Expected={}".format( | |
| model, | |
| [ | |
| "amazon.titan-embed-image-v1", | |
| "amazon.titan-embed-text-v1", | |
| "amazon.titan-embed-text-v2:0", | |
| ], | |
| ) | |
| ) | |
| batch_data.append(transformed_request) | |
| ### SET RUNTIME ENDPOINT ### | |
| endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( | |
| api_base=api_base, | |
| aws_bedrock_runtime_endpoint=optional_params.pop( | |
| "aws_bedrock_runtime_endpoint", None | |
| ), | |
| aws_region_name=aws_region_name, | |
| ) | |
| endpoint_url = f"{endpoint_url}/model/{modelId}/invoke" | |
| if batch_data is not None: | |
| if aembedding: | |
| return self._async_single_func_embeddings( # type: ignore | |
| client=( | |
| client | |
| if client is not None and isinstance(client, AsyncHTTPHandler) | |
| else None | |
| ), | |
| timeout=timeout, | |
| batch_data=batch_data, | |
| credentials=credentials, | |
| extra_headers=extra_headers, | |
| endpoint_url=endpoint_url, | |
| aws_region_name=aws_region_name, | |
| model=model, | |
| logging_obj=logging_obj, | |
| ) | |
| return self._single_func_embeddings( | |
| client=( | |
| client | |
| if client is not None and isinstance(client, HTTPHandler) | |
| else None | |
| ), | |
| timeout=timeout, | |
| batch_data=batch_data, | |
| credentials=credentials, | |
| extra_headers=extra_headers, | |
| endpoint_url=endpoint_url, | |
| aws_region_name=aws_region_name, | |
| model=model, | |
| logging_obj=logging_obj, | |
| ) | |
| elif data is None: | |
| raise Exception("Unable to map Bedrock request to provider") | |
| sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) | |
| headers = {"Content-Type": "application/json"} | |
| if extra_headers is not None: | |
| headers = {"Content-Type": "application/json", **extra_headers} | |
| request = AWSRequest( | |
| method="POST", url=endpoint_url, data=json.dumps(data), headers=headers | |
| ) | |
| sigv4.add_auth(request) | |
| if ( | |
| extra_headers is not None and "Authorization" in extra_headers | |
| ): # prevent sigv4 from overwriting the auth header | |
| request.headers["Authorization"] = extra_headers["Authorization"] | |
| prepped = request.prepare() | |
| ## ROUTING ## | |
| return cohere_embedding( | |
| model=model, | |
| input=input, | |
| model_response=model_response, | |
| logging_obj=logging_obj, | |
| optional_params=optional_params, | |
| encoding=encoding, | |
| data=data, # type: ignore | |
| complete_api_base=prepped.url, | |
| api_key=None, | |
| aembedding=aembedding, | |
| timeout=timeout, | |
| client=client, | |
| headers=prepped.headers, # type: ignore | |
| ) | |