diff --git a/fieldExtraction/src/claude_funcs.py b/fieldExtraction/src/claude_funcs.py index 108937b..9964975 100644 --- a/fieldExtraction/src/claude_funcs.py +++ b/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)