test_models.py

# What this tests ?
## Tests /models and /model/* endpoints

importpytest
importasyncio
importaiohttp
importos
importdotenv
fromdotenvimportload_dotenv

load_dotenv()

asyncdefgenerate_key(session, models=[]):
url="http://0.0.0.0:4000/key/generate"
headers= {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data= {
"models": models,
"duration": None,
 }

asyncwithsession.post(url, headers=headers, json=data) asresponse:
status=response.status
response_text=awaitresponse.text()

print(response_text)
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")
returnawaitresponse.json()


asyncdefget_models(session, key):
url="http://0.0.0.0:4000/models"
headers= {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
 }

asyncwithsession.get(url, headers=headers) asresponse:
status=response.status
response_text=awaitresponse.text()
print("response from /models")
print(response_text)
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")
returnawaitresponse.json()


@pytest.mark.asyncio
asyncdeftest_get_models():
asyncwithaiohttp.ClientSession() assession:
key_gen=awaitgenerate_key(session=session)
key=key_gen["key"]
awaitget_models(session=session, key=key)


asyncdefadd_models(session, model_id="123", model_name="azure-gpt-3.5", key="sk-1234", team_id=None):
url="http://0.0.0.0:4000/model/new"
headers= {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
 }

data= {
"model_name": model_name,
"litellm_params": {
"model": "azure/chatgpt-v-3",
"api_key": "os.environ/AZURE_API_KEY",
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"api_version": "2023-05-15",
 },
"model_info": {"id": model_id},
 }

ifteam_id:
data["model_info"]["team_id"] =team_id

asyncwithsession.post(url, headers=headers, json=data) asresponse:
status=response.status
response_text=awaitresponse.text()
print(f"Add models {response_text}")
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")

response_json=awaitresponse.json()
returnresponse_json

asyncdefupdate_model(session, model_id="123", model_name="azure-gpt-3.5", key="sk-1234"):
url="http://0.0.0.0:4000/model/update"
headers= {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
 }

data= {
"model_name": model_name,
"litellm_params": {
"model": "azure/chatgpt-v-3",
"api_key": "os.environ/AZURE_API_KEY",
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"api_version": "2023-05-15",
 },
"model_info": {"id": model_id},
 }

asyncwithsession.post(url, headers=headers, json=data) asresponse:
status=response.status
response_text=awaitresponse.text()
print(f"Add models {response_text}")
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")

response_json=awaitresponse.json()
returnresponse_json


asyncdefget_model_info(session, key, litellm_model_id=None):
"""
 Make sure only models user has access to are returned
 """
iflitellm_model_id:
url=f"http://0.0.0.0:4000/model/info?litellm_model_id={litellm_model_id}"
else:
url="http://0.0.0.0:4000/model/info"
headers= {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
 }

asyncwithsession.get(url, headers=headers) asresponse:
status=response.status
response_text=awaitresponse.text()
print(response_text)
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")
returnawaitresponse.json()


asyncdefget_model_group_info(session, key):
url="http://0.0.0.0:4000/model_group/info"
headers= {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
 }

asyncwithsession.get(url, headers=headers) asresponse:
status=response.status
response_text=awaitresponse.text()
print(response_text)
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")
returnawaitresponse.json()


asyncdefchat_completion(session, key, model="azure-gpt-3.5"):
url="http://0.0.0.0:4000/chat/completions"
headers= {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
 }
data= {
"model": model,
"messages": [
 {"role": "system", "content": "You are a helpful assistant."},
 {"role": "user", "content": "Hello!"},
 ],
 }

asyncwithsession.post(url, headers=headers, json=data) asresponse:
status=response.status
response_text=awaitresponse.text()

print(response_text)
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")


@pytest.mark.asyncio
asyncdeftest_get_models():
"""
 Get models user has access to
 """
asyncwithaiohttp.ClientSession() assession:
key_gen=awaitgenerate_key(session=session, models=["gpt-4"])
key=key_gen["key"]
response=awaitget_model_info(session=session, key=key)
models= [m["model_name"] forminresponse["data"]]
forminmodels:
assertm=="gpt-4"


@pytest.mark.asyncio
asyncdeftest_get_specific_model():
"""
 Return specific model info

 Ensure value of model_info is same as on `/model/info` (no id set)
 """
asyncwithaiohttp.ClientSession() assession:
key_gen=awaitgenerate_key(session=session, models=["gpt-4"])
key=key_gen["key"]
response=awaitget_model_info(session=session, key=key)
models= [m["model_name"] forminresponse["data"]]
model_specific_info=None
foridx, minenumerate(models):
assertm=="gpt-4"
litellm_model_id=response["data"][idx]["model_info"]["id"]
model_specific_info=response["data"][idx]
assertlitellm_model_idisnotNone
response=awaitget_model_info(
session=session, key=key, litellm_model_id=litellm_model_id
 )
assertresponse["data"][0]["model_info"]["id"] ==litellm_model_id
assert (
response["data"][0] ==model_specific_info
 ), "Model info is not the same. Got={}, Expected={}".format(
response["data"][0], model_specific_info
 )


asyncdefdelete_model(session, model_id="123", key="sk-1234"):
"""
 Make sure only models user has access to are returned
 """
url="http://0.0.0.0:4000/model/delete"
headers= {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
 }
data= {"id": model_id}

asyncwithsession.post(url, headers=headers, json=data) asresponse:
status=response.status
response_text=awaitresponse.text()
print(response_text)
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")
returnawaitresponse.json()


@pytest.mark.asyncio
asyncdeftest_add_and_delete_models():
"""
 - Add model
 - Call new model -> expect to pass
 - Delete model
 - Call model -> expect to fail
 """
importuuid

asyncwithaiohttp.ClientSession() assession:
key_gen=awaitgenerate_key(session=session)
key=key_gen["key"]
model_id=f"12345_{uuid.uuid4()}"
model_name=f"{uuid.uuid4()}"
response=awaitadd_models(
session=session, model_id=model_id, model_name=model_name
 )
assertresponse["model_id"] ==model_id
awaitasyncio.sleep(10)
awaitchat_completion(session=session, key=key, model=model_name)
awaitdelete_model(session=session, model_id=model_id)
try:
awaitchat_completion(session=session, key=key, model=model_name)
pytest.fail(f"Expected call to fail.")
exceptException:
pass


asyncdefadd_model_for_health_checking(session, model_id="123"):
url="http://0.0.0.0:4000/model/new"
headers= {
"Authorization": f"Bearer sk-1234",
"Content-Type": "application/json",
 }

data= {
"model_name": f"azure-model-health-check-{model_id}",
"litellm_params": {
"model": "azure/chatgpt-v-3",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"api_version": "2023-05-15",
 },
"model_info": {"id": model_id},
 }

asyncwithsession.post(url, headers=headers, json=data) asresponse:
status=response.status
response_text=awaitresponse.text()

print(f"Add models {response_text}")
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")


asyncdefget_model_info_v2(session, key):
url="http://0.0.0.0:4000/v2/model/info"
headers= {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
 }

asyncwithsession.get(url, headers=headers) asresponse:
status=response.status
response_text=awaitresponse.text()
print("response from v2/model/info")
print(response_text)
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")


asyncdefget_specific_model_info_v2(session, key, model_name):
url="http://0.0.0.0:4000/v2/model/info?debug=True&model="+model_name
print("running /model/info check for model=", model_name)

headers= {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
 }

asyncwithsession.get(url, headers=headers) asresponse:
status=response.status
response_text=awaitresponse.text()
print("response from v2/model/info")
print(response_text)
print()

_json_response=awaitresponse.json()
print("JSON response from /v2/model/info?model=", model_name, _json_response)

_model_info=_json_response["data"]
assertlen(_model_info) ==1, f"Expected 1 model, got {len(_model_info)}"

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")
return_model_info[0]


asyncdefget_model_health(session, key, model_name):
url="http://0.0.0.0:4000/health?model="+model_name
headers= {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
 }

asyncwithsession.get(url, headers=headers) asresponse:
status=response.status
response_text=awaitresponse.json()
print("response from /health?model=", model_name)
print(response_text)
print()

ifstatus!=200:
raiseException(f"Request did not return a 200 status code: {status}")
returnresponse_text


@pytest.mark.asyncio
asyncdeftest_add_model_run_health():
"""
 Add model
 Call /model/info and v2/model/info
 -> Admin UI calls v2/model/info
 Call /chat/completions
 Call /health
 -> Ensure the health check for the endpoint is working as expected
 """
importuuid

asyncwithaiohttp.ClientSession() assession:
key_gen=awaitgenerate_key(session=session)
key=key_gen["key"]
master_key="sk-1234"
model_id=str(uuid.uuid4())
model_name=f"azure-model-health-check-{model_id}"
print("adding model", model_name)
awaitadd_model_for_health_checking(session=session, model_id=model_id)
_old_model_info=awaitget_specific_model_info_v2(
session=session, key=key, model_name=model_name
 )
print("model info before test", _old_model_info)

awaitasyncio.sleep(30)
print("calling /model/info")
awaitget_model_info(session=session, key=key)
print("calling v2/model/info")
awaitget_model_info_v2(session=session, key=key)

print("calling /chat/completions -> expect to work")
awaitchat_completion(session=session, key=key, model=model_name)

print("calling /health?model=", model_name)
_health_info=awaitget_model_health(
session=session, key=master_key, model_name=model_name
 )
_healthy_endpooint=_health_info["healthy_endpoints"][0]

assert_health_info["healthy_count"] ==1
assert (
_healthy_endpooint["model"] =="azure/chatgpt-v-3"
 ) # this is the model that got added

# assert httpx client is is unchanges

awaitasyncio.sleep(10)

_model_info_after_test=awaitget_specific_model_info_v2(
session=session, key=key, model_name=model_name
 )

print("model info after test", _model_info_after_test)
old_openai_client=_old_model_info["openai_client"]
new_openai_client=_model_info_after_test["openai_client"]
print("old openai client", old_openai_client)
print("new openai client", new_openai_client)

"""
 PROD TEST - This is extremly important 
 The OpenAI client used should be the same after 30 seconds
 It is a serious bug if the openai client does not match here
 """
assert (
old_openai_client==new_openai_client
 ), "OpenAI client does not match for the same model after 30 seconds"

# cleanup
awaitdelete_model(session=session, model_id=model_id)

@pytest.mark.asyncio
asyncdeftest_get_personal_models_for_user():
"""
 Test /models endpoint with team
 """
fromtest_usersimportnew_user
asyncwithaiohttp.ClientSession() assession:
# Creat a user
user_data=awaitnew_user(session=session, i=0, models=["gpt-3.5-turbo"])
user_id=user_data["user_id"]
user_api_key=user_data["key"]

model_group_info=awaitget_model_group_info(session=session, key=user_api_key)
print(model_group_info)

assertlen(model_group_info["data"]) ==1
assertmodel_group_info["data"][0]["model_group"] =="gpt-3.5-turbo"

@pytest.mark.asyncio
asyncdeftest_model_group_info_e2e():
"""
 Test /model/group/info endpoint
 """
asyncwithaiohttp.ClientSession() assession:
models=awaitget_models(session=session, key="sk-1234")
print(models)

expected_models= [
"anthropic/claude-3-5-haiku-20241022",
"anthropic/claude-3-opus-20240229",
 ]

model_group_info=awaitget_model_group_info(session=session, key="sk-1234")
print(model_group_info)

has_anthropic_claude_3_5_haiku=False
has_anthropic_claude_3_opus=False
formodelinmodel_group_info["data"]:
ifmodel["model_group"] =="anthropic/claude-3-5-haiku-20241022":
has_anthropic_claude_3_5_haiku=True
ifmodel["model_group"] =="anthropic/claude-3-opus-20240229":
has_anthropic_claude_3_opus=True

asserthas_anthropic_claude_3_5_haikuandhas_anthropic_claude_3_opus


@pytest.mark.asyncio
asyncdeftest_team_model_e2e():
"""
 Test team model e2e

 - create team
 - create user
 - add user to team as admin
 - add model to team
 - update model
 - delete model
 """
fromtest_usersimportnew_user
fromtest_teamimportnew_team
importuuid
asyncwithaiohttp.ClientSession() assession:
# Creat a user
user_data=awaitnew_user(session=session, i=0)
user_id=user_data["user_id"]
user_api_key=user_data["key"]

# Create a team
member_list= [
 {"role": "admin", "user_id": user_id},
 ]
team_data=awaitnew_team(session=session, member_list=member_list, i=0)
team_id=team_data["team_id"]

model_id=str(uuid.uuid4())
model_name="my-test-model"
# Add model to team
model_data=awaitadd_models(session=session, model_id=model_id, model_name=model_name, key=user_api_key, team_id=team_id)
model_id=model_data["model_id"]

# Update model
model_data=awaitupdate_model(session=session, model_id=model_id, model_name=model_name, key=user_api_key)
model_id=model_data["model_id"]

# Delete model
awaitdelete_model(session=session, model_id=model_id, key=user_api_key)