From de9bd9de61131bd1d60b967e3ec60ccaa71dc9e1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 14 Sep 2024 18:18:21 -0700 Subject: [PATCH] refactor(bedrock/chat): separate converse api and invoke api + isolate converse api transformation logic Make it easier to see how requests are transformed for /converse --- litellm/__init__.py | 2 +- litellm/llms/bedrock/chat/__init__.py | 2 + litellm/llms/bedrock/chat/converse.py | 228 ------- litellm/llms/bedrock/chat/converse_handler.py | 407 ++++++++++++ .../{chat.py => chat/invoke_handler.py} | 606 +----------------- .../google_ai_studio_endpoints.py | 22 +- litellm/tests/test_bedrock_completion.py | 2 +- litellm/tests/test_secret_manager.py | 2 +- 8 files changed, 431 insertions(+), 840 deletions(-) create mode 100644 litellm/llms/bedrock/chat/__init__.py delete mode 100644 litellm/llms/bedrock/chat/converse.py create mode 100644 litellm/llms/bedrock/chat/converse_handler.py rename litellm/llms/bedrock/{chat.py => chat/invoke_handler.py} (69%) diff --git a/litellm/__init__.py b/litellm/__init__.py index 047927dd9fd..b884dbc9120 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -911,7 +911,7 @@ from .llms.sagemaker.sagemaker import SagemakerConfig from .llms.ollama import OllamaConfig from .llms.ollama_chat import OllamaChatConfig from .llms.maritalk import MaritTalkConfig -from .llms.bedrock.chat import ( +from .llms.bedrock.chat.invoke_handler import ( AmazonCohereChatConfig, AmazonConverseConfig, BEDROCK_CONVERSE_MODELS, diff --git a/litellm/llms/bedrock/chat/__init__.py b/litellm/llms/bedrock/chat/__init__.py new file mode 100644 index 00000000000..c3f6aef6d23 --- /dev/null +++ b/litellm/llms/bedrock/chat/__init__.py @@ -0,0 +1,2 @@ +from .converse_handler import BedrockConverseLLM +from .invoke_handler import BedrockLLM diff --git a/litellm/llms/bedrock/chat/converse.py b/litellm/llms/bedrock/chat/converse.py deleted file mode 100644 index 23f1932277c..00000000000 --- a/litellm/llms/bedrock/chat/converse.py +++ /dev/null @@ -1,228 +0,0 @@ -import json -from typing import Callable, Optional, Union - -import httpx - -import litellm -from litellm.llms.custom_httpx.http_handler import ( - AsyncHTTPHandler, - HTTPHandler, - get_async_httpx_client, -) -from litellm.types.utils import ModelResponse -from litellm.utils import CustomStreamWrapper, get_secret - -from ...base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError - -# class FakeBedrockConverseLLM(BaseAWSLLM): -# def __init__(self) -> None: -# super().__init__() - -# async def async_completion( -# self, -# model: str, -# messages: list, -# api_base: str, -# model_response: ModelResponse, -# print_verbose: Callable, -# data: str, -# timeout: Optional[Union[float, httpx.Timeout]], -# encoding, -# logging_obj, -# stream, -# optional_params: dict, -# litellm_params=None, -# logger_fn=None, -# headers={}, -# client: Optional[AsyncHTTPHandler] = None, -# ) -> Union[ModelResponse, CustomStreamWrapper]: -# 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 # type: ignore - -# try: -# response = await client.post(url=api_base, headers=headers, data=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 as e: -# raise BedrockError(status_code=408, message="Timeout error occurred.") - -# return litellm.AmazonConverseConfig()._transform_response( -# model=model, -# response=response, -# model_response=model_response, -# stream=stream if isinstance(stream, bool) else False, -# logging_obj=logging_obj, -# api_key="", -# data=data, -# messages=messages, -# print_verbose=print_verbose, -# optional_params=optional_params, -# encoding=encoding, -# ) - -# def completion( -# self, -# model: str, -# messages: list, -# api_base: Optional[str], -# custom_prompt_dict: dict, -# model_response: ModelResponse, -# print_verbose: Callable, -# encoding, -# logging_obj, -# optional_params: dict, -# acompletion: bool, -# timeout: Optional[Union[float, httpx.Timeout]], -# litellm_params: dict, -# logger_fn=None, -# extra_headers: Optional[dict] = None, -# client: Optional[Union[AsyncHTTPHandler, HTTPHandler]] = None, -# ): -# try: -# import boto3 -# from botocore.auth import SigV4Auth -# from botocore.awsrequest import AWSRequest -# from botocore.credentials import Credentials -# except ImportError: -# raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - -# ## SETUP ## -# stream = optional_params.pop("stream", None) -# modelId = optional_params.pop("model_id", None) -# if modelId is not None: -# modelId = self.encode_model_id(model_id=modelId) -# else: -# modelId = model - -# provider = model.split(".")[0] - -# # CREDENTIALS -# # pop aws_secret_access_key, aws_access_key_id, 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_bedrock_runtime_endpoint = optional_params.pop( -# "aws_bedrock_runtime_endpoint", None -# ) # https://bedrock-runtime.{region_name}.amazonaws.com -# aws_web_identity_token = optional_params.pop("aws_web_identity_token", None) -# aws_sts_endpoint = optional_params.pop("aws_sts_endpoint", None) - -# # 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" - -# if aws_region_name is None: -# aws_region_name = "us-west-2" - -# credentials = super().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, -# ) -# # sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) - -# ### SET RUNTIME ENDPOINT ### -# endpoint_url, proxy_endpoint_url = super().get_runtime_endpoint( -# api_base=api_base, -# aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, -# aws_region_name=aws_region_name, -# ) - -# if (stream is not None and stream is True) and provider != "ai21": -# endpoint_url = f"{endpoint_url}/model/{modelId}/converse-stream" -# proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/converse-stream" -# else: -# endpoint_url = f"{endpoint_url}/model/{modelId}/converse" -# proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/converse" - -# _data = litellm.AmazonConverseConfig()._transform_request( -# model=model, -# messages=messages, -# optional_params=optional_params, -# litellm_params=litellm_params, -# ) - -# data = json.dumps(_data) - -# ## COMPLETION CALL - -# # headers = {"Content-Type": "application/json"} -# # if extra_headers is not None: -# # headers = {"Content-Type": "application/json", **extra_headers} -# # print(f"endpoint_url: {endpoint_url}, data: {data}, headers: {headers}") -# # request = AWSRequest( -# # method="POST", url=endpoint_url, data=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=messages, -# # api_key="", -# # additional_args={ -# # "complete_input_dict": data, -# # "api_base": proxy_endpoint_url, -# # "headers": prepped.headers, -# # }, -# # ) - -# return self.async_completion( -# model=model, -# messages=messages, -# data=data, -# api_base=proxy_endpoint_url, -# model_response=model_response, -# print_verbose=print_verbose, -# encoding=encoding, -# logging_obj=logging_obj, -# optional_params=optional_params, -# stream=False, -# litellm_params=litellm_params, -# logger_fn=logger_fn, -# # headers=prepped.headers, -# headers={"Authorization": "my-fake-key"}, -# timeout=timeout, -# client=client, -# ) # type: ignore diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py new file mode 100644 index 00000000000..4b14f1016b5 --- /dev/null +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -0,0 +1,407 @@ +import json +import urllib +from typing import Any, Callable, Optional, Union + +import httpx + +import litellm +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + _get_httpx_client, + get_async_httpx_client, +) +from litellm.types.utils import ModelResponse +from litellm.utils import CustomStreamWrapper, get_secret + +from ...base_aws_llm import BaseAWSLLM +from ..common_utils import BedrockError +from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call + + +def make_sync_call( + client: Optional[HTTPHandler], + api_base: str, + headers: dict, + data: str, + model: str, + messages: list, + logging_obj, +): + if client is None: + client = _get_httpx_client() # Create a new client if none provided + + response = client.post( + api_base, + headers=headers, + data=data, + stream=True if "ai21" not in api_base else False, + ) + + if response.status_code != 200: + raise BedrockError(status_code=response.status_code, message=response.read()) + + if "ai21" in api_base: + aws_bedrock_process_response = BedrockConverseLLM() + model_response: ModelResponse = aws_bedrock_process_response.process_response( + model=model, + response=response, + model_response=litellm.ModelResponse(), + stream=True, + logging_obj=logging_obj, + optional_params={}, + api_key="", + data=data, + messages=messages, + print_verbose=litellm.print_verbose, + encoding=litellm.encoding, + ) # type: ignore + completion_stream: Any = MockResponseIterator(model_response=model_response) + else: + decoder = AWSEventStreamDecoder(model=model) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + + # LOGGING + logging_obj.post_call( + input=messages, + api_key="", + original_response="first stream response received", + additional_args={"complete_input_dict": data}, + ) + + return completion_stream + + +class BedrockConverseLLM(BaseAWSLLM): + def __init__(self) -> None: + super().__init__() + + def encode_model_id(self, model_id: str) -> str: + """ + Double encode the model ID to ensure it matches the expected double-encoded format. + Args: + model_id (str): The model ID to encode. + Returns: + str: The double-encoded model ID. + """ + return urllib.parse.quote(model_id, safe="") # type: ignore + + async def async_streaming( + self, + model: str, + messages: list, + api_base: str, + model_response: ModelResponse, + print_verbose: Callable, + data: str, + timeout: Optional[Union[float, httpx.Timeout]], + encoding, + logging_obj, + stream, + optional_params: dict, + litellm_params=None, + logger_fn=None, + headers={}, + client: Optional[AsyncHTTPHandler] = None, + ) -> CustomStreamWrapper: + + completion_stream = await make_call( + client=client, + api_base=api_base, + headers=headers, + data=data, + model=model, + messages=messages, + logging_obj=logging_obj, + ) + streaming_response = CustomStreamWrapper( + completion_stream=completion_stream, + model=model, + custom_llm_provider="bedrock", + logging_obj=logging_obj, + ) + return streaming_response + + async def async_completion( + self, + model: str, + messages: list, + api_base: str, + model_response: ModelResponse, + print_verbose: Callable, + data: str, + timeout: Optional[Union[float, httpx.Timeout]], + encoding, + logging_obj, + stream, + optional_params: dict, + litellm_params=None, + logger_fn=None, + headers={}, + client: Optional[AsyncHTTPHandler] = None, + ) -> Union[ModelResponse, CustomStreamWrapper]: + 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 # type: ignore + + try: + response = await client.post(url=api_base, headers=headers, data=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 as e: + raise BedrockError(status_code=408, message="Timeout error occurred.") + + return litellm.AmazonConverseConfig()._transform_response( + model=model, + response=response, + model_response=model_response, + stream=stream if isinstance(stream, bool) else False, + logging_obj=logging_obj, + api_key="", + data=data, + messages=messages, + print_verbose=print_verbose, + optional_params=optional_params, + encoding=encoding, + ) + + def completion( + self, + model: str, + messages: list, + api_base: Optional[str], + custom_prompt_dict: dict, + model_response: ModelResponse, + print_verbose: Callable, + encoding, + logging_obj, + optional_params: dict, + acompletion: bool, + timeout: Optional[Union[float, httpx.Timeout]], + litellm_params: dict, + logger_fn=None, + extra_headers: Optional[dict] = None, + client: Optional[Union[AsyncHTTPHandler, HTTPHandler]] = None, + ): + try: + import boto3 + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + from botocore.credentials import Credentials + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + + ## SETUP ## + stream = optional_params.pop("stream", None) + modelId = optional_params.pop("model_id", None) + if modelId is not None: + modelId = self.encode_model_id(model_id=modelId) + else: + modelId = model + + provider = model.split(".")[0] + + ## CREDENTIALS ## + # pop aws_secret_access_key, aws_access_key_id, 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_bedrock_runtime_endpoint = optional_params.pop( + "aws_bedrock_runtime_endpoint", None + ) # https://bedrock-runtime.{region_name}.amazonaws.com + 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, + ) + + ### SET RUNTIME ENDPOINT ### + endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( + api_base=api_base, + aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, + aws_region_name=aws_region_name, + ) + if (stream is not None and stream is True) and provider != "ai21": + endpoint_url = f"{endpoint_url}/model/{modelId}/converse-stream" + proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/converse-stream" + else: + endpoint_url = f"{endpoint_url}/model/{modelId}/converse" + proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/converse" + + sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) + + ## TRANSFORMATION ## + + _data = litellm.AmazonConverseConfig()._transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + ) + data = json.dumps(_data) + ## COMPLETION CALL + + 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=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=messages, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": proxy_endpoint_url, + "headers": prepped.headers, + }, + ) + + ### ROUTING (ASYNC, STREAMING, SYNC) + if acompletion: + if isinstance(client, HTTPHandler): + client = None + if stream is True: + return self.async_streaming( + model=model, + messages=messages, + data=data, + api_base=proxy_endpoint_url, + model_response=model_response, + print_verbose=print_verbose, + encoding=encoding, + logging_obj=logging_obj, + optional_params=optional_params, + stream=True, + litellm_params=litellm_params, + logger_fn=logger_fn, + headers=prepped.headers, + timeout=timeout, + client=client, + ) # type: ignore + ### ASYNC COMPLETION + return self.async_completion( + model=model, + messages=messages, + data=data, + api_base=proxy_endpoint_url, + model_response=model_response, + print_verbose=print_verbose, + encoding=encoding, + logging_obj=logging_obj, + optional_params=optional_params, + stream=stream, # type: ignore + litellm_params=litellm_params, + logger_fn=logger_fn, + headers=prepped.headers, + timeout=timeout, + client=client, + ) # type: ignore + + if client is None or 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_httpx_client(_params) # type: ignore + else: + client = client + + if stream is not None and stream is True: + completion_stream = make_sync_call( + client=( + client + if client is not None and isinstance(client, HTTPHandler) + else None + ), + api_base=proxy_endpoint_url, + headers=prepped.headers, # type: ignore + data=data, + model=model, + messages=messages, + logging_obj=logging_obj, + ) + streaming_response = CustomStreamWrapper( + completion_stream=completion_stream, + model=model, + custom_llm_provider="bedrock", + logging_obj=logging_obj, + ) + + return streaming_response + + ### COMPLETION + + try: + response = client.post(url=proxy_endpoint_url, headers=prepped.headers, data=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 litellm.AmazonConverseConfig()._transform_response( + model=model, + response=response, + model_response=model_response, + stream=stream if isinstance(stream, bool) else False, + logging_obj=logging_obj, + api_key="", + data=data, + messages=messages, + print_verbose=print_verbose, + optional_params=optional_params, + encoding=encoding, + ) diff --git a/litellm/llms/bedrock/chat.py b/litellm/llms/bedrock/chat/invoke_handler.py similarity index 69% rename from litellm/llms/bedrock/chat.py rename to litellm/llms/bedrock/chat/invoke_handler.py index 8068e5b03d3..fbad0940ce8 100644 --- a/litellm/llms/bedrock/chat.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -52,8 +52,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import GenericStreamingChunk as GChunk from litellm.utils import CustomStreamWrapper, ModelResponse, Usage, get_secret -from ..base_aws_llm import BaseAWSLLM -from ..prompt_templates.factory import ( +from ...base_aws_llm import BaseAWSLLM +from ...prompt_templates.factory import ( _bedrock_converse_messages_pt, _bedrock_tools_pt, cohere_message_pt, @@ -64,8 +64,8 @@ from ..prompt_templates.factory import ( parse_xml_params, prompt_factory, ) -from .common_utils import BedrockError, ModelResponseIterator, get_bedrock_tool_name -from .converse_chat.converse_transformation import AmazonConverseConfig +from ..common_utils import BedrockError, ModelResponseIterator, get_bedrock_tool_name +from .converse_transformation import AmazonConverseConfig BEDROCK_CONVERSE_MODELS = [ "anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -225,10 +225,9 @@ async def make_call( raise BedrockError(status_code=response.status_code, message=response.text) if "ai21" in api_base: - aws_bedrock_process_response = BedrockConverseLLM() model_response: ( ModelResponse - ) = aws_bedrock_process_response.process_response( + ) = litellm.AmazonConverseConfig()._transform_response( model=model, response=response, model_response=litellm.ModelResponse(), @@ -266,59 +265,6 @@ async def make_call( raise BedrockError(status_code=500, message=str(e)) -def make_sync_call( - client: Optional[HTTPHandler], - api_base: str, - headers: dict, - data: str, - model: str, - messages: list, - logging_obj, -): - if client is None: - client = _get_httpx_client() # Create a new client if none provided - - response = client.post( - api_base, - headers=headers, - data=data, - stream=True if "ai21" not in api_base else False, - ) - - if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=response.read()) - - if "ai21" in api_base: - aws_bedrock_process_response = BedrockConverseLLM() - model_response: ModelResponse = aws_bedrock_process_response.process_response( - model=model, - response=response, - model_response=litellm.ModelResponse(), - stream=True, - logging_obj=logging_obj, - optional_params={}, - api_key="", - data=data, - messages=messages, - print_verbose=litellm.print_verbose, - encoding=litellm.encoding, - ) # type: ignore - completion_stream: Any = MockResponseIterator(model_response=model_response) - else: - decoder = AWSEventStreamDecoder(model=model) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) - - # LOGGING - logging_obj.post_call( - input=messages, - api_key="", - original_response="first stream response received", - additional_args={"complete_input_dict": data}, - ) - - return completion_stream - - class BedrockLLM(BaseAWSLLM): """ Example call @@ -1118,548 +1064,6 @@ class BedrockLLM(BaseAWSLLM): return super().embedding(*args, **kwargs) -class BedrockConverseLLM(BaseAWSLLM): - def __init__(self) -> None: - super().__init__() - - def process_response( - self, - model: str, - response: Union[requests.Response, httpx.Response], - model_response: ModelResponse, - stream: bool, - logging_obj: Optional[Logging], - optional_params: dict, - api_key: str, - data: Union[dict, str], - messages: List, - print_verbose, - encoding, - ) -> Union[ModelResponse, CustomStreamWrapper]: - - ## LOGGING - if logging_obj is not None: - logging_obj.post_call( - input=messages, - api_key=api_key, - original_response=response.text, - additional_args={"complete_input_dict": data}, - ) - print_verbose(f"raw model_response: {response.text}") - json_mode: Optional[bool] = optional_params.pop("json_mode", None) - ## RESPONSE OBJECT - try: - completion_response = ConverseResponseBlock(**response.json()) # type: ignore - except Exception as e: - raise BedrockError( - message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( - response.text, str(e) - ), - status_code=422, - ) - - """ - Bedrock Response Object has optional message block - - completion_response["output"].get("message", None) - - A message block looks like this (Example 1): - "output": { - "message": { - "role": "assistant", - "content": [ - { - "text": "Is there anything else you'd like to talk about? Perhaps I can help with some economic questions or provide some information about economic concepts?" - } - ] - } - }, - (Example 2): - "output": { - "message": { - "role": "assistant", - "content": [ - { - "toolUse": { - "toolUseId": "tooluse_hbTgdi0CSLq_hM4P8csZJA", - "name": "top_song", - "input": { - "sign": "WZPZ" - } - } - } - ] - } - } - - """ - message: Optional[MessageBlock] = completion_response["output"]["message"] - chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} - content_str = "" - tools: List[ChatCompletionToolCallChunk] = [] - if message is not None: - for idx, content in enumerate(message["content"]): - """ - - Content is either a tool response or text - """ - if "text" in content: - content_str += content["text"] - if "toolUse" in content: - - ## check tool name was formatted by litellm - _response_tool_name = content["toolUse"]["name"] - response_tool_name = get_bedrock_tool_name( - response_tool_name=_response_tool_name - ) - _function_chunk = ChatCompletionToolCallFunctionChunk( - name=response_tool_name, - arguments=json.dumps(content["toolUse"]["input"]), - ) - - _tool_response_chunk = ChatCompletionToolCallChunk( - id=content["toolUse"]["toolUseId"], - type="function", - function=_function_chunk, - index=idx, - ) - tools.append(_tool_response_chunk) - chat_completion_message["content"] = content_str - - if json_mode is True and tools is not None and len(tools) == 1: - # to support 'json_schema' logic on bedrock models - json_mode_content_str: Optional[str] = tools[0]["function"].get("arguments") - if json_mode_content_str is not None: - chat_completion_message["content"] = json_mode_content_str - else: - chat_completion_message["tool_calls"] = tools - - ## CALCULATING USAGE - bedrock returns usage in the headers - input_tokens = completion_response["usage"]["inputTokens"] - output_tokens = completion_response["usage"]["outputTokens"] - total_tokens = completion_response["usage"]["totalTokens"] - - model_response.choices = [ - litellm.Choices( - finish_reason=map_finish_reason(completion_response["stopReason"]), - index=0, - message=litellm.Message(**chat_completion_message), - ) - ] - model_response.created = int(time.time()) - model_response.model = model - usage = Usage( - prompt_tokens=input_tokens, - completion_tokens=output_tokens, - total_tokens=total_tokens, - ) - setattr(model_response, "usage", usage) - - # Add "trace" from Bedrock guardrails - if user has opted in to returning it - if "trace" in completion_response: - setattr(model_response, "trace", completion_response["trace"]) - - return model_response - - def encode_model_id(self, model_id: str) -> str: - """ - Double encode the model ID to ensure it matches the expected double-encoded format. - Args: - model_id (str): The model ID to encode. - Returns: - str: The double-encoded model ID. - """ - return urllib.parse.quote(model_id, safe="") - - async def async_streaming( - self, - model: str, - messages: list, - api_base: str, - model_response: ModelResponse, - print_verbose: Callable, - data: str, - timeout: Optional[Union[float, httpx.Timeout]], - encoding, - logging_obj, - stream, - optional_params: dict, - litellm_params=None, - logger_fn=None, - headers={}, - client: Optional[AsyncHTTPHandler] = None, - ) -> CustomStreamWrapper: - streaming_response = CustomStreamWrapper( - completion_stream=None, - make_call=partial( - make_call, - client=client, - api_base=api_base, - headers=headers, - data=data, - model=model, - messages=messages, - logging_obj=logging_obj, - ), - model=model, - custom_llm_provider="bedrock", - logging_obj=logging_obj, - ) - return streaming_response - - async def async_completion( - self, - model: str, - messages: list, - api_base: str, - model_response: ModelResponse, - print_verbose: Callable, - data: str, - timeout: Optional[Union[float, httpx.Timeout]], - encoding, - logging_obj, - stream, - optional_params: dict, - litellm_params=None, - logger_fn=None, - headers={}, - client: Optional[AsyncHTTPHandler] = None, - ) -> Union[ModelResponse, CustomStreamWrapper]: - 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 # type: ignore - - try: - response = await client.post(url=api_base, headers=headers, data=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 as e: - raise BedrockError(status_code=408, message="Timeout error occurred.") - - return self.process_response( - model=model, - response=response, - model_response=model_response, - stream=stream if isinstance(stream, bool) else False, - logging_obj=logging_obj, - api_key="", - data=data, - messages=messages, - print_verbose=print_verbose, - optional_params=optional_params, - encoding=encoding, - ) - - def completion( - self, - model: str, - messages: list, - api_base: Optional[str], - custom_prompt_dict: dict, - model_response: ModelResponse, - print_verbose: Callable, - encoding, - logging_obj, - optional_params: dict, - acompletion: bool, - timeout: Optional[Union[float, httpx.Timeout]], - litellm_params: dict, - logger_fn=None, - extra_headers: Optional[dict] = None, - client: Optional[Union[AsyncHTTPHandler, HTTPHandler]] = None, - ): - try: - import boto3 - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - - ## SETUP ## - stream = optional_params.pop("stream", None) - modelId = optional_params.pop("model_id", None) - if modelId is not None: - modelId = self.encode_model_id(model_id=modelId) - else: - modelId = model - - provider = model.split(".")[0] - - ## CREDENTIALS ## - # pop aws_secret_access_key, aws_access_key_id, 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_bedrock_runtime_endpoint = optional_params.pop( - "aws_bedrock_runtime_endpoint", None - ) # https://bedrock-runtime.{region_name}.amazonaws.com - 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, - ) - - ### SET RUNTIME ENDPOINT ### - endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( - api_base=api_base, - aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, - aws_region_name=aws_region_name, - ) - if (stream is not None and stream is True) and provider != "ai21": - endpoint_url = f"{endpoint_url}/model/{modelId}/converse-stream" - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/converse-stream" - else: - endpoint_url = f"{endpoint_url}/model/{modelId}/converse" - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/converse" - - sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) - - # Separate system prompt from rest of message - system_prompt_indices = [] - system_content_blocks: List[SystemContentBlock] = [] - for idx, message in enumerate(messages): - if message["role"] == "system": - _system_content_block: Optional[SystemContentBlock] = None - if isinstance(message["content"], str) and len(message["content"]) > 0: - _system_content_block = SystemContentBlock(text=message["content"]) - elif isinstance(message["content"], list): - for m in message["content"]: - if m.get("type", "") == "text" and len(m["text"]) > 0: - _system_content_block = SystemContentBlock(text=m["text"]) - if _system_content_block is not None: - system_content_blocks.append(_system_content_block) - system_prompt_indices.append(idx) - if len(system_prompt_indices) > 0: - for idx in reversed(system_prompt_indices): - messages.pop(idx) - - inference_params = copy.deepcopy(optional_params) - additional_request_keys = [] - additional_request_params = {} - supported_converse_params = AmazonConverseConfig.__annotations__.keys() - supported_tool_call_params = ["tools", "tool_choice"] - supported_guardrail_params = ["guardrailConfig"] - json_mode: Optional[bool] = inference_params.pop( - "json_mode", None - ) # used for handling json_schema - ## TRANSFORMATION ## - - bedrock_messages: List[MessageBlock] = _bedrock_converse_messages_pt( - messages=messages, - model=model, - llm_provider="bedrock_converse", - user_continue_message=litellm_params.pop("user_continue_message", None), - ) - - # send all model-specific params in 'additional_request_params' - for k, v in inference_params.items(): - if ( - k not in supported_converse_params - and k not in supported_tool_call_params - and k not in supported_guardrail_params - ): - additional_request_params[k] = v - additional_request_keys.append(k) - for key in additional_request_keys: - inference_params.pop(key, None) - - bedrock_tools: List[ToolBlock] = _bedrock_tools_pt( - inference_params.pop("tools", []) - ) - bedrock_tool_config: Optional[ToolConfigBlock] = None - if len(bedrock_tools) > 0: - tool_choice_values: ToolChoiceValuesBlock = inference_params.pop( - "tool_choice", None - ) - bedrock_tool_config = ToolConfigBlock( - tools=bedrock_tools, - ) - if tool_choice_values is not None: - bedrock_tool_config["toolChoice"] = tool_choice_values - - _data: RequestObject = { - "messages": bedrock_messages, - "additionalModelRequestFields": additional_request_params, - "system": system_content_blocks, - "inferenceConfig": InferenceConfig(**inference_params), - } - - # Guardrail Config - guardrail_config: Optional[GuardrailConfigBlock] = None - request_guardrails_config = inference_params.pop("guardrailConfig", None) - if request_guardrails_config is not None: - guardrail_config = GuardrailConfigBlock(**request_guardrails_config) - _data["guardrailConfig"] = guardrail_config - - # Tool Config - if bedrock_tool_config is not None: - _data["toolConfig"] = bedrock_tool_config - - data = json.dumps(_data) - ## COMPLETION CALL - - 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=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=messages, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": proxy_endpoint_url, - "headers": prepped.headers, - }, - ) - - ### ROUTING (ASYNC, STREAMING, SYNC) - if acompletion: - if isinstance(client, HTTPHandler): - client = None - if stream is True: - return self.async_streaming( - model=model, - messages=messages, - data=data, - api_base=proxy_endpoint_url, - model_response=model_response, - print_verbose=print_verbose, - encoding=encoding, - logging_obj=logging_obj, - optional_params=optional_params, - stream=True, - litellm_params=litellm_params, - logger_fn=logger_fn, - headers=prepped.headers, - timeout=timeout, - client=client, - ) # type: ignore - ### ASYNC COMPLETION - return self.async_completion( - model=model, - messages=messages, - data=data, - api_base=proxy_endpoint_url, - model_response=model_response, - print_verbose=print_verbose, - encoding=encoding, - logging_obj=logging_obj, - optional_params=optional_params, - stream=stream, # type: ignore - litellm_params=litellm_params, - logger_fn=logger_fn, - headers=prepped.headers, - timeout=timeout, - client=client, - ) # type: ignore - - if stream is not None and stream is True: - - streaming_response = CustomStreamWrapper( - completion_stream=None, - make_call=partial( - make_sync_call, - client=None, - api_base=proxy_endpoint_url, - headers=prepped.headers, # type: ignore - data=data, - model=model, - messages=messages, - logging_obj=logging_obj, - ), - model=model, - custom_llm_provider="bedrock", - logging_obj=logging_obj, - ) - - return streaming_response - ### COMPLETION - - if client is None or 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_httpx_client(_params) # type: ignore - else: - client = client - try: - response = client.post(url=proxy_endpoint_url, headers=prepped.headers, data=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 self.process_response( - model=model, - response=response, - model_response=model_response, - stream=stream if isinstance(stream, bool) else False, - logging_obj=logging_obj, - optional_params=optional_params, - api_key="", - data=data, - messages=messages, - print_verbose=print_verbose, - encoding=encoding, - ) - - def get_response_stream_shape(): global _response_stream_shape_cache if _response_stream_shape_cache is None: diff --git a/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py b/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py index a272af46b3e..c4c71820b93 100644 --- a/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py @@ -80,7 +80,13 @@ async def gemini_proxy_route( updated_url = base_url.copy_with(path=encoded_endpoint) # Add or update query parameters - gemini_api_key = litellm.utils.get_secret(secret_name="GEMINI_API_KEY") + gemini_api_key: Optional[str] = litellm.utils.get_secret( # type: ignore + secret_name="GEMINI_API_KEY" + ) + if gemini_api_key is None: + raise Exception( + "Required 'GEMINI_API_KEY' in environment to make pass-through calls to Google AI Studio." + ) # Merge query parameters, giving precedence to those in updated_url merged_params = dict(request.query_params) merged_params.update({"key": gemini_api_key}) @@ -99,8 +105,8 @@ async def gemini_proxy_route( request, fastapi_response, user_api_key_dict, - query_params=merged_params, - stream=is_streaming_request, + query_params=merged_params, # type: ignore + stream=is_streaming_request, # type: ignore ) return received_value @@ -142,7 +148,7 @@ async def cohere_proxy_route( request, fastapi_response, user_api_key_dict, - stream=is_streaming_request, + stream=is_streaming_request, # type: ignore ) return received_value @@ -208,15 +214,15 @@ async def bedrock_proxy_route( endpoint_func = create_pass_through_route( endpoint=endpoint, target=str(prepped.url), - custom_headers=prepped.headers, + custom_headers=dict(prepped.headers), ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, fastapi_response, user_api_key_dict, - stream=is_streaming_request, - custom_body=data, - query_params={}, + stream=is_streaming_request, # type: ignore + custom_body=data, # type: ignore + query_params={}, # type: ignore ) return received_value diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index e949d1ee7d2..f80ace9fa4f 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -27,7 +27,7 @@ from litellm import ( completion_cost, embedding, ) -from litellm.llms.bedrock.chat import BedrockLLM, ToolBlock +from litellm.llms.bedrock.chat.invoke_handler import BedrockLLM, ToolBlock from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.prompt_templates.factory import _bedrock_tools_pt diff --git a/litellm/tests/test_secret_manager.py b/litellm/tests/test_secret_manager.py index 397128ecb06..0ca3a6cff22 100644 --- a/litellm/tests/test_secret_manager.py +++ b/litellm/tests/test_secret_manager.py @@ -17,7 +17,7 @@ sys.path.insert( import pytest from litellm.llms.AzureOpenAI.azure import get_azure_ad_token_from_oidc -from litellm.llms.bedrock.chat import BedrockConverseLLM, BedrockLLM +from litellm.llms.bedrock.chat.invoke_handler import BedrockConverseLLM, BedrockLLM from litellm.secret_managers.aws_secret_manager import load_aws_secret_manager from litellm.secret_managers.main import get_secret