doczyai-pipelines/fieldExtraction/src/investment/dynamic_funcs.py

import logging
import concurrent.futures

from src.investment import prompt_calls, preprocessing_funcs
from src.utils import string_utils
from constants.constants import Constants
from src import config
from src.prompts import prompt_templates
from src.prompts.fieldset import Field, FieldSet
from src.investment.exhibit_funcs import Exhibit


def dynamic_primary(
    exhibit_text: str,
    exhibit_level_answers: dict,
    constants: Constants,
    filename: str,
    dynamic_primary_fields: FieldSet,
):
    """
    Processes dynamic primary fields one-by-one from exhibit header and chunk text.

    Args:
        exhibit_chunk (str): The main text chunk of the exhibit
        exhibit_header (str): The header text of the exhibit
        filename (str): Name of the file being processed
        dynamic_fields (FieldSet): FieldSet containing the dynamic fields to process

    Returns:
        tuple: (dict, FieldSet) containing:
            - exhibit_level_answer_dict: Dictionary of exhibit-level answers
            - reimbursement_level_fields: FieldSet of fields assigned to reimbursement level
    """
    if not dynamic_primary_fields.contains_fields():
        return {}, FieldSet()

    dynamic_reimbursement_fields = FieldSet()
    exhibit_level_answer_dict = {}
    
    for field in dynamic_primary_fields.fields:
        exhibit_text_answer = prompt_calls.prompt_dynamic_primary(
            exhibit_text,
            field,
            constants,
            filename,
            prompt_templates.DYNAMIC_PRIMARY_TEXT,
        )
        # If there is at least ONE Non-N/A answers in the Exhibit
        if not string_utils.is_empty(exhibit_text_answer):
            field.update_valid_values(exhibit_text_answer + " | N/A")
            dynamic_reimbursement_fields.add_field(field)
            dynamic_primary_fields.remove_field(field.field_name)
        # If the field is not found, add to Exhibit-Level answers (the answer will be N/A or similar)
        else:
            exhibit_level_answer_dict[field.field_name] = exhibit_text_answer

    # Update
    exhibit_level_answers.update(exhibit_level_answer_dict)

    return exhibit_level_answers, dynamic_reimbursement_fields


def dynamic(
    exhibit_text: str,
    exhibit_header: str,
    exhibit_level_answers: dict,
    constants: Constants,
    filename: str,
    dynamic_fields: FieldSet,
    dynamic_reimbursement_fields: FieldSet,
):
    """
    Processes dynamic (code and provider info) fields from exhibit header and chunk text.

    Args:
        exhibit_chunk (str): The main text chunk of the exhibit
        exhibit_header (str): The header text of the exhibit
        filename (str): Name of the file being processed
        dynamic_fields (FieldSet): FieldSet containing the dynamic fields to process

    Returns:
        tuple: (dict, FieldSet) containing:
            - exhibit_level_answer_dict: Dictionary of exhibit-level answers
            - reimbursement_level_fields: FieldSet of fields assigned to reimbursement level
    """
    if not dynamic_fields.contains_fields():
        return {}, FieldSet()

    dynamic_to_reimbursement_level_fields = FieldSet()
    exhibit_level_answer_dict = {}

    if not string_utils.is_empty(exhibit_header):
        # Check Exhibit Header
        exhibit_header_results = prompt_calls.prompt_dynamic(
            text=exhibit_header,
            field_prompts=dynamic_fields.print_prompt_dict(constants),
            filename=filename,
        )
        # Process Exhibit Header Results - update dynamic fields and add to reimbursement level fields
        for field_name, answer in exhibit_header_results.items():
            if (
                not string_utils.is_empty(answer)
                and field_name in dynamic_fields.list_fields()
            ):
                field = dynamic_fields.get_field(field_name)
                field.update_valid_values(answer)
                field.prompt = (
                    field.prompt
                    + ". Ensure ALL values that apply to the specific reimbursement term are included."
                )
                dynamic_reimbursement_fields.add_field(field)
                dynamic_fields.remove_field(field_name)

    # Check Exhibit Text
    if dynamic_fields.contains_fields():
        exhibit_chunk_results = prompt_calls.prompt_dynamic(
            text=exhibit_text,
            field_prompts=dynamic_fields.print_prompt_dict(constants),
            filename=filename,
        )
        # Process Exhibit Chunk Results - update dynamic fields and add to reimbursement level fields
        for field_name, answer in exhibit_chunk_results.items():
            if (
                not string_utils.is_empty(answer)
                and field_name in dynamic_fields.list_fields()
            ):
                field = dynamic_fields.get_field(field_name)
                field.update_valid_values(answer)
                dynamic_to_reimbursement_level_fields.add_field(field)
                dynamic_fields.remove_field(field_name)
            # If the field is not found, add to Exhibit-Level answers (the answer will be N/A or similar)
            else:
                exhibit_level_answer_dict[field_name] = answer

    # Update
    exhibit_level_answers.update(exhibit_level_answer_dict)
    dynamic_reimbursement_fields.combine(
        dynamic_to_reimbursement_level_fields, inplace=True
    )

    return exhibit_level_answers, dynamic_reimbursement_fields


