Spaces:

ka1kuk
/

litellm

Running

App Files Files Community

litellm / litellm /tests /test_dynamodb_logs.py

ka1kuk

Upload 235 files

7db0ae4 verified 5 months ago

raw history blame contribute delete

No virus

3.78 kB

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