Siddhant Medar
ab71bd8dae
Task/unit tests investment and client * move io_utils.py * add __init__.py * refactor: move string_funcs and update import paths to use utils namespace * rename string_funcs to string_utils * move and rename table_utils.py * refactor: rename claude_funcs to llm_utils in multiple files * update unit tests * added unit tests for src/string_utils.py * added unit tests for src/table_utils.py * added unit tests for src/llm_utils.py * added unit tests for src/io_utils.py * Merge branch 'main' into task/unit-tests-investment-and-client; TODO: fix unit tests * Unit tests to fix * fix unit tests * fix invalid references Approved-by: Alex Galarce
167 lines
7.4 KiB
Python
167 lines
7.4 KiB
Python
import pytest
|
|
from src.utils.llm_utils import (
|
|
count_tokens,
|
|
create_cost_log_entry,
|
|
update_stats,
|
|
invoke_claude,
|
|
)
|
|
from src import config
|
|
|
|
class TestLLMUtils:
|
|
# Fixture to reset global stats before each test
|
|
@pytest.fixture(autouse=True)
|
|
def reset_global_stats(self):
|
|
"""Fixture to reset global stats before each test."""
|
|
config.MODEL_STATS = {
|
|
"file1.txt": {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0},
|
|
"file2.txt": {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0},
|
|
"file3.txt": {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0},
|
|
"file4.txt": {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0},
|
|
}
|
|
config.GLOBAL_STATS = {
|
|
"total_requests": 0,
|
|
"total_tokens_sent": 0,
|
|
"total_tokens_received": 0,
|
|
"total_time": 0,
|
|
}
|
|
|
|
# Test cases for count_tokens
|
|
@pytest.mark.parametrize("input_str, expected_tokens", [
|
|
("hello world", 2), # 11 characters / 6 = 1.83 → ceil to 2
|
|
("", 0), # Empty string
|
|
("a" * 12, 2), # 12 characters / 6 = 2
|
|
("a" * 13, 3), # 13 characters / 6 = 2.16 → ceil to 3
|
|
("😊", 1), # Unicode characters
|
|
])
|
|
def test_count_tokens(self, input_str, expected_tokens):
|
|
"""Test the count_tokens function."""
|
|
assert count_tokens(input_str) == expected_tokens
|
|
|
|
# Test cases for create_cost_log_entry
|
|
@pytest.mark.parametrize("filename, caller_name, prompt, response, input_cost_per_1k, output_cost_per_1k, expected_output", [
|
|
(
|
|
"file1.txt", "func1", "hello", "world", 0.1, 0.2,
|
|
{
|
|
"Filename": "file1.txt",
|
|
"Caller Function": "func1",
|
|
"Input Prompt Length": 5,
|
|
"Input Prompt Tokens": 1, # 5 chars / 6 = 0.83 → ceil to 1
|
|
"Input Cost": 0.1 * 1 / 1000,
|
|
"Output Response Length": 5,
|
|
"Output Response Tokens": 1, # 5 chars / 6 = 0.83 → ceil to 1
|
|
"Output Cost": 0.2 * 1 / 1000,
|
|
"Total Cost": (0.1 * 1 / 1000) + (0.2 * 1 / 1000),
|
|
}
|
|
),
|
|
(
|
|
"file2.txt", "func2", "", "", 0.1, 0.2,
|
|
{
|
|
"Filename": "file2.txt",
|
|
"Caller Function": "func2",
|
|
"Input Prompt Length": 0,
|
|
"Input Prompt Tokens": 0, # Empty string
|
|
"Input Cost": 0.1 * 0 / 1000,
|
|
"Output Response Length": 0,
|
|
"Output Response Tokens": 0, # Empty string
|
|
"Output Cost": 0.2 * 0 / 1000,
|
|
"Total Cost": (0.1 * 0 / 1000) + (0.2 * 0 / 1000),
|
|
}
|
|
),
|
|
(
|
|
"file3.txt", "func3", "a" * 12, "b" * 13, 0.1, 0.2,
|
|
{
|
|
"Filename": "file3.txt",
|
|
"Caller Function": "func3",
|
|
"Input Prompt Length": 12,
|
|
"Input Prompt Tokens": 2, # 12 chars / 6 = 2
|
|
"Input Cost": 0.1 * 2 / 1000,
|
|
"Output Response Length": 13,
|
|
"Output Response Tokens": 3, # 13 chars / 6 = 2.16 → ceil to 3
|
|
"Output Cost": 0.2 * 3 / 1000,
|
|
"Total Cost": (0.1 * 2 / 1000) + (0.2 * 3 / 1000),
|
|
}
|
|
),
|
|
(
|
|
"file4.txt", "func4", "😊", "😊😊", 0.1, 0.2,
|
|
{
|
|
"Filename": "file4.txt",
|
|
"Caller Function": "func4",
|
|
"Input Prompt Length": 1,
|
|
"Input Prompt Tokens": 1, # Unicode character
|
|
"Input Cost": 0.1 * 1 / 1000,
|
|
"Output Response Length": 2,
|
|
"Output Response Tokens": 1, # Unicode characters
|
|
"Output Cost": 0.2 * 1 / 1000,
|
|
"Total Cost": (0.1 * 1 / 1000) + (0.2 * 1 / 1000),
|
|
}
|
|
),
|
|
])
|
|
|
|
def test_create_cost_log_entry(self, filename, caller_name, prompt, response, input_cost_per_1k, output_cost_per_1k, expected_output):
|
|
"""Test the create_cost_log_entry function."""
|
|
result = create_cost_log_entry(filename, caller_name, prompt, response, input_cost_per_1k, output_cost_per_1k)
|
|
assert result == expected_output
|
|
|
|
# Test cases for update_stats
|
|
def test_update_stats(self):
|
|
"""Test the update_stats function."""
|
|
# Initialize the global state to zeros
|
|
config.MODEL_STATS = {"file1.txt": {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0}}
|
|
config.GLOBAL_STATS = {"total_requests": 0, "total_tokens_sent": 0, "total_tokens_received": 0, "total_time": 0}
|
|
|
|
# Verify initial state
|
|
assert config.MODEL_STATS["file1.txt"]["requests"] == 0
|
|
assert config.GLOBAL_STATS["total_requests"] == 0
|
|
|
|
# Update stats
|
|
update_stats("file1.txt", 10, 20, 5.0, "claude2")
|
|
|
|
# Check updated state
|
|
assert config.MODEL_STATS["file1.txt"]["requests"] == 1
|
|
assert config.MODEL_STATS["file1.txt"]["tokens_sent"] == 10
|
|
assert config.MODEL_STATS["file1.txt"]["tokens_received"] == 20
|
|
assert config.MODEL_STATS["file1.txt"]["total_time"] == 5.0
|
|
assert config.GLOBAL_STATS["total_requests"] == 1
|
|
assert config.GLOBAL_STATS["total_tokens_sent"] == 10
|
|
assert config.GLOBAL_STATS["total_tokens_received"] == 20
|
|
assert config.GLOBAL_STATS["total_time"] == 5.0
|
|
|
|
# Test cases for invoke_claude
|
|
@pytest.mark.parametrize("run_mode, model_id, expected_response", [
|
|
("local", config.MODEL_ID_CLAUDE2, "mocked response from claude 2 (local)"),
|
|
("local", config.MODEL_ID_CLAUDE3_HAIKU, "mocked response from claude 3 (local)"),
|
|
("local", config.MODEL_ID_CLAUDE35_SONNET, "mocked response from claude 3.5 (local)"),
|
|
("ec2", config.MODEL_ID_CLAUDE2, "mocked response from claude 2 (ec2)"),
|
|
("ec2", config.MODEL_ID_CLAUDE3_HAIKU, "mocked response from claude 3 (ec2)"),
|
|
("ec2", config.MODEL_ID_CLAUDE35_SONNET, "mocked response from claude 3.5 (ec2)"),
|
|
])
|
|
|
|
def test_invoke_claude(self, run_mode, model_id, expected_response, mocker):
|
|
"""Test the invoke_claude function."""
|
|
# Map model IDs to their corresponding mock functions
|
|
mock_functions = {
|
|
("local", config.MODEL_ID_CLAUDE2): "src.utils.llm_utils.local_claude_2",
|
|
("local", config.MODEL_ID_CLAUDE3_HAIKU): "src.utils.llm_utils.local_claude_3",
|
|
("local", config.MODEL_ID_CLAUDE35_SONNET): "src.utils.llm_utils.local_claude_3",
|
|
("ec2", config.MODEL_ID_CLAUDE2): "src.utils.llm_utils.ec2_claude2",
|
|
("ec2", config.MODEL_ID_CLAUDE3_HAIKU): "src.utils.llm_utils.ec2_claude_3",
|
|
("ec2", config.MODEL_ID_CLAUDE35_SONNET): "src.utils.llm_utils.ec2_claude_35",
|
|
}
|
|
|
|
# Mock the appropriate function based on run_mode and model_id
|
|
mock_function_path = mock_functions.get((run_mode, model_id))
|
|
if mock_function_path:
|
|
mocker.patch(mock_function_path, return_value=expected_response)
|
|
|
|
# Mock config.RUN_MODE to simulate the run mode
|
|
mocker.patch("src.utils.llm_utils.config.RUN_MODE", run_mode)
|
|
|
|
# Mock CSV writer to avoid writing to disk
|
|
mocker.patch("csv.DictWriter.writerow")
|
|
|
|
# Call the function
|
|
response = invoke_claude("test prompt", model_id, "file1.txt")
|
|
|
|
# Assertions
|
|
assert response == expected_response
|