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

from datetime import datetime
import json
import re
import warnings

import numpy as np
import pandas as pd

from src.prompts import preprocessing_prompts
import src.utils.llm_utils as llm_utils
from src import config
from src.enums.delimiters import Delimiter
from src.regex.regex_patterns import (BACKTICK_PATTERN, PIPE_PATTERN,
                                      TRIPLE_BACKTICK_PATTERN)


def extract_text_from_delimiters(
    raw_output: str, delimiter: Delimiter, match_index: int = -1
) -> str:
    """
    Extracts a match enclosed in specified delimiters from the raw output.
    This function searches for text enclosed by a specified delimiter in the given raw output string.
    It returns the match specified by the `which_match` index. If no match is found, it returns "N/A".

    - raw_output (str): The raw output string to search within.
    - which_match (int, optional): The index of the match to return. Defaults to -1, which returns the last match. If the index is out of range, the last match is returned.
    - str: The extracted text or "N/A" if no match is found.

    Raises:
        - TypeError: If `raw_output` is not a string or `delimiter` is not a Delimiter enum value.
        - ValueError: If an unsupported delimiter is provided.

    Returns:
        - str: The extracted answer or "N/A" if no match is found.
    """
    if not isinstance(raw_output, str):
        raise TypeError(
            "Expected a string for raw_output, got {0}.".format(
                type(raw_output).__name__
            )
        )

    if not isinstance(delimiter, Delimiter):
        raise TypeError(
            "Expected a Delimiter enum value, got {0}.".format(type(delimiter).__name__)
        )

    # Define a pattern based on the delimiter type
    if delimiter == Delimiter.PIPE:
        pattern = PIPE_PATTERN
    elif delimiter == Delimiter.BACKTICK:
        pattern = BACKTICK_PATTERN
    elif delimiter == Delimiter.TRIPLE_BACKTICK:
        pattern = TRIPLE_BACKTICK_PATTERN
    else:
        raise ValueError("Unsupported delimiter. Use one of the Delimiter enum values.")

    # Find all matches based on the pattern
    matches = re.findall(pattern=pattern, string=raw_output)

    if len(matches) == 0:
        return "N/A"

    # Adjust for negative indices
    if match_index < 0:
        match_index += len(matches)

    if not 0 <= match_index < len(matches):
        warnings.warn("Index out of range. Returning the last match by default.")
        match_index = -1

    return matches[match_index]

def json_parsing_search(response_text: str, field_list: list[str]) -> dict[str, str]:
    """
    Parses a JSON-like string to extract values for fields.

    This function attempts to clean and format the input string to resemble a valid JSON structure.
    It then searches for the specified fields within the string and extracts their corresponding values.

    Parameters:
        response_text (str): The JSON-like string to parse.
        field_list (list[str]): A list of field names to search for in the response_text.

    Returns:
        dict[str, str]: A dictionary where keys are field names and values are the extracted values from the response_text.
    """
    try:
        if response_text.split("{", 1)[1].strip()[0] == '"':
            response_text = "{" + response_text.split("{", 1)[1]
        else:
            response_text = "{" + response_text
    except:
        response_text = "{" + response_text

    if len(response_text.split("}", 1)) > 1:
        if response_text.rsplit("}", 1)[0].strip()[-1] == '"':
            response_text = response_text.rsplit("}", 1)[0] + "}"
        elif response_text.rsplit("}", 1)[0].strip()[-1] == "}":
            response_text = response_text.rsplit("}", 1)[0]
    else:
        response_text = response_text.rstrip(",") + "}"
    
    field_l = []
    answer_l = []
    position_dict = {}
    
    for f in field_list:
        location = response_text.find('"' + f + '"')
        if location != -1:
            position_dict[location] = f
    
    field_list = list(dict(sorted(position_dict.items())).values())
    
    for f in field_list:
        if f in response_text:
            field_l.append(f)
            value = response_text.split('"' + f + '"')[0]
            response_text = response_text.split('"' + f + '"')[1]
            if f != field_list[0] and f != field_list[-1]:
                answer_l.append(
                    value.strip("\n")
                    .strip('"')
                    .strip(":")
                    .strip(" ")
                    .strip("\n")
                    .strip(" ")
                    .strip(",")
                    .strip('"')
                )
            elif f == field_list[-1]:
                answer_l.append(
                    value.strip("\n")
                    .strip('"')
                    .strip(":")
                    .strip(" ")
                    .strip("\n")
                    .strip(" ")
                    .strip(",")
                    .strip('"')
                )
                answer_l.append(
                    response_text.strip("\n")
                    .strip('"')
                    .strip(":")
                    .strip(" ")
                    .strip("}")
                    .strip("\n")
                    .strip(" ")
                    .strip('"')
                )

    return dict(zip(field_l, answer_l))


def secondary_string_to_dict(dict_string: str, filename: str) -> dict:
    """
    Converts a string representation of a dictionary into an actual dictionary object, handling potential formatting issues.

    This function cleans the string by removing specific unwanted characters and markers that may interfere with JSON parsing.
    It then locates the substring that correctly forms a dictionary format, attempts to parse it as JSON, and handles common
    parsing errors by replacing problematic single quotes with double quotes before re-parsing.

    Parameters:
        dict_string (str): The string containing the dictionary-like content, potentially surrounded by extra text or characters.
        filename (str): The filename associated with these entries, used to tag each resulting dictionary.

    Returns:
        dict: The dictionary obtained from parsing the cleaned and corrected string.
    """
    start_index = dict_string.find("{")
    end_index = dict_string.rfind("}") + 1
    dict_substring = dict_string[start_index:end_index]
    try:
        result_dict = json.loads(dict_substring)
    except:
        try:
            dict_substring = dict_substring.replace("'", '"')
            result_dict = json.loads(dict_substring)
        except:
            try:
                prompt = preprocessing_prompts.FIX_JSON(dict_substring)
                dict_substring = llm_utils.invoke_claude(
                    prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename
                )
                result_dict = json.loads(dict_substring)
            except Exception as e:
                print(e)
                result_dict = {}
    return result_dict


def primary_string_to_dict(string_dict, filename):
    """
    Converts a dictionary of strings, where each string represents multiple dictionary entries, into a list of dictionaries,
    augmenting each with metadata such as page number and filename.

    This function iterates over each page number's string data, extracting and converting string representations of
    dictionaries into actual dictionary objects. It handles and cleans the string format to properly parse it into dictionaries.
    All dictionaries are then augmented with their respective page number and the filename before being compiled into a list.

    Parameters:
        string_dict (dict): A dictionary where keys are page numbers and values are strings containing multiple dictionary entries.
        filename (str): The filename associated with these entries, used to tag each resulting dictionary.

    Returns:
        list of dict: A list of dictionaries, each representing data extracted and converted from the input string, tagged with
                      their page number and filename.
    """
    data = []
    pattern = r"\{.*?\}"
    for page_num in string_dict.keys():
        primary_list = string_dict[page_num]
        dicts = primary_list.split("[")[1]  # Strip front
        dicts = dicts.split("]")[0]  # Strip back
        dicts = dicts.replace("\n", "")  # Remove new lines
        dicts = dicts.replace("<<<", "")
        dicts = dicts.replace(">>>", "")
        dicts = re.sub(r"(?<!\\)'([^']*?)'(?<!\\):", r'"\1":', dicts)
        dict_list = re.findall(pattern, dicts)
        dict_list = [secondary_string_to_dict(d, filename) for d in dict_list]
        for dict_ in dict_list:
            dict_["page_num"] = page_num
            data.append(dict_)
    return data

def universal_json_load(string_dict):
    """Matches json dicts and list of dicts - NOT simple lists"""
    match = re.search(r'\{.*\}|\[\s*?\{.*\}\s*?\]', string_dict, re.DOTALL) # Updated regex
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass
        raise ValueError("No valid JSON object found in the input string")
    else:
        return {}

