diff --git a/ml/billing_agent.py b/ml/billing_agent.py index a4ae12e..8f9cf1d 100644 --- a/ml/billing_agent.py +++ b/ml/billing_agent.py @@ -1,4 +1,6 @@ import os +import datetime +from google.cloud import bigquery from google.cloud import billing class BillingAgent: @@ -6,41 +8,100 @@ class BillingAgent: self.project_id = os.environ.get("GOOGLE_CLOUD_PROJECT") if not self.project_id: raise ValueError("GOOGLE_CLOUD_PROJECT environment variable not set.") + + self.billing_account_id = os.environ.get("BILLING_ACCOUNT_ID") + if not self.billing_account_id: + raise ValueError("BILLING_ACCOUNT_ID environment variable not set.") + + self.bq_client = bigquery.Client(project=self.project_id) self.billing_client = billing.CloudBillingClient() + self.billing_table = f"{self.project_id}.billing_export.gcp_billing_export_v1_*" def get_summary(self): """ Retrieves the billing summary for the last 30 days. """ - # This is a simplified example. The Cloud Billing API does not have a direct - # "summary" endpoint. You would typically query for cost data over a time range. - # This implementation will be a placeholder. - return { - "summary": "Billing summary data would be here.", - "project_id": self.project_id - } + query = f""" + SELECT + service.description as service, + SUM(cost) as total_cost + FROM `{self.billing_table}` + WHERE _PARTITIONTIME >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 30 DAY) + GROUP BY 1 + ORDER BY total_cost DESC + """ + query_job = self.bq_client.query(query) + results = query_job.result() + + summary = [] + for row in results: + summary.append({"service": row.service, "total_cost": row.total_cost}) + + return {"summary": summary} def get_forecast(self): """ Retrieves a billing forecast based on the last 7 days. """ - # The Cloud Billing API does not provide a direct forecast API. - # This would require custom logic to calculate. - # This implementation will be a placeholder. + # Get total cost for the last 7 days + query_last_7_days = f""" + SELECT SUM(cost) as total_cost + FROM `{self.billing_table}` + WHERE _PARTITIONTIME >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 7 DAY) + """ + query_job = self.bq_client.query(query_last_7_days) + results = query_job.result() + + total_cost_last_7_days = 0 + for row in results: + total_cost_last_7_days = row.total_cost or 0 + + daily_average = total_cost_last_7_days / 7 + + # Calculate remaining days in the month + today = datetime.date.today() + last_day_of_month = datetime.date(today.year, today.month, 1) + datetime.timedelta(days=32) + last_day_of_month = last_day_of_month.replace(day=1) - datetime.timedelta(days=1) + remaining_days = (last_day_of_month - today).days + + forecasted_cost = daily_average * remaining_days + + # Get current month-to-date cost + query_mtd = f""" + SELECT SUM(cost) as total_cost + FROM `{self.billing_table}` + WHERE EXTRACT(MONTH FROM _PARTITIONTIME) = EXTRACT(MONTH FROM CURRENT_DATE()) + AND EXTRACT(YEAR FROM _PARTITIONTIME) = EXTRACT(YEAR FROM CURRENT_DATE()) + """ + query_job_mtd = self.bq_client.query(query_mtd) + results_mtd = query_job_mtd.result() + mtd_cost = 0 + for row in results_mtd: + mtd_cost = row.total_cost or 0 + + total_forecast = mtd_cost + forecasted_cost + return { - "forecast": "Billing forecast data would be here.", - "project_id": self.project_id + "daily_average_last_7_days": daily_average, + "month_to_date_cost": mtd_cost, + "forecasted_remaining_cost": forecasted_cost, + "total_monthly_forecast": total_forecast, } def get_credits(self): """ Retrieves active billing credits. """ - # This is a simplified example. You would need to query for billing account - # information and then check for credits. + billing_account_name = f"billingAccounts/{self.billing_account_id}" + billing_info = self.billing_client.get_billing_account(name=billing_account_name) + + # This is a simplified representation. The actual credit information might + # be structured differently. This is a placeholder for the structure. return { - "credits": "Billing credits information would be here.", - "project_id": self.project_id + "billing_account_name": billing_info.name, + "display_name": billing_info.display_name, + "open": billing_info.open, + "credits": "Detailed credit information would require further API calls and parsing." } if __name__ == '__main__':