Katon Minhas
2535400124
Feature/claude cache * Add Bill Type * cache for llm_utils invoke_claude * Merged main into feature/claude-cache * LRUCache class * Move to OrderedDict cache * Remove lru cache call Approved-by: Alex Galarce
455 lines
16 KiB
Python
455 lines
16 KiB
Python
import csv
|
|
import inspect
|
|
import json
|
|
import math
|
|
import random
|
|
import time
|
|
|
|
import anthropic
|
|
from botocore.exceptions import ClientError, ReadTimeoutError
|
|
|
|
import src.utils.string_utils as string_utils
|
|
from src import config
|
|
from collections import OrderedDict
|
|
|
|
import hashlib
|
|
import inspect
|
|
import csv
|
|
|
|
|
|
def update_stats(filename, tokens_sent, tokens_received, elapsed_time, model_type):
|
|
# Update per-file statistics
|
|
config.MODEL_STATS[filename]["requests"] += 1
|
|
config.MODEL_STATS[filename]["tokens_sent"] += tokens_sent
|
|
config.MODEL_STATS[filename]["tokens_received"] += tokens_received
|
|
config.MODEL_STATS[filename]["total_time"] += elapsed_time
|
|
|
|
# Update per-file model-specific request count
|
|
model_key = f"{model_type}_requests"
|
|
config.MODEL_STATS[filename][model_key] = (
|
|
config.MODEL_STATS[filename].get(model_key, 0) + 1
|
|
)
|
|
|
|
# Update global statistics
|
|
config.GLOBAL_STATS["total_requests"] = (
|
|
config.GLOBAL_STATS.get("total_requests", 0) + 1
|
|
)
|
|
config.GLOBAL_STATS["total_tokens_sent"] = (
|
|
config.GLOBAL_STATS.get("total_tokens_sent", 0) + tokens_sent
|
|
)
|
|
config.GLOBAL_STATS["total_tokens_received"] = (
|
|
config.GLOBAL_STATS.get("total_tokens_received", 0) + tokens_received
|
|
)
|
|
config.GLOBAL_STATS["total_time"] = (
|
|
config.GLOBAL_STATS.get("total_time", 0) + elapsed_time
|
|
)
|
|
|
|
# Update global model-specific request count
|
|
global_model_key = f"total_{model_type}_requests"
|
|
config.GLOBAL_STATS[global_model_key] = (
|
|
config.GLOBAL_STATS.get(global_model_key, 0) + 1
|
|
)
|
|
|
|
def count_tokens(s: str) -> int:
|
|
"""Count number of tokens in a string
|
|
Currently uses very simplistic "divide characters by 6" to estimate token
|
|
https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
|
|
|
|
Args:
|
|
s (str): Input string
|
|
|
|
Returns:
|
|
int: Token count
|
|
"""
|
|
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
|
|
|
|
def get_cache_key(prompt, model_id):
|
|
"""Generates a unique hash key for caching."""
|
|
return hashlib.sha256(f"{model_id}:{prompt}".encode()).hexdigest()
|
|
|
|
def invoke_claude(prompt, model_id, filename, max_tokens=4096):
|
|
cache_key = get_cache_key(prompt, model_id)
|
|
|
|
# Check cache first - use get() method from LRUCache
|
|
if cache_key in claude_cache:
|
|
claude_cache.move_to_end(cache_key) # Mark as recently used
|
|
return claude_cache[cache_key]
|
|
|
|
# Select execution mode
|
|
if config.RUN_MODE == "local":
|
|
if model_id == config.MODEL_ID_CLAUDE2:
|
|
response = local_claude_2(prompt, model_id, filename, max_tokens)
|
|
input_cost_per_1k, output_cost_per_1k = 0.008, 0.024
|
|
elif model_id == config.MODEL_ID_CLAUDE3_HAIKU:
|
|
response = local_claude_3(prompt, model_id, filename, max_tokens)
|
|
input_cost_per_1k, output_cost_per_1k = 0.00025, 0.00125
|
|
elif model_id == config.MODEL_ID_CLAUDE35_SONNET:
|
|
response = local_claude_3(prompt, model_id, filename, max_tokens)
|
|
input_cost_per_1k, output_cost_per_1k = 0.003, 0.015
|
|
elif config.RUN_MODE == "ec2":
|
|
if model_id == config.MODEL_ID_CLAUDE2:
|
|
response = ec2_claude2(prompt, model_id, filename, max_tokens)
|
|
input_cost_per_1k, output_cost_per_1k = 0.008, 0.024
|
|
elif model_id == config.MODEL_ID_CLAUDE3_HAIKU:
|
|
response = ec2_claude_3(prompt, model_id, filename, max_tokens)
|
|
input_cost_per_1k, output_cost_per_1k = 0.00025, 0.00125
|
|
elif model_id == config.MODEL_ID_CLAUDE35_SONNET:
|
|
response = ec2_claude_35(prompt, model_id, filename, max_tokens)
|
|
input_cost_per_1k, output_cost_per_1k = 0.003, 0.015
|
|
else:
|
|
print("Usage: python local_main.py <-local> OR python main.py <ec2>")
|
|
raise
|
|
|
|
# Store response in cache with LRU eviction
|
|
if len(claude_cache) >= CACHE_LIMIT:
|
|
claude_cache.popitem(last=False) # Remove the oldest (least recently used) item
|
|
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()
|
|
caller_frame = current_frame.f_back
|
|
caller_name = caller_frame.f_code.co_name
|
|
del current_frame, caller_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
|
|
|
|
|
|
# Claude calls
|
|
def local_claude_2(prompt, model_id, filename, max_tokens):
|
|
|
|
if config.ENABLE_ANSWER_TRACKING and prompt in config.ANSWER_TRACKING_DICT.keys():
|
|
response_text = config.ANSWER_TRACKING_DICT[prompt]
|
|
|
|
else:
|
|
body = json.dumps(
|
|
{
|
|
"prompt": anthropic.HUMAN_PROMPT + prompt + anthropic.AI_PROMPT,
|
|
"max_tokens_to_sample": max_tokens,
|
|
"temperature": 0.0,
|
|
"top_p": 1,
|
|
"top_k": 250,
|
|
"stop_sequences": [anthropic.HUMAN_PROMPT],
|
|
}
|
|
)
|
|
response = config.BEDROCK_RUNTIME.invoke_model(
|
|
body=body,
|
|
modelId=model_id,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
response_body = json.loads(response.get("body").read())
|
|
response_text = response_body["completion"]
|
|
|
|
if config.ENABLE_ANSWER_TRACKING:
|
|
config.ANSWER_TRACKING_DICT[prompt] = response_text
|
|
|
|
return response_text
|
|
|
|
|
|
def local_claude_3(
|
|
prompt, model_id, filename, max_tokens, max_retries=5, initial_delay=1
|
|
):
|
|
|
|
if config.ENABLE_ANSWER_TRACKING and prompt in config.ANSWER_TRACKING_DICT.keys():
|
|
return config.ANSWER_TRACKING_DICT[prompt]
|
|
|
|
body = json.dumps(
|
|
{
|
|
"anthropic_version": "bedrock-2023-05-31",
|
|
"max_tokens": max_tokens,
|
|
"temperature": 0.0,
|
|
"messages": [
|
|
{"role": "user", "content": [{"type": "text", "text": prompt}]}
|
|
],
|
|
}
|
|
)
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
response = config.BEDROCK_RUNTIME.invoke_model(
|
|
body=body,
|
|
modelId=model_id,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
response_body = json.loads(response.get("body").read())
|
|
response_text = response_body["content"][0]["text"]
|
|
|
|
if config.ENABLE_ANSWER_TRACKING:
|
|
config.ANSWER_TRACKING_DICT[prompt] = response_text
|
|
|
|
return response_text
|
|
|
|
except ClientError as e:
|
|
if e.response["Error"]["Code"] in [
|
|
"ThrottlingException",
|
|
"TooManyRequestsException",
|
|
]:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
time.sleep(delay)
|
|
else:
|
|
raise
|
|
else:
|
|
raise
|
|
|
|
raise Exception("Max retries exceeded")
|
|
|
|
|
|
def ec2_claude2(prompt, model_id, filename, max_tokens):
|
|
start_time = time.time()
|
|
body = json.dumps(
|
|
{
|
|
"prompt": anthropic.HUMAN_PROMPT + prompt + anthropic.AI_PROMPT,
|
|
"max_tokens_to_sample": max_tokens,
|
|
"temperature": 0.0,
|
|
"top_p": 1,
|
|
"top_k": 250,
|
|
"stop_sequences": [anthropic.HUMAN_PROMPT],
|
|
}
|
|
)
|
|
response = config.BEDROCK_RUNTIME.invoke_model(
|
|
body=body,
|
|
modelId=config.MODEL_ID_CLAUDE2,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
response_body = json.loads(response.get("body").read())
|
|
response_text = response_body["completion"]
|
|
elapsed_time = time.time() - start_time
|
|
tokens_sent = len(prompt.split())
|
|
tokens_received = len(response_text.split())
|
|
|
|
update_stats(filename, tokens_sent, tokens_received, elapsed_time, "claude2")
|
|
return response_text
|
|
|
|
|
|
def ec2_claude_3(
|
|
prompt, model_id, filename, max_tokens, max_retries=5, initial_delay=1
|
|
):
|
|
start_time = time.time()
|
|
body = json.dumps(
|
|
{
|
|
"anthropic_version": "bedrock-2023-05-31",
|
|
"max_tokens": max_tokens,
|
|
"temperature": 0.0,
|
|
"messages": [
|
|
{"role": "user", "content": [{"type": "text", "text": prompt}]}
|
|
],
|
|
}
|
|
)
|
|
dev_runtime = config.assume_role_and_get_runtime(
|
|
config.dev_account_id, config.dev_role_name
|
|
)
|
|
prod_runtime = config.assume_role_and_get_runtime(
|
|
config.prod_account_id, config.prod_role_name
|
|
)
|
|
runtimes = {
|
|
"UAT": config.BEDROCK_RUNTIME, # This is the default runtime (UAT)
|
|
"DEV": dev_runtime,
|
|
"PROD": prod_runtime,
|
|
}
|
|
for runtime_key in runtimes.keys():
|
|
current_runtime = runtimes[runtime_key]
|
|
for attempt in range(max_retries):
|
|
try:
|
|
response = current_runtime.invoke_model(
|
|
body=body,
|
|
modelId=model_id,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
response_body = json.loads(response.get("body").read())
|
|
response_text = response_body["content"][0]["text"]
|
|
|
|
elapsed_time = time.time() - start_time
|
|
tokens_sent = len(prompt.split())
|
|
tokens_received = len(response_text.split())
|
|
|
|
if filename:
|
|
if model_id == config.MODEL_ID_CLAUDE35_SONNET:
|
|
update_stats(
|
|
filename,
|
|
tokens_sent,
|
|
tokens_received,
|
|
elapsed_time,
|
|
"claude35",
|
|
)
|
|
elif model_id == config.MODEL_ID_CLAUDE3_HAIKU:
|
|
update_stats(
|
|
filename,
|
|
tokens_sent,
|
|
tokens_received,
|
|
elapsed_time,
|
|
"claude3",
|
|
)
|
|
|
|
return response_text
|
|
|
|
except ClientError as e:
|
|
error_code = e.response["Error"]["Code"]
|
|
if error_code in ["ThrottlingException", "TooManyRequestsException"]:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
print(
|
|
f"Attempt {attempt + 1} failed. Retrying in {delay:.2f} seconds..."
|
|
)
|
|
time.sleep(delay)
|
|
else:
|
|
print(
|
|
f"Switching to next runtime due to {error_code}: {str(e)}."
|
|
)
|
|
break # Switch to the next runtime
|
|
else:
|
|
print(f"Unexpected error: {str(e)}")
|
|
raise # Reraise other exceptions
|
|
|
|
except Exception as e:
|
|
print(f"Unexpected error: {str(e)}")
|
|
raise
|
|
|
|
raise Exception("Max retries reached. Unable to get a response from the model.")
|
|
|
|
|
|
def ec2_claude_35(
|
|
prompt, model_id, filename, max_tokens, max_retries=5, initial_delay=1
|
|
):
|
|
start_time = time.time()
|
|
body = json.dumps(
|
|
{
|
|
"anthropic_version": "bedrock-2023-05-31",
|
|
"max_tokens": max_tokens,
|
|
"temperature": 0.0,
|
|
"messages": [
|
|
{"role": "user", "content": [{"type": "text", "text": prompt}]}
|
|
],
|
|
}
|
|
)
|
|
dev_runtime = config.assume_role_and_get_runtime(
|
|
config.dev_account_id, config.dev_role_name
|
|
)
|
|
prod_runtime = config.assume_role_and_get_runtime(
|
|
config.prod_account_id, config.prod_role_name
|
|
)
|
|
runtimes = {
|
|
"UAT": config.BEDROCK_RUNTIME, # This is the default runtime (UAT)
|
|
"DEV": dev_runtime,
|
|
"PROD": prod_runtime,
|
|
}
|
|
for runtime_key in runtimes.keys():
|
|
current_runtime = runtimes[runtime_key]
|
|
for attempt in range(max_retries):
|
|
try:
|
|
response = current_runtime.invoke_model(
|
|
body=body,
|
|
modelId=model_id,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
response_body = json.loads(response.get("body").read())
|
|
response_text = response_body["content"][0]["text"]
|
|
|
|
elapsed_time = time.time() - start_time
|
|
tokens_sent = len(prompt.split())
|
|
tokens_received = len(response_text.split())
|
|
|
|
if filename:
|
|
update_stats(
|
|
filename, tokens_sent, tokens_received, elapsed_time, "claude35"
|
|
)
|
|
|
|
return response_text
|
|
|
|
except (ReadTimeoutError, ClientError) as e:
|
|
error_code = e.response["Error"]["Code"]
|
|
if error_code in ["ThrottlingException", "TooManyRequestsException"]:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
print(
|
|
f"Attempt {attempt + 1} failed. Retrying in {delay:.2f} seconds..."
|
|
)
|
|
time.sleep(delay)
|
|
else:
|
|
print(
|
|
f"Switching to next runtime due to {error_code}: {str(e)}."
|
|
)
|
|
break # Switch to the next runtime
|
|
else:
|
|
print(f"Unexpected error: {str(e)}")
|
|
raise # Reraise other exceptions
|
|
|
|
except Exception as e:
|
|
print(f"Unexpected error: {str(e)}")
|
|
raise
|
|
|
|
raise Exception("Max retries reached. Unable to get a response from the model.")
|
|
|
|
# AC Answer - for investment
|
|
def get_one_to_one_answer(prompt, filename, fields):
|
|
response_text = invoke_claude(
|
|
prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 8192
|
|
)
|
|
response_dict = string_utils.json_parsing_search(response_text, fields)
|
|
return response_dict
|
|
|
|
|
|
def get_ac_answer(prompt, filename, fields):
|
|
response_text = invoke_claude(
|
|
prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 8192
|
|
)
|
|
response_text = response_text.replace("_DATE", "_DT")
|
|
# response_dict = dict_operations.secondary_string_to_dict(response_text, filename)
|
|
response_dict = string_utils.json_parsing_search(response_text, fields)
|
|
return response_dict
|