dtSpark 1.0.4__py3-none-any.whl

dtSpark/aws/costs.py

@@ -0,0 +1,318 @@
+"""
+Cost tracking module for AWS Bedrock usage.
+
+This module provides functionality for:
+- Retrieving Bedrock usage costs from AWS Cost Explorer
+- Breaking down costs by model/usage type
+- Reporting on monthly and daily costs
+"""
+
+import logging
+from datetime import datetime, timedelta
+from typing import Dict, Optional, List
+from botocore.exceptions import ClientError
+
+
+class CostTracker:
+    """Tracks AWS Bedrock costs using Cost Explorer API."""
+
+    def __init__(self, cost_explorer_client):
+        """
+        Initialise the cost tracker.
+
+        Args:
+            cost_explorer_client: Boto3 Cost Explorer client
+        """
+        self.ce_client = cost_explorer_client
+        self.has_permissions = None  # Cache permission check
+
+    def check_permissions(self) -> bool:
+        """
+        Check if the authenticated user has Cost Explorer permissions.
+
+        Returns:
+            True if user has permissions, False otherwise
+        """
+        if self.has_permissions is not None:
+            return self.has_permissions
+
+        try:
+            # Try a minimal query to check permissions
+            end_date = datetime.now().date()
+            start_date = end_date - timedelta(days=1)
+
+            self.ce_client.get_cost_and_usage(
+                TimePeriod={
+                    'Start': start_date.strftime('%Y-%m-%d'),
+                    'End': end_date.strftime('%Y-%m-%d')
+                },
+                Granularity='DAILY',
+                Metrics=['UnblendedCost']
+            )
+
+            self.has_permissions = True
+            logging.info("Cost Explorer permissions verified")
+            return True
+
+        except ClientError as e:
+            error_code = e.response.get('Error', {}).get('Code', '')
+            if error_code in ['AccessDeniedException', 'UnauthorizedException']:
+                logging.info("User does not have Cost Explorer permissions")
+                self.has_permissions = False
+                return False
+            else:
+                logging.warning(f"Error checking Cost Explorer permissions: {e}")
+                self.has_permissions = False
+                return False
+        except Exception as e:
+            logging.warning(f"Unexpected error checking permissions: {e}")
+            self.has_permissions = False
+            return False
+
+    def get_bedrock_costs(self) -> Optional[Dict]:
+        """
+        Retrieve Bedrock costs for current month, last month and last 24 hours.
+
+        Returns:
+            Dictionary with cost information, or None if unavailable
+        """
+        if not self.check_permissions():
+            return None
+
+        try:
+            # Calculate date ranges
+            today = datetime.now().date()
+            yesterday = today - timedelta(days=1)
+
+            # Current month (from 1st of this month to today)
+            first_day_this_month = today.replace(day=1)
+
+            # Last month (from 1st to last day of previous month)
+            last_day_last_month = first_day_this_month - timedelta(days=1)
+            first_day_last_month = last_day_last_month.replace(day=1)
+
+            # Get current month's costs (month-to-date)
+            current_month_costs = self._get_costs_for_period(
+                first_day_this_month,
+                today + timedelta(days=1),  # End date is exclusive, so add 1 day to include today
+                'Current Month'
+            )
+
+            # Get last month's costs
+            last_month_costs = self._get_costs_for_period(
+                first_day_last_month,
+                first_day_this_month,  # End date is exclusive
+                'Last Month'
+            )
+
+            # Get last 24 hours costs
+            last_24h_costs = self._get_costs_for_period(
+                yesterday,
+                today,
+                'Last 24 Hours'
+            )
+
+            return {
+                'current_month': current_month_costs,
+                'last_month': last_month_costs,
+                'last_24h': last_24h_costs,
+                'currency': 'USD'  # Cost Explorer returns USD by default
+            }
+
+        except Exception as e:
+            logging.error(f"Error retrieving Bedrock costs: {e}")
+            return None
+
+    def _get_costs_for_period(self, start_date, end_date, period_name: str) -> Dict:
+        """
+        Get Bedrock costs for a specific time period with model breakdown.
+
+        Args:
+            start_date: Start date (inclusive)
+            end_date: End date (exclusive)
+            period_name: Human-readable period name for logging
+
+        Returns:
+            Dictionary with total cost and breakdown by usage type
+        """
+        try:
+            # Query 1: Get total Bedrock costs
+            # Note: AWS Cost Explorer uses model-specific service names for Bedrock
+            # We search for services containing "Amazon Bedrock Edition"
+            response = self.ce_client.get_cost_and_usage(
+                TimePeriod={
+                    'Start': start_date.strftime('%Y-%m-%d'),
+                    'End': end_date.strftime('%Y-%m-%d')
+                },
+                Granularity='MONTHLY' if period_name in ['Last Month', 'Current Month'] else 'DAILY',
+                Metrics=['UnblendedCost'],
+                GroupBy=[{
+                    'Type': 'DIMENSION',
+                    'Key': 'SERVICE'
+                }]
+            )
+
+            # Extract total cost and breakdown by filtering for Bedrock services
+            # AWS uses service names like "Claude 3.5 Sonnet (Amazon Bedrock Edition)"
+            total_cost = 0.0
+            breakdown = {}
+
+            if response.get('ResultsByTime'):
+                for result in response['ResultsByTime']:
+                    for group in result.get('Groups', []):
+                        service_name = group.get('Keys', [''])[0]
+
+                        # Filter for Bedrock services (contain "Amazon Bedrock Edition")
+                        if 'Amazon Bedrock Edition' in service_name or 'Bedrock' in service_name:
+                            amount = float(group.get('Metrics', {}).get('UnblendedCost', {}).get('Amount', '0'))
+
+                            if amount > 0:
+                                total_cost += amount
+
+                                # Extract model name from service name
+                                # Example: "Claude 3.5 Sonnet (Amazon Bedrock Edition)" -> "Claude 3.5 Sonnet"
+                                model_name = service_name.replace(' (Amazon Bedrock Edition)', '').strip()
+
+                                if model_name in breakdown:
+                                    breakdown[model_name] += amount
+                                else:
+                                    breakdown[model_name] = amount
+
+            logging.info(f"{period_name} Bedrock costs: ${total_cost:.4f}")
+            if breakdown:
+                logging.debug(f"{period_name} breakdown: {breakdown}")
+
+            return {
+                'total': total_cost,
+                'breakdown': breakdown,
+                'period': period_name,
+                'start_date': start_date.strftime('%Y-%m-%d'),
+                'end_date': end_date.strftime('%Y-%m-%d')
+            }
+
+        except Exception as e:
+            logging.error(f"Error getting costs for {period_name}: {e}")
+            return {
+                'total': 0.0,
+                'breakdown': {},
+                'period': period_name,
+                'error': str(e)
+            }
+
+    def _parse_model_from_usage_type(self, usage_type: str) -> str:
+        """
+        Parse model name from AWS usage type string.
+
+        Usage types typically look like:
+        - "APS2-ModelInference-Claude-3-5-Sonnet-v2"
+        - "USE1-ModelInference-Titan-Text-Express"
+        - "APS2-OnDemand-Throughput"
+
+        Args:
+            usage_type: AWS usage type string
+
+        Returns:
+            Human-readable model name
+        """
+        # Remove region prefix (e.g., "APS2-", "USE1-")
+        parts = usage_type.split('-')
+
+        if 'ModelInference' in usage_type:
+            # Find the index of ModelInference and take everything after it
+            try:
+                idx = parts.index('ModelInference')
+                model_parts = parts[idx + 1:]
+
+                # Join and clean up
+                model_name = ' '.join(model_parts)
+
+                # Common model name mappings for readability
+                replacements = {
+                    'Claude 3 5 Sonnet v2': 'Claude 3.5 Sonnet v2',
+                    'Claude 3 5 Sonnet': 'Claude 3.5 Sonnet',
+                    'Claude 3 Opus': 'Claude 3 Opus',
+                    'Claude 3 Haiku': 'Claude 3 Haiku',
+                    'Titan Text Express': 'Titan Text Express',
+                    'Titan Text Lite': 'Titan Text Lite',
+                }
+
+                for old, new in replacements.items():
+                    if old in model_name:
+                        model_name = new
+                        break
+
+                return model_name
+            except (ValueError, IndexError):
+                pass
+
+        # If we can't parse it, return the usage type
+        return usage_type
+
+    def format_cost_report(self, costs: Optional[Dict]) -> List[str]:
+        """
+        Format cost information into human-readable lines.
+
+        Args:
+            costs: Cost dictionary from get_bedrock_costs()
+
+        Returns:
+            List of formatted strings for display
+        """
+        if not costs:
+            return ["Financial tracking not available with current user permissions"]
+
+        lines = []
+        currency = costs.get('currency', 'USD')
+
+        # Current Month
+        current_month = costs.get('current_month', {})
+        if current_month:
+            total = current_month.get('total', 0.0)
+            breakdown = current_month.get('breakdown', {})
+            period = current_month.get('period', 'Current Month')
+
+            lines.append(f"{period}: ${total:.2f} {currency}")
+
+            if breakdown and total > 0:
+                # Sort by cost (descending)
+                sorted_breakdown = sorted(breakdown.items(), key=lambda x: x[1], reverse=True)
+                for model, cost in sorted_breakdown:
+                    percentage = (cost / total * 100) if total > 0 else 0
+                    lines.append(f"  • {model}: ${cost:.2f} ({percentage:.1f}%)")
+
+        # Last Month
+        last_month = costs.get('last_month', {})
+        if last_month:
+            total = last_month.get('total', 0.0)
+            breakdown = last_month.get('breakdown', {})
+            period = last_month.get('period', 'Last Month')
+
+            lines.append(f"{period}: ${total:.2f} {currency}")
+
+            if breakdown and total > 0:
+                # Sort by cost (descending)
+                sorted_breakdown = sorted(breakdown.items(), key=lambda x: x[1], reverse=True)
+                for model, cost in sorted_breakdown:
+                    percentage = (cost / total * 100) if total > 0 else 0
+                    lines.append(f"  • {model}: ${cost:.2f} ({percentage:.1f}%)")
+
+        # Last 24 Hours
+        last_24h = costs.get('last_24h', {})
+        if last_24h:
+            total = last_24h.get('total', 0.0)
+            breakdown = last_24h.get('breakdown', {})
+            period = last_24h.get('period', 'Last 24 Hours')
+
+            lines.append(f"{period}: ${total:.4f} {currency}")
+
+            if breakdown and total > 0:
+                # Sort by cost (descending)
+                sorted_breakdown = sorted(breakdown.items(), key=lambda x: x[1], reverse=True)
+                for model, cost in sorted_breakdown:
+                    percentage = (cost / total * 100) if total > 0 else 0
+                    lines.append(f"  • {model}: ${cost:.4f} ({percentage:.1f}%)")
+
+        if not lines:
+            lines.append("No Bedrock usage costs found for the specified periods")
+
+        return lines