reimbursement_strings = [ # These are for `method='keyword'`
    "%",
    "$",
    "percent",
    "billed charges"
]
reimb_regex = r"(?<![$%])(?:\$\d+|\d+[$%])(?![$%])"

# TODO: check should this always be checking just page 1?  Bc in file_processing and table_utils it is but in preprocessing_funcs it's checking specific pages
def contains_reimbursement(text, page="1", method="keyword"):
    """
    Checks if the given text contains any reimbursement-related keywords or patterns.

    Args:
        text (str or dict): The text to check. If a dictionary is provided, it should have page numbers as keys.
        page (str, optional): The page number to check in the dictionary. Defaults to "1".
        method (str, optional): The method to use for checking. Options are "keyword"
         and "regex". Defaults to "keyword".

    Returns:
        bool: True if any reimbursement-related keyword or pattern is found, False otherwise.
    """
    if method == "keyword":
        if isinstance(text, dict):
            text_to_check = text.get(page, "").lower()
        elif isinstance(text, str):
            text_to_check = text.lower()
        else:
            print("contains_reimbursement - Invalid data type")
            return False
        return any(keyword in text_to_check for keyword in reimbursement_strings)
    elif method == "regex":
        return bool(re.search(reimb_regex, text))
    else:
        raise ValueError("Invalid method. Choose 'keyword' or 'regex'.")

def count_reimbursements_in_exhibit(exhibit_text: str) -> int: #JUST by regex
    """Counts reimbursements in an exhibit text.  Reimbursements are detected by a regex
    search as defined by `reimb_regex`.

    Args:
        exhibit_text (str): Input exhibit text

    Returns:
        int: Number of reimbursements detected
    """
    return len(re.findall(reimb_regex, exhibit_text))

def is_empty(value, pd_mask=True):
    """
    Checks if a value is considered empty or invalid.

    Args:
        value: The value to check.
        Can be a string, list, or pandas Series.
        pd_mask (bool): Only relevant for Series inputs. 
            If True, returns a mask for empty values in a pandas Series. 
            Otherwise, the function returns True iff the entire Series is empty.

    Returns:
        bool: True if the value is empty or invalid, False otherwise.
    """
    empty_values = [None, "", "N/A", "NA", "null", "none", "NaN", np.nan, "nan", "None"]
    
    if isinstance(value, list):  # Handle list inputs
        return not value or all(is_empty(v) for v in value)
    
    # if it's a pd.Series, return the mask (True for empty values)
    if isinstance(value, pd.Series):
        return value.isna() | (value.isin(empty_values)) if pd_mask else value.isna().all()

    if pd.isna(value):
        return True
    else:
        return value in empty_values


def get_exhibit_chunk(text_dict: dict, 
                      exhibit_chunk_mapping: dict, 
                      exhibit_page: str) -> str:
    """
    Compiles text from a specific exhibit based on its starting page number.

    This function aggregates the text of all pages that are mapped to a specific exhibit
    start page in a document. It uses the exhibit chunk mapping to determine which pages
    belong to the exhibit defined by the given start page number, and concatenates their
    texts into a single string.

    Parameters:
        text_dict (dict): A dictionary where keys are page numbers and values are the text on those pages.
        exhibit_chunk_mapping (dict): A dictionary mapping each page number to the starting page number of the exhibit it belongs to.
        exhibit_page (str): The page number that marks the beginning of the exhibit to compile text for.

    Returns:
        str: A string that combines all the text from the pages belonging to the specified exhibit, separated by newlines.

    Notes:
        - The function assumes that `exhibit_chunk_mapping` correctly maps all page numbers in `text_dict` to their respective exhibit starting pages.
        - It is important that `exhibit_page` exists as a key in the `exhibit_chunk_mapping` dictionary and corresponds to the starting page of an exhibit.
    """
    return '\n'.join([page_text for page_num, page_text in text_dict.items() if exhibit_chunk_mapping[page_num] == exhibit_page])


def datetime_str():
    return datetime.now().strftime("[%Y-%m-%d %H:%M:%S]")