(test) dynamo db writes

This commit is contained in:
ishaan-jaff 2023-12-15 16:43:33 +05:30
parent 6b0f61f29b
commit 15b1e758ab
3 changed files with 70 additions and 4 deletions

1
.gitignore vendored
View file

@ -22,4 +22,5 @@ litellm/tests/config_*.yaml
litellm/tests/langfuse.log litellm/tests/langfuse.log
litellm/tests/test_custom_logger.py litellm/tests/test_custom_logger.py
litellm/tests/langfuse.log litellm/tests/langfuse.log
litellm/tests/dynamo.log
.vscode/settings.json .vscode/settings.json

View file

@ -27,6 +27,7 @@ class DyanmoDBLogger:
def check_table_exists(self): def check_table_exists(self):
existing_tables = self.dynamodb.meta.client.list_tables()['TableNames'] existing_tables = self.dynamodb.meta.client.list_tables()['TableNames']
print_verbose(f"Dynamo DB: Existing Tables= {existing_tables}")
return self.table_name in existing_tables return self.table_name in existing_tables
@ -66,7 +67,6 @@ class DyanmoDBLogger:
id = response_obj.get("id", str(uuid.uuid4())) id = response_obj.get("id", str(uuid.uuid4()))
# convert all optional params to str # convert all optional params to str
for param, value in optional_params.items(): for param, value in optional_params.items():
try: try:
@ -95,7 +95,8 @@ class DyanmoDBLogger:
table = self.dynamodb.Table(self.table_name) table = self.dynamodb.Table(self.table_name)
# Assuming log_data is a dictionary with log information # Assuming log_data is a dictionary with log information
response = table.put_item(Item=payload) response = table.put_item(Item=payload)
print(f'Log data added to {self.table_name} successfully:', response)
print_verbose(f"Response from DynamoDB:{str(response)}")
print_verbose( print_verbose(
f"DynamoDB Layer Logging - final response object: {response_obj}" f"DynamoDB Layer Logging - final response object: {response_obj}"

View file

@ -10,12 +10,66 @@ import litellm
litellm.num_retries = 3 litellm.num_retries = 3
litellm.success_callback = ["dynamodb"] litellm.success_callback = ["dynamodb"]
litellm.set_verbose = True
import time import time
import pytest import pytest
def verify_dynamo_logs():
num_requests = 2
pass
def pre_request():
log_file = open("dynamo.log", "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
import re
def verify_log_file(log_file_path):
with open(log_file_path, 'r') as log_file:
log_content = log_file.read()
# 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 == 5
def test_dynamo_logging_async(): def test_dynamo_logging_async():
try: try:
litellm.set_verbose = True # pre
original_stdout, log_file = pre_request()
async def _test(): async def _test():
return await litellm.acompletion( return await litellm.acompletion(
model="gpt-3.5-turbo", model="gpt-3.5-turbo",
@ -26,12 +80,22 @@ def test_dynamo_logging_async():
) )
response = asyncio.run(_test()) response = asyncio.run(_test())
print(f"response: {response}") print(f"response: {response}")
time.sleep(1)
except litellm.Timeout as e: except litellm.Timeout as e:
pass pass
except Exception as e: except Exception as e:
pytest.fail(f"An exception occurred - {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("dynamo.log")
# test_dynamo_logging_async()
test_dynamo_logging_async()
def test_dynamo_logging_async_stream(): def test_dynamo_logging_async_stream():