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.

"""

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)

                # For drilling down
                # 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-25'    # 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})

    # 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})

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