From 18ef244230e6a05086acf6e9ce35a9e37e763027 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 2 Jan 2024 10:23:28 +0530 Subject: [PATCH] (test) bedrock-test passing boto3 client --- litellm/tests/test_bedrock_completion.py | 111 +++++++++++++++++++++++ litellm/tests/test_completion.py | 81 +---------------- 2 files changed, 112 insertions(+), 80 deletions(-) create mode 100644 litellm/tests/test_bedrock_completion.py diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py new file mode 100644 index 000000000..004ae86f1 --- /dev/null +++ b/litellm/tests/test_bedrock_completion.py @@ -0,0 +1,111 @@ +import sys, os +import traceback +from dotenv import load_dotenv + +load_dotenv() +import os, io + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import pytest +import litellm +from litellm import embedding, completion, completion_cost, Timeout +from litellm import RateLimitError + +# litellm.num_retries = 3 +litellm.cache = None +litellm.success_callback = [] +user_message = "Write a short poem about the sky" +messages = [{"content": user_message, "role": "user"}] + + +@pytest.fixture(autouse=True) +def reset_callbacks(): + print("\npytest fixture - resetting callbacks") + litellm.success_callback = [] + litellm._async_success_callback = [] + litellm.failure_callback = [] + litellm.callbacks = [] + + +def test_completion_bedrock_claude_completion_auth(): + print("calling bedrock claude completion params auth") + import os + + aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"] + aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"] + aws_region_name = os.environ["AWS_REGION_NAME"] + + os.environ.pop("AWS_ACCESS_KEY_ID", None) + os.environ.pop("AWS_SECRET_ACCESS_KEY", None) + os.environ.pop("AWS_REGION_NAME", None) + + try: + response = completion( + model="bedrock/anthropic.claude-instant-v1", + messages=messages, + max_tokens=10, + temperature=0.1, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_region_name=aws_region_name, + ) + # Add any assertions here to check the response + print(response) + + os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id + os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key + os.environ["AWS_REGION_NAME"] = aws_region_name + except RateLimitError: + pass + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + +test_completion_bedrock_claude_completion_auth() + + +def test_completion_bedrock_claude_external_client_auth(): + print("\ncalling bedrock claude external client auth") + import os + + aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"] + aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"] + aws_region_name = os.environ["AWS_REGION_NAME"] + + os.environ.pop("AWS_ACCESS_KEY_ID", None) + os.environ.pop("AWS_SECRET_ACCESS_KEY", None) + os.environ.pop("AWS_REGION_NAME", None) + + try: + import boto3 + + bedrock = 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, + endpoint_url=f"https://bedrock-runtime.{aws_region_name}.amazonaws.com", + ) + + response = completion( + model="bedrock/anthropic.claude-instant-v1", + messages=messages, + max_tokens=10, + temperature=0.1, + aws_bedrock_client=bedrock, + ) + # Add any assertions here to check the response + print(response) + + os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id + os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key + os.environ["AWS_REGION_NAME"] = aws_region_name + except RateLimitError: + pass + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + +test_completion_bedrock_claude_external_client_auth() diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 056a6e408..2ddb5fa13 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -1490,85 +1490,6 @@ def test_completion_bedrock_cohere(): # test_completion_bedrock_cohere() -def test_completion_bedrock_claude_completion_auth(): - print("calling bedrock claude completion params auth") - import os - - aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"] - aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"] - aws_region_name = os.environ["AWS_REGION_NAME"] - - os.environ["AWS_ACCESS_KEY_ID"] = "" - os.environ["AWS_SECRET_ACCESS_KEY"] = "" - os.environ["AWS_REGION_NAME"] = "" - - try: - response = completion( - model="bedrock/anthropic.claude-instant-v1", - messages=messages, - max_tokens=10, - temperature=0.1, - logger_fn=logger_fn, - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_region_name=aws_region_name, - ) - # Add any assertions here to check the response - print(response) - - os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id - os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key - os.environ["AWS_REGION_NAME"] = aws_region_name - except RateLimitError: - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - -# test_completion_bedrock_claude_completion_auth() - -# def test_completion_bedrock_claude_external_client_auth(): -# print("calling bedrock claude external client auth") -# import os - -# aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"] -# aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"] -# aws_region_name = os.environ["AWS_REGION_NAME"] - -# os.environ["AWS_ACCESS_KEY_ID"] = "" -# os.environ["AWS_SECRET_ACCESS_KEY"] = "" -# os.environ["AWS_REGION_NAME"] = "" - -# try: -# import boto3 -# bedrock = 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, -# endpoint_url=f"https://bedrock-runtime.{aws_region_name}.amazonaws.com" -# ) - -# response = completion( -# model="bedrock/anthropic.claude-instant-v1", -# messages=messages, -# max_tokens=10, -# temperature=0.1, -# logger_fn=logger_fn, -# aws_bedrock_client=bedrock, -# ) -# # Add any assertions here to check the response -# print(response) - -# os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id -# os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key -# os.environ["AWS_REGION_NAME"] = aws_region_name -# except RateLimitError: -# pass -# except Exception as e: -# pytest.fail(f"Error occurred: {e}") -# test_completion_bedrock_claude_external_client_auth() - # def test_completion_bedrock_claude_stream(): # print("calling claude") # litellm.set_verbose = False @@ -2017,7 +1938,7 @@ def test_completion_cloudflare(): pytest.fail(f"Error occurred: {e}") -test_completion_cloudflare() +# test_completion_cloudflare() def test_moderation():