doczyai-pipelines/src/pipelines/shared/extraction/dynamic_funcs.py

import logging
import concurrent.futures
from typing import TYPE_CHECKING, Optional

from src.pipelines.saas.prompts import prompt_calls
from src.pipelines.shared.preprocessing import preprocessing_funcs
from src.utils import string_utils
from src.constants.constants import Constants
from src import config
from src.prompts import prompt_templates
from src.prompts.fieldset import Field, FieldSet
from src.pipelines.shared.extraction.exhibit_funcs import Exhibit

if TYPE_CHECKING:
    from src.pipelines.shared.extraction.page_funcs import Page


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,
        )
        # If there is at least ONE Non-N/A answers in the Exhibit
        if not string_utils.is_empty(exhibit_text_answer):
            # Special handling for LOB: If both Medicare and Medicaid are detected, add Duals
            if field.field_name == "LOB":
                # Parse the discovered values (pipe-delimited or comma-separated)
                values_list = [
                    val.strip()
                    for val in exhibit_text_answer.replace(",", "|").split("|")
                    if val.strip()
                ]
                values_lower = [val.lower() for val in values_list]

                # Check if both Medicare and Medicaid are present (case-insensitive)
                has_medicare = any("medicare" in val for val in values_lower)
                has_medicaid = any("medicaid" in val for val in values_lower)

                # If both are present and Duals is not already in the list, add it
                if has_medicare and has_medicaid:
                    if "duals" not in values_lower:
                        exhibit_text_answer = (
                            exhibit_text_answer + " | Duals"
                            if exhibit_text_answer
                            else "Duals"
                        )
                        logging.debug(
                            f"Added 'Duals' to LOB valid values because both Medicare and Medicaid were detected"
                        )

            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,
    pages_dict: Optional[dict[str, "Page"]] = None,
    text_dict: Optional[dict[str, str]] = None,
    exhibit: Optional[Exhibit] = None,
    constants: Optional[Constants] = None,
    filename: Optional[str] = None,
) -> list[dict[str, str]]:
    """
    Enrich reimbursement rows with dynamic fields in parallel.

    Args:
        reimbursement_primary_answers: List of reimbursement answer dictionaries
        dynamic_reimbursement_fields: FieldSet of dynamic fields to assign
        pages_dict: Dictionary mapping page numbers to Page objects (preferred)
        text_dict: Dictionary mapping page numbers to text strings (for backward compatibility)
        exhibit: Exhibit object (preferred - uses exhibit.pages)
        constants: Constants object
        filename: Name of the file being processed
    """
    # Check for exhibit first (required parameter)
    if exhibit is None:
        raise ValueError("exhibit must be provided")

    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(
            pages_dict=pages_dict,
            text_dict=text_dict,
            exhibit=exhibit,
            current_page_num=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

    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()
        )

    # Special instructions for dynamic primary fields when passed to 1:1
    # Add specific template-language guidance, then append general full context instructions
    if field_to_add.field_name in ["LOB", "PRODUCT", "PROGRAM", "NETWORK"]:
        field_to_add.prompt = (
            field_to_add.prompt
            + prompt_templates.FULL_CONTEXT_DYNAMIC_PRIMARY_INSTRUCTION()
        )

    # Apply general full context instructions to all fields
    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
            )
        else:
            # Field found in some rows - keep in 1:N
            pass

    return one_to_one_fields