Merged in feature/self-repair-main (pull request #313)

Feature/self repair main * fixes * fixes * fixes * fixes * fixes * fixes * fixes * fixes * new testing * new testing * added unit tests under fieldExtraction/tests/main-unit-tests * unit test fixes * unit test fixes * merge into main * modified consolidate_output and slight changes to qa_qc_helpers to accomodate qaqc functionality w faizan and michael's changes * merging self repair into main * synced with main now * fixes * provider billing removed * provider billing removed * provider billing fixed * commented out load dotenv in config * Merged main into feature/self-repair-main * modified check_s3_bucket_exists function by returning False if credentials not found for pipeline * Merge branch 'feature/self-repair-main' of https://bitbucket.org/aarete/doczy.ai into feature/self-repair-main Approved-by: Michael McGuinness
2024-12-17 16:34:59 +00:00
parent c210052952
commit 9a736fa80d
10 changed files with 873 additions and 270 deletions
@@ -16,10 +16,32 @@ load_dotenv()


 ######################################## UTILS ###################################################
+# def check_s3_bucket_exists(s3_client, bucket_name: str):
+#     try:
+#         s3_client.head_bucket(Bucket=bucket_name)
+#         print(f"{bucket_name} exists")
+#     except s3_client.exceptions.ClientError as e:
+#         error_code = e.response["Error"]["Code"]
+#         if error_code == "404":
+#             print(f"{bucket_name} does not exist")
+#             raise
+#         else:
+#             print(f"Error {error_code}")
+#             raise
+#     except Exception as e:
+#         raise
+
+#     return
+
+from botocore.exceptions import NoCredentialsError
+
 def check_s3_bucket_exists(s3_client, bucket_name: str):
    try:
        s3_client.head_bucket(Bucket=bucket_name)
        print(f"{bucket_name} exists")
+    except NoCredentialsError:
+        print("AWS credentials not found")
+        return False
    except s3_client.exceptions.ClientError as e:
        error_code = e.response["Error"]["Code"]
        if error_code == "404":
@@ -30,9 +52,7 @@ def check_s3_bucket_exists(s3_client, bucket_name: str):
            raise
    except Exception as e:
        raise
-
-    return
-
+    return True

 ######################################## GENERAL SETTINGS ########################################
 TEST = False  # True to run test prompt - just for testing model connection
@@ -108,7 +128,7 @@ OUTPUT_DIRECTORY = get_arg_value(
    "output_dir", "output_individual"
 )  # Valid: any valid directory
 TRACKING_FILE = get_arg_value("tracking_file", f"{BATCH_ID}_tracking.csv")
-WRITE_TO_S3 = get_arg_value("write_to_s3", False) == "True"  # Valid: True or False
+WRITE_TO_S3 = get_arg_value("write_to_s3", "True") == "True"
 S3_OUTPUT_BUCKET = get_arg_value("s3_output_bucket", default="doczy-output")

 # Filtering args
@@ -1,31 +1,18 @@
 import os 
 import pandas as pd
+from pathlib import Path
+from datetime import datetime

 import config
 import costlog_funcs
 import valid
-
-from pathlib import Path
-from datetime import datetime
-from openpyxl import Workbook
-from openpyxl.utils.dataframe import dataframe_to_rows
-import config
-import valid
-import sys
-from utils import read_input
+import tracking
 from qa_qc_helpers import (table_stats_qaqc, perform_qc_qa_tests_both, 
                   perform_qc_qa_tests_ac,    
                   perform_qc_qa_tests_b, perform_qc_qa, save_qcqa_wb)

-
-
 def upload_tracking_files(run_timestamp):
-    """
-    Upload tracking files to S3 during batch processing.
-    
-    Args:
-        run_timestamp (str): Timestamp string for the current processing run
-    """
+    """Upload tracking files and aggregated cost logs to S3 during batch processing."""
    if not config.WRITE_TO_S3:
        return
        
@@ -36,41 +23,235 @@ def upload_tracking_files(run_timestamp):
    }

    try:
+        # Create aggregated costlog files
        costlog_caller_fn = "costlog_agg_by_caller.csv"
        costlog_filename_fn = "costlog_agg_by_filename.csv"

-        # Read the costlog, and aggregate it.
-        df_costlog = pd.read_csv(config.COST_LOG_NAME)
-        df_caller, df_filename = costlog_funcs.aggregate_costlog(df_costlog)
+        # Read and aggregate costlog
+        if os.path.exists(config.COST_LOG_NAME):
+            df_costlog = pd.read_csv(config.COST_LOG_NAME)
+            df_caller, df_filename = costlog_funcs.aggregate_costlog(df_costlog)

-        # Write the aggregated costlogs to .csv
-        df_caller.to_csv(costlog_caller_fn, index=False)
-        df_filename.to_csv(costlog_filename_fn, index=False)
+            # Write aggregated costlogs
+            df_caller.to_csv(costlog_caller_fn, index=False)
+            df_filename.to_csv(costlog_filename_fn, index=False)

-        tracking_files['costlog_by_caller'] = costlog_caller_fn
-        tracking_files['costlog_by_filename'] = costlog_filename_fn
+            # Add aggregated files to tracking
+            tracking_files['costlog_by_caller'] = costlog_caller_fn
+            tracking_files['costlog_by_filename'] = costlog_filename_fn

    except Exception as e:
        print(f"Costlog aggregation failed: {str(e)}")
    
+    # Upload all tracking files
    for file_type, file_name in tracking_files.items():
        if os.path.exists(file_name):
            try:
-                s3_key = os.path.join(run_timestamp, file_name).replace("\\", "/")
+                s3_key = f"{config.BATCH_ID}/{run_timestamp}/tracking/{file_name}"
                config.S3_CLIENT.upload_file(file_name, BUCKET_NAME, s3_key)
+                print(f"Uploaded {file_type} to {s3_key}")
            except Exception as e:
                print(f"Error uploading {file_name}: {str(e)}")
                raise

-
-def consolidate_atscale() -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
-    Path(config.CONSOLIDATED_OUTPUT_DIRECTORY).mkdir(parents=True, exist_ok=True)
+def get_all_run_folders(batch_id):
+    """Get all run folders for a given batch ID from S3"""
+    prefix = f"{batch_id}/"
+    runs = set()
    
