From 04d268a2f868944642adf7c48e69b06d2af0c9d5 Mon Sep 17 00:00:00 2001 From: Siddhant Medar Date: Fri, 17 Jan 2025 17:15:32 +0000 Subject: [PATCH] Merged in task/unit-tests-investment-and-client (pull request #354) Task/unit tests investment and client * move io_utils.py * add __init__.py * refactor: move string_funcs and update import paths to use utils namespace * rename string_funcs to string_utils * move and rename table_utils.py * refactor: rename claude_funcs to llm_utils in multiple files * update unit tests * added unit tests for src/string_utils.py * added unit tests for src/table_utils.py * added unit tests for src/llm_utils.py * added unit tests for src/io_utils.py * Merge branch 'main' into task/unit-tests-investment-and-client; TODO: fix unit tests * Unit tests to fix * fix unit tests * fix invalid references * add unit tests for preprocessing_funcs.py * move pytest-mock to test dependencies * Merge remote-tracking branch 'origin/main' into task/unit-tests-investment-and-client * add tests/string_utils_test.py * removed filename parameters for a few functions * save postprocessing_funcs_test midway * save postprocessing_funcs_test * postprocessing_funcs unit tests complete * Merge branch 'main' into task/unit-tests-investment-and-client Approved-by: Katon Minhas --- fieldExtraction/poetry.lock | 28 +- fieldExtraction/pyproject.toml | 1 + fieldExtraction/src/postprocessing_funcs.py | 443 +++++++++--- fieldExtraction/src/preprocessing_funcs.py | 36 +- .../tests/postprocessing_funcs_test.py | 642 ++++++++++++++++-- .../tests/preprocessing_funcs_test.py | 157 +++++ fieldExtraction/tests/string_utils_test.py | 10 +- 7 files changed, 1120 insertions(+), 197 deletions(-) create mode 100644 fieldExtraction/tests/preprocessing_funcs_test.py diff --git a/fieldExtraction/poetry.lock b/fieldExtraction/poetry.lock index e34b131..16679a4 100644 --- a/fieldExtraction/poetry.lock +++ b/fieldExtraction/poetry.lock @@ -288,17 +288,17 @@ css = ["tinycss2 (>=1.1.0,<1.5)"] [[package]] name = "boto3" -version = "1.35.94" +version = "1.35.97" description = "The AWS SDK for Python" optional = false python-versions = ">=3.8" files = [ - {file = "boto3-1.35.94-py3-none-any.whl", hash = "sha256:516c514fb447d6f216833d06a0781c003fcf43099a4ca2f5a363a8afe0942070"}, - {file = "boto3-1.35.94.tar.gz", hash = "sha256:5aa606239f0fe0dca0506e0ad6bbe4c589048e7e6c2486cee5ec22b6aa7ec2f8"}, + {file = "boto3-1.35.97-py3-none-any.whl", hash = "sha256:8e49416216a6e3a62c2a0c44fba4dd2852c85472e7b702516605b1363867d220"}, + {file = "boto3-1.35.97.tar.gz", hash = "sha256:7d398f66a11e67777c189d1f58c0a75d9d60f98d0ee51b8817e828930bf19e4e"}, ] [package.dependencies] -botocore = ">=1.35.94,<1.36.0" +botocore = ">=1.35.97,<1.36.0" jmespath = ">=0.7.1,<2.0.0" s3transfer = ">=0.10.0,<0.11.0" @@ -307,13 +307,13 @@ crt = ["botocore[crt] (>=1.21.0,<2.0a0)"] [[package]] name = "botocore" -version = "1.35.94" +version = "1.35.97" description = "Low-level, data-driven core of boto 3." optional = false python-versions = ">=3.8" files = [ - {file = "botocore-1.35.94-py3-none-any.whl", hash = "sha256:d784d944865d8279c79d2301fc09ac28b5221d4e7328fb4e23c642c253b9932c"}, - {file = "botocore-1.35.94.tar.gz", hash = "sha256:2b3309b356541faa4d88bb957dcac1d8004aa44953c0b7d4521a6cc5d3d5d6ba"}, + {file = "botocore-1.35.97-py3-none-any.whl", hash = "sha256:fed4f156b1a9b8ece53738f702ba5851b8c6216b4952de326547f349cc494f14"}, + {file = "botocore-1.35.97.tar.gz", hash = "sha256:88f2fab29192ffe2f2115d5bafbbd823ff4b6eb2774296e03ec8b5b0fe074f61"}, ] [package.dependencies] @@ -2076,13 +2076,13 @@ files = [ [[package]] name = "pydantic" -version = "2.10.4" +version = "2.10.5" description = "Data validation using Python type hints" optional = false python-versions = ">=3.8" files = [ - {file = "pydantic-2.10.4-py3-none-any.whl", hash = "sha256:597e135ea68be3a37552fb524bc7d0d66dcf93d395acd93a00682f1efcb8ee3d"}, - {file = "pydantic-2.10.4.tar.gz", hash = "sha256:82f12e9723da6de4fe2ba888b5971157b3be7ad914267dea8f05f82b28254f06"}, + {file = "pydantic-2.10.5-py3-none-any.whl", hash = "sha256:4dd4e322dbe55472cb7ca7e73f4b63574eecccf2835ffa2af9021ce113c83c53"}, + {file = "pydantic-2.10.5.tar.gz", hash = "sha256:278b38dbbaec562011d659ee05f63346951b3a248a6f3642e1bc68894ea2b4ff"}, ] [package.dependencies] @@ -2854,13 +2854,13 @@ win32 = ["pywin32"] [[package]] name = "setuptools" -version = "75.7.0" +version = "75.8.0" description = "Easily download, build, install, upgrade, and uninstall Python packages" optional = false python-versions = ">=3.9" files = [ - {file = "setuptools-75.7.0-py3-none-any.whl", hash = "sha256:84fb203f278ebcf5cd08f97d3fb96d3fbed4b629d500b29ad60d11e00769b183"}, - {file = "setuptools-75.7.0.tar.gz", hash = "sha256:886ff7b16cd342f1d1defc16fc98c9ce3fde69e087a4e1983d7ab634e5f41f4f"}, + {file = "setuptools-75.8.0-py3-none-any.whl", hash = "sha256:e3982f444617239225d675215d51f6ba05f845d4eec313da4418fdbb56fb27e3"}, + {file = "setuptools-75.8.0.tar.gz", hash = "sha256:c5afc8f407c626b8313a86e10311dd3f661c6cd9c09d4bf8c15c0e11f9f2b0e6"}, ] [package.extras] @@ -3178,4 +3178,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = "^3.12" -content-hash = "5847f5d08074f70030a8bcadea25052f725af8dd9fc47e8cb385acd3163c125e" +content-hash = "977386812e414d6b13d5237092e9da24d33cb5400a34b6a241898379eed1fba9" diff --git a/fieldExtraction/pyproject.toml b/fieldExtraction/pyproject.toml index cd7edae..63cf31e 100644 --- a/fieldExtraction/pyproject.toml +++ b/fieldExtraction/pyproject.toml @@ -26,6 +26,7 @@ isort = "^5.13.2" [tool.poetry.group.test.dependencies] pytest = "^8.3.3" +pytest-mock = "^3.14.0" [build-system] requires = ["poetry-core"] diff --git a/fieldExtraction/src/postprocessing_funcs.py b/fieldExtraction/src/postprocessing_funcs.py index 742c644..671a470 100644 --- a/fieldExtraction/src/postprocessing_funcs.py +++ b/fieldExtraction/src/postprocessing_funcs.py @@ -14,21 +14,27 @@ from src.constants.valid import ( keywords_to_detect_invalid_state_name, ) from src.utils.string_utils import is_empty +from typing import Any logging.basicConfig( level=logging.ERROR, format="%(asctime)s - %(levelname)s - %(message)s" ) - -def filter_service_column(answer_dicts): +def filter_service_column(answer_dicts: list[dict[str, str]]) -> list[dict[str, str]]: """ - Filters the answer_dicts based on the presence of specific keywords. + Filters the answer_dicts based on the presence of specific keywords in the 'FULL_SERVICE' field. + + Args: + answer_dicts (list[dict[str, str]]): A list of dictionaries containing service information. + + Returns: + list[dict[str, str]]: A filtered list of dictionaries where the 'FULL_SERVICE' field does not contain any of the specified keywords. """ filtered_list = [] for d in answer_dicts: clean_dict = True for keyword in valid.SERVICE_FILTER: - if ( + if "FULL_SERVICE" in d and ( keyword.upper() in d["FULL_SERVICE"].upper() or d["FULL_SERVICE"].upper() in keyword.upper() ): @@ -37,13 +43,21 @@ def filter_service_column(answer_dicts): filtered_list.append(d) return filtered_list +def filter_methodology_column(answer_dicts: list[dict[str, str]]) -> list[dict[str, str]]: + """ + Filters the answer_dicts based on the presence of specific keywords in the 'FULL_METHODOLOGY' field. -def filter_methodology_column(answer_dicts): + Args: + answer_dicts (list[dict[str, str]]): A list of dictionaries containing methodology information. + + Returns: + list[dict[str, str]]: A filtered list of dictionaries where the 'FULL_METHODOLOGY' field does not contain any of the specified keywords. + """ filtered_list = [] for d in answer_dicts: clean_dict = True for keyword in valid.METHODOLOGY_FILTER: - if ( + if "FULL_METHODOLOGY" in d and ( keyword.upper() in d["FULL_METHODOLOGY"].upper() or d["FULL_METHODOLOGY"].upper() in keyword.upper() ): @@ -52,8 +66,16 @@ def filter_methodology_column(answer_dicts): filtered_list.append(d) return filtered_list +def clean_td(td: list[dict[str, Any]]) -> list[dict[str,Any]]: + """ + Cleans the input list of dictionaries by transforming certain fields into lists and handling specific cases. -def clean_td(td): + Args: + td (list[dict[str, any]]): A list of dictionaries containing various fields. + + Returns: + list[dict[str, any]]: A cleaned list of dictionaries with transformed fields. + """ td_clean = [] for d in td: new_d = {} @@ -89,7 +111,16 @@ def get_parent_agreement_code(filename): return "N/A" -def consolidate_subheader(dict_list): +def consolidate_subheader(dict_list: list[dict[str, str]]) -> list[dict[str, str]]: + """ + Consolidates the 'SUBHEADER' field into the 'FULL_SERVICE' field for each dictionary in the list. + + Args: + dict_list (list[dict[str, str]]): A list of dictionaries containing service information. + + Returns: + list[dict[str, str]]: A modified list of dictionaries where the 'SUBHEADER' field is concatenated to the 'FULL_SERVICE' field and then removed. + """ modified_list = [] for d in dict_list: if "FULL_SERVICE" in d and "SUBHEADER" in d: @@ -100,8 +131,23 @@ def consolidate_subheader(dict_list): modified_list.append(d) return modified_list +def clean_msr_lesser(df: pd.DataFrame, ls: list[str] = valid.VALID_MSR) -> pd.DataFrame: + """ + Cleans the 'LESSER' and 'LESSER_RATE' columns in the given DataFrame based on the 'FULL_SERVICE' and 'EXHIBIT' columns. -def clean_msr_lesser(df, ls=valid.VALID_MSR): + This function performs the following operations: + 1. Identifies rows where the 'FULL_SERVICE' column contains any of the specified keywords. + 2. For each identified row, checks other rows with the same 'EXHIBIT' value. + 3. If any of these rows have 'LESSER' set to 'Y', updates the 'LESSER' value of the original row to 'Y'. + 4. Moves the first 'LESSER_RATE' value where 'LESSER' is 'Y' to the original row. + + Args: + df (pd.DataFrame): The input DataFrame containing the columns 'FULL_SERVICE', 'EXHIBIT', 'LESSER', and 'LESSER_RATE'. + ls (list[str], optional): A list of keywords to search for in the 'FULL_SERVICE' column. Defaults to valid.VALID_MSR. + + Returns: + pd.DataFrame: The modified DataFrame with updated 'LESSER' and 'LESSER_RATE' values. + """ pattern = "|".join(re.escape(item) for item in ls) mask = df["FULL_SERVICE"].str.contains(pattern, case=False, na=False) target_rows = df[mask] @@ -117,14 +163,25 @@ def clean_msr_lesser(df, ls=valid.VALID_MSR): df.at[index, "LESSER"] = "Y" # Move first LESSER_RATE where LESSER==Y to MSR row - lesser_rate = exhibit_rows[exhibit_rows["LESSER"] == "Y"][ - "LESSER_RATE" - ].iloc[0] + lesser_rate = exhibit_rows[exhibit_rows["LESSER"] == "Y"]["LESSER_RATE"].iloc[0] df.at[index, "LESSER_RATE"] = lesser_rate return df +def clean_default_term(df: pd.DataFrame) -> pd.DataFrame: + """ + Cleans the 'DEFAULT_TERM' and 'DEFAULT_RATE' columns in the given DataFrame. -def clean_default_term(df): + This function performs the following operations: + 1. Removes any invalid default terms based on a predefined pattern. + 2. Replaces the 'DEFAULT_RATE' values with acronyms as specified in the valid module. + 3. Ensures that 'DEFAULT_TERM' values match the 'IP_OP' column values appropriately. + + Args: + df (pd.DataFrame): The input DataFrame containing the columns 'DEFAULT_TERM', 'DEFAULT_RATE', and 'IP_OP'. + + Returns: + pd.DataFrame: The cleaned DataFrame with updated 'DEFAULT_TERM' and 'DEFAULT_RATE' values. + """ # Remove any invalid Defaults pattern = "|".join(re.escape(term) for term in valid.INVALID_DEFAULT) df["DEFAULT_TERM"] = df["DEFAULT_TERM"].astype(str) @@ -150,7 +207,19 @@ def clean_default_term(df): return df -def clean_lesser_rate(df): +def clean_lesser_rate(df: pd.DataFrame) -> pd.DataFrame: + """ + Cleans the 'LESSER_RATE' column in the given DataFrame by updating it with the 'RATE_STANDARD' value if certain conditions are met. + + This function iterates over each row in the DataFrame and checks if the 'LESSER' column contains 'Y', + the 'LESSER_RATE' column is empty, and the 'RATE_STANDARD' column is not empty. If these conditions are met, it updates the 'LESSER_RATE' column with the value from the 'RATE_STANDARD' column and sets the 'RATE_STANDARD' column to 'N/A'. + + Args: + df (pd.DataFrame): The input DataFrame containing the columns 'LESSER', 'LESSER_RATE', and 'RATE_STANDARD'. + + Returns: + pd.DataFrame: The modified DataFrame with updated 'LESSER_RATE' and 'RATE_STANDARD' values. + """ for index, row in df.iterrows(): lesser = str(row["LESSER"]) if not pd.isna(row["LESSER"]) else "" lesser_rate = row["LESSER_RATE"] @@ -161,8 +230,23 @@ def clean_lesser_rate(df): df.at[index, "RATE_STANDARD"] = "N/A" return df +def clean_prov_2(df: pd.DataFrame) -> pd.DataFrame: + """ + Cleans the 'PROV_TYPE_LEVEL_2' column in the given DataFrame by finding and assigning valid values based on other columns. -def clean_prov_2(df): + The function performs the following steps: + 1. Selects valid prov 2 values using 'PROV_TYPE' column. + 2. Identifies rows where 'PROV_TYPE_LEVEL_2' is empty. + 3. Attempts to find exact matches for 'PROV_TYPE_LEVEL_2' from 'FULL_SERVICE' and 'EXHIBIT' columns. + 4. If no match is found, it checks other rows with the same 'EXHIBIT' value for potential matches. + 5. Logs the matching process and any invalid assignments. + + Args: + df (pd.DataFrame): The input DataFrame containing columns 'PROV_TYPE', 'PROV_TYPE_LEVEL_2', 'FULL_SERVICE', and 'EXHIBIT'. + + Returns: + pd.DataFrame: The DataFrame with cleaned 'PROV_TYPE_LEVEL_2' values. + """ valid_types = valid.select_valid_prov_2(df["PROV_TYPE"]) target_rows = df[df["PROV_TYPE_LEVEL_2"].apply(is_empty)] @@ -236,8 +320,19 @@ def clean_prov_2(df): return df +def clean_term_clause(final_df: pd.DataFrame) -> pd.DataFrame: + """ + Cleans the 'TERM_CLAUSE' column in the given DataFrame by setting specific patterns to 'N/A'. -def clean_term_clause(final_df): + This function identifies rows where the 'TERM_CLAUSE' column starts with or contains certain phrases + and sets the value to 'N/A' for those rows. + + Args: + final_df (pd.DataFrame): The input DataFrame containing the 'TERM_CLAUSE' column. + + Returns: + pd.DataFrame: The cleaned DataFrame with updated 'TERM_CLAUSE' values. + """ if "TERM_CLAUSE" in final_df: final_df.loc[ final_df["TERM_CLAUSE"].str.startswith("IL-4 Termination"), "TERM_CLAUSE" @@ -297,9 +392,17 @@ def clean_auto_renewal_ind(final_df): return final_df -def clean_npi(final_df): +def clean_npi(final_df: pd.DataFrame) -> pd.DataFrame: + """ + Cleans the 'PROV_GROUP_NPI' column in the given DataFrame by replacing strings with a digit count not equal to 10 with "N/A". + + Args: + final_df (pd.DataFrame): The input DataFrame containing the 'PROV_GROUP_NPI' column. + + Returns: + pd.DataFrame: The modified DataFrame with cleaned 'PROV_GROUP_NPI' values. + """ if "PROV_GROUP_NPI" in final_df: - # Replace strings with a digit count not equal to 10 with "N/A" final_df["PROV_GROUP_NPI"] = final_df["PROV_GROUP_NPI"].apply( lambda x: ( f"N/A - (model detected: {x})" @@ -309,13 +412,20 @@ def clean_npi(final_df): ) return final_df +def clean_network_access_fees(final_df: pd.DataFrame) -> pd.DataFrame: + """ + Cleans the 'NETWORK_ACCESS_FEES_IND' column in the given DataFrame. -def clean_network_access_fees(final_df): + This function sets the 'NETWORK_ACCESS_FEES_IND' to "Yes" for rows where the value is not empty. + + Args: + final_df (pd.DataFrame): The input DataFrame containing the 'NETWORK_ACCESS_FEES_IND' column. + + Returns: + pd.DataFrame: The cleaned DataFrame with updated 'NETWORK_ACCESS_FEES_IND' values. + """ if "NETWORK_ACCESS_FEES_IND" in final_df: - final_df.loc[ - not is_empty(final_df["NETWORK_ACCESS_FEES_IND"]), - "NETWORK_ACCESS_FEES_IND", - ] = "Yes" + final_df.loc[~final_df["NETWORK_ACCESS_FEES_IND"].apply(is_empty), "NETWORK_ACCESS_FEES_IND"] = "Yes" return final_df @@ -363,14 +473,15 @@ def get_health_plan_state(text_dict: dict[str, str]) -> str: for page_num, page_content in text_dict.items(): state_matches = re.findall(pattern, page_content, flags=re.IGNORECASE) - state_matches = [state.strip().title() for state in state_matches] + state_matches = [state.strip().title() if len(state) > 2 else state.strip().upper() for state in state_matches] states_set.update(state_matches) - - states_list = list(states_set) + # Expand any abbreviations. If the state is already its full name, we're just gonna # default to that. E.g., "CT" -> "Connecticut", "Connecticut" -> "Connecticut" - states_list = [STATE_MAP.get(state, state) for state in states_list] + states_set = set([STATE_MAP.get(state, state) for state in states_set]) + states_list = list(states_set) + if len(states_list) > 1: state = "Multiple" elif len(states_list) == 1: @@ -395,6 +506,7 @@ def state_check(states: str | list[str], text_dict: dict[str, str]) -> str: states, str ): # If it's a string, turn it to a list (even a singleton) if "," in states: + states = states.replace("[", "").replace("]", "") states = states.split(",") else: states = [states] @@ -432,6 +544,8 @@ def state_check(states: str | list[str], text_dict: dict[str, str]) -> str: "N/A" ) # Discard any N/As (so if `result_states` is ['Illinois', 'N/A'], it returns 'Illinois' instead of 'Multiple') + print("Processed states", result_states) + if len(result_states) == 1: return list(result_states)[0] @@ -450,11 +564,25 @@ def state_check(states: str | list[str], text_dict: dict[str, str]) -> str: def clean_health_plan_state_from_hotfix( - df: pd.DataFrame, + df: pd.DataFrame, contract_name: str, text_dict: dict[str, str], ) -> pd.DataFrame: + """ + Cleans the 'HEALTH_PLAN_STATE' column in the given DataFrame using a hotfix. + Args: + df (pd.DataFrame): The DataFrame containing the data to be cleaned. + contract_name (str): The name of the contract. + text_dict (dict[str, str]): A dictionary containing text mappings for state checking. + Returns: + pd.DataFrame: The cleaned DataFrame. + Raises: + TypeError: If the input 'df' is not a pandas DataFrame. + ValueError: If the 'contract_name' is None. + KeyError: If the 'HEALTH_PLAN_STATE' column is not found in the DataFrame. + """ + if not isinstance(df, pd.DataFrame): # raise TypeError(f"Expected a dataframe, got {type(df).__name__}") return df @@ -472,32 +600,55 @@ def clean_health_plan_state_from_hotfix( return df health_plan_state = df["HEALTH_PLAN_STATE"].unique().tolist()[0] - + + print("Sending to state_check", health_plan_state) answer = state_check(health_plan_state, text_dict) df["HEALTH_PLAN_STATE"] = answer return df +def clean_notice_provider_name_and_address(final_df: pd.DataFrame) -> pd.DataFrame: + """ + Cleans the 'NOTICE_PROVIDER_NAME' and 'NOTICE_PROVIDER_ADDRESS' columns in the given DataFrame. -def clean_notice_provider_name_and_address(final_df): + This function sets the 'NOTICE_PROVIDER_NAME' and 'NOTICE_PROVIDER_ADDRESS' to NaN + for rows where the 'NOTICE_PROVIDER_NAME' contains "Superior HealthPlan". + + Args: + final_df (pd.DataFrame): The input DataFrame containing the columns to be cleaned. + + Returns: + pd.DataFrame: The cleaned DataFrame with updated 'NOTICE_PROVIDER_NAME' and 'NOTICE_PROVIDER_ADDRESS' values. + """ if "NOTICE_PROVIDER_NAME" in final_df and "NOTICE_PROVIDER_ADDRESS" in final_df: - final_df.loc[ - final_df["NOTICE_PROVIDER_NAME"].str.contains( - "Superior HealthPlan", na=False - ), - "NOTICE_PROVIDER_NAME", - ] = np.nan - final_df.loc[ - final_df["NOTICE_PROVIDER_NAME"].str.contains( - "Superior HealthPlan", na=False - ), - "NOTICE_PROVIDER_ADDRESS", - ] = np.nan + # Create a mask to identify rows where "Superior HealthPlan" is found + mask = final_df["NOTICE_PROVIDER_NAME"].str.contains("Superior HealthPlan", na=False) + + # Set "NOTICE_PROVIDER_NAME" to NaN where the condition is true + final_df.loc[mask, "NOTICE_PROVIDER_NAME"] = np.nan + + # Set "NOTICE_PROVIDER_ADDRESS" to NaN where the condition is true + final_df.loc[mask, "NOTICE_PROVIDER_ADDRESS"] = np.nan + return final_df +def get_tin_from_filename(final_df: pd.DataFrame) -> pd.DataFrame: + """ + Extracts TIN (Taxpayer Identification Number) from the filename or contract name + and updates the 'PROV_GROUP_TIN' column in the given DataFrame. -def get_tin_from_filename(final_df): + The function checks if the DataFrame contains the 'Contract Name' or 'Filename' column. + If found, it extracts the first 10 characters from the respective column, removes any hyphens, + and assigns this value to a temporary column 'temp_filename'. If 'PROV_GROUP_TIN' is empty or + contains only spaces, and 'temp_filename' is numeric, it updates 'PROV_GROUP_TIN' with the value from 'temp_filename'. Finally, it drops the 'temp_filename' column. + + Args: + final_df (pd.DataFrame): The input DataFrame containing the columns 'Contract Name' or 'Filename' and 'PROV_GROUP_TIN'. + + Returns: + pd.DataFrame: The updated DataFrame with 'PROV_GROUP_TIN' potentially modified based on the filename or contract name. + """ if "Contract Name" in final_df.columns: final_df["temp_filename"] = ( final_df["Contract Name"].str[:10].str.replace("-", "") @@ -536,34 +687,66 @@ def clean_policies_and_procedures(final_df): ] = "N/A" return final_df +# Function to format and validate TIN numbers +def format_tin(x: str) -> str: + """ + Formats a given string as a Taxpayer Identification Number (TIN). -def clean_tin(final_df): + The function extracts all digits from the input string and formats them as '##-#######' + if there are exactly 9 digits. If the input does not contain exactly 9 digits, it returns + 'N/A' or 'N/A - (model detected: {x})' depending on the input value. + + Args: + x (str): The input string to be formatted. + + Returns: + str: The formatted TIN or an appropriate 'N/A' message. + """ + digits = "".join(filter(str.isdigit, x)) + if len(digits) == 9: + # Format as '##-#######' + return f"{digits[:2]}-{digits[2:]}" + else: + # Retain prior handling for invalid entries + if x != "N/A": + return f"N/A - (model detected: {x})" + else: + return "N/A" + +def clean_tin(final_df: pd.DataFrame) -> pd.DataFrame: + """ + Cleans the 'PROV_GROUP_TIN' column in the given DataFrame by replacing + SSN-like strings with "N/A" and applying a formatting function. + + Args: + final_df (pd.DataFrame): The DataFrame containing the 'PROV_GROUP_TIN' column. + + Returns: + pd.DataFrame: The modified DataFrame with cleaned 'PROV_GROUP_TIN' values. + """ if "PROV_GROUP_TIN" in final_df: # Replace SSN-like strings with "N/A" final_df["PROV_GROUP_TIN"] = final_df["PROV_GROUP_TIN"].replace( r"\b\d{3}-\d{2}-\d{4}\b", "N/A", regex=True ) - # Function to format and validate TIN numbers - def format_tin(x): - digits = "".join(filter(str.isdigit, x)) - if len(digits) == 9: - # Format as '##-#######' - return f"{digits[:2]}-{digits[2:]}" - else: - # Retain prior handling for invalid entries - if x != "N/A": - return f"N/A - (model detected: {x})" - else: - return "N/A" - # Apply the formatting function to the 'PROV_GROUP_TIN' column final_df["PROV_GROUP_TIN"] = final_df["PROV_GROUP_TIN"].apply(format_tin) return final_df -def clean_tin_npi_other(final_df): +def clean_tin_npi_other(final_df: pd.DataFrame) -> pd.DataFrame: + """ + Cleans the 'PROV_TIN_OTHER' and 'PROV_NPI_OTHER' columns in the given DataFrame by removing + square brackets from the string values. + + Args: + final_df (pd.DataFrame): The DataFrame containing the columns to be cleaned. + + Returns: + pd.DataFrame: The cleaned DataFrame with square brackets removed from the specified columns. + """ if "PROV_TIN_OTHER" in final_df: final_df["PROV_TIN_OTHER"] = ( final_df["PROV_TIN_OTHER"] @@ -579,30 +762,93 @@ def clean_tin_npi_other(final_df): return final_df +def replace_quotes(value: Any) -> str: + """ + Replaces escaped double quotes in a string with actual double quotes. -def replace_quotes(value): + Args: + value (Any): The input value which will be converted to a string and processed. + + Returns: + str: The processed string with escaped double quotes replaced by actual double quotes. + """ try: return str(value).replace(r"\"", '"') except: return value +def replace_null_terms(value: Any) -> str: + """ + Replace null terms in the given value with "N/A". -def replace_null_terms(value): + This function converts the input value to a string and checks if any term from + the NULL_ANSWER_TERMS list (defined in the valid module) is present in the value. + If a match is found, it returns "N/A". Otherwise, it returns the original value. + + Args: + value (Any): The input value to be checked and potentially replaced. + + Returns: + str: "N/A" if a null term is found in the value, otherwise the original value. + """ # Convert value to string and check if any term from NULL_ANSWER_TERMS is in the value if any(term.lower() in str(value).lower() for term in valid.NULL_ANSWER_TERMS): return "N/A" return value +def filter_dict(d: dict, pattern: 're.Pattern') -> dict: + """ + Filters a dictionary by removing entries where the value matches a given pattern. -def filter_dict(d, pattern): - final_dict = {} - for page_num, answer in d.items(): - if not pattern.search(answer): - final_dict[page_num] = answer - return final_dict + Args: + d (dict): The dictionary to filter, where keys are page numbers and values are answers. + pattern (re.Pattern): The compiled regular expression pattern to search for in the answers. + Returns: + dict: A new dictionary with entries that do not match the pattern. + """ + final_dict = {} + for page_num, answer in d.items(): + if not pattern.search(answer): + final_dict[page_num] = answer + return final_dict -def filter_add_ons(df): +def check_add_on(group: pd.DataFrame) -> pd.DataFrame: + """ + Checks if any 'FULL_SERVICE' contains 'Add-on' or 'Add On' and updates the group accordingly. + + If any row in the group has 'FULL_SERVICE' containing 'Add-on' or 'Add On', this function sets + the 'ADD_ON_REIMBURSEMENT_LANGUAGE' to 'N/A' and 'ADD_ON_REIMBURSEMENT_IND' to 'N' for the entire group. + + Args: + group (pd.DataFrame): The input DataFrame group to be checked and updated. + + Returns: + pd.DataFrame: The updated DataFrame group with modified 'ADD_ON_REIMBURSEMENT_LANGUAGE' and 'ADD_ON_REIMBURSEMENT_IND' values. + """ + if ( + group["FULL_SERVICE"] + .str.contains("Add-on|Add On", regex=True, case=False, na=False) + .any() + ): + group["ADD_ON_REIMBURSEMENT_LANGUAGE"] = "N/A" + group["ADD_ON_REIMBURSEMENT_IND"] = "N" + return group + +def filter_add_ons(df: pd.DataFrame) -> pd.DataFrame: + """ + Filters and processes add-on reimbursement language in the given DataFrame. + + This function performs the following operations: + 1. Identifies rows where the 'ADD_ON_REIMBURSEMENT_LANGUAGE' column matches a specific pattern and sets the value to 'N/A' and 'ADD_ON_REIMBURSEMENT_IND' to 'N'. + 2. Groups the DataFrame by 'Contract Name' and 'EXHIBIT' or 'Filename' and 'EXHIBIT', and checks if any 'FULL_SERVICE' contains 'Add-on' or 'Add On'. If found, sets the 'ADD_ON_REIMBURSEMENT_LANGUAGE' to 'N/A' and 'ADD_ON_REIMBURSEMENT_IND' to 'N' for the whole group. + + Args: + df (pd.DataFrame): The input DataFrame containing the columns 'ADD_ON_REIMBURSEMENT_LANGUAGE', 'ADD_ON_REIMBURSEMENT_IND', 'FULL_SERVICE', and either 'Contract Name' or 'Filename', and 'EXHIBIT'. + + Returns: + pd.DataFrame: The processed DataFrame with updated 'ADD_ON_REIMBURSEMENT_LANGUAGE' and 'ADD_ON_REIMBURSEMENT_IND' values. + """ pattern = r"(in no event).+(includ.?)|(forward).+(payments)" mask = df["ADD_ON_REIMBURSEMENT_LANGUAGE"].str.contains( @@ -611,19 +857,7 @@ def filter_add_ons(df): df.loc[mask, "ADD_ON_REIMBURSEMENT_LANGUAGE"] = "N/A" df.loc[mask, "ADD_ON_REIMBURSEMENT_IND"] = "N" - def check_add_on(group): - # Check if any 'SERVICE_TYPE' contains 'Add-on' or 'Add On' - if ( - group["FULL_SERVICE"] - .str.contains("Add-on|Add On", regex=True, case=False, na=False) - .any() - ): - group["ADD_ON_REIMBURSEMENT_LANGUAGE"] = ( - "N/A" # Set 'ADD_ON' to 'N/A' for the whole group - ) - group["ADD_ON_REIMBURSEMENT_IND"] = "N" - return group - + # Group by 'Contract Name' and 'EXHIBIT', then apply the check_add_on function if "Contract Name" in df.columns: df = df.groupby(["Contract Name", "EXHIBIT"]).apply(check_add_on) @@ -633,10 +867,32 @@ def filter_add_ons(df): # TODO : If original add on value is actually an Exclusion, and the Exclusion value for the row is invalid, then move the Add On value to the Exclusions column return df +def count_dollar_values(s: str) -> int: + """ + Count the number of dollar signs ('$') and percentage signs ('%') in a given string. -def add_scmr(df): - def count_dollar_values(s): - return str(s).count("$") + str(s).count("%") + Args: + s (str): The input string to be analyzed. + + Returns: + int: The total count of dollar signs and percentage signs in the input string. + """ + return str(s).count("$") + str(s).count("%") + +def add_scmr(df: pd.DataFrame) -> pd.DataFrame: + """ + Adds columns to the DataFrame to indicate if a row has multiple rates for a single code. + This function checks for the presence of specific columns ("FLAT_FEE_STANDARD", "FLAT_FEE", "FLAT FEE") + in the DataFrame and uses the first one it finds. It then counts the number of dollar values in the + selected column and adds two new columns to the DataFrame: + - "Single Code Multiple Rates (Language)": Contains the value from the selected column if there are + multiple dollar values, otherwise "N/A". + - "Single Code Multiple Rates (Y/N)": Contains "Y" if there are multiple dollar values, otherwise "N". + Parameters: + df (pd.DataFrame): The input DataFrame containing the fee columns. + Returns: + pd.DataFrame: The modified DataFrame with the new columns added. + """ if "FLAT_FEE_STANDARD" in df.columns: col = "FLAT_FEE_STANDARD" @@ -733,23 +989,30 @@ def methodology_short_fix(df: pd.DataFrame) -> pd.DataFrame: return df - def yn_fixes(df: pd.DataFrame) -> pd.DataFrame: + """ + Fixes columns with (Y/N) values in the DataFrame by converting 'Yes' to 'Y' and 'No' to 'N'. + Also ensures that any other values are set to 'N'. Additionally, clears language columns + associated with 'Y/N' columns if the value is not 'Y'. + Args: + df (pd.DataFrame): The input DataFrame containing columns with (Y/N) values. + + Returns: + pd.DataFrame: The modified DataFrame with fixed (Y/N) values and cleared language columns. + """ for col in df.columns: if "(Y/N)" in col: df.loc[df[col] == "Yes", col] = "Y" df.loc[df[col] == "No", col] = "N" df.loc[(df[col] != "Y") & (df[col] != "N"), col] = "N" - for yn in valid.AC_IND_LANG_COLS.keys(): - if yn in df.columns and valid.AC_IND_LANG_COLS[yn] in df.columns: - lang = valid.AC_IND_LANG_COLS[yn] + for yn, lang in valid.AC_IND_LANG_COLS.items(): + if yn in df.columns and lang in df.columns: df.loc[df[yn] != "Y", lang] = "" - for yn in valid.B_IND_LANG_COLS.keys(): - if yn in df.columns and valid.B_IND_LANG_COLS[yn] in df.columns: - lang = valid.B_IND_LANG_COLS[yn] + for yn, lang in valid.B_IND_LANG_COLS.items(): + if yn in df.columns and lang in df.columns: df.loc[df[yn] != "Y", lang] = "" return df diff --git a/fieldExtraction/src/preprocessing_funcs.py b/fieldExtraction/src/preprocessing_funcs.py index e4af2cf..40e3d67 100644 --- a/fieldExtraction/src/preprocessing_funcs.py +++ b/fieldExtraction/src/preprocessing_funcs.py @@ -5,9 +5,9 @@ from itertools import chain, count, groupby from src.prompts import preprocessing_prompts import src.utils.llm_utils as llm_utils import src.utils.string_utils as string_utils -from src import config, keywords, smart_chunking_funcs -from src.regex.regex_patterns import PIPE_PATTERN +from src import config, keywords, smart_chunking_funcs +from src.enums.delimiters import Delimiter def clean_newlines(contract_text: str) -> str: """Clean textract output by removing newlines and page number indicators. @@ -22,8 +22,9 @@ def clean_newlines(contract_text: str) -> str: Returns: str: cleaned text with newlines and page number indicators removed """ - cleaned_text = re.sub(r"(? dict[str, str]: Returns: dict[str, str]: A dictionary, keyed by the string page number and valued by the page text. """ - temp_list = text.split("Start of Page No. = ") - text_list = re.split(r"Start of Page No. = [0-9]+\n", text) + + if isinstance(text, str): + if not text: + return {} + + temp_list = text.split("Start of Page No. = ") + text_list = re.split(r"Start of Page No. = [0-9]+\n", text) - text_dict = {} - for i in range(len(text_list)): - text_dict[temp_list[i].split()[0]] = text_list[i] + text_dict = {} + for i in range(len(text_list)): + text_dict[temp_list[i].split()[0]] = text_list[i] # splits on whitespace characters, which includes spaces, tabs, and newline characters - return {k: v for k, v in text_dict.items() if k != "Document"} + return {k: v for k, v in text_dict.items() if k != "Document"} def clean_law_symbols(contract_text): contract_text = contract_text.replace("$$", "$") - contract_text = re.sub(r"(U\.?S\.?C\.?) \$", r"\1§", contract_text) - contract_text = re.sub(r"(C\.?F\.?R\.?) \$", r"\1§", contract_text) + contract_text = re.sub(r"(U\.?S\.?C\.?) \$", r"\1§", contract_text) # replaces $ with § when it follows abbreviations like "U.S.C." (United States Code). + contract_text = re.sub(r"(C\.?F\.?R\.?) \$", r"\1§", contract_text) # replaces $ with § when it follows abbreviations like "C.F.R." (Code of Federal Regulations). # Second correction: Replace '$' with '§' when followed by a number with three decimal places contract_text = re.sub(r"\$(?=\d+\.\d{3})", "§", contract_text) @@ -247,7 +253,7 @@ def smart_chunk_ac( raise ValueError( "methodology must be 'hierarchy' or 'or' - using entire contract as chunk" ) - if field_group == "initial_page_group": + if field_group == "initial_page_group": #TODO: this needs to be removed? as it's not present in GROUPED_KEYWORD_MAPPINGS keyword_page_dict["initial_pages"] = [1, 2] # chunks[field_group] = concatenated page texts @@ -303,7 +309,9 @@ def get_exhibit_pages(text_dict, filename): claude_answer_raw = llm_utils.invoke_claude( prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=10 ) - claude_answer_extracted = re.findall(pattern=PIPE_PATTERN, string=claude_answer_raw)[-1] + claude_answer_extracted = string_utils.extract_text_from_delimiters( + claude_answer_raw, Delimiter.PIPE + ) if "Y" in claude_answer_extracted: exhibit_pages.append(page_num) return exhibit_pages diff --git a/fieldExtraction/tests/postprocessing_funcs_test.py b/fieldExtraction/tests/postprocessing_funcs_test.py index 5fc9d68..dc83ad0 100644 --- a/fieldExtraction/tests/postprocessing_funcs_test.py +++ b/fieldExtraction/tests/postprocessing_funcs_test.py @@ -1,13 +1,41 @@ + import pytest import pandas as pd import numpy as np +from src.postprocessing_funcs import check_add_on from src.postprocessing_funcs import ( InvalidDateException, yn_fixes, convert_to_us_date_format, + filter_service_column, + filter_methodology_column, + clean_td, + get_parent_agreement_code, + consolidate_subheader, + clean_msr_lesser, + clean_default_term, + clean_lesser_rate, + clean_prov_2, + clean_term_clause, + clean_auto_renewal_ind, + clean_npi, + clean_network_access_fees, + clean_health_plan_state, + get_health_plan_state, + state_check, + clean_health_plan_state_from_hotfix, + clean_notice_provider_name_and_address, + get_tin_from_filename, + clean_policies_and_procedures, + clean_tin, + clean_tin_npi_other, + filter_add_ons, + add_scmr, + add_hyphen_if_needed, + remove_unnamed_columns, + format_tin, + count_dollar_values ) - - class TestPostprocessingFuncs: @pytest.fixture def valid_dates(self): @@ -16,23 +44,18 @@ class TestPostprocessingFuncs: ("2023-01-12", "01/12/2023"), ("2023-13-01", "01/13/2023"), ] - @pytest.fixture def invalid_dates(self): return ["2023-13-32", "2023-00-01", "2023-12-00", "2023-12-32", "2023-13-13"] - @pytest.fixture def empty_dates(self): return ["", "N/A", "null", "none"] - @pytest.fixture def invalid_format_dates(self): return ["2023/12/01", "12-01-2023", "2023.12.01"] - @pytest.fixture def invalid_type_dates(self): return [None] - def test_convert_to_us_date_format( self, valid_dates, @@ -43,83 +66,560 @@ class TestPostprocessingFuncs: ): for input_date, expected_output in valid_dates: assert convert_to_us_date_format(input_date) == expected_output - for input_date in invalid_dates: with pytest.raises(InvalidDateException): convert_to_us_date_format(input_date) - for input_date in empty_dates: with pytest.raises(InvalidDateException): convert_to_us_date_format(input_date) - for input_date in invalid_format_dates: with pytest.raises(InvalidDateException): convert_to_us_date_format(input_date) - for input_date in invalid_type_dates: with pytest.raises(TypeError): convert_to_us_date_format(input_date) - + @pytest.fixture def sample_answer_dicts(self): return [ - { - "FULL_SERVICE": "This is a test statement containing keyword 'Lease'", - "FULL_METHODOLOGY": "PMPM", - }, - {"FULL_SERVICE": "unspecified", "FULL_METHODOLOGY": "valid"}, - { - "FULL_SERVICE": "This is an example of clean dict", - "FULL_METHODOLOGY": "Will be effective for 30 days", - }, + {"FULL_SERVICE": "Service Risk", "FULL_METHODOLOGY": "Methodology interest"}, + {"FULL_SERVICE": "Service B", "FULL_METHODOLOGY": "PMPM"}, + {"FULL_SERVICE": "damages", "FULL_METHODOLOGY": "Invalid Methodology"}, ] - - @pytest.fixture - def sample_df(self): - data = { - "FULL_SERVICE": ["Service A", "Service B", "Invalid Service"], - "FULL_METHODOLOGY": [ - "Methodology X", - "Methodology Y", - "Invalid Methodology", - ], - "LESSER": ["N", "Y", "N"], - "LESSER_RATE": [100, 200, 300], - "RATE_STANDARD": [150, 250, 350], - "DEFAULT_TERM": ["Outpatient", "Inpatient", "Invalid"], - "DEFAULT_RATE": ["AC", "BC", "Invalid"], - "IP_OP": ["OP", "IP", "OP"], - "PROV_TYPE_LEVEL_2": ["", "Type A", "Type B"], - "TERM_CLAUSE": [ - "IL-4 Termination", - "bonus payment shall be effective", - "No term or termination", - ], - "CONTRACT_AUTO_RENEWAL_IND": ["Y", "N", "Y"], - "CONTRACT_TERMINATION_DT": [np.nan, "2023-12-31", np.nan], - "PROV_GROUP_NPI": ["1234567890", "123", "N/A"], - "NETWORK_ACCESS_FEES_IND": ["Yes", "No", "Yes"], - "PAYER_NAME": ["Illini Health", "Superior HealthPlan", "Other Payer"], - "HEALTH_PLAN_STATE": ["IL", "TX", "CA"], - "NOTICE_PROVIDER_NAME": ["Superior HealthPlan", "Provider A", "Provider B"], - "NOTICE_PROVIDER_ADDRESS": ["Address A", "Address B", "Address C"], - "POLICIES_AND_PROCEDURES": ["Policy A", "Procedure B", "Invalid"], - "PROV_GROUP_TIN": ["12-3456789", "123456789", "N/A"], - "ADD_ON_REIMBURSEMENT_LANGUAGE": [ - "in no event includ", - "forward payments", - "N/A", - ], - "ADD_ON_REIMBURSEMENT_IND": ["Y", "N", "Y"], - "FLAT_FEE_STANDARD": ["$100", "$200", "$300"], - "CONTRACT_LOB": ["LOB A", "LOB B", "LOB C"], - "Parent Agreement Code": ["1234", "5678", "91011"], - "RATE_SHORT": ["$100", "$200", "$300"], - "DUMMY_COL(Y/N)": ["Yes", "No", ""], - } - return pd.DataFrame(data) - - def test_yn_fixes(self, sample_df): - cleaned_df = yn_fixes(sample_df) - assert cleaned_df.loc[1, "ADD_ON_REIMBURSEMENT_LANGUAGE"] == "" - assert cleaned_df["DUMMY_COL(Y/N)"].tolist() == ["Y", "N", "N"] + + def test_filter_service_column(self, sample_answer_dicts): + filtered = filter_service_column(sample_answer_dicts) + assert len(filtered) == 1 + assert all(d["FULL_SERVICE"] == "Service B" for d in filtered) + def test_filter_methodology_column(self, sample_answer_dicts): + filtered = filter_methodology_column(sample_answer_dicts) + assert len(filtered) == 1 + assert all(d["FULL_METHODOLOGY"] == "Invalid Methodology" for d in filtered) + @pytest.mark.parametrize("input_td, expected_output", [ + ( + {"DATE": "01/01/2023", "Health Plan State": "Chicago", "page_num": 1, "Filename": "file1.txt"}, + {"DATE": ["01/01/2023"], "Health Plan State": ["Chicago"], "page_num": 1, "Filename": "file1.txt"} + ), + ( + {"DATE": "01/01/2023", "Health Plan State": ["New York", "Michigan"], "page_num": 1, "Filename": "file1.txt"}, + {"DATE": ["01/01/2023"], "Health Plan State": ["New York", "Michigan"], "page_num": 1, "Filename": "file1.txt"} + ), + ( + {"DATE": "N/A", "Health Plan State": "N/A", "page_num": 2, "Filename": "file2.txt"}, + {"DATE": ["N/A"], "Health Plan State": [], "page_num": 2, "Filename": "file2.txt"} + ), + ( + {"DATE": "02/01/2023", "Health Plan State": "California, Texas", "page_num": 3, "Filename": "file3.txt"}, + {"DATE": ["02/01/2023"], "Health Plan State": ["California", "Texas"], "page_num": 3, "Filename": "file3.txt"} + ), + ]) + def test_clean_td(self, input_td, expected_output): + cleaned = clean_td([input_td]) + assert cleaned[0] == expected_output + @pytest.mark.parametrize("filename, expected_code", [ + ("contract_1234(1).txt", "1234"), + ("contract_5678.txt", "5678"), + ("invalid_filename.txt", "N/A"), + ]) + def test_get_parent_agreement_code(self, filename, expected_code): + code = get_parent_agreement_code(filename) + assert code == expected_code + + @pytest.mark.parametrize("input_dict, expected_output", [ + ({"FULL_SERVICE": "Service A", "SUBHEADER": "Header A"}, {"FULL_SERVICE": "Header A - Service A"}), + ({"FULL_SERVICE": "Service B", "SUBHEADER": ""}, {"FULL_SERVICE": " - Service B"}), + ({"FULL_SERVICE": "Service C", "SUBHEADER": "N/A"}, {"FULL_SERVICE": "Service C"}), + ]) + def test_consolidate_subheader(self, input_dict, expected_output): + consolidated = consolidate_subheader([input_dict]) + assert consolidated[0] == expected_output + @pytest.mark.parametrize("input_df, expected_lesser, expected_lesser_rate", [ + ( + pd.DataFrame({ + "FULL_SERVICE": ["Multiple Outpatient", "Multiple Outpatient", "Not multiple"], + "EXHIBIT": ["Type A", "Type A", "Type B"], + "LESSER": ["", "Y", "N"], + "LESSER_RATE": [100, 200, 300] + }), + ["Y", "Y", "N"], + [200, 200, 300] + ), + ( + pd.DataFrame({ + "FULL_SERVICE": ["Multiple Outpatient", "Multiple Outpatient", "Not multiple"], + "EXHIBIT": ["Type A", "Type A", "Type B"], + "LESSER": ["Y", "Y", "N"], + "LESSER_RATE": [100, 200, 300] + }), + ["Y", "Y", "N"], + [100, 100, 300] + ), + ]) + def test_clean_msr_lesser(self, input_df, expected_lesser, expected_lesser_rate): + cleaned_df = clean_msr_lesser(input_df) + assert cleaned_df["LESSER"].tolist() == expected_lesser + assert cleaned_df["LESSER_RATE"].tolist() == expected_lesser_rate + @pytest.mark.parametrize("input_df, expected_df", [ + ( + pd.DataFrame({ + "DEFAULT_TERM": ["Except as otherwise provided"], + "DEFAULT_RATE": ["Invalid"], + "IP_OP": ["IP"] + }), + pd.DataFrame({ + "DEFAULT_TERM": ["N/A"], + "DEFAULT_RATE": ["N/A"], + "IP_OP": ["IP"] + }) + ), + ( + pd.DataFrame({ + "DEFAULT_TERM": ["Medicare Outpatient", "Medicaid Inpatient", "Invalid"], + "DEFAULT_RATE": ["AC", "BC", "Amount Payable by Medicaid:NA"], + "IP_OP": ["IP", "OP", "IP"] + }), + pd.DataFrame({ + "DEFAULT_TERM": ["N/A", "N/A", "Invalid"], + "DEFAULT_RATE": ["N/A", "N/A", "Amount Payable by MCD:NA"], + "IP_OP": ["IP", "OP", "IP"] + }) + ), + ]) + def test_clean_default_term(self, input_df, expected_df): + cleaned_df = clean_default_term(input_df) + pd.testing.assert_frame_equal(cleaned_df, expected_df) + @pytest.mark.parametrize("input_df, expected_output", [ + ( + pd.DataFrame({ + "LESSER_RATE": [np.nan, "200", "300"], + "RATE_STANDARD": ["150", "250", "350"], + "LESSER": ["Y", "Y", "N"] + }), + pd.DataFrame({ + "LESSER_RATE": ["150", "200", "300"], + "RATE_STANDARD": ["N/A", "250", "350"], + "LESSER": ["Y", "Y", "N"] + }) + ) + ]) + def test_clean_lesser_rate(self, input_df, expected_output): + cleaned_df = clean_lesser_rate(input_df) + pd.testing.assert_frame_equal(cleaned_df, expected_output) + + @pytest.mark.parametrize("input_df, expected_output", [ + ( + pd.DataFrame({ + "PROV_TYPE_LEVEL_2": ["", "Type A", "Type B"], + "FULL_SERVICE": ["Specialist", "ASC", "Invalid Service - Skilled Nursing Facility"], + "EXHIBIT": ["Exhibit A", "Exhibit B", "Exhibit C"], + "PROV_TYPE": ["Professional", "Ancilliary","Invalid"] + }), + ["Specialist", "ASC", "Skilled Nursing Facility"], + ), + ( + pd.DataFrame({ + "PROV_TYPE_LEVEL_2": [""], + "FULL_SERVICE": [""], + "EXHIBIT": ["This has keyword, Primary Care Provider"], + "PROV_TYPE": ["Invalid"] + }), + ["Primary Care Provider"], + ), + ( + pd.DataFrame({ + "PROV_TYPE_LEVEL_2": ["", "Lab","",""], + "FULL_SERVICE": ["", "ASC","","Hospital"], + "EXHIBIT": ["Invalid", "Invalid","Invalid2","Invalid2"], + "PROV_TYPE": ["", "Ancilliary","","Facility"] + }), + ["Lab", "ASC", "Hospital", "Hospital"], + ) + ]) + def test_clean_prov_2(self, input_df, expected_output): + cleaned_df = clean_prov_2(input_df) + cleaned_df['PROV_TYPE_LEVEL_2'].tolist() == expected_output + @pytest.mark.parametrize("input_df, expected_output", [ + ( + pd.DataFrame({ + "TERM_CLAUSE": ["IL-4 Termination", "bonus payment shall be effective", "No term or termination", "does not contain", "valid term clause"] + }), + pd.DataFrame({ + "TERM_CLAUSE": ["N/A", "N/A", "N/A", "N/A", "valid term clause"] + }) + ) + ]) + def test_clean_term_clause(self, input_df, expected_output): + cleaned_df = clean_term_clause(input_df) + pd.testing.assert_frame_equal(cleaned_df, expected_output) + @pytest.mark.parametrize("input_df, column, expected_output", [ + ( + pd.DataFrame({ + "TERM_CLAUSE": ["Contract automatically renews", "coterminous with the agreement", "N/A", " "], + "CONTRACT_AUTO_RENEWAL_IND": ["N", "N", "N", "N"]}), + "CONTRACT_AUTO_RENEWAL_IND", + ["Y", "Y", "N", "N"] + ), + ( + pd.DataFrame({ + "TERM_CLAUSE": ["N/A", ""], + "CONTRACT_TERMINATION_DT": ["Y", "Y"]}), + "CONTRACT_TERMINATION_DT", + [np.nan, np.nan] + ), + ( + pd.DataFrame({ + "CONTRACT_AUTO_RENEWAL_IND": ["Yes", "Y", "N"], + "CONTRACT_TERMINATION_DT": ["", "N/A", "12/31/2022"], + }), + "CONTRACT_TERMINATION_DT", + [np.nan, np.nan, "12/31/2022"] + ), + ( + pd.DataFrame({ + "CONTRACT_AUTO_RENEWAL_IND": ["Y",""], + }), + "CONTRACT_AUTO_RENEWAL_IND", + ["Y", "N"] + ), + ( + pd.DataFrame({ + "Contract Name": ["1234_contract.txt"]}), + "CONTRACT_AUTO_RENEWAL_IND", + ["N"] + ) + ]) + def test_clean_auto_renewal_ind(self, input_df, column, expected_output): + cleaned_df = clean_auto_renewal_ind(input_df) + assert cleaned_df[column].tolist() == expected_output + @pytest.mark.parametrize("input_df, expected_npi", [ + ( + pd.DataFrame({"PROV_GROUP_NPI": ["1234567890", "123", "N/A"]}), + ["1234567890", "N/A - (model detected: 123)", "N/A"] + ) + ]) + def test_clean_npi(self, input_df, expected_npi): + cleaned_df = clean_npi(input_df) + assert cleaned_df["PROV_GROUP_NPI"].tolist() == expected_npi + @pytest.mark.parametrize("input_df, expected_output", [ + ( + pd.DataFrame({ + "NETWORK_ACCESS_FEES_IND": ["Yes", "N/A", np.nan, None, "nan","Test", "No"] + }), + ["Yes", "N/A", np.nan, None, "nan", "Yes", "Yes"] + ) + ]) + def test_clean_network_access_fees(self, input_df, expected_output): + cleaned_df = clean_network_access_fees(input_df) + assert cleaned_df["NETWORK_ACCESS_FEES_IND"].tolist() == expected_output + @pytest.mark.parametrize("input_df, expected_state", [ + (pd.DataFrame({"HEALTH_PLAN_STATE": ["IL", "TX", "CA"], + "PAYER_NAME":["Illinios Payer","",""]}), ["Illinois","Texas","California"]), + (pd.DataFrame({"HEALTH_PLAN_STATE": ["TX", "TX", "CA"]}), ["Texas","Texas","California"]), + (pd.DataFrame({"HEALTH_PLAN_STATE": ["CA", "TX", "CA"]}), ["California","Texas","California"]), + ]) + def test_clean_health_plan_state(self, input_df, expected_state): + cleaned_df = clean_health_plan_state(input_df) + assert cleaned_df["HEALTH_PLAN_STATE"].tolist() == expected_state + @pytest.mark.parametrize("text_dict, expected_state", [ + ({"page1": "This contract is for Illinois and Texas."}, "Multiple"), + ({"page1": "This contract is for IL."}, "Illinois"), + ({"page1": "This contract is for TX and Texas."}, "Texas"), + ({"page1": "This contract is for California."}, "California"), + ({"page1": "This contract is for New York."}, "New York"), + ({"page1": "This contract is for Illinois, Texas, and California."}, "Multiple"), + ({"page1": "This contract is for Illinois and California."}, "Multiple"), + ({"page1": "This contract is for Texas and New York."}, "Multiple"), + ({"page1": "This contract is for."}, "N/A"), + ({"page1": ""}, "N/A"), + ]) + def test_get_health_plan_state(self, text_dict, expected_state): + state = get_health_plan_state(text_dict) + assert state == expected_state + @pytest.mark.parametrize("state_abbr, text_dict, expected_state", [ + ("IL", {"page1": "This contract is for Illinois and Texas."}, "Illinois"), + ("TX", {"page1": "This contract is for Illinois and Texas."}, "Texas"), + ("CA", {"page1": "This contract is for California."}, "California"), + ("NY", {"page1": "This contract is for New York."}, "New York"), + ("FL", {"page1": "This contract is for Hlorida."}, "Florida"), + ("", {"page1": "This contract is for Illinois and Texas."}, "Multiple"), + ]) + def test_state_check(self, state_abbr, text_dict, expected_state): + state = state_check(state_abbr, text_dict) + assert state == expected_state + + @pytest.mark.parametrize("input_df, contract_name, text_dict, expected_state", [ + ( + pd.DataFrame({"HEALTH_PLAN_STATE": [["IL", "TX", "CA"]]}), + "contract_1234_v1.txt", + {"page1": "This is a sample contract contract"}, + "Multiple" + ), + ( + pd.DataFrame({"HEALTH_PLAN_STATE": ["IL,"]}), + "contract_1234_v1.txt", + {"page1": "This is a sample contract"}, + "Illinois" + ), + ( + pd.DataFrame({"HEALTH_PLAN_STATE": [""]}), + "contract_1234_v1.txt", + {"page1": "This contract is for Illinois and Texas."}, + "Multiple" + ) + ]) + def test_clean_health_plan_state_from_hotfix(self, input_df, contract_name, text_dict, expected_state): + input_df["HEALTH_PLAN_STATE"] = input_df["HEALTH_PLAN_STATE"].apply( + lambda x: f"[{','.join(x)}]" if isinstance(x, list) else x + ) + cleaned_df = clean_health_plan_state_from_hotfix(input_df, contract_name, text_dict) + assert cleaned_df.loc[0, "HEALTH_PLAN_STATE"] == expected_state + @pytest.mark.parametrize("input_df, expected_name, expected_address", [ + ( + pd.DataFrame({ + "NOTICE_PROVIDER_NAME": ["Superior HealthPlan", "Provider A", "Provider B"], + "NOTICE_PROVIDER_ADDRESS": ["Address A", "Address B", "Address C"] + }), + [np.nan, "Provider A", "Provider B"], + [np.nan, "Address B", "Address C"] + ), + ]) + def test_clean_notice_provider_name_and_address(self, input_df, expected_name, expected_address): + cleaned_df = clean_notice_provider_name_and_address(input_df) + assert cleaned_df["NOTICE_PROVIDER_NAME"].tolist() == expected_name + assert cleaned_df["NOTICE_PROVIDER_ADDRESS"].tolist() == expected_address + + @pytest.mark.parametrize("input_df, expected_tin", [ + (pd.DataFrame({"Contract Name":["12-3456789_contract(1)", "12-3456789_contract(2)"], + "PROV_GROUP_TIN": ["", "123456789"]} + ), ["123456789", "123456789"] + ), + ]) + def test_get_tin_from_filename(self, input_df, expected_tin): + cleaned_df = get_tin_from_filename(input_df) + assert cleaned_df["PROV_GROUP_TIN"].tolist() == expected_tin + @pytest.mark.parametrize("input_df, expected_output", [ + ( + pd.DataFrame({ + "POLICIES_AND_PROCEDURES": ["Policy A", "Procedure B", "Invalid", ""] + }), + pd.DataFrame({ + "POLICIES_AND_PROCEDURES": ["Policy A", "Procedure B", "N/A", "N/A"] + }) + ), + ]) + def test_clean_policies_and_procedures(self, input_df, expected_output): + cleaned_df = clean_policies_and_procedures(input_df) + pd.testing.assert_frame_equal(cleaned_df, expected_output) + @pytest.mark.parametrize("input_df, expected_tin", [ + (pd.DataFrame({"PROV_GROUP_TIN": ["123-45-6789"]}), "N/A"), + (pd.DataFrame({"PROV_GROUP_TIN": ["12-34-56"]}), "N/A - (model detected: 12-34-56)"), + (pd.DataFrame({"PROV_GROUP_TIN": ["12-3456789"]}), "12-3456789"), + (pd.DataFrame({"PROV_GROUP_TIN": ["N/A"]}), "N/A"), + ]) + def test_clean_tin(self, input_df, expected_tin): + cleaned_df = clean_tin(input_df) + assert cleaned_df.loc[0, "PROV_GROUP_TIN"] == expected_tin + @pytest.mark.parametrize("input_df, column, expected_tin", [ + (pd.DataFrame({"PROV_TIN_OTHER": [["123456789", "123456789"], "N/A"]}), "PROV_TIN_OTHER", ["123456789,123456789","N/A"]), + (pd.DataFrame({"PROV_NPI_OTHER": ["12-3456789"]}), "PROV_NPI_OTHER", ["12-3456789"]), + (pd.DataFrame({"PROV_TIN_OTHER": [""]}), "PROV_TIN_OTHER", [""]), + (pd.DataFrame({"PROV_NPI_OTHER": [""]}), "PROV_NPI_OTHER", [""]), + ]) + + def test_clean_tin_npi_other(self, input_df, column, expected_tin): + input_df[column] = input_df[column].apply(lambda x: f"[{','.join(x)}]" if isinstance(x, list) else x) + cleaned_df = clean_tin_npi_other(input_df) + assert cleaned_df[column].tolist() == expected_tin + @pytest.mark.parametrize("input_df, expected_output", [ + ( + pd.DataFrame({ + "ADD_ON_REIMBURSEMENT_LANGUAGE": ["in no event", "forward", "Valid"], + "ADD_ON_REIMBURSEMENT_IND": ["Y", "N", "Y"], + "FULL_SERVICE": ["Add-on service", "No add-on", "Valid service"], + "Contract Name": ["Contract A", "Contract A", "Contract B"], + "EXHIBIT": ["Exhibit 1", "Exhibit 1", "Exhibit 2"] + }), + pd.DataFrame({ + "ADD_ON_REIMBURSEMENT_LANGUAGE": ["N/A", "N/A", "Valid"], + "ADD_ON_REIMBURSEMENT_IND": ["N", "N", "Y"], + "FULL_SERVICE": ["Add-on service", "No add-on", "Valid service"], + "Contract Name": ["Contract A", "Contract A", "Contract B"], + "EXHIBIT": ["Exhibit 1", "Exhibit 1", "Exhibit 2"] + }) + ), + ( + pd.DataFrame({ + "ADD_ON_REIMBURSEMENT_LANGUAGE": ["", "additional fees", "N/A"], + "ADD_ON_REIMBURSEMENT_IND": ["Y", "Y", "N"], + "FULL_SERVICE": ["No add-on", "Add On service", "N/A"], + "Filename": ["File1", "File1", "File2"], + "EXHIBIT": ["Exhibit 1", "Exhibit 1", "Exhibit 2"] + }), + pd.DataFrame({ + "ADD_ON_REIMBURSEMENT_LANGUAGE": ["N/A", "N/A", "N/A"], + "ADD_ON_REIMBURSEMENT_IND": ["N", "N", "N"], + "FULL_SERVICE": ["No add-on", "Add On service", "N/A"], + "Filename": ["File1", "File1", "File2"], + "EXHIBIT": ["Exhibit 1", "Exhibit 1", "Exhibit 2"] + }) + ) + ]) + + @pytest.mark.filterwarnings("ignore::UserWarning") + @pytest.mark.filterwarnings("ignore::DeprecationWarning") + def test_filter_add_ons(self, input_df, expected_output): + # Call the function and reset the index + cleaned_df = filter_add_ons(input_df).reset_index(drop=True) + + # Assert that the cleaned DataFrame matches the expected output + pd.testing.assert_frame_equal(cleaned_df, expected_output) + + @pytest.mark.parametrize("input_df, expected_df", [ + (pd.DataFrame({"FLAT_FEE_STANDARD": ["$100", "$200 30%", "$400"]}), + pd.DataFrame({ + "Single Code Multiple Rates (Language)": ["N/A", "$200 30%", "N/A"], + "Single Code Multiple Rates (Y/N)": ["N", "Y", "N"] + })) + ]) + + def test_add_scmr(self, input_df, expected_df): + cleaned_df = add_scmr(input_df) + + pd.testing.assert_series_equal( + cleaned_df["Single Code Multiple Rates (Language)"], + expected_df["Single Code Multiple Rates (Language)"], + ) + pd.testing.assert_series_equal( + cleaned_df["Single Code Multiple Rates (Y/N)"], + expected_df["Single Code Multiple Rates (Y/N)"], + ) + @pytest.mark.parametrize("input_df, expected_df", [ + ( + pd.DataFrame({ + "DUMMY_COL(Y/N)": ["Yes", "No", ""], + "SINGLE_CODE_MULTIPLE_RATES_IND": ["Y", "N", "Y"], + "SINGLE_CODE_MULTIPLE_RATES_LANGUAGE": ["Language 1", "Language 2", "N/A"], + "RATE_ESCALATOR_IND": ["Y", "N", "N"], + "RATE_ESCALATOR_DT": ["12/1/2022", "12/11/2022", "01/05/2018"], + }), + pd.DataFrame({ + "DUMMY_COL(Y/N)": ["Y", "N", "N"], + "SINGLE_CODE_MULTIPLE_RATES_IND": ["Y", "N", "Y"], + "SINGLE_CODE_MULTIPLE_RATES_LANGUAGE": ["Language 1", "", "N/A"], + "RATE_ESCALATOR_IND": ["Y", "N", "N"], + "RATE_ESCALATOR_DT": ["12/1/2022", "", ""], + }) + ), + + ]) + def test_yn_fixes(self, input_df, expected_df): + cleaned_df = yn_fixes(input_df) + pd.testing.assert_frame_equal(cleaned_df, expected_df) + + @pytest.mark.parametrize("input_tin, expected_output", [ + ("12345678", "01-2345678"), + ("123456789", "12-3456789"), + ("1-2345678", "01-2345678"), + ("123-45", None), + ("", None), + ]) + + def test_add_hyphen_if_needed(self, input_tin, expected_output): + assert add_hyphen_if_needed(input_tin) == expected_output + @pytest.mark.parametrize("input_date, expected_output", [ + ("2023-12-31", "12/31/2023"), + ("2023-01-12", "01/12/2023"), + ]) + def test_convert_to_us_date_format_valid(self, input_date, expected_output): + assert convert_to_us_date_format(input_date) == expected_output + @pytest.mark.parametrize("input_date", [ + "N/A", + "2023-13-32", + "2023-00-01", + "2023-12-00", + "2023-12-32", + "2023-13-13", + "", + "null", + "none", + "2023/12/01", + "12-01-2023", + "2023.12.01", + None, + ]) + def test_convert_to_us_date_format_invalid(self, input_date): + with pytest.raises((InvalidDateException, TypeError)): + convert_to_us_date_format(input_date) + @pytest.mark.parametrize("input_df, expected_columns", [ + (pd.DataFrame({"Unnamed: 0": [1, 2], "A": [3, 4]}), ["A"]), + (pd.DataFrame({"Unnamed: 0": [1, 2], "Unnamed: 1": [3, 4], "B": [5, 6]}), ["B"]), + (pd.DataFrame({"A": [1, 2], "B": [3, 4]}), ["A", "B"]), + ]) + def test_remove_unnamed_columns(self, input_df, expected_columns): + cleaned_df = remove_unnamed_columns(input_df) + assert list(cleaned_df.columns) == expected_columns + @pytest.mark.parametrize("input_tin, expected_output", [ + ("123456789", "12-3456789"), + ("12-3456789", "12-3456789"), + ("123-45-6789", "12-3456789"), + ("N/A", "N/A"), + ("invalidTIN", "N/A - (model detected: invalidTIN)"), + ("", "N/A - (model detected: )"), + ]) + def test_format_tin(self, input_tin, expected_output): + assert format_tin(input_tin) == expected_output + @pytest.mark.parametrize("input_str, expected_count", [ + ("$100", 1), + ("100%", 1), + ("$100 and 50%", 2), + ("No dollar or percentage signs", 0), + ("$$$ and %%%", 6), + ("$100 $200 $300", 3), + ("100% 200% 300%", 3), + ("", 0), + ("$%$%", 4), + ]) + def test_count_dollar_values(self, input_str, expected_count): + assert count_dollar_values(input_str) == expected_count + + @pytest.mark.parametrize("input_df, expected_df", [ + ( + pd.DataFrame({ + "FULL_SERVICE": ["Add-on Service", "Add-on Service"], + "ADD_ON_REIMBURSEMENT_LANGUAGE": ["Language A", "Language B"], + "ADD_ON_REIMBURSEMENT_IND": ["Y", "Y"] + }), + pd.DataFrame({ + "FULL_SERVICE": ["Add-on Service", "Add-on Service"], + "ADD_ON_REIMBURSEMENT_LANGUAGE": ["N/A", "N/A"], + "ADD_ON_REIMBURSEMENT_IND": ["N", "N"] + }) + ), + ( + pd.DataFrame({ + "FULL_SERVICE": ["Add On Service", "Add On Service"], + "ADD_ON_REIMBURSEMENT_LANGUAGE": ["Language A", "Language B"], + "ADD_ON_REIMBURSEMENT_IND": ["Y", "Y"] + }), + pd.DataFrame({ + "FULL_SERVICE": ["Add On Service", "Add On Service"], + "ADD_ON_REIMBURSEMENT_LANGUAGE": ["N/A", "N/A"], + "ADD_ON_REIMBURSEMENT_IND": ["N", "N"] + }) + ), + ( + pd.DataFrame({ + "FULL_SERVICE": ["Regular Service", "Regular Service"], + "ADD_ON_REIMBURSEMENT_LANGUAGE": ["Language A", "Language B"], + "ADD_ON_REIMBURSEMENT_IND": ["Y", "Y"] + }), + pd.DataFrame({ + "FULL_SERVICE": ["Regular Service", "Regular Service"], + "ADD_ON_REIMBURSEMENT_LANGUAGE": ["Language A", "Language B"], + "ADD_ON_REIMBURSEMENT_IND": ["Y", "Y"] + }) + ) + ]) + def test_check_add_on(self, input_df, expected_df): + cleaned_df = check_add_on(input_df) + pd.testing.assert_frame_equal(cleaned_df, expected_df) diff --git a/fieldExtraction/tests/preprocessing_funcs_test.py b/fieldExtraction/tests/preprocessing_funcs_test.py new file mode 100644 index 0000000..16f422b --- /dev/null +++ b/fieldExtraction/tests/preprocessing_funcs_test.py @@ -0,0 +1,157 @@ +import pytest +from src.preprocessing_funcs import ( + clean_newlines, + split_text, + clean_law_symbols, + smart_chunk_ac, + filter_quick_review, + get_exhibit_pages, +) + +class TestPreprocessingFuncs: + # Test cases for clean_newlines + @pytest.mark.parametrize("input_text, expected_output", [ + # Removes single newlines + ("This is a test.\nThis is only a test.", "This is a test. This is only a test."), + # Removes page number indicators + ("Page 1 of 10\n\nThis is a test.", " This is a test."), + # Handles multiple newlines + ("Line 1\n\nLine 2\n\nLine 3", "Line 1\n\nLine 2\n\nLine 3"), + # Handles no newlines + ("This is a single line.", "This is a single line."), + # Handles mixed newlines and page numbers + ("Page 1 of 10\n\nLine 1\nLine 2\nPage 2 of 10\n\nLine 3", " Line 1 Line 2 Line 3"), + ]) + def test_clean_newlines(self, input_text, expected_output): + assert clean_newlines(input_text) == expected_output + + # Test cases for split_text + @pytest.mark.parametrize("input_text, expected_output", [ + # Splits on page markers + ( + "Title Start of Page No. = 1\nContent1\nStart of Page No. = 2\nContent2", + {"Title" : "Title ", "1": "Content1\n", "2": "Content2"} + ), + # Ignores "Document" key + ( + "Document\nStart of Page No. = 1\nContent1", + {"1": "Content1"} + ), + # Handles empty input + ("", {}), + # Handles no page markers + ("This is a single page.", {"This": "This is a single page."}), + # Handles multiple page markers + ( + "Introduction\nStart of Page No. = 1\nContent1\nStart of Page No. = 2\nContent2\nStart of Page No. = 3\nContent3", + {'Introduction': 'Introduction\n', '1': 'Content1\n', '2': 'Content2\n', '3': 'Content3'} + ), + ]) + def test_split_text(self, input_text, expected_output): + assert split_text(input_text) == expected_output + + # Test cases for clean_law_symbols + @pytest.mark.parametrize("input_text, expected_output", [ + # Replaces double dollar signs + ("This is a $$ test.", "This is a $ test."), + # Replaces U.S.C. symbol + ("U.S.C. $1234", "U.S.C.§1234"), + # Replaces C.F.R. symbol + ("C.F.R. $1234", "C.F.R.§1234"), + # Replaces $ followed by three decimal places + ("$123.456", "§123.456"), + # Handles no replacements + ("This is a normal text.", "This is a normal text."), + # Handles mixed replacements + ("U.S.C. $1234 and C.F.R. $5678 and $123.456", "U.S.C.§1234 and C.F.R.§5678 and §123.456"), + ]) + def test_clean_law_symbols(self, input_text, expected_output): + assert clean_law_symbols(input_text) == expected_output + + # Test cases for smart_chunk_ac + @pytest.mark.parametrize("methodology, mock_return, chunking_fn_name, expected_chunk, expected_pages", [ + # OR methodology + ("or", [1, 2], "or", "Content1\nContent2", {"test": [1, 2]}), + # Hierarchy methodology + ("hierarchy", [1, 2], "hierarchical", "Content1\nContent2", {"test": [1, 2]}), + # AND methodology + ("and", [1], "and", "Content1", {"test": [1]}), + # Regex methodology + ("regex", [1, 2], "regex", "Content1\nContent2", {"test": [1, 2]}), + # Include/Exclude methodology + ("include_exclude", [1], "with_include_exclude_keywords", "Content1", {"test": [1]}), + ]) + + def test_smart_chunk_ac(self, mocker, methodology, mock_return, chunking_fn_name, expected_chunk, expected_pages): + # Mock the chunking functions + mock_chunk_func = mocker.patch(f"src.smart_chunking_funcs.chunk_{chunking_fn_name}", return_value=mock_return) + mock_keyword_search = mocker.patch("src.smart_chunking_funcs.keyword_search", return_value={"test": mock_return}) + + text_dict = {"1": "Content1", "2": "Content2"} + keyword_mappings = {"group1": {"methodology": methodology, "keywords": ["test"], "case_sensitive": False}} + + if methodology == "regex": keyword_mappings["group1"]["regex"] = "Content" + + if methodology == "include_exclude": + del keyword_mappings["group1"]["keywords"] + keyword_mappings["group1"]["included_keywords"] = ["Content"] + keyword_mappings["group1"]["excluded_keywords"] = [] + + result = smart_chunk_ac(text_dict, "Content1\nContent2", keyword_mappings) + + assert result["group1"] == expected_chunk + assert result["group1_pages"] == expected_pages + mock_chunk_func.assert_called_once() + mock_keyword_search.assert_called_once() + + # Test cases for filter_quick_review + @pytest.mark.parametrize("input_dict, expected_contract, expected_quick_review", [ + # Filters quick review pages + ( + {"1": "QUICK REVIEW", "2": "Normal content", "3": "COVER SHEET"}, + {"2": "Normal content"}, + {"1": "QUICK REVIEW", "3": "COVER SHEET"} + ), + # No quick review pages + ( + {"1": "Normal content", "2": "More content"}, + {"1": "Normal content", "2": "More content"}, + {} + ), + # All quick review pages + ( + {"1": "QUICK REVIEW", "2": "TOP SHEET"}, + {}, + {"1": "QUICK REVIEW", "2": "TOP SHEET"} + ), + # Case-insensitive matching + ( + {"1": "quick review", "2": "Normal content"}, + {"2": "Normal content"}, + {"1": "quick review"} + ), + ]) + def test_filter_quick_review(self, input_dict, expected_contract, expected_quick_review): + contract, quick_review = filter_quick_review(input_dict) + assert contract == expected_contract + assert quick_review == expected_quick_review + + # Test cases for get_exhibit_pages + @pytest.mark.parametrize("mock_response, expected_pages, text_dict", [ + # Exhibit page detected + (["|Y|","|N|"], ["1"], {"1": "Exhibit content", "2": "Non-exhibit content"}), + # No exhibit page detected + ([], [], {}), + # Multiple exhibit pages + (["|Y|","|Y|"], ["1", "2"], {"1": "Exhibit content", "2": "Attachment"}), + # Invalid response + (["Invalid"], [], {"1": "Non-exhibit content"}), + ]) + def test_get_exhibit_pages(self, mocker, mock_response, expected_pages, text_dict): + # Mock the LLM response + mocker_invoke_claude = mocker.patch("src.utils.llm_utils.invoke_claude", side_effect=mock_response) + filename = "test.pdf" + assert get_exhibit_pages(text_dict, filename) == expected_pages + + # # verify calls made + # assert mocker_invoke_claude.call_count == len(text_dict) diff --git a/fieldExtraction/tests/string_utils_test.py b/fieldExtraction/tests/string_utils_test.py index 6effc02..32edaca 100644 --- a/fieldExtraction/tests/string_utils_test.py +++ b/fieldExtraction/tests/string_utils_test.py @@ -5,10 +5,9 @@ import src.utils.llm_utils as llm_utils import numpy as np import pandas as pd - from src.enums.delimiters import Delimiter -class TestExtractTextFromDelimiters: +class TestStringUtils: @pytest.mark.parametrize("raw_text, delimiter, match_index, expected", [ ("Here is the answer |this is a pipe answer| and |more text.|", Delimiter.PIPE, 0, "this is a pipe answer"), ("Here is the answer |this is a pipe answer| and |more text.|", Delimiter.PIPE, -1, "more text."), @@ -24,7 +23,6 @@ class TestExtractTextFromDelimiters: result = extract_text_from_delimiters(raw_text, delimiter, match_index) assert result == expected -class TestJsonParsingSearch: @pytest.mark.parametrize("response_text, field_list, expected", [ ('{"name": "John", "age": "30", "city": "New York"}', ["name", "age", "city"], {"name": "John", "age": "30", "city": "New York"}), ('{"name": "John", "city": "New York"}', ["name", "age", "city"], {"name": "John", "city": "New York"}), @@ -36,7 +34,6 @@ class TestJsonParsingSearch: def test_json_parsing_search(self, response_text, field_list, expected): assert json_parsing_search(response_text, field_list) == expected -class TestSecondaryStringToDict: @pytest.fixture def mock_invoke_claude(self, mocker): return mocker.patch.object(llm_utils, 'invoke_claude') @@ -79,7 +76,6 @@ class TestSecondaryStringToDict: assert result == {} mock_invoke_claude.assert_called_once() -class TestPrimaryStringToDict: @pytest.fixture def mock_secondary_string_to_dict(self, mocker): return mocker.patch.object(utils.string_utils, "secondary_string_to_dict") @@ -164,7 +160,6 @@ class TestPrimaryStringToDict: assert result == expected_output assert mock_secondary_string_to_dict.call_count == 4 -class TestContainsReimbursement: @pytest.mark.parametrize("text, page, expected", [ ({"1": "This page contains a 10% reimbursement schedule.", "2": "No relevant content here."}, "1", True), ({"1": "This page has no relevant content.", "2": "Still no relevant content here."}, "1", False), @@ -184,7 +179,6 @@ class TestContainsReimbursement: captured = capsys.readouterr() assert "contains_reimbursement - Invalid data type" in captured.out -class TestIsEmpty: @pytest.mark.parametrize("value, expected", [ (None, True), ("", True), @@ -199,4 +193,4 @@ class TestIsEmpty: ]) def test_is_empty(self, value, expected): result = is_empty(value) - assert result == expected \ No newline at end of file + assert result == expected