OSVauco/ml/billing_agent.py

import os
import datetime
from google.cloud import bigquery

class BillingAgent:
    def __init__(self):
        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_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)

    def get_summary(self):
        """
        Retrieves the billing summary for the last 30 days.
        """
        query = f"""
            SELECT
                service.description as service,
                SUM(cost) + SUM(IFNULL((SELECT SUM(c.amount) FROM UNNEST(credits) c), 0)) as total_cost
            FROM `{self.billing_table}`
            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()
        
        summary = []
        for row in results:
            summary.append({"service": row.service, "total_cost": row.total_cost})
            
        if not summary:
            return {
                "onboarding_status": {
                    "state": "awaiting_data",
                    "message": "Første fakturaeksport fra BigQuery kan ta 24-48 timer. Data er på vei."
                }
            }
            
        return {"summary": summary}

    def get_forecast(self):
        """
        Retrieves a billing forecast based on the last 7 days.
        """
        query_last_7_days = f"""
            SELECT SUM(cost) + SUM(IFNULL((SELECT SUM(c.amount) FROM UNNEST(credits) c), 0)) 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

        today = datetime.date.today()
        if today.month == 12:
            next_month_first_day = datetime.date(today.year + 1, 1, 1)
        else:
            next_month_first_day = datetime.date(today.year, today.month + 1, 1)
        last_day_of_month = next_month_first_day - datetime.timedelta(days=1)
        remaining_days = (last_day_of_month - today).days

        forecasted_cost = daily_average * remaining_days
        
        query_mtd = f"""
            SELECT SUM(cost) + SUM(IFNULL((SELECT SUM(c.amount) FROM UNNEST(credits) c), 0)) 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 {
            "daily_average_last_7_days": daily_average,
            "month_to_date_cost": mtd_cost,
            "forecasted_remaining_cost": forecasted_cost,
            "total_monthly_forecast": total_forecast,
            "data_note": "Prognose basert på siste 7 dager. Fakturadata fra BigQuery kan ha 24-48 timers forsinkelse."
        }

    def get_anomalies(self):
        """
        Detects anomalies in billing data by comparing today's cost to the 7-day average.
        """
        query = f"""
            WITH daily_costs AS (
                SELECT
                    service.description AS service,
                    DATE(_PARTITIONTIME) AS usage_date,
                    SUM(cost) + SUM(IFNULL((SELECT SUM(c.amount) FROM UNNEST(credits) c), 0)) AS daily_cost
                FROM `{self.billing_table}`
                WHERE _PARTITIONTIME >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 14 DAY)
                GROUP BY 1, 2
            ),
            costs_with_avg AS (
                SELECT
                    service,
                    usage_date,
                    daily_cost,
                    AVG(daily_cost) OVER (PARTITION BY service ORDER BY usage_date ROWS BETWEEN 7 PRECEDING AND 1 PRECEDING) AS avg_7day
                FROM daily_costs
            )
            SELECT
                service,
                daily_cost AS today_cost,
                avg_7day,
                (daily_cost / avg_7day) AS ratio
            FROM costs_with_avg
            WHERE usage_date = CURRENT_DATE()
              AND avg_7day > 0
              AND daily_cost > (2.0 * avg_7day)
        """
        query_job = self.bq_client.query(query)
        results = query_job.result()

        anomalies = []
        for row in results:
            anomalies.append({
                "service": row.service,
                "today_cost": row.today_cost,
                "avg_7d": row.avg_7day,
                "ratio": row.ratio
            })

        return {"anomalies": anomalies}

if __name__ == '__main__':
    agent = BillingAgent()
    print("Summary:", agent.get_summary())
    print("Forecast:", agent.get_forecast())
    print("Anomalies:", agent.get_anomalies())