From 04aa44783151c7dead615e9cffcbdb18e762b22f Mon Sep 17 00:00:00 2001 From: Alex Galarce Date: Thu, 21 Aug 2025 17:45:18 +0000 Subject: [PATCH] Merged in feature/remove-costlog (pull request #675) 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 --- fieldExtraction/src/config.py | 20 +---- fieldExtraction/src/tracking/costlog_utils.py | 72 ------------------ .../src/tracking/tracking_utils.py | 1 - fieldExtraction/src/utils/llm_utils.py | 74 ------------------- fieldExtraction/tests/test_costlog_funcs.py | 9 --- fieldExtraction/tests/test_llm_utils.py | 19 +---- 6 files changed, 2 insertions(+), 193 deletions(-) delete mode 100644 fieldExtraction/src/tracking/costlog_utils.py delete mode 100644 fieldExtraction/tests/test_costlog_funcs.py diff --git a/fieldExtraction/src/config.py b/fieldExtraction/src/config.py index 221a110..12be7b9 100644 --- a/fieldExtraction/src/config.py +++ b/fieldExtraction/src/config.py @@ -59,25 +59,7 @@ def check_s3_bucket_exists(s3_client, bucket_name: str): TEST = False # True to run test prompt - just for testing model connection TODAY = datetime.now().strftime("%Y%m%d") -######################################## COST LOG ######################################## -COST_LOG_NAME = "costlog.csv" - -COST_LOG_FIELDS = [ - "Filename", - "Caller Function", - "Input Prompt Length", - "Input Prompt Tokens", - "Input Cost", - "Output Response Length", - "Output Response Tokens", - "Output Cost", - "Total Cost", -] - -if not os.path.exists(COST_LOG_NAME): - with open(COST_LOG_NAME, "w", newline="", encoding="utf-8") as f: - writer = csv.DictWriter(f, fieldnames=COST_LOG_FIELDS) - writer.writeheader() +######################################## TRACKING ######################################## UPLOAD_FREQUENCY = 10 diff --git a/fieldExtraction/src/tracking/costlog_utils.py b/fieldExtraction/src/tracking/costlog_utils.py deleted file mode 100644 index b3f8369..0000000 --- a/fieldExtraction/src/tracking/costlog_utils.py +++ /dev/null @@ -1,72 +0,0 @@ -import pandas as pd - -from src import config - - -def aggregate_costlog(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]: - """This function aggregates the costlog in two ways: - - by caller function - - by filename - And returns two dataframes, corresponding to each of these aggregation methods - - Args: - df (pd.DataFrame): Input costlog. Each row should be one call to the LLM. - The columns are determined by `config.COST_LOG_FIELDS` - - Returns: - tuple[pd.DataFrame, pd.DataFrame]: Aggregated dataframes by caller and filename - respectively. - """ - # check required input columns - missing_columns = set(config.COST_LOG_FIELDS) - set(df.columns) - if missing_columns: - raise ValueError(f"Missing required columns: {', '.join(missing_columns)}") - - group_by_caller_function = ( - df.groupby("Caller Function") - .agg( - count_invocations=pd.NamedAgg(column="Filename", aggfunc="count"), - unique_filenames=pd.NamedAgg(column="Filename", aggfunc="nunique"), - sum_input_prompt_length=pd.NamedAgg( - column="Input Prompt Length", aggfunc="sum" - ), - sum_input_prompt_tokens=pd.NamedAgg( - column="Input Prompt Tokens", aggfunc="sum" - ), - sum_input_cost=pd.NamedAgg(column="Input Cost", aggfunc="sum"), - sum_output_response_length=pd.NamedAgg( - column="Output Response Length", aggfunc="sum" - ), - sum_output_response_tokens=pd.NamedAgg( - column="Output Response Tokens", aggfunc="sum" - ), - sum_output_cost=pd.NamedAgg(column="Output Cost", aggfunc="sum"), - sum_total_cost=pd.NamedAgg(column="Total Cost", aggfunc="sum"), - ) - .reset_index() - ) - - group_by_filename = ( - df.groupby("Filename") - .agg( - count_llm_calls=pd.NamedAgg(column="Caller Function", aggfunc="count"), - sum_input_prompt_length=pd.NamedAgg( - column="Input Prompt Length", aggfunc="sum" - ), - sum_input_prompt_tokens=pd.NamedAgg( - column="Input Prompt Tokens", aggfunc="sum" - ), - sum_input_cost=pd.NamedAgg(column="Input Cost", aggfunc="sum"), - sum_output_response_length=pd.NamedAgg( - column="Output Response Length", aggfunc="sum" - ), - sum_output_response_tokens=pd.NamedAgg( - column="Output Response Tokens", aggfunc="sum" - ), - sum_output_cost=pd.NamedAgg(column="Output Cost", aggfunc="sum"), - sum_total_cost=pd.NamedAgg(column="Total Cost", aggfunc="sum"), - ) - .reset_index() - ) - - return group_by_caller_function, group_by_filename diff --git a/fieldExtraction/src/tracking/tracking_utils.py b/fieldExtraction/src/tracking/tracking_utils.py index e6fe9f7..0473d34 100644 --- a/fieldExtraction/src/tracking/tracking_utils.py +++ b/fieldExtraction/src/tracking/tracking_utils.py @@ -250,7 +250,6 @@ def upload_tracking_logs(run_timestamp, batch_id): write_stats_to_csv() tracking_files = { - "cost_log": config.COST_LOG_NAME, "file_log": config.FILE_LOG_NAME, "tracking": config.TRACKING_FILE, } diff --git a/fieldExtraction/src/utils/llm_utils.py b/fieldExtraction/src/utils/llm_utils.py index 60445bc..0544279 100644 --- a/fieldExtraction/src/utils/llm_utils.py +++ b/fieldExtraction/src/utils/llm_utils.py @@ -1,7 +1,5 @@ import base64 -import csv import hashlib -import inspect import json import logging import math @@ -64,46 +62,6 @@ def count_tokens(s: str) -> int: return math.ceil(len(s) / 6) -def create_cost_log_entry( - filename: str, - caller_name: str, - prompt: str, - response: str, - input_cost_per_1k: float, - output_cost_per_1k: float, -) -> dict: - """Generates an entry for the running cost log. - - Args: - filename (str): Name of file being operated upon - caller_name (str): Caller function to log - prompt (str): Input prompt to LLM - response (str): Output response from LLM - input_cost_per_1k (float): Cost per 1k tokens of input - output_cost_per_1k (float): Cost per 1k tokens of output - - Returns: - dict: Cost log engry - """ - input_tokens = count_tokens(prompt) - input_cost = input_cost_per_1k * input_tokens / 1000 - output_tokens = count_tokens(response) - output_cost = output_cost_per_1k * output_tokens / 1000 - - cost_log_entry = { - "Filename": filename, - "Caller Function": caller_name, - "Input Prompt Length": len(prompt), - "Input Prompt Tokens": input_tokens, - "Input Cost": input_cost, - "Output Response Length": len(response), - "Output Response Tokens": count_tokens(response), - "Output Cost": output_cost, - "Total Cost": input_cost + output_cost, - } - return cost_log_entry - - # In-memory cache claude_cache = OrderedDict() CACHE_LIMIT = 20000 @@ -212,22 +170,6 @@ def invoke_claude( claude_cache[cache_key] = response # Insert new response claude_cache.move_to_end(cache_key) # Mark as recently used - # Cost logging - current_frame = inspect.currentframe() - if current_frame is not None and current_frame.f_back is not None: - caller_name = current_frame.f_back.f_code.co_name - else: - caller_name = "unknown" - - del current_frame # Avoid reference cycles - - cost_log_entry = create_cost_log_entry( - filename, caller_name, prompt, response, input_cost_per_1k, output_cost_per_1k - ) - with open(config.COST_LOG_NAME, "a", newline="", encoding="utf-8") as f: - writer = csv.DictWriter(f, fieldnames=config.COST_LOG_FIELDS) - writer.writerow(cost_log_entry) - return response @@ -285,22 +227,6 @@ def invoke_claude_with_image_array( claude_cache[cache_key] = response # Insert new response claude_cache.move_to_end(cache_key) # Mark as recently used - # Cost logging - current_frame = inspect.currentframe() - if current_frame is not None and current_frame.f_back is not None: - caller_name = current_frame.f_back.f_code.co_name - else: - caller_name = "unknown" - - del current_frame # Avoid reference cycles - - cost_log_entry = create_cost_log_entry( - filename, caller_name, prompt, response, input_cost_per_1k, output_cost_per_1k - ) - with open(config.COST_LOG_NAME, "a", newline="", encoding="utf-8") as f: - writer = csv.DictWriter(f, fieldnames=config.COST_LOG_FIELDS) - writer.writerow(cost_log_entry) - return response diff --git a/fieldExtraction/tests/test_costlog_funcs.py b/fieldExtraction/tests/test_costlog_funcs.py deleted file mode 100644 index a23c065..0000000 --- a/fieldExtraction/tests/test_costlog_funcs.py +++ /dev/null @@ -1,9 +0,0 @@ -import pandas as pd -import pytest - -from src.tracking.costlog_utils import aggregate_costlog - - -def test_missing_columns(): - with pytest.raises(ValueError): - aggregate_costlog(pd.DataFrame()) diff --git a/fieldExtraction/tests/test_llm_utils.py b/fieldExtraction/tests/test_llm_utils.py index a3ff622..4aac239 100644 --- a/fieldExtraction/tests/test_llm_utils.py +++ b/fieldExtraction/tests/test_llm_utils.py @@ -43,36 +43,19 @@ class TestLLMUtils(unittest.TestCase): 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): + 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) - mock_csv_writer.assert_called() if __name__ == "__main__":