mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
feat(aws_base_llm.py): prevents recreating boto3 credentials during high traffic
Leads to 100ms perf boost in local testing
This commit is contained in:
Krrish Dholakia
parent
dad1ad2077
commit
3f08382591
8 changed files with 801 additions and 426 deletions
|
|
@ -1,5 +1,6 @@
|
|||
import hashlib
|
||||
import json
|
||||
from typing import List, Optional
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -28,6 +29,14 @@ class BaseAWSLLM(BaseLLM):
|
|||
self.iam_cache = DualCache()
|
||||
super().__init__()
|
||||
|
||||
def get_cache_key(self, credential_args: Dict[str, Optional[str]]) -> str:
|
||||
"""
|
||||
Generate a unique cache key based on the credential arguments.
|
||||
"""
|
||||
# Convert credential arguments to a JSON string and hash it to create a unique key
|
||||
credential_str = json.dumps(credential_args, sort_keys=True)
|
||||
return hashlib.sha256(credential_str.encode()).hexdigest()
|
||||
|
||||
def get_credentials(
|
||||
self,
|
||||
aws_access_key_id: Optional[str] = None,
|
||||
|
|
@ -43,7 +52,12 @@ class BaseAWSLLM(BaseLLM):
|
|||
"""
|
||||
Return a boto3.Credentials object
|
||||
"""
|
||||
args = locals()
|
||||
args.pop("self")
|
||||
cache_key = self.get_cache_key(args)
|
||||
|
||||
import boto3
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
## CHECK IS 'os.environ/' passed in
|
||||
params_to_check: List[Optional[str]] = [
|
||||
|
|
@ -186,7 +200,6 @@ class BaseAWSLLM(BaseLLM):
|
|||
|
||||
# Extract the credentials from the response and convert to Session Credentials
|
||||
sts_credentials = sts_response["Credentials"]
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
credentials = Credentials(
|
||||
access_key=sts_credentials["AccessKeyId"],
|
||||
|
|
@ -211,12 +224,59 @@ class BaseAWSLLM(BaseLLM):
|
|||
secret_key=aws_secret_access_key,
|
||||
token=aws_session_token,
|
||||
)
|
||||
|
||||
return credentials
|
||||
else:
|
||||
|
||||
# Check if credentials are already in cache. These credentials have no expiry time.
|
||||
cached_credentials: Optional[Credentials] = self.iam_cache.get_cache(
|
||||
cache_key
|
||||
)
|
||||
if cached_credentials:
|
||||
return cached_credentials
|
||||
|
||||
session = boto3.Session(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
region_name=aws_region_name,
|
||||
)
|
||||
|
||||
return session.get_credentials()
|
||||
credentials = session.get_credentials()
|
||||
|
||||
self.iam_cache.set_cache(cache_key, credentials, ttl=3600 - 60)
|
||||
|
||||
return credentials
|
||||
|
||||
def get_runtime_endpoint(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
aws_bedrock_runtime_endpoint: Optional[str],
|
||||
aws_region_name: str,
|
||||
) -> Tuple[str, str]:
|
||||
env_aws_bedrock_runtime_endpoint = get_secret("AWS_BEDROCK_RUNTIME_ENDPOINT")
|
||||
if api_base is not None:
|
||||
endpoint_url = api_base
|
||||
elif aws_bedrock_runtime_endpoint is not None and isinstance(
|
||||
aws_bedrock_runtime_endpoint, str
|
||||
):
|
||||
endpoint_url = aws_bedrock_runtime_endpoint
|
||||
elif env_aws_bedrock_runtime_endpoint and isinstance(
|
||||
env_aws_bedrock_runtime_endpoint, str
|
||||
):
|
||||
endpoint_url = env_aws_bedrock_runtime_endpoint
|
||||
else:
|
||||
endpoint_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com"
|
||||
|
||||
# Determine proxy_endpoint_url
|
||||
if env_aws_bedrock_runtime_endpoint and isinstance(
|
||||
env_aws_bedrock_runtime_endpoint, str
|
||||
):
|
||||
proxy_endpoint_url = env_aws_bedrock_runtime_endpoint
|
||||
elif aws_bedrock_runtime_endpoint is not None and isinstance(
|
||||
aws_bedrock_runtime_endpoint, str
|
||||
):
|
||||
proxy_endpoint_url = aws_bedrock_runtime_endpoint
|
||||
else:
|
||||
proxy_endpoint_url = endpoint_url
|
||||
|
||||
return endpoint_url, proxy_endpoint_url
|
||||
|
|
|
|||
|
|
@ -64,7 +64,8 @@ from ..prompt_templates.factory import (
|
|||
parse_xml_params,
|
||||
prompt_factory,
|
||||
)
|
||||
from .common_utils import BedrockError, ModelResponseIterator, get_runtime_endpoint
|
||||
from .common_utils import BedrockError, ModelResponseIterator, get_bedrock_tool_name
|
||||
from .converse_chat.converse_transformation import AmazonConverseConfig
|
||||
|
||||
BEDROCK_CONVERSE_MODELS = [
|
||||
"anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
|
|
@ -416,6 +417,7 @@ class BedrockLLM(BaseAWSLLM):
|
|||
except:
|
||||
raise BedrockError(message=response.text, status_code=422)
|
||||
|
||||
outputText: Optional[str] = None
|
||||
try:
|
||||
if provider == "cohere":
|
||||
if "text" in completion_response:
|
||||
|
|
@ -565,23 +567,27 @@ class BedrockLLM(BaseAWSLLM):
|
|||
|
||||
try:
|
||||
if (
|
||||
len(outputText) > 0
|
||||
outputText is not None
|
||||
and len(outputText) > 0
|
||||
and hasattr(model_response.choices[0], "message")
|
||||
and getattr(model_response.choices[0].message, "tool_calls", None)
|
||||
and getattr(model_response.choices[0].message, "tool_calls", None) # type: ignore
|
||||
is None
|
||||
):
|
||||
model_response.choices[0].message.content = outputText
|
||||
model_response.choices[0].message.content = outputText # type: ignore
|
||||
elif (
|
||||
hasattr(model_response.choices[0], "message")
|
||||
and getattr(model_response.choices[0].message, "tool_calls", None)
|
||||
and getattr(model_response.choices[0].message, "tool_calls", None) # type: ignore
|
||||
is not None
|
||||
):
|
||||
pass
|
||||
else:
|
||||
raise Exception()
|
||||
except:
|
||||
except Exception as e:
|
||||
raise BedrockError(
|
||||
message=json.dumps(outputText), status_code=response.status_code
|
||||
message="Error parsing received text={}.\nError-{}".format(
|
||||
outputText, str(e)
|
||||
),
|
||||
status_code=response.status_code,
|
||||
)
|
||||
|
||||
if stream and provider == "ai21":
|
||||
|
|
@ -593,8 +599,8 @@ class BedrockLLM(BaseAWSLLM):
|
|||
streaming_choice = litellm.utils.StreamingChoices()
|
||||
streaming_choice.index = model_response.choices[0].index
|
||||
delta_obj = litellm.utils.Delta(
|
||||
content=getattr(model_response.choices[0].message, "content", None),
|
||||
role=model_response.choices[0].message.role,
|
||||
content=getattr(model_response.choices[0].message, "content", None), # type: ignore
|
||||
role=model_response.choices[0].message.role, # type: ignore
|
||||
)
|
||||
streaming_choice.delta = delta_obj
|
||||
streaming_model_response.choices = [streaming_choice]
|
||||
|
|
@ -730,7 +736,7 @@ class BedrockLLM(BaseAWSLLM):
|
|||
)
|
||||
|
||||
### SET RUNTIME ENDPOINT ###
|
||||
endpoint_url, proxy_endpoint_url = get_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,
|
||||
|
|
@ -1001,7 +1007,7 @@ class BedrockLLM(BaseAWSLLM):
|
|||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
error_code = err.response.status_code
|
||||
raise BedrockError(status_code=error_code, message=response.text)
|
||||
raise BedrockError(status_code=error_code, message=err.response.text)
|
||||
except httpx.TimeoutException as e:
|
||||
raise BedrockError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
|
|
@ -1112,182 +1118,6 @@ class BedrockLLM(BaseAWSLLM):
|
|||
return super().embedding(*args, **kwargs)
|
||||
|
||||
|
||||
class AmazonConverseConfig:
|
||||
"""
|
||||
Reference - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
|
||||
#2 - https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html#conversation-inference-supported-models-features
|
||||
"""
|
||||
|
||||
maxTokens: Optional[int]
|
||||
stopSequences: Optional[List[str]]
|
||||
temperature: Optional[int]
|
||||
topP: Optional[int]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
maxTokens: Optional[int] = None,
|
||||
stopSequences: Optional[List[str]] = None,
|
||||
temperature: Optional[int] = None,
|
||||
topP: Optional[int] = None,
|
||||
) -> None:
|
||||
locals_ = locals()
|
||||
for key, value in locals_.items():
|
||||
if key != "self" and value is not None:
|
||||
setattr(self.__class__, key, value)
|
||||
|
||||
@classmethod
|
||||
def get_config(cls):
|
||||
return {
|
||||
k: v
|
||||
for k, v in cls.__dict__.items()
|
||||
if not k.startswith("__")
|
||||
and not isinstance(
|
||||
v,
|
||||
(
|
||||
types.FunctionType,
|
||||
types.BuiltinFunctionType,
|
||||
classmethod,
|
||||
staticmethod,
|
||||
),
|
||||
)
|
||||
and v is not None
|
||||
}
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> List[str]:
|
||||
supported_params = [
|
||||
"max_tokens",
|
||||
"stream",
|
||||
"stream_options",
|
||||
"stop",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"extra_headers",
|
||||
"response_format",
|
||||
]
|
||||
|
||||
if (
|
||||
model.startswith("anthropic")
|
||||
or model.startswith("mistral")
|
||||
or model.startswith("cohere")
|
||||
or model.startswith("meta.llama3-1")
|
||||
):
|
||||
supported_params.append("tools")
|
||||
|
||||
if model.startswith("anthropic") or model.startswith("mistral"):
|
||||
# only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
|
||||
supported_params.append("tool_choice")
|
||||
|
||||
return supported_params
|
||||
|
||||
def map_tool_choice_values(
|
||||
self, model: str, tool_choice: Union[str, dict], drop_params: bool
|
||||
) -> Optional[ToolChoiceValuesBlock]:
|
||||
if tool_choice == "none":
|
||||
if litellm.drop_params is True or drop_params is True:
|
||||
return None
|
||||
else:
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message="Bedrock doesn't support tool_choice={}. To drop it from the call, set `litellm.drop_params = True.".format(
|
||||
tool_choice
|
||||
),
|
||||
status_code=400,
|
||||
)
|
||||
elif tool_choice == "required":
|
||||
return ToolChoiceValuesBlock(any={})
|
||||
elif tool_choice == "auto":
|
||||
return ToolChoiceValuesBlock(auto={})
|
||||
elif isinstance(tool_choice, dict):
|
||||
# only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
|
||||
specific_tool = SpecificToolChoiceBlock(
|
||||
name=tool_choice.get("function", {}).get("name", "")
|
||||
)
|
||||
return ToolChoiceValuesBlock(tool=specific_tool)
|
||||
else:
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message="Bedrock doesn't support tool_choice={}. Supported tool_choice values=['auto', 'required', json object]. To drop it from the call, set `litellm.drop_params = True.".format(
|
||||
tool_choice
|
||||
),
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
def get_supported_image_types(self) -> List[str]:
|
||||
return ["png", "jpeg", "gif", "webp"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
model: str,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
for param, value in non_default_params.items():
|
||||
if param == "response_format":
|
||||
json_schema: Optional[dict] = None
|
||||
schema_name: str = ""
|
||||
if "response_schema" in value:
|
||||
json_schema = value["response_schema"]
|
||||
schema_name = "json_tool_call"
|
||||
elif "json_schema" in value:
|
||||
json_schema = value["json_schema"]["schema"]
|
||||
schema_name = value["json_schema"]["name"]
|
||||
"""
|
||||
Follow similar approach to anthropic - translate to a single tool call.
|
||||
|
||||
When using tools in this way: - https://docs.anthropic.com/en/docs/build-with-claude/tool-use#json-mode
|
||||
- You usually want to provide a single tool
|
||||
- You should set tool_choice (see Forcing tool use) to instruct the model to explicitly use that tool
|
||||
- Remember that the model will pass the input to the tool, so the name of the tool and description should be from the model’s perspective.
|
||||
"""
|
||||
if json_schema is not None:
|
||||
_tool_choice = self.map_tool_choice_values(
|
||||
model=model, tool_choice="required", drop_params=drop_params # type: ignore
|
||||
)
|
||||
|
||||
_tool = ChatCompletionToolParam(
|
||||
type="function",
|
||||
function=ChatCompletionToolParamFunctionChunk(
|
||||
name=schema_name, parameters=json_schema
|
||||
),
|
||||
)
|
||||
|
||||
optional_params["tools"] = [_tool]
|
||||
optional_params["tool_choice"] = _tool_choice
|
||||
optional_params["json_mode"] = True
|
||||
else:
|
||||
if litellm.drop_params is True or drop_params is True:
|
||||
pass
|
||||
else:
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message="Bedrock doesn't support response_format={}. To drop it from the call, set `litellm.drop_params = True.".format(
|
||||
value
|
||||
),
|
||||
status_code=400,
|
||||
)
|
||||
if param == "max_tokens":
|
||||
optional_params["maxTokens"] = value
|
||||
if param == "stream":
|
||||
optional_params["stream"] = value
|
||||
if param == "stop":
|
||||
if isinstance(value, str):
|
||||
if len(value) == 0: # converse raises error for empty strings
|
||||
continue
|
||||
value = [value]
|
||||
optional_params["stopSequences"] = value
|
||||
if param == "temperature":
|
||||
optional_params["temperature"] = value
|
||||
if param == "top_p":
|
||||
optional_params["topP"] = value
|
||||
if param == "tools":
|
||||
optional_params["tools"] = value
|
||||
if param == "tool_choice":
|
||||
_tool_choice_value = self.map_tool_choice_values(
|
||||
model=model, tool_choice=value, drop_params=drop_params # type: ignore
|
||||
)
|
||||
if _tool_choice_value is not None:
|
||||
optional_params["tool_choice"] = _tool_choice_value
|
||||
return optional_params
|
||||
|
||||
|
||||
class BedrockConverseLLM(BaseAWSLLM):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
|
@ -1612,7 +1442,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
)
|
||||
|
||||
### SET RUNTIME ENDPOINT ###
|
||||
endpoint_url, proxy_endpoint_url = get_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,
|
||||
|
|
@ -1811,7 +1641,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
error_code = err.response.status_code
|
||||
raise BedrockError(status_code=error_code, message=response.text)
|
||||
raise BedrockError(status_code=error_code, message=err.response.text)
|
||||
except httpx.TimeoutException:
|
||||
raise BedrockError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
|
|
@ -1845,24 +1675,6 @@ def get_response_stream_shape():
|
|||
return _response_stream_shape_cache
|
||||
|
||||
|
||||
def get_bedrock_tool_name(response_tool_name: str) -> str:
|
||||
"""
|
||||
If litellm formatted the input tool name, we need to convert it back to the original name.
|
||||
|
||||
Args:
|
||||
response_tool_name (str): The name of the tool as received from the response.
|
||||
|
||||
Returns:
|
||||
str: The original name of the tool.
|
||||
"""
|
||||
|
||||
if response_tool_name in litellm.bedrock_tool_name_mappings.cache_dict:
|
||||
response_tool_name = litellm.bedrock_tool_name_mappings.cache_dict[
|
||||
response_tool_name
|
||||
]
|
||||
return response_tool_name
|
||||
|
||||
|
||||
class AWSEventStreamDecoder:
|
||||
def __init__(self, model: str) -> None:
|
||||
from botocore.parsers import EventStreamJSONParser
|
||||
|
|
|
|||
228
litellm/llms/bedrock/chat/converse.py
Normal file
228
litellm/llms/bedrock/chat/converse.py
Normal file
|
|
@ -0,0 +1,228 @@
|
|||
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
|
||||
431
litellm/llms/bedrock/chat/converse_transformation.py
Normal file
431
litellm/llms/bedrock/chat/converse_transformation.py
Normal file
|
|
@ -0,0 +1,431 @@
|
|||
"""
|
||||
Translating between OpenAI's `/chat/completion` format and Amazon's `/converse` format
|
||||
"""
|
||||
|
||||
import copy
|
||||
import time
|
||||
import types
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.llms.bedrock import *
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionResponseMessage,
|
||||
ChatCompletionToolCallChunk,
|
||||
ChatCompletionToolCallFunctionChunk,
|
||||
ChatCompletionToolParam,
|
||||
ChatCompletionToolParamFunctionChunk,
|
||||
)
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
from ...prompt_templates.factory import _bedrock_converse_messages_pt, _bedrock_tools_pt
|
||||
from ..common_utils import BedrockError, get_bedrock_tool_name
|
||||
|
||||
|
||||
class AmazonConverseConfig:
|
||||
"""
|
||||
Reference - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
|
||||
#2 - https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html#conversation-inference-supported-models-features
|
||||
"""
|
||||
|
||||
maxTokens: Optional[int]
|
||||
stopSequences: Optional[List[str]]
|
||||
temperature: Optional[int]
|
||||
topP: Optional[int]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
maxTokens: Optional[int] = None,
|
||||
stopSequences: Optional[List[str]] = None,
|
||||
temperature: Optional[int] = None,
|
||||
topP: Optional[int] = None,
|
||||
) -> None:
|
||||
locals_ = locals()
|
||||
for key, value in locals_.items():
|
||||
if key != "self" and value is not None:
|
||||
setattr(self.__class__, key, value)
|
||||
|
||||
@classmethod
|
||||
def get_config(cls):
|
||||
return {
|
||||
k: v
|
||||
for k, v in cls.__dict__.items()
|
||||
if not k.startswith("__")
|
||||
and not isinstance(
|
||||
v,
|
||||
(
|
||||
types.FunctionType,
|
||||
types.BuiltinFunctionType,
|
||||
classmethod,
|
||||
staticmethod,
|
||||
),
|
||||
)
|
||||
and v is not None
|
||||
}
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> List[str]:
|
||||
supported_params = [
|
||||
"max_tokens",
|
||||
"stream",
|
||||
"stream_options",
|
||||
"stop",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"extra_headers",
|
||||
"response_format",
|
||||
]
|
||||
|
||||
if (
|
||||
model.startswith("anthropic")
|
||||
or model.startswith("mistral")
|
||||
or model.startswith("cohere")
|
||||
or model.startswith("meta.llama3-1")
|
||||
):
|
||||
supported_params.append("tools")
|
||||
|
||||
if model.startswith("anthropic") or model.startswith("mistral"):
|
||||
# only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
|
||||
supported_params.append("tool_choice")
|
||||
|
||||
return supported_params
|
||||
|
||||
def map_tool_choice_values(
|
||||
self, model: str, tool_choice: Union[str, dict], drop_params: bool
|
||||
) -> Optional[ToolChoiceValuesBlock]:
|
||||
if tool_choice == "none":
|
||||
if litellm.drop_params is True or drop_params is True:
|
||||
return None
|
||||
else:
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message="Bedrock doesn't support tool_choice={}. To drop it from the call, set `litellm.drop_params = True.".format(
|
||||
tool_choice
|
||||
),
|
||||
status_code=400,
|
||||
)
|
||||
elif tool_choice == "required":
|
||||
return ToolChoiceValuesBlock(any={})
|
||||
elif tool_choice == "auto":
|
||||
return ToolChoiceValuesBlock(auto={})
|
||||
elif isinstance(tool_choice, dict):
|
||||
# only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
|
||||
specific_tool = SpecificToolChoiceBlock(
|
||||
name=tool_choice.get("function", {}).get("name", "")
|
||||
)
|
||||
return ToolChoiceValuesBlock(tool=specific_tool)
|
||||
else:
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message="Bedrock doesn't support tool_choice={}. Supported tool_choice values=['auto', 'required', json object]. To drop it from the call, set `litellm.drop_params = True.".format(
|
||||
tool_choice
|
||||
),
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
def get_supported_image_types(self) -> List[str]:
|
||||
return ["png", "jpeg", "gif", "webp"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
model: str,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
for param, value in non_default_params.items():
|
||||
if param == "response_format":
|
||||
json_schema: Optional[dict] = None
|
||||
schema_name: str = ""
|
||||
if "response_schema" in value:
|
||||
json_schema = value["response_schema"]
|
||||
schema_name = "json_tool_call"
|
||||
elif "json_schema" in value:
|
||||
json_schema = value["json_schema"]["schema"]
|
||||
schema_name = value["json_schema"]["name"]
|
||||
"""
|
||||
Follow similar approach to anthropic - translate to a single tool call.
|
||||
|
||||
When using tools in this way: - https://docs.anthropic.com/en/docs/build-with-claude/tool-use#json-mode
|
||||
- You usually want to provide a single tool
|
||||
- You should set tool_choice (see Forcing tool use) to instruct the model to explicitly use that tool
|
||||
- Remember that the model will pass the input to the tool, so the name of the tool and description should be from the model’s perspective.
|
||||
"""
|
||||
if json_schema is not None:
|
||||
_tool_choice = self.map_tool_choice_values(
|
||||
model=model, tool_choice="required", drop_params=drop_params # type: ignore
|
||||
)
|
||||
|
||||
_tool = ChatCompletionToolParam(
|
||||
type="function",
|
||||
function=ChatCompletionToolParamFunctionChunk(
|
||||
name=schema_name, parameters=json_schema
|
||||
),
|
||||
)
|
||||
|
||||
optional_params["tools"] = [_tool]
|
||||
optional_params["tool_choice"] = _tool_choice
|
||||
optional_params["json_mode"] = True
|
||||
else:
|
||||
if litellm.drop_params is True or drop_params is True:
|
||||
pass
|
||||
else:
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message="Bedrock doesn't support response_format={}. To drop it from the call, set `litellm.drop_params = True.".format(
|
||||
value
|
||||
),
|
||||
status_code=400,
|
||||
)
|
||||
if param == "max_tokens":
|
||||
optional_params["maxTokens"] = value
|
||||
if param == "stream":
|
||||
optional_params["stream"] = value
|
||||
if param == "stop":
|
||||
if isinstance(value, str):
|
||||
if len(value) == 0: # converse raises error for empty strings
|
||||
continue
|
||||
value = [value]
|
||||
optional_params["stopSequences"] = value
|
||||
if param == "temperature":
|
||||
optional_params["temperature"] = value
|
||||
if param == "top_p":
|
||||
optional_params["topP"] = value
|
||||
if param == "tools":
|
||||
optional_params["tools"] = value
|
||||
if param == "tool_choice":
|
||||
_tool_choice_value = self.map_tool_choice_values(
|
||||
model=model, tool_choice=value, drop_params=drop_params # type: ignore
|
||||
)
|
||||
if _tool_choice_value is not None:
|
||||
optional_params["tool_choice"] = _tool_choice_value
|
||||
return optional_params
|
||||
|
||||
def _transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> RequestObject:
|
||||
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
|
||||
|
||||
return _data
|
||||
|
||||
def _transform_response(
|
||||
self,
|
||||
model: str,
|
||||
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
|
||||
|
|
@ -5,7 +5,7 @@ Common utilities used across bedrock chat/embedding/image generation
|
|||
import os
|
||||
import types
|
||||
from enum import Enum
|
||||
from typing import List, Optional, Union, Tuple
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -575,7 +575,7 @@ def init_bedrock_client(
|
|||
# Iterate over parameters and update if needed
|
||||
for i, param in enumerate(params_to_check):
|
||||
if param and param.startswith("os.environ/"):
|
||||
params_to_check[i] = get_secret(param)
|
||||
params_to_check[i] = get_secret(param) # type: ignore
|
||||
# Assign updated values back to parameters
|
||||
(
|
||||
aws_access_key_id,
|
||||
|
|
@ -618,13 +618,13 @@ def init_bedrock_client(
|
|||
import boto3
|
||||
|
||||
if isinstance(timeout, float):
|
||||
config = boto3.session.Config(connect_timeout=timeout, read_timeout=timeout)
|
||||
config = boto3.session.Config(connect_timeout=timeout, read_timeout=timeout) # type: ignore
|
||||
elif isinstance(timeout, httpx.Timeout):
|
||||
config = boto3.session.Config(
|
||||
config = boto3.session.Config( # type: ignore
|
||||
connect_timeout=timeout.connect, read_timeout=timeout.read
|
||||
)
|
||||
else:
|
||||
config = boto3.session.Config()
|
||||
config = boto3.session.Config() # type: ignore
|
||||
|
||||
### CHECK STS ###
|
||||
if (
|
||||
|
|
@ -725,40 +725,6 @@ def init_bedrock_client(
|
|||
return client
|
||||
|
||||
|
||||
def get_runtime_endpoint(
|
||||
api_base: Optional[str],
|
||||
aws_bedrock_runtime_endpoint: Optional[str],
|
||||
aws_region_name: str,
|
||||
) -> Tuple[str, str]:
|
||||
env_aws_bedrock_runtime_endpoint = get_secret("AWS_BEDROCK_RUNTIME_ENDPOINT")
|
||||
if api_base is not None:
|
||||
endpoint_url = api_base
|
||||
elif aws_bedrock_runtime_endpoint is not None and isinstance(
|
||||
aws_bedrock_runtime_endpoint, str
|
||||
):
|
||||
endpoint_url = aws_bedrock_runtime_endpoint
|
||||
elif env_aws_bedrock_runtime_endpoint and isinstance(
|
||||
env_aws_bedrock_runtime_endpoint, str
|
||||
):
|
||||
endpoint_url = env_aws_bedrock_runtime_endpoint
|
||||
else:
|
||||
endpoint_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com"
|
||||
|
||||
# Determine proxy_endpoint_url
|
||||
if env_aws_bedrock_runtime_endpoint and isinstance(
|
||||
env_aws_bedrock_runtime_endpoint, str
|
||||
):
|
||||
proxy_endpoint_url = env_aws_bedrock_runtime_endpoint
|
||||
elif aws_bedrock_runtime_endpoint is not None and isinstance(
|
||||
aws_bedrock_runtime_endpoint, str
|
||||
):
|
||||
proxy_endpoint_url = aws_bedrock_runtime_endpoint
|
||||
else:
|
||||
proxy_endpoint_url = endpoint_url
|
||||
|
||||
return endpoint_url, proxy_endpoint_url
|
||||
|
||||
|
||||
class ModelResponseIterator:
|
||||
def __init__(self, model_response):
|
||||
self.model_response = model_response
|
||||
|
|
@ -783,3 +749,21 @@ class ModelResponseIterator:
|
|||
raise StopAsyncIteration
|
||||
self.is_done = True
|
||||
return self.model_response
|
||||
|
||||
|
||||
def get_bedrock_tool_name(response_tool_name: str) -> str:
|
||||
"""
|
||||
If litellm formatted the input tool name, we need to convert it back to the original name.
|
||||
|
||||
Args:
|
||||
response_tool_name (str): The name of the tool as received from the response.
|
||||
|
||||
Returns:
|
||||
str: The original name of the tool.
|
||||
"""
|
||||
|
||||
if response_tool_name in litellm.bedrock_tool_name_mappings.cache_dict:
|
||||
response_tool_name = litellm.bedrock_tool_name_mappings.cache_dict[
|
||||
response_tool_name
|
||||
]
|
||||
return response_tool_name
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ from litellm.types.llms.bedrock import AmazonEmbeddingRequest, CohereEmbeddingRe
|
|||
from litellm.types.utils import Embedding, EmbeddingResponse, Usage
|
||||
|
||||
from ...base_aws_llm import BaseAWSLLM
|
||||
from ..common_utils import BedrockError, get_runtime_endpoint
|
||||
from ..common_utils import BedrockError
|
||||
from .amazon_titan_g1_transformation import AmazonTitanG1Config
|
||||
from .amazon_titan_multimodal_transformation import (
|
||||
AmazonTitanMultimodalEmbeddingG1Config,
|
||||
|
|
@ -141,7 +141,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
error_code = err.response.status_code
|
||||
raise BedrockError(status_code=error_code, message=response.text)
|
||||
raise BedrockError(status_code=error_code, message=err.response.text)
|
||||
except httpx.TimeoutException:
|
||||
raise BedrockError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
|
|
@ -197,7 +197,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
client=client,
|
||||
timeout=timeout,
|
||||
api_base=prepped.url,
|
||||
headers=prepped.headers,
|
||||
headers=dict(prepped.headers),
|
||||
data=data,
|
||||
)
|
||||
|
||||
|
|
@ -288,7 +288,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
client=client,
|
||||
timeout=timeout,
|
||||
api_base=prepped.url,
|
||||
headers=prepped.headers,
|
||||
headers=dict(prepped.headers),
|
||||
data=data,
|
||||
)
|
||||
|
||||
|
|
@ -342,8 +342,8 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
aembedding: Optional[bool],
|
||||
extra_headers: Optional[dict],
|
||||
optional_params=None,
|
||||
litellm_params=None,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> EmbeddingResponse:
|
||||
try:
|
||||
import boto3
|
||||
|
|
@ -392,10 +392,21 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
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 = get_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
|
||||
|
|
@ -467,170 +478,5 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
aembedding=aembedding,
|
||||
timeout=timeout,
|
||||
client=client,
|
||||
headers=prepped.headers,
|
||||
headers=dict(prepped.headers),
|
||||
)
|
||||
|
||||
# def _embedding_func_single(
|
||||
# model: str,
|
||||
# input: str,
|
||||
# client: Any,
|
||||
# optional_params=None,
|
||||
# encoding=None,
|
||||
# logging_obj=None,
|
||||
# ):
|
||||
# if isinstance(input, str) is False:
|
||||
# raise BedrockError(
|
||||
# message="Bedrock Embedding API input must be type str | List[str]",
|
||||
# status_code=400,
|
||||
# )
|
||||
# # logic for parsing in - calling - parsing out model embedding calls
|
||||
# ## FORMAT EMBEDDING INPUT ##
|
||||
# provider = model.split(".")[0]
|
||||
# inference_params = copy.deepcopy(optional_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
|
||||
# if provider == "amazon":
|
||||
# input = input.replace(os.linesep, " ")
|
||||
# data = {"inputText": input, **inference_params}
|
||||
# # data = json.dumps(data)
|
||||
# elif provider == "cohere":
|
||||
# inference_params["input_type"] = inference_params.get(
|
||||
# "input_type", "search_document"
|
||||
# ) # aws bedrock example default - https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=cohere.embed-english-v3
|
||||
# data = {"texts": [input], **inference_params} # type: ignore
|
||||
# body = json.dumps(data).encode("utf-8") # type: ignore
|
||||
# ## LOGGING
|
||||
# request_str = f"""
|
||||
# response = client.invoke_model(
|
||||
# body={body},
|
||||
# modelId={modelId},
|
||||
# accept="*/*",
|
||||
# contentType="application/json",
|
||||
# )""" # type: ignore
|
||||
# logging_obj.pre_call(
|
||||
# input=input,
|
||||
# api_key="", # boto3 is used for init.
|
||||
# additional_args={
|
||||
# "complete_input_dict": {"model": modelId, "texts": input},
|
||||
# "request_str": request_str,
|
||||
# },
|
||||
# )
|
||||
# try:
|
||||
# response = client.invoke_model(
|
||||
# body=body,
|
||||
# modelId=modelId,
|
||||
# accept="*/*",
|
||||
# contentType="application/json",
|
||||
# )
|
||||
# response_body = json.loads(response.get("body").read())
|
||||
# ## LOGGING
|
||||
# logging_obj.post_call(
|
||||
# input=input,
|
||||
# api_key="",
|
||||
# additional_args={"complete_input_dict": data},
|
||||
# original_response=json.dumps(response_body),
|
||||
# )
|
||||
# if provider == "cohere":
|
||||
# response = response_body.get("embeddings")
|
||||
# # flatten list
|
||||
# response = [item for sublist in response for item in sublist]
|
||||
# return response
|
||||
# elif provider == "amazon":
|
||||
# return response_body.get("embedding")
|
||||
# except Exception as e:
|
||||
# raise BedrockError(
|
||||
# message=f"Embedding Error with model {model}: {e}", status_code=500
|
||||
# )
|
||||
|
||||
# def embedding(
|
||||
# model: str,
|
||||
# input: Union[list, str],
|
||||
# model_response: litellm.EmbeddingResponse,
|
||||
# api_key: Optional[str] = None,
|
||||
# logging_obj=None,
|
||||
# optional_params=None,
|
||||
# encoding=None,
|
||||
# ):
|
||||
# ### BOTO3 INIT ###
|
||||
# # 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_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_bedrock_runtime_endpoint = optional_params.pop(
|
||||
# "aws_bedrock_runtime_endpoint", None
|
||||
# )
|
||||
# aws_web_identity_token = optional_params.pop("aws_web_identity_token", None)
|
||||
|
||||
# # use passed in BedrockRuntime.Client if provided, otherwise create a new one
|
||||
# client = init_bedrock_client(
|
||||
# aws_access_key_id=aws_access_key_id,
|
||||
# aws_secret_access_key=aws_secret_access_key,
|
||||
# aws_region_name=aws_region_name,
|
||||
# aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
|
||||
# aws_web_identity_token=aws_web_identity_token,
|
||||
# aws_role_name=aws_role_name,
|
||||
# aws_session_name=aws_session_name,
|
||||
# )
|
||||
# if isinstance(input, str):
|
||||
# ## Embedding Call
|
||||
# embeddings = [
|
||||
# _embedding_func_single(
|
||||
# model,
|
||||
# input,
|
||||
# optional_params=optional_params,
|
||||
# client=client,
|
||||
# logging_obj=logging_obj,
|
||||
# )
|
||||
# ]
|
||||
# elif isinstance(input, list):
|
||||
# ## Embedding Call - assuming this is a List[str]
|
||||
# embeddings = [
|
||||
# _embedding_func_single(
|
||||
# model,
|
||||
# i,
|
||||
# optional_params=optional_params,
|
||||
# client=client,
|
||||
# logging_obj=logging_obj,
|
||||
# )
|
||||
# for i in input
|
||||
# ] # [TODO]: make these parallel calls
|
||||
# else:
|
||||
# # enters this branch if input = int, ex. input=2
|
||||
# raise BedrockError(
|
||||
# message="Bedrock Embedding API input must be type str | List[str]",
|
||||
# status_code=400,
|
||||
# )
|
||||
|
||||
# ## Populate OpenAI compliant dictionary
|
||||
# embedding_response = []
|
||||
# for idx, embedding in enumerate(embeddings):
|
||||
# embedding_response.append(
|
||||
# {
|
||||
# "object": "embedding",
|
||||
# "index": idx,
|
||||
# "embedding": embedding,
|
||||
# }
|
||||
# )
|
||||
# model_response.object = "list"
|
||||
# model_response.data = embedding_response
|
||||
# model_response.model = model
|
||||
# input_tokens = 0
|
||||
|
||||
# input_str = "".join(input)
|
||||
|
||||
# input_tokens += len(encoding.encode(input_str))
|
||||
|
||||
# usage = Usage(
|
||||
# prompt_tokens=input_tokens,
|
||||
# completion_tokens=0,
|
||||
# total_tokens=input_tokens + 0,
|
||||
# )
|
||||
# model_response.usage = usage
|
||||
|
||||
# return model_response
|
||||
|
|
|
|||
|
|
@ -2367,6 +2367,7 @@ def completion(
|
|||
)
|
||||
|
||||
if model in litellm.BEDROCK_CONVERSE_MODELS:
|
||||
|
||||
response = bedrock_converse_chat_completion.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
|
|
|
|||
|
|
@ -1,7 +1,20 @@
|
|||
model_list:
|
||||
- model_name: gpt-3.5-turbo
|
||||
- model_name: fake-claude-endpoint
|
||||
litellm_params:
|
||||
model: anthropic.claude-3-sonnet-20240229-v1:0
|
||||
api_base: https://exampleopenaiendpoint-production.up.railway.app
|
||||
# aws_session_token: "IQoJb3JpZ2luX2VjELj//////////wEaCXVzLXdlc3QtMiJHMEUCIQDatCRVkIZERLcrR6P7Qd1vNfZ8r8xB/LUeaVaTW/lBTwIgAgmHSBe41d65GVRKSkpgVonjsCmOmAS7s/yklM9NsZcq3AEI4P//////////ARABGgw4ODg2MDIyMjM0MjgiDJrio0/CHYEfyt5EqyqwAfyWO4t3bFVWAOIwTyZ1N6lszeJKfMNus2hzVc+r73hia2Anv88uwPxNg2uqnXQNJumEo0DcBt30ZwOw03Isboy0d5l05h8gjb4nl9feyeKmKAnRdcqElrEWtCC1Qcefv78jQv53AbUipH1ssa5NPvptqZZpZYDPMlBEnV3YdvJJiuE23u2yOkCt+EoUJLaOYjOryoRyrSfbWB+JaUsB68R3rNTHzReeN3Nob/9Ic4HrMMmzmLcGOpgBZxclO4w8Z7i6TcVqbCwDOskxuR6bZaiFxKFG+9tDrWS7jaQKpq/YP9HUT0YwYpZplaBEEZR5sbIndg5yb4dRZrSHplblqKz8XLaUf5tuuyRJmwr96PTpw/dyEVk9gicFX6JfLBEv0v5rN2Z0JMFLdfIP4kC1U2PjcPOWoglWO3fLmJ4Lol2a3c5XDSMwMxjcJXq+c8Ue1v0="
|
||||
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
|
||||
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
|
||||
- model_name: gemini-vision
|
||||
litellm_params:
|
||||
model: vertex_ai/gemini-1.0-pro-vision-001
|
||||
api_base: https://exampleopenaiendpoint-production.up.railway.app/v1/projects/adroit-crow-413218/locations/us-central1/publishers/google/models/gemini-1.0-pro-vision-001
|
||||
vertex_project: "adroit-crow-413218"
|
||||
vertex_location: "us-central1"
|
||||
|
||||
- model_name: fake-openai-endpoint
|
||||
litellm_params:
|
||||
model: gpt-3.5-turbo
|
||||
|
||||
router_settings:
|
||||
model_group_alias: {"gpt-4": {"model": "gpt-3.5-turbo", "hidden": false}}
|
||||
api_base: https://exampleopenaiendpoint-production.up.railway.app
|
||||
|
||||
Loading…
Add table
Reference in a new issue