File size: 3,780 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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
import sys
import os
import io, asyncio
# import logging
# logging.basicConfig(level=logging.DEBUG)
sys.path.insert(0, os.path.abspath("../.."))
from litellm import completion
import litellm
litellm.num_retries = 3
import time, random
import pytest
def pre_request():
file_name = f"dynamo.log"
log_file = open(file_name, "a+")
# Clear the contents of the file by truncating it
log_file.truncate(0)
# Save the original stdout so that we can restore it later
original_stdout = sys.stdout
# Redirect stdout to the file
sys.stdout = log_file
return original_stdout, log_file, file_name
import re
def verify_log_file(log_file_path):
with open(log_file_path, "r") as log_file:
log_content = log_file.read()
print(
f"\nVerifying DynamoDB file = {log_file_path}. File content=", log_content
)
# Define the pattern to search for in the log file
pattern = r"Response from DynamoDB:{.*?}"
# Find all matches in the log content
matches = re.findall(pattern, log_content)
# Print the DynamoDB success log matches
print("DynamoDB Success Log Matches:")
for match in matches:
print(match)
# Print the total count of lines containing the specified response
print(f"Total occurrences of specified response: {len(matches)}")
# Count the occurrences of successful responses (status code 200 or 201)
success_count = sum(
1
for match in matches
if "'HTTPStatusCode': 200" in match or "'HTTPStatusCode': 201" in match
)
# Print the count of successful responses
print(f"Count of successful responses from DynamoDB: {success_count}")
assert success_count == 3 # Expect 3 success logs from dynamoDB
def test_dynamo_logging():
# all dynamodb requests need to be in one test function
# since we are modifying stdout, and pytests runs tests in parallel
try:
# pre
# redirect stdout to log_file
litellm.success_callback = ["dynamodb"]
litellm.dynamodb_table_name = "litellm-logs-1"
litellm.set_verbose = True
original_stdout, log_file, file_name = pre_request()
print("Testing async dynamoDB logging")
async def _test():
return await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "This is a test"}],
max_tokens=100,
temperature=0.7,
user="ishaan-2",
)
response = asyncio.run(_test())
print(f"response: {response}")
# streaming + async
async def _test2():
response = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "This is a test"}],
max_tokens=10,
temperature=0.7,
user="ishaan-2",
stream=True,
)
async for chunk in response:
pass
asyncio.run(_test2())
# aembedding()
async def _test3():
return await litellm.aembedding(
model="text-embedding-ada-002", input=["hi"], user="ishaan-2"
)
response = asyncio.run(_test3())
time.sleep(1)
except Exception as e:
pytest.fail(f"An exception occurred - {e}")
finally:
# post, close log file and verify
# Reset stdout to the original value
sys.stdout = original_stdout
# Close the file
log_file.close()
verify_log_file(file_name)
print("Passed! Testing async dynamoDB logging")
# test_dynamo_logging_async()
|