def dynamic_assignment(reimbursement_primary_answers: list[dict[str, str]], 
                       dynamic_reimbursement_fields: FieldSet, 
                       text_dict: dict[str, str],
                       exhibit: Exhibit,
                       constants: Constants, 
                       filename: str
                       ) -> list[dict[str, str]]:
    """
    Enrich reimbursement rows with dynamic fields in parallel.
    """
    if not reimbursement_primary_answers or not dynamic_reimbursement_fields.fields:
        return reimbursement_primary_answers

    def _assign_single(answer_dict: dict[str, str]) -> dict[str, str]:
        # copy to avoid mutating caller-provided dicts when running in threads
        updated = answer_dict.copy()
        page_num = updated["REIMB_PAGE"]
        exhibit_text_simplified = preprocessing_funcs.simplify_exhibit(
            text_dict, exhibit.exhibit_page_nums, page_num
        )
        service_term = updated['SERVICE_TERM']
        reimb_term = updated['REIMB_TERM']
        for dynamic_field in dynamic_reimbursement_fields.fields:
            dynamic_field_answer = prompt_calls.prompt_dynamic_assignment(
                service_term, reimb_term, dynamic_field, exhibit_text_simplified, page_num, constants, filename
            )
            updated.update(dynamic_field_answer)
        return updated

    max_workers = min(len(reimbursement_primary_answers), 8)
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        return list(executor.map(_assign_single, reimbursement_primary_answers))


def add_one_to_one_field(one_to_one_fields, field_to_add, answer_dicts, constants: Constants):
    field_to_add.field_type = "smart_chunked"
    field_to_add.relationship = "one_to_one"

    # Special Handling for Program, Product, and Network:
    # If any LOB value has been found, skip PROGRAM, PRODUCT, and NETWORK
    # Only search for these fields when NO LOB has been found
    # NOTE: Check raw LOB field, not AARETE_DERIVED_LOB, because AARETE_DERIVED_LOB
    # is populated later via crosswalk and may be derived from PROGRAM/PRODUCT
    if field_to_add.field_name in ["PROGRAM", "PRODUCT", "NETWORK"]:
        lob_values = [
            answer_dict.get("LOB", "") for answer_dict in answer_dicts
        ]
        unique_lobs = set([lob for lob in lob_values if not string_utils.is_empty(lob)])
        if len(unique_lobs) > 0:
            return one_to_one_fields

    if field_to_add.field_name == "CLAIM_TYPE_CD":
        field_to_add.prompt = (
            field_to_add.prompt
            + prompt_templates.FULL_CONTEXT_CLAIM_TYPES_ADDITIONAL_INSTRUCTION()
        )

    field_to_add.prompt = (
        field_to_add.prompt
        + prompt_templates.FULL_CONTEXT_ADDITIONAL_INSTRUCTIONS(field_to_add.allow_na)
    )

    if field_to_add.valid_values:
        field_to_add.keywords = field_to_add.get_valid_values(constants)

    one_to_one_fields.add_field(field_to_add)
    return one_to_one_fields


def get_dynamic_one_to_one_fields(answer_dicts: list[dict[str, str]], constants: Constants):
    """
    Returns a FieldSet object of dynamic fields that have are empty for at least one dict in answer_dicts
    """
    # These fields get passed to 1:1 if ALL of the answers are empty
    all_empty_fields = FieldSet(
        config.FIELD_JSON_PATH, relationship="one_to_n", base_field=True
    )

    one_to_one_fields = FieldSet()  # Empty FieldSet to populate if criteria are met

    # Check if any LOB value has been found - if so, skip PROGRAM, PRODUCT, NETWORK
    # Only search for PROGRAM, PRODUCT, NETWORK when NO LOB has been found
    # NOTE: Check raw LOB field, not AARETE_DERIVED_LOB, because AARETE_DERIVED_LOB
    # is populated later via crosswalk and may be derived from PROGRAM/PRODUCT
    lob_values = [
        answer_dict.get("LOB", "") for answer_dict in answer_dicts
    ]
    unique_lobs = set([lob for lob in lob_values if not string_utils.is_empty(lob)])
    has_lob = len(unique_lobs) > 0

    # Handle ALL empty fields
    for field in all_empty_fields.fields:
        field_name = field.field_name
        if "REIMB" in field_name:  # Don't do this for the REIMB_ fields
            continue
        
        empty_count = sum(
            1 for answer_dict in answer_dicts 
            if string_utils.is_empty(answer_dict.get(field_name))
        )
        total_count = len(answer_dicts)
        
        if empty_count == total_count:
            # Skip PROGRAM, PRODUCT, NETWORK if any LOB has been found
            # Only search for these fields when NO LOB has been found
            if has_lob and field_name in ["PROGRAM", "PRODUCT", "NETWORK"]:
                continue
            
            field_to_add = Field.load_from_file(
                file_path=config.FIELD_JSON_PATH,
                field_name=field.base_field if field.base_field else field.field_name 
            ) # Try base_field first, fall back to field_name if base_field is None/empty
            one_to_one_fields = add_one_to_one_field(
                one_to_one_fields, field_to_add, answer_dicts, constants
            )
    
    return one_to_one_fields