# @pytest.mark.skip(reason="AWS Suspended Account") 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, ModelResponse 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_2_1_completion_auth(): print("calling bedrock claude 2.1 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-v2:1", 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_2_1_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 litellm.set_verbose = True 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() @pytest.mark.skip(reason="Expired token, need to renew") def test_completion_bedrock_claude_sts_client_auth(): print("\ncalling bedrock claude external client auth") import os aws_access_key_id = os.environ["AWS_TEMP_ACCESS_KEY_ID"] aws_secret_access_key = os.environ["AWS_TEMP_SECRET_ACCESS_KEY"] aws_region_name = os.environ["AWS_REGION_NAME"] aws_role_name = os.environ["AWS_TEMP_ROLE_NAME"] try: import boto3 litellm.set_verbose = True response = completion( model="bedrock/anthropic.claude-instant-v1", 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_role_name=aws_role_name, aws_session_name="my-test-session", ) response = embedding( model="cohere.embed-multilingual-v3", input=["hello world"], aws_region_name="us-east-1", aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_role_name=aws_role_name, aws_session_name="my-test-session", ) response = completion( model="gpt-3.5-turbo", messages=messages, aws_region_name="us-east-1", aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_role_name=aws_role_name, aws_session_name="my-test-session", ) # Add any assertions here to check the response print(response) except RateLimitError: pass except Exception as e: pytest.fail(f"Error occurred: {e}") # test_completion_bedrock_claude_sts_client_auth() def test_bedrock_claude_3(): try: litellm.set_verbose = True response: ModelResponse = completion( model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", messages=messages, max_tokens=10, temperature=0.78, ) # Add any assertions here to check the response assert len(response.choices) > 0 assert len(response.choices[0].message.content) > 0 except RateLimitError: pass except Exception as e: pytest.fail(f"Error occurred: {e}") def test_bedrock_claude_3_tool_calling(): try: litellm.set_verbose = True tools = [ { "type": "function", "function": { "name": "get_current_weather", "description": "Get the current weather in a given location", "parameters": { "type": "object", "properties": { "location": { "type": "string", "description": "The city and state, e.g. San Francisco, CA", }, "unit": { "type": "string", "enum": ["celsius", "fahrenheit"], }, }, "required": ["location"], }, }, } ] messages = [ {"role": "user", "content": "What's the weather like in Boston today?"} ] response: ModelResponse = completion( model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", messages=messages, tools=tools, tool_choice="auto", ) print(f"response: {response}") # Add any assertions here to check the response assert isinstance(response.choices[0].message.tool_calls[0].function.name, str) assert isinstance( response.choices[0].message.tool_calls[0].function.arguments, str ) except RateLimitError: pass except Exception as e: pytest.fail(f"Error occurred: {e}") def encode_image(image_path): import base64 with open(image_path, "rb") as image_file: return base64.b64encode(image_file.read()).decode("utf-8") @pytest.mark.skip( reason="we already test claude-3, this is just another way to pass images" ) def test_completion_claude_3_base64(): try: litellm.set_verbose = True litellm.num_retries = 3 image_path = "../proxy/cached_logo.jpg" # Getting the base64 string base64_image = encode_image(image_path) resp = litellm.completion( model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", messages=[ { "role": "user", "content": [ {"type": "text", "text": "Whats in this image?"}, { "type": "image_url", "image_url": { "url": "data:image/jpeg;base64," + base64_image }, }, ], } ], ) prompt_tokens = resp.usage.prompt_tokens raise Exception("it worked!") except Exception as e: if "500 Internal error encountered.'" in str(e): pass else: pytest.fail(f"An exception occurred - {str(e)}") def test_provisioned_throughput(): try: litellm.set_verbose = True import botocore, json, io import botocore.session from botocore.stub import Stubber bedrock_client = botocore.session.get_session().create_client( "bedrock-runtime", region_name="us-east-1" ) expected_params = { "accept": "application/json", "body": '{"prompt": "\\n\\nHuman: Hello, how are you?\\n\\nAssistant: ", ' '"max_tokens_to_sample": 256}', "contentType": "application/json", "modelId": "provisioned-model-arn", } response_from_bedrock = { "body": io.StringIO( json.dumps( { "completion": " Here is a short poem about the sky:", "stop_reason": "max_tokens", "stop": None, } ) ), "contentType": "contentType", "ResponseMetadata": {"HTTPStatusCode": 200}, } with Stubber(bedrock_client) as stubber: stubber.add_response( "invoke_model", service_response=response_from_bedrock, expected_params=expected_params, ) response = litellm.completion( model="bedrock/anthropic.claude-instant-v1", model_id="provisioned-model-arn", messages=[{"content": "Hello, how are you?", "role": "user"}], aws_bedrock_client=bedrock_client, ) print("response stubbed", response) except Exception as e: pytest.fail(f"Error occurred: {e}") # test_provisioned_throughput() def test_completion_bedrock_mistral_completion_auth(): print("calling bedrock mistral 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: ModelResponse = completion( model="bedrock/mistral.mistral-7b-instruct-v0:2", messages=messages, max_tokens=10, temperature=0.1, ) # Add any assertions here to check the response assert len(response.choices) > 0 assert len(response.choices[0].message.content) > 0 # 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_mistral_completion_auth()