(test) s3 logs for /chat/completions

This commit is contained in:
ishaan-jaff
2024-01-11 09:16:06 +05:30
parent df0f689027
commit f263cf51ea
+43 -20
View File
@@ -18,6 +18,7 @@ import pytest
def test_s3_logging():
# all s3 requests need to be in one test function
# since we are modifying stdout, and pytests runs tests in parallel
# on circle ci - we only test litellm.acompletion()
try:
# pre
# redirect stdout to log_file
@@ -25,7 +26,9 @@ def test_s3_logging():
litellm.success_callback = ["s3"]
litellm.set_verbose = True
print("Testing async dynamoDB logging")
print("Testing async s3 logging")
expected_keys = []
async def _test():
return await litellm.acompletion(
@@ -38,30 +41,50 @@ def test_s3_logging():
response = asyncio.run(_test())
print(f"response: {response}")
expected_keys.append(response.id)
# streaming + async
async def _test2():
response = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "what llm are u"}],
max_tokens=10,
temperature=0.7,
user="ishaan-2",
stream=True,
)
async for chunk in response:
pass
# # streaming + async
# async def _test2():
# response = await litellm.acompletion(
# model="gpt-3.5-turbo",
# messages=[{"role": "user", "content": "what llm are u"}],
# max_tokens=10,
# temperature=0.7,
# user="ishaan-2",
# stream=True,
# )
# async for chunk in response:
# pass
asyncio.run(_test2())
# asyncio.run(_test2())
# aembedding()
async def _test3():
return await litellm.aembedding(
model="text-embedding-ada-002", input=["hi"], user="ishaan-2"
)
# async def _test3():
# return await litellm.aembedding(
# model="text-embedding-ada-002", input=["hi"], user="ishaan-2"
# )
response = asyncio.run(_test3())
time.sleep(1)
# response = asyncio.run(_test3())
# expected_keys.append(response.id)
# time.sleep(1)
import boto3
s3 = boto3.client("s3")
bucket_name = "litellm-logs"
# List objects in the bucket
response = s3.list_objects(Bucket=bucket_name)
# Sort the objects based on the LastModified timestamp
objects = sorted(
response["Contents"], key=lambda x: x["LastModified"], reverse=True
)
# Get the keys of the most recent objects
most_recent_keys = [obj["Key"] for obj in objects]
print("\n most recent keys", most_recent_keys)
print("\n Expected keys: ", expected_keys)
for key in expected_keys:
assert key in most_recent_keys
except Exception as e:
pytest.fail(f"An exception occurred - {e}")
finally: