doczyai-pipelines/fieldExtraction/src/utils/llm_utils.py

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


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


def invoke_claude(prompt, model_id, filename, max_tokens=4096):
    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 = 0.008
            output_cost_per_1k = 0.024
        elif model_id == config.MODEL_ID_CLAUDE3_HAIKU:
            response = local_claude_3(prompt, model_id, filename, max_tokens)
            input_cost_per_1k = 0.00025
            output_cost_per_1k = 0.00125
        elif model_id == config.MODEL_ID_CLAUDE35_SONNET:
            response = local_claude_3(prompt, model_id, filename, max_tokens)
            input_cost_per_1k = 0.003
            output_cost_per_1k = 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 = 0.008
            output_cost_per_1k = 0.024
        elif model_id == config.MODEL_ID_CLAUDE3_HAIKU:
            response = ec2_claude_3(prompt, model_id, filename, max_tokens)
            input_cost_per_1k = 0.00025
            output_cost_per_1k = 0.00125
        elif model_id == config.MODEL_ID_CLAUDE35_SONNET:
            # TODO: make these `ec2_*` and `local_*` return response, options
            # with options making the costs
            response = ec2_claude_35(prompt, model_id, filename, max_tokens)
            input_cost_per_1k = 0.003
            output_cost_per_1k = 0.015
    else:
        print("Usage: python local_main.py <-local> OR python main.py <ec2>")
        raise

    # Cost logging
    ## Get current frame
    current_frame = inspect.currentframe()
    ## Get caller's frame (previous frame)
    caller_frame = current_frame.f_back
    ## Get name of caller's function from its frame
    caller_name = caller_frame.f_code.co_name

    ## Clean up to prevent memory leaks / reference cycles
    del current_frame
    del caller_frame

    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.")


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