doczyai-pipelines/ops_scripts/cost_analysis/aws_cost_automation.py

import boto3
import pandas as pd
from datetime import datetime, timedelta

"""
This script is a work in progress. 
TODO:
1. Read the DS batch tracker and get the total document counts per client and hence calculate the total bedrock cost per client. 

This script is used to automate the process of fetching AWS costs for specific services and additional charges, as well as counting the number of records in each client's S3 bucket. 
The script performs the following main steps:

1. **Get Service Costs**: Fetch costs for specified AWS services using the AWS Cost Explorer API.
2. **Get Costs by Record Type**: Fetch additional charges (e.g., Tax, Support) using the AWS Cost Explorer API.
3. **Process Bucket**: Count the number of records in each client's S3 bucket for a specified date range.
4. **Export Data**: Export the cost data and bucket record counts to CSV files. (This is currently disabled by default.)

"""

def get_service_costs(start_date, end_date, services, profile_name):
    client = boto3.Session(profile_name=profile_name).client('ce')  # Cost Explorer client
    response = client.get_cost_and_usage(
        TimePeriod={
            'Start': start_date,
            'End': end_date  # AWS Cost Explorer 'End' date is exclusive
        },
        Granularity='DAILY',
        Metrics=['UnblendedCost'],
        GroupBy=[{'Type': 'DIMENSION', 'Key': 'SERVICE'}],
        Filter={
            'Dimensions': {
                'Key': 'SERVICE',
                'Values': services
            }
        }
    )

    costs = {}
    for result in response['ResultsByTime']:
        for group in result.get('Groups', []):
            service = group['Keys'][0]
            amount = float(group['Metrics']['UnblendedCost']['Amount'])
            costs[service] = costs.get(service, 0) + amount
    return costs

def get_costs_by_record_type(start_date, end_date, record_types, profile_name):
    client = boto3.Session(profile_name=profile_name).client('ce')
    response = client.get_cost_and_usage(
        TimePeriod={'Start': start_date, 'End': end_date},
        Granularity='DAILY',
        Metrics=['UnblendedCost'],
        GroupBy=[{'Type': 'DIMENSION', 'Key': 'RECORD_TYPE'}],
        Filter={
            'Dimensions': {
                'Key': 'RECORD_TYPE',
                'Values': record_types
            }
        }
    )

    costs = {}
    for result in response['ResultsByTime']:
        for group in result.get('Groups', []):
            record_type = group['Keys'][0]
            amount = float(group['Metrics']['UnblendedCost']['Amount'])
            costs[record_type] = costs.get(record_type, 0) + amount
    return costs

def process_bucket(bucket_name, start_date_obj, end_date_obj, profile_name):
    s3 = boto3.Session(profile_name=profile_name).client('s3')
    prefix = 'processed_batch_ids/'
    paginator = s3.get_paginator('list_objects_v2')
    page_iterator = paginator.paginate(Bucket=bucket_name, Prefix=prefix)

    total_records = 0
    for page in page_iterator:
        for obj in page.get('Contents', []):
            key = obj['Key']
            last_modified = obj['LastModified']
            if key.endswith('.csv') and start_date_obj <= last_modified.date() <= end_date_obj:
                # Read the CSV file into a dataframe
                csv_obj = s3.get_object(Bucket=bucket_name, Key=key)
                df = pd.read_csv(csv_obj['Body'])
                total_records += len(df)
                # print(f"Number of records in {key}: {len(df)}")
    return total_records

