doczyai-pipelines/src/pipelines/shared/postprocessing/aarete_derived.py

import logging
import src.config as config
import src.utils.string_utils as string_utils
from src.constants.constants import Constants
from src.crosswalk.crosswalk_builder import CrosswalkBuilder
from src.prompts.fieldset import FieldSet
from src.utils.crosswalk_utils import apply_crosswalk


def fill_na_from_field(results_dict, to_field, from_field, crosswalk_path):
    from_str = results_dict.get(from_field)
    if string_utils.is_empty(from_str):
        return results_dict.get(to_field)
    MAPPING_DICT = CrosswalkBuilder().from_json(crosswalk_path).mapping
    to_value = apply_crosswalk(from_str, MAPPING_DICT)
    return to_value if to_value else results_dict.get(to_field)


def fill_na_mapping(answer_dicts):
    """
    Fill in missing values in fields using mappings from other fields.
    Always checks all crosswalks (PROGRAM and PRODUCT) and merges unique values.

    Args:
        answer_dicts (list[dict]): List of answer dictionaries

    Returns:
        list[dict]: List of answer dictionaries with filled/merged values
    """
    for answer_dict in answer_dicts:
        # Collect all AARETE_DERIVED_LOB values from different sources
        all_lob_values = set()

        # Get existing AARETE_DERIVED_LOB values (if any)
        existing_lob_list = answer_dict.get("AARETE_DERIVED_LOB", "")

        if not string_utils.is_empty(existing_lob_list):
            for lob_val in existing_lob_list:
                if "|" in existing_lob_list:
                    all_lob_values.update(
                        v.strip()
                        for v in lob_val.split("|")
                        if v.strip() and v.strip() != "N/A"
                    )
                elif lob_val.strip() and lob_val.strip() != "N/A":
                    all_lob_values.add(lob_val)

        # Get AARETE_DERIVED_LOB from AARETE_DERIVED_PROGRAM crosswalk (always check if PROGRAM exists)
        program_lob_values = set()
        if not string_utils.is_empty(answer_dict.get("AARETE_DERIVED_PROGRAM")):
            program_filled_value_list = fill_na_from_field(
                answer_dict,
                "AARETE_DERIVED_LOB",
                "AARETE_DERIVED_PROGRAM",
                "src/constants/mappings/crosswalk_program_lob.json",
            )
            if (
                not string_utils.is_empty(program_filled_value_list)
                and "N/A" not in program_filled_value_list
            ):
                for program_lob_val in program_filled_value_list:
                    if "|" in program_lob_val:
                        program_lob_values = set(
                            v.strip()
                            for v in program_lob_val.split("|")
                            if v.strip() and v.strip() != "N/A"
                        )
                    elif program_lob_val.strip() and program_lob_val.strip() != "N/A":
                        program_lob_values = {program_lob_val.strip()}
                    all_lob_values.update(program_lob_values)

        # Get AARETE_DERIVED_LOB from PRODUCT crosswalk (always check if PRODUCT exists)
        product_lob_values = set()
        if not string_utils.is_empty(answer_dict.get("PRODUCT")):
            product_filled_value_list = fill_na_from_field(
                answer_dict,
                "AARETE_DERIVED_LOB",
                "PRODUCT",
                "src/constants/mappings/crosswalk_product_lob.json",
            )
            if (
                not string_utils.is_empty(product_filled_value_list)
                and "N/A" not in product_filled_value_list
            ):
                for product_lob_val in product_filled_value_list:
                    if "|" in product_lob_val:
                        product_lob_values = set(
                            v.strip()
                            for v in product_filled_value_list.split("|")
                            if v.strip() and v.strip() != "N/A"
                        )
                    elif product_lob_val.strip() and product_lob_val.strip() != "N/A":
                        product_lob_values = {product_lob_val.strip()}
                    all_lob_values.update(product_lob_values)

        # Set the merged, deduplicated value
        if all_lob_values:
            answer_dict["AARETE_DERIVED_LOB"] = sorted(all_lob_values)
            # Set relationship flags if values came from PROGRAM or PRODUCT
            # Only set if field is empty or "N/A" (preserve LLM-determined Inclusive/Exclusive)
            if program_lob_values:
                existing_relationship = answer_dict.get("LOB_PROGRAM_RELATIONSHIP", "")
                # Original code (CHANGED - now preserves LLM values):
                # answer_dict["LOB_PROGRAM_RELATIONSHIP"] = "Exclusive"
                if string_utils.is_empty(existing_relationship):
                    answer_dict["LOB_PROGRAM_RELATIONSHIP"] = "Exclusive"
            if product_lob_values:
                existing_relationship = answer_dict.get("LOB_PRODUCT_RELATIONSHIP", "")
                # Original code (CHANGED - now preserves LLM values):
                # answer_dict["LOB_PRODUCT_RELATIONSHIP"] = "Exclusive"
                if string_utils.is_empty(existing_relationship):
                    answer_dict["LOB_PRODUCT_RELATIONSHIP"] = "Exclusive"
        elif string_utils.is_empty(answer_dict.get("AARETE_DERIVED_LOB")):
            # If still empty after all attempts, keep it as is (will be N/A or empty)
            pass

        # Handle cases where PROGRAM/PRODUCT is NA/blank (LLM was bypassed)
        # Set relationship to "Exclusive" as default when PROGRAM/PRODUCT is missing and relationship is not set
        # This covers both cases: LOB exists but PROGRAM/PRODUCT is NA, or both LOB and PROGRAM/PRODUCT are NA
        if string_utils.is_empty(answer_dict.get("AARETE_DERIVED_PROGRAM")):
            if string_utils.is_empty(answer_dict.get("LOB_PROGRAM_RELATIONSHIP")):
                answer_dict["LOB_PROGRAM_RELATIONSHIP"] = "Exclusive"

        if string_utils.is_empty(answer_dict.get("AARETE_DERIVED_PRODUCT")):
            if string_utils.is_empty(answer_dict.get("LOB_PRODUCT_RELATIONSHIP")):
                answer_dict["LOB_PRODUCT_RELATIONSHIP"] = "Exclusive"

    return answer_dicts


