run mock tests for test_completion_sagemaker

This commit is contained in:
Ishaan Jaff 2024-08-15 14:49:21 -07:00
parent bfa3278cde
commit d0bdafa49c

View file

@ -3337,33 +3337,6 @@ def test_customprompt_together_ai():
# test_customprompt_together_ai()
@pytest.mark.skip(reason="AWS Suspended Account")
def test_completion_sagemaker():
try:
litellm.set_verbose = True
print("testing sagemaker")
response = completion(
model="sagemaker/jumpstart-dft-hf-llm-mistral-7b-ins-20240329-150233",
model_id="huggingface-llm-mistral-7b-instruct-20240329-150233",
messages=messages,
temperature=0.2,
max_tokens=80,
aws_region_name=os.getenv("AWS_REGION_NAME_2"),
aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID_2"),
aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY_2"),
input_cost_per_second=0.000420,
)
# Add any assertions here to check the response
print(response)
cost = completion_cost(completion_response=response)
print("calculated cost", cost)
assert (
cost > 0.0 and cost < 1.0
) # should never be > $1 for a single completion call
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_sagemaker()