From 5a6588342cab071322737d69043a3731e34d0468 Mon Sep 17 00:00:00 2001 From: Brian Schultheiss Date: Sun, 23 Jun 2024 15:19:54 -0700 Subject: [PATCH] added test for change --- docs/my-website/docs/providers/bedrock.md | 8 +- litellm/tests/test_bedrock_completion.py | 181 ++++++++++++++++++++++ 2 files changed, 185 insertions(+), 4 deletions(-) diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index 7f9b21b96b7..1b073ad25eb 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -474,10 +474,10 @@ from litellm import completion response = completion( model="bedrock/anthropic.claude-instant-v1", messages=[{ "content": "Hello, how are you?","role": "user"}], + aws_region_name="", aws_access_key_id="", aws_secret_access_key="", - aws_session_token="", - aws_region_name="", + aws_session_token=None, ) ``` @@ -553,7 +553,7 @@ This is a deprecated flow. Boto3 is not async. And boto3.client does not let us Experimental - 2024-Jun-23: - aws_access_key_id, aws_secret_access_key=, and aws_session_token will be extracted from boto3.client and be passed onto the httpx client + `aws_access_key_id`, `aws_secret_access_key`, and `aws_session_token` will be extracted from boto3.client and be passed onto the httpx client ::: @@ -569,7 +569,7 @@ bedrock = boto3.client( region_name="us-east-1", aws_access_key_id="", aws_secret_access_key="", - aws_session_token="", + aws_session_token=None, ) response = completion( diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index b953ca2a3a6..614d1b76fad 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -15,6 +15,7 @@ from litellm import embedding, completion, completion_cost, Timeout, ModelRespon from litellm import RateLimitError from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler from unittest.mock import patch, AsyncMock, Mock +from litellm.llms.bedrock_httpx import BedrockLLM # litellm.num_retries = 3 litellm.cache = None @@ -205,6 +206,186 @@ def test_completion_bedrock_claude_sts_client_auth(): except Exception as e: pytest.fail(f"Error occurred: {e}") +@pytest.fixture() +def bedrock_session_token_creds(): + print("\ncalling oidc auto to get aws_session_token credentials") + import os + + aws_region_name = os.environ["AWS_REGION_NAME"] + aws_session_token = os.environ.get("AWS_SESSION_TOKEN") + + bllm = BedrockLLM() + if aws_session_token is not None: + # For local testing + creds = bllm.get_credentials( + aws_region_name=aws_region_name, + aws_access_key_id=os.environ['AWS_ACCESS_KEY_ID'], + aws_secret_access_key=os.environ['AWS_SECRET_ACCESS_KEY'], + aws_session_token=aws_session_token + ) + else: + # For circle-ci testing + # aws_role_name = os.environ["AWS_TEMP_ROLE_NAME"] + # TODO: This is using ai.moda's IAM role, we should use LiteLLM's IAM role eventually + aws_role_name = "arn:aws:iam::335785316107:role/litellm-github-unit-tests-circleci" + aws_web_identity_token = "oidc/circleci_v2/" + + creds = bllm.get_credentials( + aws_region_name=aws_region_name, + aws_web_identity_token=aws_web_identity_token, + aws_role_name=aws_role_name, + aws_session_name="my-test-session", + ) + return creds + +@pytest.mark.skipif( + os.environ.get("CIRCLE_OIDC_TOKEN_V2") is None, + reason="Cannot run without being in CircleCI Runner", +) +def test_completion_bedrock_claude_aws_session_token(bedrock_session_token_creds): + print("\ncalling bedrock claude with aws_session_token auth") + + import os + aws_region_name = os.environ["AWS_REGION_NAME"] + aws_access_key_id = bedrock_session_token_creds.access_key + aws_secret_access_key = bedrock_session_token_creds.secret_key + aws_session_token = bedrock_session_token_creds.token + + try: + litellm.set_verbose = True + + response_1 = completion( + model="bedrock/anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + max_tokens=10, + temperature=0.1, + aws_region_name=aws_region_name, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + ) + print(response_1) + assert len(response_1.choices) > 0 + assert len(response_1.choices[0].message.content) > 0 + + # This second call is to verify that the cache isn't breaking anything + response_2 = completion( + model="bedrock/anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + max_tokens=5, + temperature=0.2, + aws_region_name=aws_region_name, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + ) + print(response_2) + assert len(response_2.choices) > 0 + assert len(response_2.choices[0].message.content) > 0 + + # This third call is to verify that the cache isn't used for a different region + response_3 = completion( + model="bedrock/anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + max_tokens=6, + temperature=0.3, + aws_region_name="us-east-1", + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + ) + print(response_3) + assert len(response_3.choices) > 0 + assert len(response_3.choices[0].message.content) > 0 + + except RateLimitError: + pass + except Exception as e: + pytest.fail(f"Error occurred: {e}") + +@pytest.mark.skipif( + os.environ.get("CIRCLE_OIDC_TOKEN_V2") is None, + reason="Cannot run without being in CircleCI Runner", +) +def test_completion_bedrock_claude_aws_bedrock_client(bedrock_session_token_creds): + print("\ncalling bedrock claude with aws_session_token auth") + + import os + import boto3 + from botocore.client import Config + + aws_region_name = os.environ["AWS_REGION_NAME"] + aws_access_key_id = bedrock_session_token_creds.access_key + aws_secret_access_key = bedrock_session_token_creds.secret_key + aws_session_token = bedrock_session_token_creds.token + + aws_bedrock_client_west = boto3.client( + service_name="bedrock-runtime", + region_name=aws_region_name, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + config= Config( + read_timeout=600 + ) + ) + + + try: + litellm.set_verbose = True + + response_1 = completion( + model="bedrock/anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + max_tokens=10, + temperature=0.1, + aws_bedrock_client=aws_bedrock_client_west, + ) + print(response_1) + assert len(response_1.choices) > 0 + assert len(response_1.choices[0].message.content) > 0 + + # This second call is to verify that the cache isn't breaking anything + response_2 = completion( + model="bedrock/anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + max_tokens=5, + temperature=0.2, + aws_bedrock_client=aws_bedrock_client_west, + ) + print(response_2) + assert len(response_2.choices) > 0 + assert len(response_2.choices[0].message.content) > 0 + + # This third call is to verify that the cache isn't used for a different region + aws_bedrock_client_east = boto3.client( + service_name="bedrock-runtime", + region_name="us-east-1", + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + config= Config( + read_timeout=600 + ) + ) + + response_3 = completion( + model="bedrock/anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + max_tokens=6, + temperature=0.3, + aws_bedrock_client=aws_bedrock_client_east, + ) + print(response_3) + assert len(response_3.choices) > 0 + assert len(response_3.choices[0].message.content) > 0 + + except RateLimitError: + pass + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + # test_completion_bedrock_claude_sts_client_auth()