-    ac_final_df = None    
-    b_final_df = None
+    paginator = config.S3_CLIENT.get_paginator('list_objects_v2')
+    for page in paginator.paginate(Bucket=config.S3_OUTPUT_BUCKET, Prefix=prefix):
+        for obj in page.get('Contents', []):
+            parts = obj['Key'].split('/')
+            if len(parts) > 1 and parts[1].startswith('run_'):
+                runs.add(parts[1])
+    
+    return sorted(list(runs))
+
+def consolidate_batch_outputs(batch_id):
+    """Consolidate outputs across all runs for a batch ID"""
+    print(f"Starting batch-level consolidation for {batch_id}")
+    
+    ac_dfs = []
+    b_dfs = []
+    
+    run_folders = get_all_run_folders(batch_id)
+    
+    for run_timestamp in run_folders:
+        ac_key = f"{batch_id}/{run_timestamp}/consolidated/{batch_id}-AC.csv"
+        b_key = f"{batch_id}/{run_timestamp}/consolidated/{batch_id}-B.csv"
+        
+        try:
+            response = config.S3_CLIENT.get_object(
+                Bucket=config.S3_OUTPUT_BUCKET,
+                Key=ac_key
+            )
+            ac_df = pd.read_csv(response['Body'])
+            ac_df['run_timestamp'] = run_timestamp
+            ac_dfs.append(ac_df)
+        except Exception as e:
+            print(f"Could not read AC file from run {run_timestamp}: {str(e)}")
+            
+        try:
+            response = config.S3_CLIENT.get_object(
+                Bucket=config.S3_OUTPUT_BUCKET,
+                Key=b_key
+            )
+            b_df = pd.read_csv(response['Body'])
+            b_df['run_timestamp'] = run_timestamp
+            b_dfs.append(b_df)
+        except Exception as e:
+            print(f"Could not read B file from run {run_timestamp}: {str(e)}")
+    
+    batch_ac_df = None
+    batch_b_df = None
+    
+    if ac_dfs:
+        batch_ac_df = pd.concat(ac_dfs, ignore_index=True)
+        batch_ac_df = batch_ac_df.sort_values('run_timestamp').drop_duplicates(
+            subset=['Contract Name'], 
+            keep='last'
+        ).drop('run_timestamp', axis=1)
+        
+        ac_output_path = f"{batch_id}/consolidated/{batch_id}-AC.csv"
+        csv_buffer = batch_ac_df.to_csv(index=False).encode()
+        config.S3_CLIENT.put_object(
+            Bucket=config.S3_OUTPUT_BUCKET,
+            Key=ac_output_path,
+            Body=csv_buffer
+        )
+        print(f"Saved batch AC file: {ac_output_path}")
+    
+    if b_dfs:
+        batch_b_df = pd.concat(b_dfs, ignore_index=True)
+        batch_b_df = batch_b_df.sort_values('run_timestamp').drop_duplicates(
+            subset=['Contract Name'], 
+            keep='last'
+        ).drop('run_timestamp', axis=1)
+        
+        b_output_path = f"{batch_id}/consolidated/{batch_id}-B.csv"
+        csv_buffer = batch_b_df.to_csv(index=False).encode()
+        config.S3_CLIENT.put_object(
+            Bucket=config.S3_OUTPUT_BUCKET,
+            Key=b_output_path,
+            Body=csv_buffer
+        )
+        print(f"Saved batch B file: {b_output_path}")
+    
+    # Create batch ABC if both exist
+    if batch_ac_df is not None and batch_b_df is not None:
+        try:
+            batch_b_df_copy = batch_b_df.copy()
+            batch_b_df_copy.drop(["Pages", "Parent Agreement Code"], axis=1, inplace=True)
+            
+            batch_abc_df = pd.merge(
+                batch_b_df_copy, 
+                batch_ac_df,
+                on="Contract Name",
+                how="right"
+            )
+            
+            batch_abc_df = batch_abc_df[
+                [col for col in valid.ABC_COLUMNS if col in batch_abc_df]
+            ]
+            
+            abc_output_path = f"{batch_id}/consolidated/{batch_id}-ABC.csv"
+            csv_buffer = batch_abc_df.to_csv(index=False).encode()
+            config.S3_CLIENT.put_object(
+                Bucket=config.S3_OUTPUT_BUCKET,
+                Key=abc_output_path,
+                Body=csv_buffer
+            )
+            print(f"Saved batch ABC file: {abc_output_path}")
+            
+            # Perform QAQC on batch consolidated files
+            perform_qc_qa(batch_ac_df, batch_b_df, batch_abc_df, f"{batch_id}_batch")
+            
+        except Exception as e:
+            print(f"Error creating batch ABC file: {str(e)}")
+    
+    print("Batch-level consolidation complete")
+    return batch_ac_df, batch_b_df, batch_abc_df
+    
+
+def upload_individual_outputs(run_timestamp):
+    """Upload individual AC and B output files to their respective S3 folders"""
+    BUCKET_NAME = config.S3_OUTPUT_BUCKET
+    
+    for subdir, dirs, files in os.walk(config.OUTPUT_DIRECTORY):
+        # Skip the consolidated directory
+        if subdir == config.CONSOLIDATED_OUTPUT_DIRECTORY:
+            continue
+            
+        basename = os.path.basename(subdir)
+        if basename:  # Skip root directory
+            filename = basename + ".txt"
+            
+            # If already fully processed and uploaded, skip
+            if tracking.check_file_processed(config.BATCH_ID, filename, 'AC') and \
+               tracking.check_file_processed(config.BATCH_ID, filename, 'B'):
+                print(f"Skipping {filename} - already processed")
+                continue
+                
+            individual_folder = f"{config.BATCH_ID}/{run_timestamp}/individual/{basename}"
+            
+            # Check for AC results
+            ac_path = os.path.join(subdir, config.AC_RESULTS_NAME)
+            if os.path.exists(ac_path):
+                s3_key = f"{individual_folder}/{config.AC_RESULTS_NAME}"
+                try:
+                    config.S3_CLIENT.upload_file(ac_path, BUCKET_NAME, s3_key)
+                    tracking.log_file_processing(filename, 'AC', s3_key, 'Completed')
+                    tracking.update_master_batch_tracking(
+                        config.BATCH_ID,
+                        run_timestamp,
+                        [filename],
+                        'AC',
+                        s3_path=s3_key,
+                        upload_status='Uploaded'
+                    )
+                except Exception as e:
+                    error_msg = str(e)
+                    print(f"Error uploading AC file {basename}: {error_msg}")
+                    tracking.log_file_processing(filename, 'AC', s3_key, 'Failed', error_msg)
+            
+            # Check for B results
+            b_path = os.path.join(subdir, config.B_RESULTS_NAME)
+            if os.path.exists(b_path):
+                s3_key = f"{individual_folder}/{config.B_RESULTS_NAME}"
+                try:
+                    config.S3_CLIENT.upload_file(b_path, BUCKET_NAME, s3_key)
+                    tracking.log_file_processing(filename, 'B', s3_key, 'Completed')
+                    tracking.update_master_batch_tracking(
+                        config.BATCH_ID,
+                        run_timestamp,
+                        [filename],
+                        'B',
+                        s3_path=s3_key,
+                        upload_status='Uploaded'
+                    )
+                except Exception as e:
+                    error_msg = str(e)
+                    print(f"Error uploading B file {basename}: {error_msg}")
+                    tracking.log_file_processing(filename, 'B', s3_key, 'Failed', error_msg)
+
+def consolidate_output(run_timestamp, ac_df = None, b_df = None):
+    """Consolidate and upload outputs"""
+    Path(config.CONSOLIDATED_OUTPUT_DIRECTORY).mkdir(parents=True, exist_ok=True)
+    processed_ac_files = []
+    processed_b_files = []
+    
+    ac_final_df = None if ac_df is None else ac_df
+    b_final_df = None if b_df is None else b_df
+    abc_final_df = None
+
+    # First upload individual outputs
+    if config.WRITE_TO_S3:
+        print(f"Uploading individual outputs...")
+        upload_individual_outputs(run_timestamp)

    # If AC only
