Alex Galarce
04aa447831
Remove cost logging functionality and related tests * Remove cost logging functionality and related tests * Remove cost log entry test and related CSV writer mock * Merged main into feature/remove-costlog Approved-by: Katon Minhas
63 lines
2.4 KiB
Python
63 lines
2.4 KiB
Python
import csv
|
|
import hashlib
|
|
import unittest
|
|
from collections import OrderedDict
|
|
from unittest.mock import mock_open, patch
|
|
|
|
import src.utils.llm_utils as llm_utils
|
|
from src import config
|
|
|
|
|
|
class TestLLMUtils(unittest.TestCase):
|
|
def setUp(self):
|
|
self.filename = "test_file.txt"
|
|
self.prompt = "This is a test prompt."
|
|
self.response = "This is a test response."
|
|
self.model_type = "test_model"
|
|
self.input_cost_per_1k = 0.01
|
|
self.output_cost_per_1k = 0.02
|
|
self.mock_stats = {
|
|
self.filename: {
|
|
"requests": 0,
|
|
"tokens_sent": 0,
|
|
"tokens_received": 0,
|
|
"total_time": 0,
|
|
}
|
|
}
|
|
config.MODEL_STATS = self.mock_stats
|
|
config.GLOBAL_STATS = {}
|
|
|
|
def test_update_stats(self):
|
|
llm_utils.update_stats(self.filename, 10, 20, 5, self.model_type)
|
|
self.assertEqual(config.MODEL_STATS[self.filename]["requests"], 1)
|
|
self.assertEqual(config.MODEL_STATS[self.filename]["tokens_sent"], 10)
|
|
self.assertEqual(config.MODEL_STATS[self.filename]["tokens_received"], 20)
|
|
self.assertEqual(config.MODEL_STATS[self.filename]["total_time"], 5)
|
|
self.assertEqual(config.GLOBAL_STATS["total_requests"], 1)
|
|
self.assertEqual(config.GLOBAL_STATS["total_tokens_sent"], 10)
|
|
self.assertEqual(config.GLOBAL_STATS["total_tokens_received"], 20)
|
|
self.assertEqual(config.GLOBAL_STATS["total_time"], 5)
|
|
|
|
def test_count_tokens(self):
|
|
self.assertEqual(llm_utils.count_tokens("123456"), 1)
|
|
self.assertEqual(llm_utils.count_tokens("1234567"), 2)
|
|
self.assertEqual(llm_utils.count_tokens(""), 0)
|
|
|
|
def test_get_cache_key(self):
|
|
model_id = "test_model"
|
|
expected_hash = hashlib.sha256(f"{model_id}:{self.prompt}".encode()).hexdigest()
|
|
self.assertEqual(llm_utils.get_cache_key(self.prompt, model_id), expected_hash)
|
|
|
|
@patch("builtins.open", new_callable=mock_open)
|
|
def test_invoke_claude_local_mode(self, mock_open_file):
|
|
config.RUN_MODE = "local"
|
|
model_id = config.MODEL_ID_CLAUDE2
|
|
|
|
with patch("src.utils.llm_utils.local_claude_2", return_value=self.response):
|
|
response = llm_utils.invoke_claude(self.prompt, model_id, self.filename)
|
|
self.assertEqual(response, self.response)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|