mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
added test for change
This commit is contained in:
Brian Schultheiss
parent
3fbb25f903
commit
5a6588342c
2 changed files with 185 additions and 4 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue