Katon Minhas
c4e519894b
Refactor/one time io * Delete client work * Streamline imports * Fix list class * Update tests * Clear code funcs test * Fix unit tests * Single read code funcs * Single-read model * Single-load embeddings * Successful E2E test * update unit tests * Update importsg * Reload poetry.lock * remove old exhibit header function * Remove aarete_derived generic function * remove align and format tables * Remove strings to dict * references * Clear code * Move preprocessing_funcs * remove keywords * refactor postprocessing_funcs * Pass unit test - remove qa_qc directory * Black and isort * remove print Approved-by: Alex Galarce
80 lines
3.0 KiB
Python
80 lines
3.0 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_create_cost_log_entry(self):
|
|
entry = llm_utils.create_cost_log_entry(
|
|
self.filename,
|
|
"caller_func",
|
|
self.prompt,
|
|
self.response,
|
|
self.input_cost_per_1k,
|
|
self.output_cost_per_1k,
|
|
)
|
|
self.assertEqual(entry["Filename"], self.filename)
|
|
self.assertEqual(entry["Caller Function"], "caller_func")
|
|
self.assertEqual(entry["Input Prompt Length"], len(self.prompt))
|
|
self.assertEqual(entry["Output Response Length"], len(self.response))
|
|
self.assertIn("Total Cost", entry)
|
|
|
|
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)
|
|
@patch("csv.DictWriter.writerow")
|
|
def test_invoke_claude_local_mode(self, mock_csv_writer, 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)
|
|
mock_csv_writer.assert_called()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|