forked from phoenix/litellm-mirror
402 lines
13 KiB
Python
402 lines
13 KiB
Python
# @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()
|