import asyncio
import sys
from loguru import logger
from swarms.utils.litellm_wrapper import LiteLLM
# Configure loguru logger
logger.remove() # Remove default handler
logger.add(
"test_litellm.log",
rotation="1 MB",
format="{time} | {level} | {message}",
level="DEBUG",
)
logger.add(sys.stdout, level="INFO")
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Retrieve detailed current weather information for a specified location, including temperature, humidity, wind speed, and atmospheric conditions.",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA, or a specific geographic coordinate in the format 'latitude,longitude'.",
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit", "kelvin"],
"description": "The unit of temperature measurement to be used in the response.",
},
"include_forecast": {
"type": "boolean",
"description": "Indicates whether to include a short-term weather forecast along with the current conditions.",
},
"time": {
"type": "string",
"format": "date-time",
"description": "Optional parameter to specify the time for which the weather data is requested, in ISO 8601 format.",
},
},
"required": [
"location",
"unit",
"include_forecast",
"time",
],
},
},
}
]
# Initialize LiteLLM with streaming enabled
llm = LiteLLM(model_name="gpt-4o-mini", tools_list_dictionary=tools)
async def main():
# When streaming is enabled, arun returns a stream of chunks
# We need to handle these chunks directly, not try to access .choices
stream = await llm.arun("What is the weather in San Francisco?")
logger.info(f"Received stream from LLM. {stream}")
if stream is not None:
logger.info(f"Stream is not None. {stream}")
else:
logger.info("Stream is None.")
def run_test_suite():
"""Run all test cases and generate a comprehensive report."""
logger.info("Starting LiteLLM Test Suite")
total_tests = 0
passed_tests = 0
failed_tests = []
def log_test_result(test_name: str, passed: bool, error=None):
nonlocal total_tests, passed_tests
total_tests += 1
if passed:
passed_tests += 1
logger.success(f"✅ {test_name} - PASSED")
else:
failed_tests.append((test_name, error))
logger.error(f"❌ {test_name} - FAILED: {error}")
# Test 1: Basic Initialization
try:
logger.info("Testing basic initialization")
llm = LiteLLM()
assert llm.model_name == "gpt-4o"
assert llm.temperature == 0.5
assert llm.max_tokens == 4000
log_test_result("Basic Initialization", True)
except Exception as e:
log_test_result("Basic Initialization", False, str(e))
# Test 2: Custom Parameters
try:
logger.info("Testing custom parameters")
llm = LiteLLM(
model_name="gpt-3.5-turbo",
temperature=0.7,
max_tokens=2000,
system_prompt="You are a helpful assistant",
)
assert llm.model_name == "gpt-3.5-turbo"
assert llm.temperature == 0.7
assert llm.max_tokens == 2000
assert llm.system_prompt == "You are a helpful assistant"
log_test_result("Custom Parameters", True)
except Exception as e:
log_test_result("Custom Parameters", False, str(e))
# Test 3: Message Preparation
try:
logger.info("Testing message preparation")
llm = LiteLLM(system_prompt="Test system prompt")
messages = llm._prepare_messages("Test task")
assert len(messages) == 2
assert messages[0]["role"] == "system"
assert messages[0]["content"] == "Test system prompt"
assert messages[1]["role"] == "user"
assert messages[1]["content"] == "Test task"
log_test_result("Message Preparation", True)
except Exception as e:
log_test_result("Message Preparation", False, str(e))
# Test 4: Basic Completion
try:
logger.info("Testing basic completion")
llm = LiteLLM()
response = llm.run("What is 2+2?")
assert isinstance(response, str)
assert len(response) > 0
log_test_result("Basic Completion", True)
except Exception as e:
log_test_result("Basic Completion", False, str(e))
try:
# tool usage
asyncio.run(main())
except Exception as e:
log_test_result("Tool Usage", False, str(e))
# Test 5: Tool Calling
try:
logger.info("Testing tool calling")
tools = [
{
"type": "function",
"function": {
"name": "test_function",
"description": "A test function",
"parameters": {
"type": "object",
"properties": {"test": {"type": "string"}},
},
},
}
]
llm = LiteLLM(
tools_list_dictionary=tools,
tool_choice="auto",
model_name="gpt-4o-mini",
)
assert llm.tools_list_dictionary == tools
assert llm.tool_choice == "auto"
log_test_result("Tool Calling Setup", True)
except Exception as e:
log_test_result("Tool Calling Setup", False, str(e))
# Test 6: Async Completion
async def test_async():
try:
logger.info("Testing async completion")
llm = LiteLLM()
response = await llm.arun("What is 3+3?")
assert isinstance(response, str)
assert len(response) > 0
log_test_result("Async Completion", True)
except Exception as e:
log_test_result("Async Completion", False, str(e))
asyncio.run(test_async())
# Test 7: Batched Run
try:
logger.info("Testing batched run")
llm = LiteLLM()
tasks = ["Task 1", "Task 2", "Task 3"]
responses = llm.batched_run(tasks, batch_size=2)
assert isinstance(responses, list)
assert len(responses) == 3
log_test_result("Batched Run", True)
except Exception as e:
log_test_result("Batched Run", False, str(e))
# Generate test report
success_rate = (passed_tests / total_tests) * 100
logger.info("\n=== Test Suite Report ===")
logger.info(f"Total Tests: {total_tests}")
logger.info(f"Passed Tests: {passed_tests}")
logger.info(f"Failed Tests: {len(failed_tests)}")
logger.info(f"Success Rate: {success_rate:.2f}%")
if failed_tests:
logger.error("\nFailed Tests Details:")
for test_name, error in failed_tests:
logger.error(f"{test_name}: {error}")
return {
"total_tests": total_tests,
"passed_tests": passed_tests,
"failed_tests": failed_tests,
"success_rate": success_rate,
}
if __name__ == "__main__":
test_results = run_test_suite()
logger.info(
"Test suite completed. Check test_litellm.log for detailed logs."
)