From 80b4af7abec2393288aaa4cc01584b34a5d0e6db Mon Sep 17 00:00:00 2001 From: Brian Schultheiss Date: Tue, 25 Jun 2024 13:29:33 -0700 Subject: [PATCH] Revert some non-essential changes --- docs/my-website/docs/providers/aws_sagemaker.md | 1 - litellm/llms/sagemaker.py | 14 ++------------ litellm/proxy/proxy_server.py | 2 -- litellm/types/router.py | 4 ---- 4 files changed, 2 insertions(+), 19 deletions(-) diff --git a/docs/my-website/docs/providers/aws_sagemaker.md b/docs/my-website/docs/providers/aws_sagemaker.md index 5793fb05ae8..2b65709e8e0 100644 --- a/docs/my-website/docs/providers/aws_sagemaker.md +++ b/docs/my-website/docs/providers/aws_sagemaker.md @@ -59,7 +59,6 @@ response = completion( messages=[{ "content": "Hello, how are you?","role": "user"}], aws_access_key_id="", aws_secret_access_key="", - aws_session_token="", aws_region_name="", ) ``` diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 7d639b7bb21..8e75428bb7c 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -162,10 +162,9 @@ def completion( ): import boto3 - # pop aws_secret_access_key, aws_access_key_id, aws_session_token, aws_region_name from kwargs, since completion calls fail with them + # 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) model_id = optional_params.pop("model_id", None) @@ -176,7 +175,6 @@ def completion( service_name="sagemaker-runtime", aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, region_name=aws_region_name, ) else: @@ -251,7 +249,6 @@ def completion( model_id=model_id, aws_secret_access_key=aws_secret_access_key, aws_access_key_id=aws_access_key_id, - aws_session_token=aws_session_token, aws_region_name=aws_region_name, ) return response @@ -284,7 +281,6 @@ def completion( model_id=model_id, aws_secret_access_key=aws_secret_access_key, aws_access_key_id=aws_access_key_id, - aws_session_token=aws_session_token, aws_region_name=aws_region_name, ) data = json.dumps({"inputs": prompt, "parameters": inference_params}).encode( @@ -418,7 +414,6 @@ async def async_streaming( aws_secret_access_key: Optional[str], aws_access_key_id: Optional[str], aws_region_name: Optional[str], - aws_session_token: Optional[str] = None, ): """ Use aioboto3 @@ -434,7 +429,6 @@ async def async_streaming( service_name="sagemaker-runtime", aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, region_name=aws_region_name, ) else: @@ -487,7 +481,6 @@ async def async_completion( aws_secret_access_key: Optional[str], aws_access_key_id: Optional[str], aws_region_name: Optional[str], - aws_session_token: Optional[str] = None, ): """ Use aioboto3 @@ -503,7 +496,6 @@ async def async_completion( service_name="sagemaker-runtime", aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, region_name=aws_region_name, ) else: @@ -647,10 +639,9 @@ def embedding( ### BOTO3 INIT import boto3 - # pop aws_secret_access_key, aws_access_key_id, aws_session_token, aws_region_name from kwargs, since completion calls fail with them + # 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) if aws_access_key_id is not None: @@ -660,7 +651,6 @@ def embedding( service_name="sagemaker-runtime", aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, region_name=aws_region_name, ) else: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9c1039f51ea..30b90abe64c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6343,7 +6343,6 @@ async def model_info_v2( _model["litellm_params"].pop("vertex_credentials", None) _model["litellm_params"].pop("aws_access_key_id", None) _model["litellm_params"].pop("aws_secret_access_key", None) - _model["litellm_params"].pop("aws_session_token", None) verbose_proxy_logger.debug("all_models: %s", all_models) return {"data": all_models} @@ -6860,7 +6859,6 @@ async def model_info_v1( model["litellm_params"].pop("vertex_credentials", None) model["litellm_params"].pop("aws_access_key_id", None) model["litellm_params"].pop("aws_secret_access_key", None) - model["litellm_params"].pop("aws_session_token", None) verbose_proxy_logger.debug("all_models: %s", all_models) return {"data": all_models} diff --git a/litellm/types/router.py b/litellm/types/router.py index 059a8620e56..e6864ffe2e9 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -145,7 +145,6 @@ class GenericLiteLLMParams(BaseModel): ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: Optional[str] = None aws_secret_access_key: Optional[str] = None - aws_session_token: Optional[str] = None aws_region_name: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None @@ -179,7 +178,6 @@ class GenericLiteLLMParams(BaseModel): ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None, - aws_session_token: Optional[str] = None, aws_region_name: Optional[str] = None, ## IBM WATSONX ## watsonx_region_name: Optional[str] = None, @@ -244,7 +242,6 @@ class LiteLLM_Params(GenericLiteLLMParams): ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None, - aws_session_token: Optional[str] = None, aws_region_name: Optional[str] = None, **params, ): @@ -310,7 +307,6 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: Optional[str] aws_secret_access_key: Optional[str] - aws_session_token: Optional[str] aws_region_name: Optional[str] ## IBM WATSONX ## watsonx_region_name: Optional[str]