doczyai-pipelines/fieldExtraction/scripts/textract_template.py

"""
This script was written to extract .txt files from pdf for adhoc client runs.
Ensure you have permissions for API gateway, S3, Lambda function and SQS to execute this (DEVELOPER & above roles in DEV & UAT for Doczy should suffice)
If you have Analyst / Test roles and need the permissions to be elevated or policies to be updated then please ask Sannan Iqbal to make the changes.
 
The flow to execute the pipeline is:
Create batch using api req
Add files to the newly created batch along with the batch id as the tag
Construct payload to trigger pipeline
Post request to trigger pipeline with payload
The output text file can be found in the client_bucket/contract_text_file/batch_123456
And the final LLM parsed outputs can be found in client_bucket/final_output/batch_123456
"""

import os
import boto3
import requests
import json
import boto3
from datetime import datetime
import os
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import get_credentials
from botocore.session import Session
import time


# Function to upload files to S3
def upload_files_to_s3(directory, bucket, batch_id):
    contract_list = []
    for filename in os.listdir(directory):
        if filename.endswith(".pdf"):
            file_path = os.path.join(directory, filename)
            s3_key = f"contracts-landing-zone/{batch_id}/{filename}"

            # Upload file to S3
            s3_client.upload_file(file_path, bucket, s3_key)
            print(f"Uploaded {filename} to S3 bucket\n")

            # Add tags to the uploaded file
            s3_client.put_object_tagging(
                Bucket=bucket,
                Key=s3_key,
                Tagging={"TagSet": [{"Key": "BatchId", "Value": batch_id}]},
            )
            print(f"Added tags to {filename}\n")
            # Add file details to contract list
            # For now we can leave it as is since only A, C have been operationalized
            contract_list.append(
                {
                    "contract_name": filename,
                    "groups": [
                        "A",  # This needs to be dynamic in UI 1 based on what group has been selected
                        "C",
                    ],
                    "contract_source_path": s3_key,
                }
            )
    return contract_list


def create_batch(client_bucket, create_batch_url):
    myobj = {"client-bucket-name": client_bucket}

    # Call create batch API endpoint
    response = requests.post(create_batch_url, json=myobj)
    if response.status_code >= 200 and response.status_code < 300:
        try:
            new_batch_id = json.loads(json.loads(response.text)["body"])["batch_id"]
            landing_zone = json.loads(json.loads(response.text)["body"])["landing_zone"]
        except:
            print(myobj)
            print(response.text)
            new_batch_id = "failed_cases"
            landing_zone = "contracts_landing_zone"
    else:
        print(response.text)
        new_batch_id = "failed_cases"
        landing_zone = "contracts_landing_zone"
    return new_batch_id, landing_zone


def list_filtered_files(bucket_name, prefix, start_date, end_date, profile_name):
    """
    List all files in an S3 bucket filtered by date range.

    Parameters:
    - bucket_name: str, the name of the S3 bucket
    - prefix: str, the prefix (folder) in the S3 bucket
    - start_date: str, the start date in ISO 8601 format
    - end_date: str, the end date in ISO 8601 format
    - profile_name: str, the AWS profile name

    Returns:
    - list of str: the keys of the filtered files
    """
    # Parse the dates
    start_date = datetime.fromisoformat(start_date.replace("Z", "+00:00"))
    end_date = datetime.fromisoformat(end_date.replace("Z", "+00:00"))

    # Initialize a session using the specified profile
    session = boto3.Session(profile_name=profile_name)
    s3_client = session.client("s3")

    # List objects in the bucket with the specified prefix
    paginator = s3_client.get_paginator("list_objects_v2")
    page_iterator = paginator.paginate(Bucket=bucket_name, Prefix=prefix)

    # Filter files by date
    filtered_files = []
    for page in page_iterator:
        if "Contents" in page:
            for obj in page["Contents"]:
                last_modified = obj["LastModified"]
                if start_date <= last_modified <= end_date:
                    filtered_files.append(obj["Key"])

    return filtered_files


def download_files(bucket_name, file_keys, profile_name, local_directory):
    """
    Download files from an S3 bucket.

    Parameters:
    - bucket_name: str, the name of the S3 bucket
    - file_keys: list of str, the keys of the files to download
    - profile_name: str, the AWS profile name
    - local_directory: str, the local directory to download files to
    """
    # Initialize a session using the specified profile
    session = boto3.Session(profile_name=profile_name)
    s3_client = session.client("s3")

    # Ensure the local directory exists
    if not os.path.exists(local_directory):
        os.makedirs(local_directory)

    # Download each file
    for key in file_keys:
        file_name = os.path.basename(key)  # Get only the file name from the key
        local_file_path = os.path.join(local_directory, file_name)
        print(f"Downloading {key} to {local_file_path}")
        s3_client.download_file(bucket_name, key, local_file_path)

    print("Download completed.")


if __name__ == "__main__":
    # Define your variables
    s3_bucket = "doczyai-use2-u-cn1-s3-textract-processing-001"
    # batch_id = 'batch_100524101551'
    client_name = "Priority Health"
    username = "ADHOC USER"

    # These endpoints are in UAT, please change them to DEV if there are access issues with UAT
    api_endpoint = (
        "https://29gm8cek03.execute-api.us-east-2.amazonaws.com/dev/trigger-pipeline"
    )
    create_batch_url = (
        "https://29gm8cek03.execute-api.us-east-2.amazonaws.com/dev/create-batch"
    )

    session = boto3.Session(
        profile_name="temp_cred"
    )  # Change the profile name to the one you have in your .aws/credentials file
    s3_client = session.client("s3")

    # Use current directory as PDF directory
    pdf_directory = "C:\\Doczy\\Priority Health\\For_Textract\\For_Textract\\UAT- test East Paris Surgical Center copy"

    # Create batch
    batch_id, landing_zone = create_batch(s3_bucket, create_batch_url)
    print(f"Batch ID: {batch_id}")
    print(f"Landing Zone: {landing_zone}")
    # batch_id = 'batch_110624213433'
    # landing_zone = 'contracts_landing_zone/batch_110624213433/'

    if batch_id == "failed_cases":
        print("Batch creation failed. Exiting...")
        exit()
    # Upload files and get contract list
    contract_list = upload_files_to_s3(pdf_directory, s3_bucket, batch_id)

    # Create JSON object
    data = {
        "s3_bucket": s3_bucket,
        "batch_id": batch_id,
        "client_name": client_name,
        "username": username,
        "contract_list": contract_list,
    }

    # Make POST request to API
    response = requests.post(api_endpoint, json=data)

    # Print response
    print(response.status_code)
    print(response.json())

    ##################
    # Get text files

    prefix = f"contract-text-file/{batch_id}/"  # if you have a specific prefix (folder) in your bucket

    # These dates are to filter the contracts in case there are older contracts in the same batch
    start_date = "2024-06-04T00:00:00Z"  # ISO 8601 format
    end_date = "2024-06-14T23:59:59Z"  # ISO 8601 format
    profile_name = "temp_cred"
    local_directory = "C:\\Doczy\\Priority Health\\For_Textract\\For_Textract\\text otuput"  # local directory to save files

    # List filtered files
    time.sleep(
        30
    )  # Waiting for the text files to be generated, this may take longer and files may not be available after 30 seconds sometimes
    filtered_files = list_filtered_files(
        s3_bucket, prefix, start_date, end_date, profile_name
    )
    print(f"Filtered files: {len(filtered_files)}")

    # Download files
    download_files(s3_bucket, filtered_files, profile_name, local_directory)