Merged in feature/reimb-primary-improvements (pull request #524)

Reimbursement Exhibit Prefiltering * Add IDENTIFY_REIMBURSEMENT_EXHIBITS function to analyze exhibit headers for reimbursement information * Add Zone.Identifier to .gitignore to prevent Windows zone information files from being tracked * Add functionality to identify reimbursement exhibits * Add functionality to identify reimbursement exhibits in run_one_to_n_prompts * Output YAML parsing * Refactor identify_reimbursement_exhibits to improve YAML parsing and error handling * New reimbursement exhibit logic * Update IDENTIFY_REIMBURSEMENT_EXHIBITS_PROMPT to return PAGE NUMBERS instead of exhibit headers * Fix identify_reimbursement_exhibits to return page numbers as strings after YAML parsing * Remove debugging print statement for reimbursement exhibits in run_one_to_n_prompts * Remove commented debugging print statements in identify_reimbursement_exhibits * Increase read_timeout in EC2 configuration from 2000 to 4000 * Merge remote-tracking branch 'origin/main' into feature/reimb-primary-improvements * Remove debugging print statements in get_reimbursement_primary function * mypy error * Fix return type in identify_reimbursement_exhibits to ensure keys are returned as a list * Merge remote-tracking branch 'origin/main' into feature/reimb-primary-improvements Approved-by: Katon Minhas
2025-05-13 20:22:43 +00:00
parent 384b3c5d31
commit 55fa4f2e79
5 changed files with 103 additions and 33 deletions
@@ -75,6 +75,7 @@ streamlit/venv
 *.PDF
 *.TXT
 *.txt
+*:Zone.Identifier
 *.tfplan
 *.pem
 textfiles/
@@ -217,7 +217,7 @@ if RUN_MODE == "local":

 elif RUN_MODE == "ec2":
    # Bedrock
-    config = Config(read_timeout=2000, max_pool_connections=60)
+    config = Config(read_timeout=4000, max_pool_connections=60) # Doubled read_timeout to 4000
    BEDROCK_RUNTIME = boto3.client(
        service_name="bedrock-runtime", region_name="us-east-1", config=config
    )
@@ -138,48 +138,57 @@ def run_one_to_one_prompts(filename, contract_text, text_dict, top_sheet_dict, d

 def run_one_to_n_prompts(filename, text_dict, exhibit_pages, exhibit_chunk_mapping, all_dataset):

+
+    ################## IDENTIFY COMPENSATION EXHIBITS ##################
+    # Get all exhibit headers first
+    all_exhibit_headers = {page: one_to_n_funcs.get_exhibit_header(string_utils.get_exhibit_chunk(text_dict, exhibit_chunk_mapping, page), filename) 
+                           for page in exhibit_pages}
+
+    # Then identify the compensation exhibits
+    reimbursement_exhibits = one_to_n_funcs.identify_reimbursement_exhibits(all_exhibit_headers, filename)
+    # print("Reimbursement Exhibits: ", reimbursement_exhibits) # Debugging line
+
    ################## RUN PROMPTS ##################
    one_to_n_results = []
-    for exhibit_page in exhibit_pages:
+    for exhibit_page in reimbursement_exhibits:
        exhibit_chunk = string_utils.get_exhibit_chunk(text_dict, exhibit_chunk_mapping, exhibit_page)
-        if string_utils.contains_reimbursement(exhibit_chunk):
            
-            ################## INITIALIZE FIELDS ##################
-            reimbursement_level_fields = FieldSet(relationship="one_to_n", field_type="reimbursement_level", file_path=FIELD_JSON_PATH)
+        ################## INITIALIZE FIELDS ##################
+        reimbursement_level_fields = FieldSet(relationship="one_to_n", field_type="reimbursement_level", file_path=FIELD_JSON_PATH)

-            ################## GET EXHIBIT HEADER ##################
-            exhibit_header = one_to_n_funcs.get_exhibit_header(exhibit_chunk, filename)
+        ################## GET EXHIBIT HEADER ##################
+        exhibit_header = one_to_n_funcs.get_exhibit_header(exhibit_chunk, filename)

-            ################## GET REIMBURSEMENT TIN/NPI ##################
-            tin_npi_answers, reimbursement_level_fields = reimbursement_tin_npi(exhibit_chunk, reimbursement_level_fields, filename)
+        ################## GET REIMBURSEMENT TIN/NPI ##################
+        tin_npi_answers, reimbursement_level_fields = reimbursement_tin_npi(exhibit_chunk, reimbursement_level_fields, filename)

-            ################## GET EXHIBIT-LEVEL ANSWERS ##################
-            exhibit_level_answers = one_to_n_funcs.get_exhibit_level_answers(exhibit_chunk, filename)
-            exhibit_level_answers['EXHIBIT_PAGE'] = exhibit_page
-            exhibit_level_answers['EXHIBIT_TITLE'] = exhibit_header
+        ################## GET EXHIBIT-LEVEL ANSWERS ##################
+        exhibit_level_answers = one_to_n_funcs.get_exhibit_level_answers(exhibit_chunk, filename)
+        exhibit_level_answers['EXHIBIT_PAGE'] = exhibit_page
+        exhibit_level_answers['EXHIBIT_TITLE'] = exhibit_header

-            ################## GET DYNAMIC-PRIMARY ANSWERS ##################
-            dynamic_to_exhibit_level_answers, dynamic_to_reimbursement_level_fields = dynamic_funcs.get_dynamic_answers(exhibit_chunk, 
-                                                                                                                        exhibit_header, 
-                                                                                                                        filename,
-                                                                                                                        FieldSet(relationship="one_to_n", field_type="dynamic", file_path=config.FIELD_JSON_PATH))
-            exhibit_level_answers.update(dynamic_to_exhibit_level_answers)
-            reimbursement_level_fields.combine(dynamic_to_reimbursement_level_fields, inplace=True)
-            
-            ################## GET DYNAMIC-CODE ANSWERS ##################
-            code_to_exhibit_level_answers, code_to_reimbursement_level_fields  = dynamic_funcs.get_dynamic_answers(exhibit_chunk,
-                                                                                                                   exhibit_header, 
-                                                                                                                   filename,
-                                                                                                                   FieldSet(relationship="one_to_n", field_type="dynamic_code", file_path=config.FIELD_JSON_PATH))
-            exhibit_level_answers.update(code_to_exhibit_level_answers)
-            reimbursement_level_fields.combine(code_to_reimbursement_level_fields, inplace=True)
+        ################## GET DYNAMIC-PRIMARY ANSWERS ##################
+        dynamic_to_exhibit_level_answers, dynamic_to_reimbursement_level_fields = dynamic_funcs.get_dynamic_answers(exhibit_chunk, 
+                                                                                                                    exhibit_header, 
+                                                                                                                    filename,
+                                                                                                                    FieldSet(relationship="one_to_n", field_type="dynamic", file_path=config.FIELD_JSON_PATH))
+        exhibit_level_answers.update(dynamic_to_exhibit_level_answers)
+        reimbursement_level_fields.combine(dynamic_to_reimbursement_level_fields, inplace=True)
+        
+        ################## GET DYNAMIC-CODE ANSWERS ##################
+        code_to_exhibit_level_answers, code_to_reimbursement_level_fields  = dynamic_funcs.get_dynamic_answers(exhibit_chunk,
+                                                                                                                exhibit_header, 
+                                                                                                                filename,
+                                                                                                                FieldSet(relationship="one_to_n", field_type="dynamic_code", file_path=config.FIELD_JSON_PATH))
+        exhibit_level_answers.update(code_to_exhibit_level_answers)
+        reimbursement_level_fields.combine(code_to_reimbursement_level_fields, inplace=True)

-            ################## GET REIMBURSEMENT-LEVEL ANSWERS (INCLUDING DYNAMIC) ##################
-            reimbursement_level_answers = reimbursement_level(exhibit_chunk, filename, reimbursement_level_fields, all_dataset, exhibit_page) # Return list of dictionaries
+        ################## GET REIMBURSEMENT-LEVEL ANSWERS (INCLUDING DYNAMIC) ##################
+        reimbursement_level_answers = reimbursement_level(exhibit_chunk, filename, reimbursement_level_fields, all_dataset, exhibit_page) # Return list of dictionaries

-            ################# COMBINE ANSWERS ##################
-            full_answer_dict = combine_one_to_n_answers(exhibit_level_answers, reimbursement_level_answers, tin_npi_answers)
-            one_to_n_results += full_answer_dict
+        ################# COMBINE ANSWERS ##################
+        full_answer_dict = combine_one_to_n_answers(exhibit_level_answers, reimbursement_level_answers, tin_npi_answers)
+        one_to_n_results += full_answer_dict
    
    ################## Crosswalk Fields ################## 
    one_to_n_results = aarete_derived.get_crosswalk_fields(one_to_n_results)
@@ -13,6 +13,8 @@ from src.enums.delimiters import Delimiter
 from src.prompts.investment_prompts import Field, FieldSet, EXHIBIT_LEVEL
 from src.config import FIELD_JSON_PATH

+import yaml # This is for experimental YAML parsing
+
 METHODOLOGY_BREAKOUT_QUESTIONS = FieldSet(
    file_path=config.FIELD_JSON_PATH, field_type="methodology_breakout"
 ).print_prompt_dict()
@@ -55,13 +57,55 @@ def get_exhibit_header(exhibit_chunk, filename):
    )
    return llm_answer_final

+def identify_reimbursement_exhibits(exhibit_headers: dict, filename: str) -> list[str]:
+    """Use LLM to identify which exhibits likely contain reimbursement terms.
+
+    Args:
+        exhibit_headers (dict): Dictionary of exhibit headers to analyze, keyed by 
+            string page number, valued by the exhibit header text.
+        filename (str): The name of the file being processed.
+
+    Returns:
+        list[str]: List of page numbers (as strings) that likely contain reimbursement terms.
+    """
+    # print(exhibit_headers) # Debugging line
+    # print(type(exhibit_headers)) # Debugging line
+    prompt = investment_prompts.IDENTIFY_REIMBURSEMENT_EXHIBITS_PROMPT(exhibit_headers, yaml_output=True) # Trying YAML output
+    # print(prompt) # Debugging line
+    llm_answer_raw = llm_utils.invoke_claude(
+        prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 2000
+    )
+
+    yaml_content = string_utils.extract_text_from_delimiters(
+        llm_answer_raw, Delimiter.TRIPLE_BACKTICK)
+    
+    # Check if the first line is just `yaml` and if so remove it
+    lines = yaml_content.strip().split('\n')
+    if lines[0].strip().lower() == 'yaml':
+        yaml_content = '\n'.join(lines[1:])
+    
+    # print(llm_answer_raw) # Debugging line
+    # print("Extracted YAML content:\n", yaml_content) # Debugging line
+    try: 
+        # Parse the YAML output into a dictionary
+        exhibits_dict = yaml.safe_load(yaml_content)
+        # YAML parsing has built-in type conversion for "NO" -> False, "YES" -> True
+        # So we'll check for boolean values and strings (just in case)
+        return [str(page) for page, value in exhibits_dict.items() # Convert back to string
+                if value is True or (isinstance(value, str) and value.upper() == "YES")]
+    except yaml.YAMLError as e:
+        print(f"Error parsing YAML: {e}")
+        print(f"Raw LLM output: {llm_answer_raw}")
+        return list(exhibit_headers.keys())  # Return all keys of the exhibit headers if parsing fails

 def get_reimbursement_primary(reimbursement_level_fields, exhibit_text, filename):
    field_prompts = reimbursement_level_fields.get_prompt_dict()
    prompt = investment_prompts.REIMBURSEMENT_LEVEL_PRIMARY(exhibit_text, field_prompts)
+    # print(prompt) # Debugging line
    llm_answer_raw = llm_utils.invoke_claude(
        prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, max_tokens=25000
    )
+    # print(llm_answer_raw) # Debugging line
    llm_answer_final = string_utils.universal_json_load(llm_answer_raw)
    return llm_answer_final

@@ -309,6 +309,22 @@ Before returning any output, ensure that all instructions for inclusion were fol
 Enclose your final answer in |pipes|.
 """

+def IDENTIFY_REIMBURSEMENT_EXHIBITS_PROMPT(exhibit_headers: dict, yaml_output: bool = False):
+    formatted_headers = [f"Page {page}: {header}" for page, header in exhibit_headers.items()]
+    return f"""Review these exhibit headers and identify which ones are likely to contain reimbursement methodologies, payment terms, or compensation schedules.
+
+Look for keywords like "compensation", "payment", "rate", "reimbursement", "fee schedule", "pricing", or "financial" in the headers.
+
+For each header, return 'YES' if it's likely to contain reimbursement information, or 'NO' if it's not.
+
+Headers to review:
+{"\n".join(formatted_headers)}
+
+Return your answer as a {"JSON object" if not yaml_output else "YAML list"} where keys are PAGE NUMBERS and values are 'YES' or 'NO'.
+Feel free to explain your reasoning, but place your final {"JSON" if not yaml_output else "YAML"} output between triple backticks.
+If you're uncertain about a header, err on the side of inclusion (YES).
+"""
+

 def EXHIBIT_LEVEL(context, fields):
    return f"""Extract attributes from a section of a contract.