def get_crosswalk_fields(answer_dicts: list, constants: Constants):

    crosswalk_fields = FieldSet(file_path=config.FIELD_JSON_PATH, crosswalk=True)
    for to_field in crosswalk_fields.fields:
        to_field_name, from_field_name = to_field.field_name, to_field.base_field
        # Find crosswalk
        crosswalk = constants.get_constant(to_field.crosswalk)
        if crosswalk:
            for answer_dict in answer_dicts:
                to_field_value = answer_dict.get(to_field_name)
                from_field_value = answer_dict.get(from_field_name)
                # If from_field_value is populated and to_field_value is not populated, perform mapping
                if not string_utils.is_empty(from_field_value) and (
                    string_utils.is_empty(to_field_value) or "N/A" in to_field_value
                ):
                    if isinstance(from_field_value, list):
                        from_field_value_list = from_field_value
                    else:
                        from_field_value_list = [from_field_value]

                    # Map each value in the list
                    to_field_answer_list = []
                    for individual_from_field_value in from_field_value_list:
                        individual_from_field_value = (
                            individual_from_field_value.strip()
                        )
                        if individual_from_field_value in crosswalk.mapping.keys():
                            # Value is a key in the mapping - map it to the target value
                            to_field_answer_list.append(
                                crosswalk.mapping.get(individual_from_field_value)
                            )
                        elif individual_from_field_value in crosswalk.mapping.values():
                            # Value is already in the target format (e.g., "CHIP" when target is "CHIP")
                            # Keep it as-is for AARETE_DERIVED fields, or use reverse mapping for others
                            if "AARETE_DERIVED" in to_field_name:
                                to_field_answer_list.append(individual_from_field_value)
                            else:
                                to_field_answer_list.append(
                                    crosswalk.create_reverse_mapping().get(
                                        individual_from_field_value
                                    )
                                )
                    answer_dict[to_field_name] = to_field_answer_list
                    # update from field value to list if not already
                    if not isinstance(from_field_value, list):
                        answer_dict[from_field_name] = from_field_value_list

    return answer_dicts