-    if "a" in config.FIELDS and "c" in config.FIELDS:
+    if "a" in config.FIELDS and "c" in config.FIELDS and ac_final_df is None:
        ac_dfs = []
        for subdir, dirs, files in os.walk(config.OUTPUT_DIRECTORY):
            ac_path = os.path.join(subdir, config.AC_RESULTS_NAME)
@@ -78,18 +259,37 @@ def consolidate_atscale() -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
                try:
                    ac_df = pd.read_csv(ac_path)
                    ac_dfs.append(ac_df)
-                except:
-                    pass
-        ac_final_df = pd.concat(ac_dfs, ignore_index=True)
-
-        ac_final_df.to_csv(
-            os.path.join(
-                config.CONSOLIDATED_OUTPUT_DIRECTORY, f"{config.BATCH_ID}-AC.csv"
+                    processed_ac_files.append(os.path.basename(subdir) + ".txt")
+                except Exception as e:
+                    print(f"Error reading AC file {ac_path}: {str(e)}")
+        
+        if ac_dfs:
+            ac_final_df = pd.concat(ac_dfs, ignore_index=True)
+            output_path = os.path.join(
+                config.CONSOLIDATED_OUTPUT_DIRECTORY, 
+                f"{config.BATCH_ID}-AC.csv"
            )
-        )
+            ac_final_df.to_csv(output_path)
+            
+            if config.WRITE_TO_S3:
+                s3_key = f"{config.BATCH_ID}/{run_timestamp}/consolidated/{config.BATCH_ID}-AC.csv"
+                try:
+                    config.S3_CLIENT.upload_file(output_path, config.S3_OUTPUT_BUCKET, s3_key)
+                    tracking.log_file_processing(f"{config.BATCH_ID}-AC.csv", 'consolidated', s3_key, 'Completed')
+                    tracking.update_master_batch_tracking(
+                        config.BATCH_ID,
+                        run_timestamp,
+                        [f"{config.BATCH_ID}-AC.csv"],
+                        'consolidated',
+                        s3_path=s3_key,
+                        upload_status='Uploaded'
+                    )
+                except Exception as e:
+                    print(f"Error uploading consolidated AC file: {str(e)}")
+                    tracking.log_file_processing(f"{config.BATCH_ID}-AC.csv", 'consolidated', s3_key, 'Failed', str(e))

    # If B only
-    if "b" in config.FIELDS:
+    if "b" in config.FIELDS and b_final_df is None:
        b_dfs = []
        for subdir, dirs, files in os.walk(config.OUTPUT_DIRECTORY):
            b_path = os.path.join(subdir, config.B_RESULTS_NAME)
@@ -97,40 +297,78 @@ def consolidate_atscale() -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
                try:
                    b_df = pd.read_csv(b_path)
                    b_dfs.append(b_df)
-                except:
-                    pass
-        b_final_df = pd.concat(b_dfs, ignore_index=True)
-        final_path = os.path.join(
-                config.CONSOLIDATED_OUTPUT_DIRECTORY, f"{config.BATCH_ID}-B.csv"
+                    processed_b_files.append(os.path.basename(subdir) + ".txt")
+                except Exception as e:
+                    print(f"Error reading B file {b_path}: {str(e)}")
+                    
+        if b_dfs:
+            b_final_df = pd.concat(b_dfs, ignore_index=True)
+            output_path = os.path.join(
+                config.CONSOLIDATED_OUTPUT_DIRECTORY,
+                f"{config.BATCH_ID}-B.csv"
            )
-        b_final_df.to_csv(
-            final_path
-        )
-        print(f"file written to {final_path}")
-
-    abc_final_df = None    
+            b_final_df.to_csv(output_path)
+            
+            if config.WRITE_TO_S3:
+                s3_key = f"{config.BATCH_ID}/{run_timestamp}/consolidated/{config.BATCH_ID}-B.csv"
+                try:
+                    config.S3_CLIENT.upload_file(output_path, config.S3_OUTPUT_BUCKET, s3_key)
+                    tracking.log_file_processing(f"{config.BATCH_ID}-B.csv", 'consolidated', s3_key, 'Completed')
+                    tracking.update_master_batch_tracking(
+                        config.BATCH_ID,
+                        run_timestamp,
+                        [f"{config.BATCH_ID}-B.csv"],
+                        'consolidated',
+                        s3_path=s3_key,
+                        upload_status='Uploaded'
+                    )
+                except Exception as e:
+                    print(f"Error uploading consolidated B file: {str(e)}")
+                    tracking.log_file_processing(f"{config.BATCH_ID}-B.csv", 'consolidated', s3_key, 'Failed', str(e))
+    
    # If ABC
-    if "a" in config.FIELDS and "c" in config.FIELDS and "b" in config.FIELDS and ac_final_df is not None and b_final_df is not None:
-        b_final_df.drop(
-            ["Pages", "Parent Agreement Code"], axis=1, inplace=True
-        )  # Drop Pages so only one in result
+    if "a" in config.FIELDS and "c" in config.FIELDS and "b" in config.FIELDS:
+        try:
+            b_final_df.drop(
+                ["Pages", "Parent Agreement Code"], axis=1, inplace=True
+            )
+            abc_path = os.path.join(
+                config.CONSOLIDATED_OUTPUT_DIRECTORY, f"{config.BATCH_ID}-ABC.csv"
+            )
+            abc_final_df = pd.merge(
+                b_final_df, ac_final_df, on="Contract Name", how="right"
+            )
+            abc_final_df = abc_final_df[
+                [col for col in valid.ABC_COLUMNS if col in abc_final_df]
+            ]
+            
+            print(
+                "Columns not found: ",
+                [f for f in valid.ABC_COLUMNS if f not in abc_final_df.columns],
+            )
+            abc_final_df.to_csv(abc_path)
+            print(f"Consolidation complete. Results written to {abc_path}")

-        abc_path = os.path.join(
-            config.CONSOLIDATED_OUTPUT_DIRECTORY, f"{config.BATCH_ID}-ABC.csv"
-        )
-        abc_final_df = pd.merge(
-            b_final_df, ac_final_df, on="Contract Name", how="right"
-        )
-        abc_final_df = abc_final_df[
-            [col for col in valid.ABC_COLUMNS if col in abc_final_df]
-        ]
+            if config.WRITE_TO_S3:
+                s3_key = f"{config.BATCH_ID}/{run_timestamp}/consolidated/{config.BATCH_ID}-ABC.csv"
+                try:
+                    config.S3_CLIENT.upload_file(abc_path, config.S3_OUTPUT_BUCKET, s3_key)
+                    tracking.log_file_processing(f"{config.BATCH_ID}-ABC.csv", 'consolidated', s3_key, 'Completed')
+                except Exception as e:
+                    error_msg = str(e)
+                    print(f"Error uploading ABC file: {error_msg}")
+                    tracking.log_file_processing(f"{config.BATCH_ID}-ABC.csv", 'consolidated', s3_key, 'Failed', error_msg)
+        except Exception as e:
+            print(f"Error creating ABC consolidated file: {str(e)}")
+
+    # Upload tracking files
+    if config.WRITE_TO_S3:
+        try:
+            upload_tracking_files(run_timestamp)
+            print(f"Upload completed to s3://{config.S3_OUTPUT_BUCKET}/{config.BATCH_ID}/{run_timestamp}")
+        except Exception as e:
+            print(f"Error uploading tracking files: {str(e)}")

-        print(
-            "Columns not found: ",
-            [f for f in valid.ABC_COLUMNS if f not in abc_final_df.columns],
-        )
-        abc_final_df.to_csv(abc_path)
-        print(f"Consolidation complete. Results written to {abc_path}")

    return ac_final_df, b_final_df, abc_final_df

@@ -1,3 +1,4 @@
+
 import concurrent.futures
 import traceback
 import os
@@ -21,89 +22,259 @@ import qa_qc_helpers

 def process_b(item):
    filename, contract_text = item
+    start_time = datetime.now()
+    
    if utils.contains_reimbursement(contract_text, 0):
        try:
-            start_time = datetime.now()
            file_processing.run_b_prompts(item)
            end_time = datetime.now()

-            batch_tracker.update_file_progress(filename, 'B', start_time, end_time)
-            return item
+            output_dir = os.path.join(config.OUTPUT_DIRECTORY, os.path.splitext(filename)[0])
+            output_file = os.path.join(output_dir, config.B_RESULTS_NAME)
+            
+            if os.path.exists(output_file):
+                batch_tracker.update_file_progress(filename, 'B', start_time, end_time)
+                tracking.log_file_processing(
+                    filename, 
+                    'B', 
+                    f"{config.BATCH_ID}/processing/{filename}",
+                    'Completed'
+                )
+                return ('B', item)  
+            else:
+                error_msg = 'B output file not created'
+                batch_tracker.update_file_progress(
+                    filename, 'B', start_time, datetime.now(),
+                    status='Failed', error=error_msg
+                )
+                tracking.log_file_processing(
+                    filename, 
+                    'B',
+                    f"{config.BATCH_ID}/processing/{filename}",
+                    'Failed',
+                    error_msg
+                )
+                return None
+                
        except Exception as e:
-            print(f"Error processing item {filename}: {e}")
+            error_msg = str(e)
+            print(f"Error processing item {filename}: {error_msg}")
            traceback.print_exc()
            batch_tracker.update_file_progress(
                filename, 'B', start_time, datetime.now(),
-                status='Failed', error=str(e)
+                status='Failed', error=error_msg
            )
+            tracking.log_file_processing(
+                filename,
+                'B',
+                f"{config.BATCH_ID}/processing/{filename}",
+                'Failed',
+                error_msg
+            )
+            return None
    else:
        print(f"No reimbursement information found in {filename}. Skipping processing.")
        batch_tracker.update_file_progress(
            filename, 'B', datetime.now(), datetime.now(),
            status='Skipped', error='No reimbursement information found'
        )
+        tracking.log_file_processing(
+            filename,
+            'B',
+            f"{config.BATCH_ID}/processing/{filename}",
+            'Skipped',
+            'No reimbursement information found'
+        )
+        return None

 def process_ac(item):
    filename, contract_text = item
+    start_time = datetime.now()
+    
    try:
-        start_time = datetime.now()
        file_processing.run_ac_prompts(item)
        end_time = datetime.now()
        
-        batch_tracker.update_file_progress(filename, 'AC', start_time, end_time)
-        return item
+        output_dir = os.path.join(config.OUTPUT_DIRECTORY, os.path.splitext(filename)[0])
+        output_file = os.path.join(output_dir, config.AC_RESULTS_NAME)
+        
+        if os.path.exists(output_file):
+            batch_tracker.update_file_progress(filename, 'AC', start_time, end_time)
+            tracking.log_file_processing(
+                filename,
+                'AC',
+                f"{config.BATCH_ID}/processing/{filename}",
+                'Completed'
+            )
+            return ('AC', item)
+        else:
+            error_msg = 'AC output file not created'
+            batch_tracker.update_file_progress(
+                filename, 'AC', start_time, datetime.now(),
+                status='Failed', error=error_msg
+            )
+            tracking.log_file_processing(
+                filename,
+                'AC',
+                f"{config.BATCH_ID}/processing/{filename}",
+                'Failed',
+                error_msg
+            )
+            return None
    except Exception as e:
-        print(f"Error processing item {filename}: {e}")
+        error_msg = str(e)
+        print(f"Error processing item {filename}: {error_msg}")
        traceback.print_exc()
        batch_tracker.update_file_progress(
            filename, 'AC', start_time, datetime.now(),
-            status='Failed', error=str(e)
+            status='Failed', error=error_msg
        )
+        tracking.log_file_processing(
+            filename,
+            'AC',
+            f"{config.BATCH_ID}/processing/{filename}",
+            'Failed',
+            error_msg
+        )
+        return None

-def process_results(result, processed_count, total_files):
-    """Process individual results and update tracking"""
-    if result and result[0]:
-        filename = result[0][0]
-        print(f"Completed processing: {filename}")
+def process_results(result, filename, input_dict_b, stats):
+    """Process individual results and update tracking stats"""
+    if not result or not result[1]:
+        return stats
+    
+    proc_type, item = result
+    filename = item[0]
+    
+    output_dir = os.path.join(config.OUTPUT_DIRECTORY, os.path.splitext(filename)[0])
+    output_file = os.path.join(output_dir, 
+                              config.B_RESULTS_NAME if proc_type == 'B' else config.AC_RESULTS_NAME)
+    
+    if os.path.exists(output_file):
+        if filename not in stats['processed_files'].get(proc_type, set()):
+            if proc_type == 'B':
+                stats['completed_b'] += 1
+            else:
+                stats['completed_ac'] += 1
+            stats['processed_files'].setdefault(proc_type, set()).add(filename)
+            
+            tracking.log_file_processing(
+                filename,
+                proc_type,
+                f"{config.BATCH_ID}/processing/{filename}",
+                'Completed'
+            )
+            
+        stats['processed_count'] += 1
        
-        # Calculate and display progress
-        progress = (processed_count / total_files) * 100
-        print(f"Progress: {progress:.2f}% ({processed_count}/{total_files})")
+    return stats
+
+
+def process_remaining_files(input_dict, input_dict_b, run_timestamp, stats):
+    """Process any remaining files in self-repair mode"""
+    print("\nEntering self-repair mode...")
+    
+    
+    completed_ac, completed_b = tracking.get_processed_files(config.BATCH_ID)
+    
+    
+    stats['processed_files']['AC'] = completed_ac
+    stats['processed_files']['B'] = completed_b
+    stats['completed_ac'] = len(completed_ac)
+    stats['completed_b'] = len(completed_b)
+    
+    # Find truly remaining files
+    remaining_ac = {k: v for k, v in input_dict.items() 
+                   if k not in completed_ac}
+    remaining_b = {k: v for k, v in input_dict_b.items() 
+                  if k not in completed_b}
+    
+    if not remaining_ac and not remaining_b:
+        print("No files need repair processing.")
+        return stats
+        
+    print(f"Found {len(remaining_ac)} AC files and {len(remaining_b)} B files requiring processing")
+    
+    with concurrent.futures.ThreadPoolExecutor(max_workers=config.MAX_WORKERS) as executor:
+        futures = []
+        if remaining_b:
+            print(f"Processing {len(remaining_b)} remaining B files...")
+            futures.extend([executor.submit(process_b, item) for item in remaining_b.items()])
+        if remaining_ac:
+            print(f"Processing {len(remaining_ac)} remaining AC files...")
+            futures.extend([executor.submit(process_ac, item) for item in remaining_ac.items()])
+            
+        for future in concurrent.futures.as_completed(futures):
+            try:
+                result = future.result()
+                if result:
+                    stats = process_results(result, result[1][0], input_dict_b, stats)
+                    print(f"Repair processing completed for: {result[1][0]}")
+            except Exception as e:
+                print(f"Error in repair processing: {e}")
+                traceback.print_exc()
+    
+    print("Self-repair processing complete")
+    print("Reconsolidating outputs after repair...")
+    consolidate_output.consolidate_output(run_timestamp)
+    
+    return stats

 def main():
    global batch_tracker
    batch_tracker = BatchTracker()
    
    try:
-        # Create single timestamp for the entire run
        run_timestamp = datetime.now().strftime(f"run_%Y%m%d_%H:%M_{config.BATCH_ID}")
        
-        # Read input first
+        # Read and process input
        input_dict = utils.read_input()
        print(f"Total Input Files : {len(input_dict)}")
-        
-        # Filter B
+
+        # Filter B - files with reimbursement
        input_dict_b = {
            k: v for k, v in input_dict.items() 
            if utils.contains_reimbursement(str(v))
        }
        
-        # Set file counts after we have both numbers
        batch_tracker.set_input_file_counts(
            total_files=len(input_dict),
            reimbursement_files=len(input_dict_b)
        )
-        
-        # Now add the batch run with the correct counts
        batch_tracker.add_batch_run()
        
        if config.FILTER_ALREADY_PROCESSED:
-            input_dict_ac, input_dict_b = utils.filter_already_processed(
-                input_dict, input_dict_b
-            )
+            # Get processed files using new tracking function
+            processed_ac, processed_b = tracking.get_processed_files(config.BATCH_ID)
+            
+            # Filter AC files
+            input_dict_ac = {
+                k: v for k, v in input_dict.items() 
+                if k not in processed_ac
+            }
+            
+            # Filter B files
+            input_dict_b = {
+                k: v for k, v in input_dict_b.items() 
+                if k not in processed_b
+            }
+            print(f"Filtered out {len(processed_ac) + len(processed_b)} already processed files")
        else:
            input_dict_ac = input_dict
-
+        
+        # Calculate totals
+        total_ac = len(input_dict_ac) if "a" in config.FIELDS.lower() and "c" in config.FIELDS.lower() else 0
+        total_b = len(input_dict_b) if "b" in config.FIELDS.lower() else 0
+        
+        # Initialize tracking stats
+        stats = {
+            'completed_ac': 0,
+            'completed_b': 0,
+            'processed_count': 0,
+            'processed_files': {'AC': set(), 'B': set()}
+        }
+        
+        # Setup processing types
        files_processing_types = {}
        if "b" in config.FIELDS.lower():
            for filename in input_dict_b:
@@ -115,28 +286,21 @@ def main():
        
        total_processing_events = sum(len(types) for types in files_processing_types.values())
        print(f"Total unique files: {len(files_processing_types)}")
-        print(f"Total processing events expected: {total_processing_events}")
+        print(f"Total AC files to process: {total_ac}")
+        print(f"Total B files to process: {total_b}")
        
-        processed_count = 0
-        completed_results = []
-        futures = []
        last_upload_count = 0
-        UPLOAD_FREQUENCY = 10  # Upload every 10 files
+        UPLOAD_FREQUENCY = 10
        
+        # Process files
        with concurrent.futures.ThreadPoolExecutor(max_workers=config.MAX_WORKERS) as executor:
-            if "b" in config.FIELDS.lower():
-                if len(input_dict_b) > 0:
-                    futures.extend([
-                        executor.submit(process_b, item) 
-                        for item in input_dict_b.items()
-                    ])
+            futures = []
+            
+            if "b" in config.FIELDS.lower() and len(input_dict_b) > 0:
+                futures.extend([executor.submit(process_b, item) for item in input_dict_b.items()])

-            if "a" in config.FIELDS.lower() and "c" in config.FIELDS.lower():
-                if len(input_dict_ac) > 0:
-                    futures.extend([
-                        executor.submit(process_ac, item) 
-                        for item in input_dict_ac.items()
-                    ])
+            if "a" in config.FIELDS.lower() and "c" in config.FIELDS.lower() and len(input_dict_ac) > 0:
+                futures.extend([executor.submit(process_ac, item) for item in input_dict_ac.items()])

            if not futures:
                raise ValueError("no input files were staged for processing")
@@ -145,15 +309,23 @@ def main():
                try:
                    result = future.result()
                    if result:
-                        processed_count += 1
-                        completed_results.append(result)
-                        process_results(result, processed_count, total_processing_events)
+                        filename = result[1][0]
+                        stats = process_results(result, filename, input_dict_b, stats)
                        
-                        # Periodic upload check - update existing files
-                        if processed_count - last_upload_count >= UPLOAD_FREQUENCY:
+                        # Calculate and display progress
+                        ac_progress, b_progress = utils.calculate_progress_stats(
+                            stats['completed_ac'], total_ac,
+                            stats['completed_b'], total_b
+                        )
+                        print(f"\nProgress:")
+                        print(f"AC: {ac_progress:.2f}% ({stats['completed_ac']}/{total_ac})")
+                        print(f"B: {b_progress:.2f}% ({stats['completed_b']}/{total_b})")
+                        print(f"Overall: {(stats['processed_count'] / total_processing_events * 100):.2f}%")
+                        
+                        if stats['processed_count'] - last_upload_count >= UPLOAD_FREQUENCY:
                            try:
-                                consolidate_output.upload_tracking_files(run_timestamp)
-                                last_upload_count = processed_count
+                                consolidate_output.consolidate_output(run_timestamp)
+                                last_upload_count = stats['processed_count']
                            except Exception as e:
                                print(f"Warning: Could not update tracking data: {str(e)}")
                                    
@@ -161,28 +333,32 @@ def main():
                    print(f"Error processing future: {e}")
                    traceback.print_exc()

-        print("\nIndividual Processing Complete")
+        print("\nInitial Processing Complete")
        
        print("Starting output consolidation and QA/QC...")
        try:
-            ac_final_df, b_final_df, abc_final_df = consolidate_output.consolidate_atscale()
+            ac_final_df, b_final_df, abc_final_df = consolidate_output.consolidate_output(run_timestamp)
            wb = qa_qc_helpers.perform_qc_qa(ac_final_df, b_final_df, abc_final_df, input_dict)
            qa_qc_helpers.save_qcqa_wb(wb, run_timestamp)
            print("Consolidation and QA/QC complete")
        except Exception as e:
            print(f"Error during consolidation: {str(e)}")
        
+        # Process remaining files and get updated stats
+        stats = process_remaining_files(input_dict, input_dict_b, run_timestamp, stats)
+        
        batch_tracker.update_batch_run(end_time=datetime.now())
        
        print(f"\nFinal Statistics:")
        print(f"Total input files: {batch_tracker.total_input_files}")
        print(f"Files with reimbursement: {batch_tracker.reimbursement_files}")
-        print(f"Unique files processed: {len(files_processing_types)}")
-        print(f"Processing events completed: {processed_count}")
+        print(f"AC files processed: {stats['completed_ac']}/{total_ac}")
+        print(f"B files processed: {stats['completed_b']}/{total_b}")
+        print(f"Total processing events completed: {stats['processed_count']}/{total_processing_events}")
        
-        if processed_count > 0 and total_processing_events > 0:
-            success_rate = (processed_count / total_processing_events) * 100
-            print(f"Success rate: {success_rate:.2f}%")
+        if stats['processed_count'] > 0 and total_processing_events > 0:
+            success_rate = (stats['processed_count'] / total_processing_events) * 100
+            print(f"Overall success rate: {success_rate:.2f}%")
        
    except Exception as e:
        print(f"Error in main processing: {e}")
@@ -134,7 +134,55 @@ def TOP_DOWN_PRIMARY(page):
 The preceding text is one page of a contract between a Payer and a provider in their network. Extract the names of certain entities listed in the contract and return a single dictionary with the values below:

 Here are the attributes to be included in each dictionary, and instructions on how to correctly answer:
-'EXHIBIT' : List any exhibit, attachment, or amendment names found on the page. Write the full name of the exhibit, including the exhibit number or letter, as well as any other subtitles describing the contents of the exhibit. Do NOT write the page number. If no Exhibit, Attachment, or Amendment is found, write 'N/A'.
+'EXHIBIT': CRITICAL - Capture the complete hierarchy of document identifiers found on the page:
+- MUST include ALL of the following when present:
+  * Full Attachment names and numbers (e.g., "Attachment C: Commercial-Exchange - EPO")
+  * Complete Exhibit numbers/letters with full titles
+  * All relevant subtitles and section names
+  * Provider/Entity names when listed with exhibit information
+- Format: List each component on a new line in the exact order found on page
+- Example format for multi-level exhibits:
+  "Attachment C: Commercial-Exchange - EPO
+   EXHIBIT 1
+   COMPENSATION SCHEDULE
+   PHYSICIAN SERVICES
+   LAILA HIRJEE MD PA"
+- Rules for inclusion:
+  * Do NOT abbreviate or summarize any part of the exhibit names
+  * Do NOT include page numbers
+  * Do NOT cherry-pick only certain parts - capture the complete exhibit hierarchy
+  * Include ALL text that appears to be part of the exhibit/attachment header
+  * Keep exact capitalization and formatting as shown in document
+- Write 'N/A' only if NO exhibit, attachment, or amendment identifiers are found
+
+Example 1:
+Document text shows:
+"Attachment C: Commercial-Exchange - EPO
+EXHIBIT 1
+COMPENSATION SCHEDULE
+PHYSICIAN SERVICES
+LAILA HIRJEE MD PA"
+Correct output: "Attachment C: Commercial-Exchange - EPO
+EXHIBIT 1
+COMPENSATION SCHEDULE
+PHYSICIAN SERVICES
+LAILA HIRJEE MD PA"
+
+Example 2:
+Document text shows:
+"ATTACHMENT A - PPO
+EXHIBIT 3
+MEDICARE ADVANTAGE"
+Correct output: "ATTACHMENT A - PPO
+EXHIBIT 3
+MEDICARE ADVANTAGE"
+
+Final check before returning result:
+1. Did you capture the COMPLETE attachment/exhibit name including all descriptive text?
+2. Did you include ALL levels of the exhibit hierarchy?
+3. Did you maintain the exact formatting and line breaks as shown in the document?
+4. Did you include any provider/entity names that were part of the exhibit header?
+
 'CONTRACT_LOB' : List the Line of Business mentioned on the page. Choose ONLY from the following: {valid.VALID_LOBS}. If there are multiple on the page, write them in a comma-separated list. If no LOB is found, write 'N/A'.
 'CONTRACT_PROGRAM' : List any Programs or Plans mentioned on the page. {valid.VALID_PROGRAMS}. If no program is found, write 'N/A'.

@@ -564,32 +564,17 @@ def save_qcqa_wb(wb, run_timestamp):
    
        for filename in consolidated_files:            
            local_path = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, filename)            
-            if os.path.exists(local_path):                
-                s3_key = f"{run_timestamp}/consolidated/{filename}"                
+            if os.path.exists(local_path):
+                if filename == report_filename:
+                    s3_key=f"{config.BATCH_ID}/{run_timestamp}/consolidated/qcqa/{filename}"                  
+                else:
+                    s3_key = f"{config.BATCH_ID}/{run_timestamp}/consolidated/{filename}"                
                try:                    
                    print(f"uploading {local_path} to s3://{BUCKET_NAME}/{s3_key}")                    
                    config.S3_CLIENT.upload_file(local_path, BUCKET_NAME, s3_key)                
                except Exception as e:                    
                    print(f"Error uploading {local_path}: {str(e)}")            
            else:                
-                print(f"File not found: {local_path}")        
-    
-    
-        # Upload tracking files        
-        tracking_files = {            
-            'cost_log': config.COST_LOG_NAME,            
-            'file_log': config.FILE_LOG_NAME        
-            }        
-    
-    
-        for file_type, file_name in tracking_files.items():            
-            if os.path.exists(file_name):                
-                try:                    
-                    print(f"uploading final {file_name}")                    
-                    s3_key = os.path.join(run_timestamp, file_name).replace("\\", "/")                    
-                    config.S3_CLIENT.upload_file(file_name, BUCKET_NAME, s3_key)                
-                
-                except Exception as e:                    
-                    print(f"Error uploading {file_name}: {str(e)}")        
+                print(f"File not found: {local_path}")         
        
-        print(f"Upload completed to s3://{BUCKET_NAME}/{run_timestamp}")
+        print(f"Upload completed to s3://{BUCKET_NAME}/{config.BATCH_ID}/{run_timestamp}")
@@ -1,30 +1,23 @@
 import config
-
 import time
 import random
 import os
 import csv
-
+import pandas as pd
+from datetime import datetime

 def write_stats_to_csv(filename=config.TRACKING_FILE):
+    """Write model and global stats to CSV file with retry logic"""
    max_retries = 5
    for attempt in range(max_retries):
        try:
            with open(filename, "w", newline="") as csvfile:
                writer = csv.writer(csvfile)
-                writer.writerow(
-                    [
-                        "Filename",
-                        "Requests",
-                        "Tokens Sent",
-                        "Tokens Received",
-                        "Total Time (s)",
-                        "Avg Time per Request (s)",
-                        "Claude 2 Requests",
-                        "Claude 3 Requests",
-                        "Claude 3.5 Requests",
-                    ]
-                )
+                writer.writerow([
+                    "Filename", "Requests", "Tokens Sent", "Tokens Received",
+                    "Total Time (s)", "Avg Time per Request (s)", 
+                    "Claude 2 Requests", "Claude 3 Requests", "Claude 3.5 Requests",
+                ])

                for filename, stats in config.MODEL_STATS.items():
                    avg_time = (
@@ -32,76 +25,169 @@ def write_stats_to_csv(filename=config.TRACKING_FILE):
                        if stats["requests"] > 0
                        else 0
                    )
-                    writer.writerow(
-                        [
-                            filename,
-                            stats["requests"],
-                            stats["tokens_sent"],
-                            stats["tokens_received"],
-                            round(stats["total_time"], 2),
-                            avg_time,
-                            stats["claude2_requests"],
-                            stats["claude3_requests"],
-                            stats["claude35_requests"],
-                        ]
-                    )
+                    writer.writerow([
+                        filename, stats["requests"], stats["tokens_sent"],
+                        stats["tokens_received"], round(stats["total_time"], 2),
+                        avg_time, stats["claude2_requests"],
+                        stats["claude3_requests"], stats["claude35_requests"],
+                    ])

                # Write global stats
                writer.writerow([])
                writer.writerow(["GLOBAL STATS", "Value"])
-                writer.writerow(
-                    ["Total Requests", config.GLOBAL_STATS["total_requests"]]
-                )
-                writer.writerow(
-                    ["Total Tokens Sent", config.GLOBAL_STATS["total_tokens_sent"]]
-                )
-                writer.writerow(
-                    [
-                        "Total Tokens Received",
-                        config.GLOBAL_STATS["total_tokens_received"],
-                    ]
-                )
-                writer.writerow(
-                    ["Total Time (s)", round(config.GLOBAL_STATS["total_time"], 2)]
-                )
-                writer.writerow(
-                    [
-                        "Avg Time per Request (s)",
-                        (
-                            round(
-                                config.GLOBAL_STATS["total_time"]
-                                / config.GLOBAL_STATS["total_requests"],
-                                2,
-                            )
-                            if config.GLOBAL_STATS["total_requests"] > 0
-                            else 0
-                        ),
-                    ]
-                )
-                writer.writerow(
-                    [
-                        "Total Claude 2 Requests",
-                        config.GLOBAL_STATS["total_claude2_requests"],
-                    ]
-                )
-                writer.writerow(
-                    [
-                        "Total Claude 3 Requests",
-                        config.GLOBAL_STATS["total_claude3_requests"],
-                    ]
-                )
-                writer.writerow(
-                    [
-                        "Total Claude 3.5 Requests",
-                        config.GLOBAL_STATS["total_claude35_requests"],
-                    ]
-                )
+                for stat_name, stat_value in config.GLOBAL_STATS.items():
+                    if "time" in stat_name.lower():
+                        stat_value = round(stat_value, 2)
+                    writer.writerow([stat_name, stat_value])
                writer.writerow(["Total Files Processed", len(config.MODEL_STATS)])
            break
        except PermissionError:
            if attempt < max_retries - 1:
                time.sleep(random.uniform(0.1, 0.5))
            else:
-                print(
-                    f"Unable to write to {filename} after {max_retries} attempts. It may be open in another program."
-                )
+                print(f"Unable to write to {filename} after {max_retries} attempts. It may be open in another program.")
+
+def get_master_batch_tracking(batch_id):
+    """
+    Get master batch tracking file from S3
+    Returns DataFrame with processed files and their status
+    """
+    s3_key = f"{batch_id}/master_batch_tracking.csv"
+    try:
+        response = config.S3_CLIENT.get_object(
+            Bucket=config.S3_OUTPUT_BUCKET,
+            Key=s3_key
+        )
+        df = pd.read_csv(response['Body'])
+        # Ensure required columns exist
+        required_cols = ['filename', 'batch_id', 'run_timestamp', 'output_type', 's3_path', 'upload_status']
+        for col in required_cols:
+            if col not in df.columns:
+                df[col] = None
+        return df
+    except config.S3_CLIENT.exceptions.NoSuchKey:
+        return pd.DataFrame(
+            columns=['filename', 'batch_id', 'run_timestamp', 'output_type', 's3_path', 'upload_status']
+        )
+
+def update_master_batch_tracking(batch_id, run_timestamp, processed_files, output_type, s3_path=None, upload_status='Uploaded'):
+    """
+    Update master batch tracking with newly processed files
+    Only updates S3 when files are successfully uploaded
+    """
+    tracking_df = get_master_batch_tracking(batch_id)
+    
+    
+    if output_type == 'ABC':
+        return tracking_df
+    
+    new_records = []
+    for filename in processed_files:
+        
+        mask = (tracking_df['filename'] == filename) & \
+               (tracking_df['output_type'] == output_type) & \
+               (tracking_df['upload_status'] == 'Uploaded')
+               
+        if not mask.any():
+            new_records.append({
+                'filename': filename,
+                'batch_id': batch_id,
+                'run_timestamp': run_timestamp,
+                'output_type': output_type,
+                's3_path': s3_path or f"{batch_id}/{run_timestamp}/{output_type}",
+                'upload_status': upload_status
+            })
+    
+    if new_records:
+        new_df = pd.DataFrame(new_records)
+        tracking_df = pd.concat([tracking_df, new_df], ignore_index=True)
+        
+        
+        if upload_status == 'Uploaded':
+            csv_buffer = tracking_df.to_csv(index=False).encode()
+            s3_key = f"{batch_id}/master_batch_tracking.csv"
+            config.S3_CLIENT.put_object(
+                Bucket=config.S3_OUTPUT_BUCKET,
+                Key=s3_key,
+                Body=csv_buffer
+            )
+    
+    return tracking_df
+
+def get_processed_files(batch_id):
+    """Get list of successfully processed and uploaded files"""
+    tracking_df = get_master_batch_tracking(batch_id)
+    
+    if tracking_df.empty:
+        return set(), set()
+    
+    # Only consider successfully uploaded files
+    successful_df = tracking_df[tracking_df['upload_status'] == 'Uploaded']
+    
+    ac_files = set(successful_df[successful_df['output_type'] == 'AC']['filename'].unique())
+    b_files = set(successful_df[successful_df['output_type'] == 'B']['filename'].unique())
+    
+    return ac_files, b_files
+
+def log_file_processing(filename, process_type, s3_path, status='Completed', error=None):
+    """Log file processing details to file_log"""
+    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+    log_entry = {
+        'timestamp': timestamp,
+        'filename': filename,
+        'process_type': process_type,
+        's3_path': s3_path,
+        'status': status,
+        'error': error or 'None'
+    }
+    
+    try:
+        if os.path.exists(config.FILE_LOG_NAME):
+            df = pd.read_csv(config.FILE_LOG_NAME)
+        else:
+            df = pd.DataFrame(columns=['timestamp', 'filename', 'process_type', 's3_path', 'status', 'error'])
+        
+        df = pd.concat([df, pd.DataFrame([log_entry])], ignore_index=True)
+        df.to_csv(config.FILE_LOG_NAME, index=False)
+    except Exception as e:
+        print(f"Error writing to file log: {str(e)}")
+
+def upload_tracking_logs(run_timestamp, batch_id):
+    """Upload all tracking logs to S3"""
+    if not config.WRITE_TO_S3:
+        return
+        
+    BUCKET_NAME = config.S3_OUTPUT_BUCKET
+    
+    
+    write_stats_to_csv()
+    
+    
+    tracking_files = {
+        'cost_log': config.COST_LOG_NAME,
+        'file_log': config.FILE_LOG_NAME,
+        'tracking': config.TRACKING_FILE
+    }
+    
+    for file_type, file_name in tracking_files.items():
+        if os.path.exists(file_name):
+            try:
+                s3_key = f"{batch_id}/{run_timestamp}/tracking/{file_name}"
+                print(f"Uploading {file_type} to s3://{BUCKET_NAME}/{s3_key}")
+                config.S3_CLIENT.upload_file(file_name, BUCKET_NAME, s3_key)
+            except Exception as e:
+                print(f"Error uploading {file_name}: {str(e)}")
+                raise
+
+def check_file_processed(batch_id, filename, output_type):
+    """Check if a specific file has been successfully processed and uploaded"""
+    tracking_df = get_master_batch_tracking(batch_id)
+    
+    if tracking_df.empty:
+        return False
+        
+    return any(
+        (tracking_df['filename'] == filename) & 
+        (tracking_df['output_type'] == output_type) & 
+        (tracking_df['upload_status'] == 'Uploaded')
+    )
@@ -2,10 +2,8 @@ import os
 import re
 import pandas as pd
 import shutil
-import time
-import sys
 import numpy as np
-from collections import defaultdict
+import tracking 
 import config
 from botocore.exceptions import ClientError
 from io import StringIO
@@ -188,6 +186,11 @@ def preprocess_text_file(file_path):
    pages = re.split(r"Start of Page No\. = \d+", text)
    return pages

+def calculate_progress_stats(completed_ac, total_ac, completed_b, total_b):
+    """Calculate and return progress statistics"""
+    ac_progress = min((completed_ac / total_ac * 100), 100.0) if total_ac > 0 else 0
+    b_progress = min((completed_b / total_b * 100), 100.0) if total_b > 0 else 0
+    return ac_progress, b_progress

 def format_td_check(td_dicts, dont_include_list):
    final_str = ""
@@ -275,32 +278,56 @@ def contains_lol(text, page='1'):


 def filter_already_processed(input_dict, input_dict_b):
-    already_processed_ac, already_processed_b = [], []
-    for folder_name in os.listdir(config.OUTPUT_DIRECTORY):
-        # AC
-        if config.AC_RESULTS_NAME in os.listdir(
-            os.path.join(config.OUTPUT_DIRECTORY, folder_name)
-        ):
-            already_processed_ac.append(folder_name + ".txt")
-        # B
-        if config.B_RESULTS_NAME in os.listdir(
-            os.path.join(config.OUTPUT_DIRECTORY, folder_name)
-        ):
-            already_processed_b.append(folder_name + ".txt")
+    """Filter out already processed files based on either local files or S3 master tracking"""
+    if config.WRITE_TO_S3:
+        # Get processed files from S3 master tracking for this specific batch
+        already_processed_ac, already_processed_b = tracking.get_already_processed_files(config.BATCH_ID)
+        
+        # Keep files that still need either AC or B processing
+        input_dict_ac = {
+            k: v for k, v in input_dict.items() 
+            if k not in already_processed_ac
+        }
+        
+        input_dict_b = {
+            k: v for k, v in input_dict_b.items() 
+            if k not in already_processed_b
+        }
+        
+        print(f"Filtering based on S3 master tracking for batch {config.BATCH_ID}:")
+        print(f"AC: Found {len(already_processed_ac)} already processed, {len(input_dict_ac)} remaining")
+        print(f"B: Found {len(already_processed_b)} already processed, {len(input_dict_b)} remaining")
+        
+    else:
+        # Original local directory checking logic
+        already_processed_ac, already_processed_b = [], []
+        for folder_name in os.listdir(config.OUTPUT_DIRECTORY):
+            folder_path = os.path.join(config.OUTPUT_DIRECTORY, folder_name)
+            if os.path.isdir(folder_path):
+                if config.AC_RESULTS_NAME in os.listdir(folder_path):
+                    already_processed_ac.append(folder_name + ".txt")
+                if config.B_RESULTS_NAME in os.listdir(folder_path):
+                    already_processed_b.append(folder_name + ".txt")
+        
+        input_dict_ac = {
+            k: v for k, v in input_dict.items()
+            if k not in already_processed_ac
+        }
+        
+        input_dict_b = {
+            k: v for k, v in input_dict_b.items()
+            if k not in already_processed_b
+        }
+        
+        print(f"Filtering based on local output directory:")
+        print(f"AC: Found {len(already_processed_ac)} already processed, {len(input_dict_ac)} remaining")
+        print(f"B: Found {len(already_processed_b)} already processed, {len(input_dict_b)} remaining")

-    input_dict_ac = {
-        key: input_dict[key]
-        for key in input_dict.keys()
-        if key not in already_processed_ac
-    }
-    input_dict_b = {
-        key: input_dict_b[key]
-        for key in input_dict_b.keys()
-        if key not in already_processed_b
-    }
    return input_dict_ac, input_dict_b


+
+
 def is_empty(value):
    if pd.isna(value):
        return True
@@ -389,7 +389,6 @@ AC_MAPPING = {
    "TEMPLATE": "Template",
    "PROVIDER_BASED_BILLING_EXCLUSION_LANGUAGE_IND": "Provider-Based Billing Exclusion (Y/N)",
    "PROVIDER_BASED_BILLING_EXCLUSION_LANGUAGE": "Provider-Based Billing Exclusion (Language)",
-    'Provider-based Billing Exclusion (Language)' : "Provider-Based Billing Exclusion (Language)",
    "NON_RENEWAL_LANGUAGE": "Non-Renewal Language",
    "NON_RENEWAL_DAYS": "Non-Renewal - Days",
    "TIMELY_FILING": "Timely Filing",
@@ -509,8 +508,8 @@ ABC_COLUMNS = [
    "Medical Necessity Language (Y/N)",
    "Medical Necessity Language (Language)",
    "Template",
-    "Provider-based Billing Exclusion (Y/N)",
-    "Provider-based Billing Exclusion (Language)",
+    "Provider-Based Billing Exclusion (Y/N)",
+    "Provider-Based Billing Exclusion (Language)",
 ]


@@ -789,9 +788,7 @@ HOTFIX_MAPPING = {
    "INVOICE_PRICING_LANGUAGE": "Invoice Pricing (Language)",
    "TEMPLATE": "Template",
    "PROVIDER_BASED_BILLING_EXCLUSION_LANGUAGE_IND": "Provider-Based Billing Exclusion (Y/N)",
-    'Provider-based Billing Exclusion (Y/N)' : "Provider-Based Billing Exclusion (Y/N)",
    "PROVIDER_BASED_BILLING_EXCLUSION_LANGUAGE": "Provider-Based Billing Exclusion (Language)",
-    "Provider-based Billing Exclusion (Language)" : "Provider-Based Billing Exclusion (Language)",
    "NON_RENEWAL_LANGUAGE": "Non-Renewal Language",
    "NON_RENEWAL_DAYS": "Non-Renewal - Days",
    "TIMELY_FILING": "Timely Filing",