test_spend_logs.py

# What this tests?
## Tests /spend endpoints.

importpytest, time, uuid, json
importasyncio
importaiohttp


asyncdefgenerate_key(session, models=[], team_id=None):
url="http://0.0.0.0:4000/key/generate"
headers= {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data= {
"models": models,
"duration": None,
 }
ifteam_idisnotNone:
data["team_id"] =team_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()


asyncdefchat_completion(session, key, model="gpt-3.5-turbo"):
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": f"Hello! {uuid.uuid4()}"},
 ],
 }

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()


asyncdefchat_completion_high_traffic(session, key, model="gpt-3.5-turbo"):
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": f"Hello! {uuid.uuid4()}"},
 ],
 }
try:
asyncwithsession.post(url, headers=headers, json=data) asresponse:
status=response.status
response_text=awaitresponse.text()

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

returnawaitresponse.json()
exceptExceptionase:
returnNone


asyncdefget_spend_logs(session, request_id=None, api_key=None):
ifapi_keyisnotNone:
url=f"http://0.0.0.0:4000/spend/logs?api_key={api_key}"
else:
url=f"http://0.0.0.0:4000/spend/logs?request_id={request_id}"
headers= {"Authorization": "Bearer sk-1234", "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()


@pytest.mark.asyncio
asyncdeftest_spend_logs():
"""
 - Create key
 - Make call (makes sure it's in spend logs)
 - Get request id from logs
 """
asyncwithaiohttp.ClientSession() assession:
key_gen=awaitgenerate_key(session=session)
key=key_gen["key"]
response=awaitchat_completion(session=session, key=key)
awaitasyncio.sleep(20)
awaitget_spend_logs(session=session, request_id=response["id"])


asyncdefgenerate_org(session: aiohttp.ClientSession) ->dict:
"""
 Generate a new organization using the API.

 Args:
 session: aiohttp client session

 Returns:
 dict: Response containing org_id
 """
url="http://0.0.0.0:4000/organization/new"
headers= {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}

request_body= {
"organization_alias": f"test-org-{uuid.uuid4()}",
 }

asyncwithsession.post(url, headers=headers, json=request_body) asresponse:
returnawaitresponse.json()


asyncdefgenerate_team(session: aiohttp.ClientSession, org_id: str) ->dict:
"""
 Generate a new team within an organization using the API.

 Args:
 session: aiohttp client session
 org_id: Organization ID to create the team in

 Returns:
 dict: Response containing team_id
 """
url="http://0.0.0.0:4000/team/new"
headers= {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data= {"organization_id": org_id}

asyncwithsession.post(url, headers=headers, json=data) asresponse:
returnawaitresponse.json()


@pytest.mark.asyncio
asyncdeftest_spend_logs_with_org_id():
"""
 - Create Organization
 - Create Team in organization
 - Create Key in organization
 - Make call (makes sure it's in spend logs)
 - Get request id from logs
 - Assert spend logs have correct org_id and team_id
 """
asyncwithaiohttp.ClientSession() assession:
org_gen=awaitgenerate_org(session=session)
print("org_gen: ", json.dumps(org_gen, indent=4, default=str))
org_id=org_gen["organization_id"]
team_gen=awaitgenerate_team(session=session, org_id=org_id)
print("team_gen: ", json.dumps(team_gen, indent=4, default=str))
team_id=team_gen["team_id"]
key_gen=awaitgenerate_key(session=session, team_id=team_id)
print("key_gen: ", json.dumps(key_gen, indent=4, default=str))
key=key_gen["key"]
response=awaitchat_completion(session=session, key=key)
awaitasyncio.sleep(20)
spend_logs_response=awaitget_spend_logs(
session=session, request_id=response["id"]
 )
print(
"spend_logs_response: ",
json.dumps(spend_logs_response, indent=4, default=str),
 )
spend_logs_response=spend_logs_response[0]
assertspend_logs_response["metadata"]["user_api_key_org_id"] ==org_id
assertspend_logs_response["metadata"]["user_api_key_team_id"] ==team_id
assertspend_logs_response["team_id"] ==team_id


asyncdefget_predict_spend_logs(session):
url="http://0.0.0.0:4000/global/predict/spend/logs"
headers= {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data= {
"data": [
 {
"date": "2024-03-09",
"spend": 200000,
"api_key": "f19bdeb945164278fc11c1020d8dfd70465bffd931ed3cb2e1efa6326225b8b7",
 }
 ]
 }

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_spend_report(session, start_date, end_date):
url="http://0.0.0.0:4000/global/spend/report"
headers= {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
asyncwithsession.get(
url, headers=headers, params={"start_date": start_date, "end_date": end_date}
 ) 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.skip(reason="datetime in ci/cd gets set weirdly")
@pytest.mark.asyncio
asyncdeftest_get_predicted_spend_logs():
"""
 - Create key
 - Make call (makes sure it's in spend logs)
 - Get request id from logs
 """
asyncwithaiohttp.ClientSession() assession:
result=awaitget_predict_spend_logs(session=session)
print(result)

assert"response"inresult
assertlen(result["response"]) >0


@pytest.mark.skip(reason="High traffic load test, meant to be run locally")
@pytest.mark.asyncio
asyncdeftest_spend_logs_high_traffic():
"""
 - Create key
 - Make 30 concurrent calls
 - Get all logs for that key
 - Wait 10s
 - Assert it's 30
 """

asyncdefretry_request(func, *args, _max_attempts=5, **kwargs):
forattemptinrange(_max_attempts):
try:
returnawaitfunc(*args, **kwargs)
except (
aiohttp.client_exceptions.ClientOSError,
aiohttp.client_exceptions.ServerDisconnectedError,
 ) ase:
ifattempt+1==_max_attempts:
raise# re-raise the last ClientOSError if all attempts failed
print(f"Attempt {attempt+1} failed, retrying...")

asyncwithaiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=600)
 ) assession:
start=time.time()
key_gen=awaitgenerate_key(session=session)
key=key_gen["key"]
n=1000
tasks= [
retry_request(
chat_completion_high_traffic,
session=session,
key=key,
model="azure-gpt-3.5",
 )
for_inrange(n)
 ]
chat_completions=awaitasyncio.gather(*tasks)
successful_completions= [cforcinchat_completionsifcisnotNone]
print(f"Num successful completions: {len(successful_completions)}")
awaitasyncio.sleep(10)
try:
response=awaitretry_request(get_spend_logs, session=session, api_key=key)
print(f"response: {response}")
print(f"len responses: {len(response)}")
assertlen(response) ==n
print(n, time.time() -start, len(response))
exceptException:
print(n, time.time() -start, 0)
raiseException("it worked!")


@pytest.mark.asyncio
asyncdeftest_spend_report_endpoint():
asyncwithaiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=600)
 ) assession:
importdatetime

todays_date=datetime.date.today() +datetime.timedelta(days=1)
todays_date=todays_date.strftime("%Y-%m-%d")

print("todays_date", todays_date)
thirty_days_ago= (
datetime.date.today() -datetime.timedelta(days=30)
 ).strftime("%Y-%m-%d")
spend_report=awaitget_spend_report(
session=session, start_date=thirty_days_ago, end_date=todays_date
 )
print("spend report", spend_report)

forrowinspend_report:
date=row["group_by_day"]
teams=row["teams"]
forteaminteams:
team_name=team["team_name"]
total_spend=team["total_spend"]
metadata=team["metadata"]

assertteam_nameisnotNone

print(f"Date: {date}")
print(f"Team: {team_name}")
print(f"Total Spend: {total_spend}")
print("Metadata: ", metadata)
print()