File size: 2,893 Bytes
7db0ae4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
#### What this tests ####
# This tests if ahealth_check() actually works
import sys, os
import traceback
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm, asyncio
@pytest.mark.asyncio
async def test_azure_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/chatgpt-v-2",
"messages": [{"role": "user", "content": "Hey, how's it going?"}],
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
}
)
print(f"response: {response}")
assert "x-ratelimit-remaining-tokens" in response
return response
# asyncio.run(test_azure_health_check())
@pytest.mark.asyncio
async def test_azure_embedding_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/azure-embedding-model",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
},
input=["test for litellm"],
mode="embedding",
)
print(f"response: {response}")
assert "x-ratelimit-remaining-tokens" in response
return response
@pytest.mark.asyncio
async def test_openai_img_gen_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "dall-e-3",
"api_key": os.getenv("OPENAI_API_KEY"),
},
mode="image_generation",
prompt="cute baby sea otter",
)
print(f"response: {response}")
assert isinstance(response, dict) and "error" not in response
return response
# asyncio.run(test_openai_img_gen_health_check())
async def test_azure_img_gen_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/",
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": "2023-06-01-preview",
},
mode="image_generation",
prompt="cute baby sea otter",
)
assert isinstance(response, dict) and "error" not in response
return response
# asyncio.run(test_azure_img_gen_health_check())
@pytest.mark.asyncio
async def test_sagemaker_embedding_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "sagemaker/berri-benchmarking-gpt-j-6b-fp16",
"messages": [{"role": "user", "content": "Hey, how's it going?"}],
},
mode="embedding",
input=["test from litellm"],
)
print(f"response: {response}")
assert isinstance(response, dict)
return response
# asyncio.run(test_sagemaker_embedding_health_check())
|