forked from phoenix/litellm-mirror
(test) dynamo db writes
This commit is contained in:
parent
6b0f61f29b
commit
15b1e758ab
3 changed files with 70 additions and 4 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
@ -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
|
||||||
|
|
|
@ -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}"
|
||||||
|
|
|
@ -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():
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue