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