diff --git a/ml/billing_agent.py b/ml/billing_agent.py index 8f9cf1d..660d4ea 100644 --- a/ml/billing_agent.py +++ b/ml/billing_agent.py @@ -1,7 +1,6 @@ import os import datetime from google.cloud import bigquery -from google.cloud import billing class BillingAgent: def __init__(self): @@ -9,13 +8,11 @@ class BillingAgent: 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.billing_table = os.environ.get("BILLING_TABLE") + if not self.billing_table: + raise ValueError("BILLING_TABLE 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): """ @@ -29,6 +26,7 @@ class BillingAgent: WHERE _PARTITIONTIME >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 30 DAY) GROUP BY 1 ORDER BY total_cost DESC + LIMIT 10 """ query_job = self.bq_client.query(query) results = query_job.result() @@ -88,24 +86,7 @@ class BillingAgent: "total_monthly_forecast": total_forecast, } - def get_credits(self): - """ - Retrieves active billing 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 { - "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__': agent = BillingAgent() print("Summary:", agent.get_summary()) print("Forecast:", agent.get_forecast()) - print("Credits:", agent.get_credits())