jacob-mathison/doczyai-pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Merged in feature/remove-costlog (pull request #675)

Browse Source 
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
This commit is contained in:
Alex Galarce
2025-08-21 17:45:18 +00:00
parent 9a3a93ce7c
commit 04aa447831
6 changed files with 2 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fieldExtraction/src/config.py
+1 -19
View File
@@ -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
fieldExtraction/src/tracking/costlog_utils.py
-72
View File
@@ -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
fieldExtraction/src/tracking/tracking_utils.py
-1
View File
@@ -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,
}
fieldExtraction/src/utils/llm_utils.py
-74
View File
@@ -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
fieldExtraction/tests/test_costlog_funcs.py
-9
View File
@@ -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())
fieldExtraction/tests/test_llm_utils.py
+1 -18
View File
@@ -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__":