Merged in refactor/tidy-claude-funcs (pull request #279)

Refactor/tidy claude funcs

* pull out create_cost_log_entry() into its own function

* add docstring

* isort

* Merged main into refactor/tidy-claude-funcs

* Merged main into refactor/tidy-claude-funcs


Approved-by: Michael McGuinness
Approved-by: Katon Minhas

fieldExtraction/src/claude_funcs.py

@@ -1,17 +1,15 @@
-### EC2 ###
-
+import csv
+import inspect
 import json
-import anthropic
-import config
-from botocore.exceptions import ReadTimeoutError, ClientError
+import math
 import random
 import time
 
-from botocore.exceptions import ClientError
+import anthropic
+from botocore.exceptions import ClientError, ReadTimeoutError
+
 import config
-import inspect
-import csv
-import math
+
 
 def update_stats(filename, tokens_sent, tokens_received, elapsed_time, model_type):
     # Update per-file statistics
@@ -59,6 +57,44 @@ 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
 
 
 def invoke_claude(prompt, model_id, filename, max_tokens=4096):
@@ -106,25 +142,14 @@ def invoke_claude(prompt, model_id, filename, max_tokens=4096):
     del current_frame
     del caller_frame
 
-    #TODO: make a `create_cost_log_entry()` function that takes in:
-    # caller_name, prompt, response, input_cost_per_1k, output_cost_per_1k
-    # and returns the below `cost_log_entry`
-
-    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,
-    }
+    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)