def main():
    # Input parameters
    start_date = '2024-11-01'  # Replace with your start date
    end_date = '2024-11-30'    # Replace with your end date
    client_buckets = ['highmark-files', 'la-care-files',
                      'centene-fidelis-files',
                      'centene-healthnet-files',
                      'centene-national-contracting-files',
                      'centene-texas-files',
                      'bcbs-files',
                      'community-health-choice-files',
                      'humana-files',
                      'scan-health-files',
                      'texas-children-files',
                      'village-care']
    profile_name = 'doczy_uat'

    # List of services to fetch costs for
    services = [
        'Amazon Textract',
        'Amazon Bedrock',
        'Claude Instant (Amazon Bedrock Edition)',
        'Claude 3 Haiku (Amazon Bedrock Edition)',
        'Claude 3.5 Sonnet (Amazon Bedrock Edition)'
    ]

    # Convert date strings to datetime objects
    start_date_obj = datetime.strptime(start_date, '%Y-%m-%d').date()
    end_date_obj = datetime.strptime(end_date, '%Y-%m-%d').date()

    # Adjust end_date for AWS Cost Explorer API (end date is exclusive)
    end_date_ce = (datetime.strptime(end_date, '%Y-%m-%d') + timedelta(days=1)).strftime('%Y-%m-%d')

    # Get costs for specified services
    service_costs = get_service_costs(start_date, end_date_ce, services, profile_name)

    # Get charges for Tax and Support
    tax_costs = get_costs_by_record_type(start_date, end_date_ce, ['Tax'], profile_name)
    support_costs = get_costs_by_record_type(start_date, end_date_ce, ['Support'], profile_name)

    # Prepare data for CSV export
    cost_data = []

    # Output costs
    print("AWS Service Costs:")
    for service in services:
        cost = service_costs.get(service, 0.0)
        print(f"  {service}: ${cost:.2f}")
        cost_data.append({'Category': 'Service', 'Item': service, 'Cost': cost})

    print("\nAdditional Charges:")
    tax_cost = tax_costs.get('Tax', 0.0)
    support_cost = support_costs.get('Support', 0.0)
    print(f"  Tax: ${tax_cost:.2f}")
    print(f"  Support: ${support_cost:.2f}")
    cost_data.append({'Category': 'Additional Charges', 'Item': 'Tax', 'Cost': tax_cost})
    cost_data.append({'Category': 'Additional Charges', 'Item': 'Support', 'Cost': support_cost})

    total_costs = sum([cost['Cost'] for cost in cost_data])
    cost_data.append({'Category': 'Total', 'Item': 'Total', 'Cost': total_costs})

    # Adding up all costs that are not Amazon Textract and other costs as bedrock costs
    # then based on the costs, we split the support and tax costs respective to the textract vs bedrock costs

    textract_costs = service_costs.get('Amazon Textract', 0.0)
    bedrock_costs = total_costs - textract_costs

    support_costs_textract = support_costs.get('Support', 0.0) * (textract_costs / total_costs)
    support_costs_bedrock = support_costs.get('Support', 0.0) * (bedrock_costs / total_costs)

    tax_costs_textract = tax_costs.get('Tax', 0.0) * (textract_costs / total_costs)
    tax_costs_bedrock = tax_costs.get('Tax', 0.0) * (bedrock_costs / total_costs)

    print(f"Pro-rated Support + Tax costs for Textract based on total costs: ${support_costs_textract + tax_costs_textract:.2f}")
    print(f"Pro-rated Support + Tax costs for Bedrock based on total costs: ${support_costs_bedrock + tax_costs_bedrock:.2f}")


    total_textract_costs = textract_costs + support_costs_textract + tax_costs_textract
    total_bedrock_costs = bedrock_costs + support_costs_bedrock + tax_costs_bedrock

    print("\nAWS Service Costs pro-rated taxes by service :")
    print(f" TOTAL Amazon Textract: ${total_textract_costs:.2f}")
    print(f" TOTAL Amazon Bedrock: ${total_bedrock_costs:.2f}")
    
    # Process each bucket and count records
    print("\nBucket-wise Record Counts:")
    bucket_data = []
    for bucket in client_buckets:
        total_records = process_bucket(bucket, start_date_obj, end_date_obj, profile_name)
        print(f"  {bucket}: {total_records} records")
        bucket_data.append({'Bucket': bucket, 'Records': total_records})

    # Combine all record counts for buckets and calcuate the percent of total for each bucket. Ingoring buckets with 0 records

    total_records = sum([record['Records'] for record in bucket_data])
    for record in bucket_data:
        if record['Records'] > 0:
            record['Percent of Total'] = (record['Records'] / total_records) * 100

    
    # Printing the record weightage for each bucket 
    print("\nBucket-wise Record Weightage:")

    for record in bucket_data:
        if record['Records'] > 0:
            print(f"  {record['Bucket']}: {record['Percent of Total']:.2f}%")
            # Calcualting the textract cost with respect to the percent of total 
            textract_cost_bucket = total_textract_costs * record['Percent of Total'] / 100
            print(f"  {record['Bucket']}: ${textract_cost_bucket:.2f}")

    print(f'Total textract documents from buckets: {total_records}') 


    # Option to export data as CSV
    export_csv = True  # Set to True to export CSV
    if export_csv:
        # Export cost data
        cost_df = pd.DataFrame(cost_data)
        # cost_df.to_csv('aws_costs.csv', index=False)
        # print("\nService costs exported to 'aws_costs.csv'.")

        # Export bucket data
        bucket_df = pd.DataFrame(bucket_data)
        # bucket_df.to_csv('bucket_records.csv', index=False)
        # print("Bucket record counts exported to 'bucket_records.csv'.")

if __name__ == '__main__':
    main()